Compare commits
16
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Sanderhoff-alt
|
e4e5f8b285 | ||
|
Nicolò Boschi
|
6582e26ef9 | ||
|
Nicolò Boschi
|
188eaa3dc7 | ||
|
Nicolò Boschi
|
39de3974ea | ||
|
Nicolò Boschi
|
efef4fa398 | ||
|
Nicolò Boschi
|
e11a59ff64 | ||
|
Nicolò Boschi
|
0de91b8b73 | ||
|
Nicolò Boschi
|
31c1aaf213 | ||
|
github-actions[bot]
|
9db22115a4 | ||
|
Nicolò Boschi
|
8b78b4ac04 | ||
|
Nicolò Boschi
|
edd0d0c5bb | ||
|
Nicolò Boschi
|
5f137bf391 | ||
|
Nicolò Boschi
|
6692e38c80 | ||
|
Nicolò Boschi
|
df8ac42b52 | ||
|
Nicolò Boschi
|
d98990b46e | ||
|
Nicolò BoschiandSanderhoff-alt
|
2e2dfe1309 |
@@ -77,6 +77,7 @@ results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
- **No direct database access in `api/http.py`** (or any API router). HTTP handlers must not build SQL, call `acquire_with_retry` / `conn.fetch` / `conn.fetchrow` / `conn.execute`, or reference `fq_table(...)`. All persistence and queries live in `MemoryEngine` (the engine layer). A handler parses/validates the request, calls an engine method, shapes the HTTP response, and maps domain results to status codes (e.g. a `None` return → 404).
|
||||
- **Authentication/tenancy is enforced inside each engine method, not assumed by the handler.** Every engine method that touches bank-scoped data must authenticate via `request_context` — typically `await self._authenticate_tenant(request_context)` (often indirectly through `get_bank_profile(...)`) — so the correct tenant schema is resolved before any query runs. Handlers must thread `request_context` through to the engine method; never query a tenant-scoped table assuming the schema is already set.
|
||||
- Engine methods return typed models (Pydantic/dataclass), not raw dicts (see Type Safety).
|
||||
- **Every list endpoint paginates, following the existing ones.** A `GET` that returns a collection whose size grows with the data (banks, documents, memories, entities, operations, webhook deliveries, audit logs, …) must take `limit`/`offset` and bound its result — an unbounded list is an unbounded payload plus unbounded per-row work (per-item counts, config resolution, embedding hydration). Copy the shape `list_documents` uses, don't invent a new one: `limit: int = Query(default=100, ge=0)` and `offset: int = Query(default=0, ge=0)` on the handler, matching keyword args on the engine method, and a response carrying the page **plus `total`, `limit`, `offset`** so a client knows when to stop. Add a `q` search param when the collection is something a user picks from in a UI — client-side filtering only ever sees the loaded page. Bounded-by-construction endpoints are the exception, not the rule: a tree/export that is whole-structure by design, or a table capped at write time (e.g. `observation_history` / `mental_model_history`, trimmed to `*_max_entries` on insert). If it isn't bounded, paginate it.
|
||||
|
||||
### Bank/Tenant Isolation in Queries
|
||||
- **Bank isolation is a hard security invariant: no query may read, count, update, or delete another bank's rows.** Tenant isolation is enforced at the schema level (the resolved `search_path` / `fq_table(...)` qualifier, gated by `_authenticate_tenant`); bank isolation is enforced *within* a schema by a `bank_id` predicate on every statement that touches a multi-bank table.
|
||||
@@ -163,6 +164,28 @@ Flag any new logic that lacks test coverage.
|
||||
|
||||
See CLAUDE.md → Key Conventions → Testing for the full pattern.
|
||||
|
||||
### 6a. Check tests assert memory state via the engine API, not raw SQL
|
||||
|
||||
Tests must verify what a retain / recall / consolidation produced by calling the public
|
||||
`MemoryEngine` read API — `list_memory_units` (units and their `metadata` / `tags`; counts via
|
||||
`total`; `document_id` / `fact_type` / `entity_id` filters), `list_entities` (canonical names,
|
||||
mention counts), `get_graph_data` (nodes/edges), `get_bank_stats`, `recall_async` — **not** by
|
||||
reaching into the memory tables (`memory_units`, `memory_links`, `unit_entities`) with raw SQL via
|
||||
`pool.acquire()` / `conn.fetch*`. Asserting on those tables couples the test to a storage-layer
|
||||
detail and checks a proxy instead of the observable property (see **General Principles** → tests
|
||||
assert the property, and the handler rule in **7b**).
|
||||
|
||||
**Flag as should fix** any added or changed test whose assertion runs a `SELECT` / `COUNT` against
|
||||
`memory_units` / `memory_links` / `unit_entities` where an engine read method returns the same
|
||||
fact. Prime tell: `async with pool.acquire() as conn:` followed by `SELECT ... FROM memory_units`
|
||||
inside a test body; a `fetchval("SELECT count(*) FROM memory_units ...")` that `list_memory_units`
|
||||
`["total"]` would return; a `canonical_name` query that `list_entities` covers.
|
||||
|
||||
Direct SQL on those tables is legitimate **only** when it forces or inspects internal state the
|
||||
public API cannot express — e.g. an `UPDATE documents SET updated_at` that forges a race, or a
|
||||
raw `memory_links` row-count that the deduped `get_graph_data` edge list cannot reproduce. Those
|
||||
must carry a comment saying why the direct access is necessary; flag any that do not.
|
||||
|
||||
### 7. Check API consistency
|
||||
|
||||
If any files in `hindsight-api-slim/hindsight_api/api/` were changed:
|
||||
@@ -199,6 +222,15 @@ For each statement against a multi-bank table (`memory_units`, `documents`, `ent
|
||||
|
||||
**Flag as a must fix** any statement filtering a multi-bank table by a caller-supplied, non-globally-unique key (`document_id`, `mental_models.id`, an entity name, …) with **no** `bank_id` predicate — construct the concrete two-bank scenario (two banks share the id; the statement reads/counts/updates/deletes the wrong bank's rows or over-reports) to confirm it's real before flagging. Prime tells: a `bank_id`-carrying sibling statement right next to a `bank_id`-less one; a `WHERE bank_id` guarded by `if bank_id:` with a `None` default; an import/transfer write that inherits a source `bank_id` instead of pinning the destination.
|
||||
|
||||
### 7d. Check list endpoints paginate
|
||||
|
||||
For every added or changed `GET` handler that returns a collection, confirm it takes `limit`/`offset` and returns `total` — see **API Layer & Data Access** above for the exact shape. Then check the fix is real end to end, since a param that nothing enforces is worse than none:
|
||||
|
||||
- **The bound reaches the work, not just the response.** Verify the page size actually limits the expensive part — the SQL `LIMIT`/`OFFSET`, or (when paging must happen after an in-process filter, as in `list_banks` where the `filter_bank_list` extension hook can drop any bank) an explicit slice with the per-item work — live store counts, `get_bank_configs`, re-embedding — done for the page only. Paging in SQL *before* a filter that can drop rows is a **must fix**: it hands back short or empty pages and a `total` that counts rows the caller can't see.
|
||||
- **Every in-repo consumer pages.** A new default `limit` silently truncates callers that used to get everything: the control plane (`src/lib/api.ts` + the `src/app/api/` proxy route + any context/selector that holds the full list), the CLI (`hindsight-cli/src/api.rs`), MCP tools, and the Zapier dynamic dropdowns. Each must either page through to completion or expose paging in its UI — flag any consumer left on a single default-sized page.
|
||||
- **Search moves server-side with it.** A picker that filtered client-side over the full list now only filters the loaded page. If the endpoint gained `q`, the UI must send it (and disable its local filtering, e.g. cmdk's `shouldFilter={false}`); if it didn't, say why the collection is small enough not to need it.
|
||||
- **Tests that look up their own row must not depend on landing on page 1** — they should search or pass an explicit `limit`, not rely on default ordering.
|
||||
|
||||
### 8. Check code comments
|
||||
|
||||
For each non-trivial change:
|
||||
|
||||
@@ -188,9 +188,23 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Vector Extension (Optional - uses pgvector by default)
|
||||
# Options: "pgvector" (default), "vchord", "pgvectorscale" (DiskANN)
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvector
|
||||
|
||||
# Let a vector index scan resume until the query's LIMIT is satisfied, instead of
|
||||
# stopping when its first candidate list drains (pgvector: hnsw.ef_search, 200) — with
|
||||
# it off, a larger recall budget cannot retrieve more rows. Needs pgvector 0.8.0+;
|
||||
# older servers reject it and it is dropped automatically. Set false and restart as a
|
||||
# quick revert to the previous retrieval depth, with no code change.
|
||||
# HINDSIGHT_API_ANN_ITERATIVE_SCAN=true
|
||||
# Ceiling on tuples one resumed scan may visit. Bounds the CPU and memory a selective
|
||||
# query can spend resuming (filters are applied after the scan, so it resumes often).
|
||||
# Lower it to trade depth back for latency. pgvector's own default is 20000.
|
||||
# HINDSIGHT_API_ANN_MAX_SCAN_TUPLES=4000
|
||||
# For Azure PostgreSQL with DiskANN:
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale # Auto-detects pg_diskann on Azure
|
||||
|
||||
# Per-bank vector indexes (pgvector / pgvectorscale / vchord only; ScaNN and Oracle use one global index)
|
||||
# HINDSIGHT_API_VECTOR_INDEX_MIN_ROWS=0 # Memories a bank needs in one fact type before that fact type gets its own vector index. 0 (default) = no minimum, every bank holding memories is indexed. Set ~10000 on deployments with thousands of banks: every index lives on the shared memory_units table and is planned against by every OTHER bank's queries, so unconditional per-bank indexes put a ceiling on bank count. Smaller banks then use exact search, which is faster AND exact.
|
||||
|
||||
# Text Search Extension (Optional - uses native PostgreSQL full-text search by default)
|
||||
# Backend options: "native" (default), "vchord", "pg_textsearch", "pgroonga", "pg_search"
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION=native
|
||||
|
||||
File diff suppressed because one or more lines are too long
|
Before Width: | Height: | Size: 41 KiB After Width: | Height: | Size: 41 KiB |
@@ -1044,6 +1044,10 @@
|
||||
{
|
||||
"date": "2026-08-18",
|
||||
"stars": 20114
|
||||
},
|
||||
{
|
||||
"date": "2026-08-19",
|
||||
"stars": 20214
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -2,13 +2,14 @@
|
||||
|
||||
|
||||
|
||||
[Documentation](https://hindsight.vectorize.io) • [Paper](https://arxiv.org/abs/2512.12818) • [Cookbook](https://hindsight.vectorize.io/cookbook) • [Hindsight Cloud](https://ui.hindsight.vectorize.io/signup)
|
||||
[Documentation](https://hindsight.vectorize.io) • [Integrations](https://hindsight.vectorize.io/integrations) • [Cookbook](https://hindsight.vectorize.io/cookbook) • [Benchmarks](https://benchmarks.hindsight.vectorize.io/) • [Paper](https://arxiv.org/abs/2512.12818) • [Hindsight Cloud](https://ui.hindsight.vectorize.io/signup)
|
||||
|
||||
[](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml)
|
||||
[](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml)
|
||||
[](https://pypi.org/project/hindsight-api/)
|
||||
[](https://pypi.org/project/hindsight-client/)
|
||||
[](https://www.npmjs.com/package/@vectorize-io/hindsight-client)
|
||||
[](https://join.slack.com/t/hindsight-space/shared_invite/zt-3nhbm4w29-LeSJ5Ixi6j8PdiYOCPlOgg)
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
|
||||
|
||||
<br/>
|
||||
|
||||
<a href="https://trendshift.io/repositories/15603" target="_blank"><img src="https://trendshift.io/api/badge/repositories/15603" alt="vectorize-io%2Fhindsight | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
|
||||
@@ -20,28 +21,33 @@
|
||||
|
||||
Hindsight™ is an agent memory system built to create smarter agents that learn over time. Most agent memory systems focus on recalling conversation history. Hindsight is focused on making agents that learn, not just remember.
|
||||
|
||||
|
||||
<video src="https://github.com/user-attachments/assets/923b798d-3581-4897-bb62-9cfa5a931682" controls></video>
|
||||
|
||||
It eliminates the shortcomings of alternative techniques such as RAG and knowledge graph and delivers state-of-the-art performance on long term memory tasks.
|
||||
|
||||
**Contents**
|
||||
|
||||
- [Memory Performance & Accuracy](#memory-performance--accuracy)
|
||||
- [Quick Start](#quick-start) — [server](#1-start-a-server) · [clients](#2-connect-a-client) · [platforms](#supported-platforms) · [embedded](#python-embedded-no-server-required)
|
||||
- [Adding Hindsight to Your Agent](#adding-hindsight-to-your-agent) — [LLM Wrapper](#llm-wrapper-2-lines-of-code) · [integrations](#integrations) · [coding agents](#coding-agents) · [MCP](#mcp-server)
|
||||
- [Core Concepts](#core-concepts) — [memory types](#memory-types) · [retain / recall / reflect](#the-three-operations) · [observations](#observations) · [mental models & knowledge pages](#mental-models--knowledge-pages) · [banks](#memory-banks)
|
||||
- [Use Cases](#use-cases)
|
||||
- [Running in Production](#running-in-production)
|
||||
- [Resources](#resources)
|
||||
|
||||
---
|
||||
|
||||
## Memory Performance & Accuracy
|
||||
|
||||
Hindsight is the most accurate agent memory system ever tested according to benchmark performance. It has achieved state-of-the-art performance on the LongMemEval benchmark, widely used to assess memory system performance across a variety of conversational AI scenarios. The current reported performance of Hindsight and other agent memory solutions as of January 2026 is shown here:
|
||||
|
||||
|
||||
|
||||
> Live, continuously updated results — including per-model accuracy, latency and cost — are published at [benchmarks.hindsight.vectorize.io](https://benchmarks.hindsight.vectorize.io/).
|
||||
|
||||
The benchmark performance data for Hindsight has been independently reproduced by research collaborators at the Virginia Tech [Sanghani Center for Artificial Intelligence and Data Analytics](https://sanghani.cs.vt.edu/) and The Washington Post. Other scores are self-reported by software vendors.
|
||||
|
||||
Hindsight is being used in production at Fortune 500 enterprises and by a growing number of AI startups.
|
||||
|
||||
## Adding Hindsight to Your AI Agents
|
||||
|
||||
The easiest way to use Hindsight with an existing agent is with the LLM Wrapper. You can add memory to your agent with 2 lines of code. That will swap your current LLM client out with the Hindsight wrapper. After that, memories will be stored and retrieved automatically as you make LLM calls.
|
||||
|
||||
If you need more control over how and when your agent stores and recalls memories, there's also a simple API you can integrate with using the SDKs or directly via HTTP.
|
||||
|
||||
|
||||
Hindsight is being used in production at Fortune 500 enterprises and by a growing number of AI startups.
|
||||
|
||||
---
|
||||
|
||||
@@ -53,10 +59,11 @@ If you need more control over how and when your agent stores and recalls memorie
|
||||
|
||||
---
|
||||
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Docker (recommended)
|
||||
### 1. Start a server
|
||||
|
||||
#### Docker (recommended)
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
@@ -70,31 +77,52 @@ docker run -it --pull always --name hindsight --restart unless-stopped -p 8888:8
|
||||
>API: http://localhost:8888
|
||||
>UI: http://localhost:9999
|
||||
|
||||
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio`, `minimax`, and `atlas` ([Atlas Cloud](https://www.atlascloud.ai/?utm_source=github&utm_medium=link&utm_campaign=hindsight)). The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
Hindsight works with **25+ LLM providers** via `HINDSIGHT_API_LLM_PROVIDER` — hosted (`openai`, `anthropic`, `gemini`, `groq`, `bedrock`, `vertexai`, `minimax`, `deepseek`, `atlas`, …), fully local (`ollama`, `lmstudio`, `llamacpp`), any OpenAI-compatible endpoint, and gateways (`litellm`, `litellmrouter`) that reach the rest. Existing subscriptions work too: `openai-codex` (ChatGPT Plus/Pro) and `claude-code` (Claude Pro/Max) need no API key. See [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
|
||||
|
||||
|
||||
### Docker (external PostgreSQL)
|
||||
#### Docker (external PostgreSQL)
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
export HINDSIGHT_DB_PASSWORD=choose-a-password
|
||||
cd docker/docker-compose
|
||||
docker compose up
|
||||
docker compose up
|
||||
```
|
||||
|
||||
> Oracle AI Database is also supported for enterprise deployments with full feature parity. See the [storage documentation](https://hindsight.vectorize.io/developer/storage) for details.
|
||||
|
||||
|
||||
>API: http://localhost:8888
|
||||
>UI: http://localhost:9999
|
||||
|
||||
### Client
|
||||
#### Bare metal (pip)
|
||||
|
||||
```bash
|
||||
pip install hindsight-client -U
|
||||
# or
|
||||
npm install @vectorize-io/hindsight-client
|
||||
pip install hindsight-api
|
||||
export HINDSIGHT_API_LLM_API_KEY=sk-xxx
|
||||
|
||||
hindsight-api
|
||||
```
|
||||
|
||||
#### Kubernetes (Helm)
|
||||
|
||||
```bash
|
||||
helm install hindsight oci://ghcr.io/vectorize-io/charts/hindsight \
|
||||
--set api.llm.provider=openai \
|
||||
--set api.llm.apiKey=sk-xxx \
|
||||
--set postgresql.enabled=true
|
||||
```
|
||||
|
||||
#### Managed (no server)
|
||||
|
||||
[Hindsight Cloud](https://vectorize.io/pricing) is the hosted option: managed infrastructure that scales automatically, plus a dashboard, backups, team collaboration and a 99.9% uptime SLA. Billing is usage-based with free credits to start — no fixed monthly or per-seat fee. Point any client at `https://api.hindsight.vectorize.io` with your API key and skip the deployment entirely.
|
||||
|
||||
[Compare self-hosted, Cloud and Enterprise →](https://vectorize.io/pricing) · [Sign up →](https://ui.hindsight.vectorize.io/signup)
|
||||
|
||||
All options, including Windows and air-gapped setups, are covered in the [installation guide](https://hindsight.vectorize.io/developer/installation).
|
||||
|
||||
### 2. Connect a client
|
||||
|
||||
```bash
|
||||
pip install hindsight-client -U # Python
|
||||
npm install @vectorize-io/hindsight-client # Node.js / TypeScript
|
||||
go get github.com/vectorize-io/hindsight/hindsight-clients/go # Go
|
||||
curl -fsSL https://hindsight.vectorize.io/get-cli | bash # CLI
|
||||
```
|
||||
|
||||
#### Python
|
||||
@@ -116,10 +144,6 @@ client.reflect(bank_id="my-bank", query="Tell me about Alice")
|
||||
|
||||
#### Node.js / TypeScript
|
||||
|
||||
```bash
|
||||
npm install @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
```javascript
|
||||
const { HindsightClient } = require('@vectorize-io/hindsight-client');
|
||||
|
||||
@@ -135,6 +159,18 @@ const main = async () => {
|
||||
main();
|
||||
```
|
||||
|
||||
Full reference: [Python](https://hindsight.vectorize.io/sdks/python) · [Node.js](https://hindsight.vectorize.io/sdks/nodejs) · [Go](https://hindsight.vectorize.io/sdks/go) · [CLI](https://hindsight.vectorize.io/sdks/cli) · [REST API](https://hindsight.vectorize.io/api-reference)
|
||||
|
||||
### Supported Platforms
|
||||
|
||||
| Platform | Docker | Bare Metal (pip) | Embedded DB (pg0) |
|
||||
|----------|--------|------------------|--------------------|
|
||||
| **Linux** (x86_64, ARM64) | ✅ | ✅ | ✅ |
|
||||
| **macOS** (Apple Silicon / arm64) | ✅ | ✅ | ✅ |
|
||||
| **macOS** (Intel / x86_64) | ✅ | ⚠️ | ✅ |
|
||||
| **Windows** (x86_64) | ✅ | ✅ | ✅ |
|
||||
|
||||
⚠️ Intel Macs: use `hindsight-all-slim` — see the [installation guide](https://hindsight.vectorize.io/developer/installation#supported-platforms) for details.
|
||||
|
||||
### Python Embedded (no server required)
|
||||
|
||||
@@ -150,7 +186,7 @@ from hindsight import HindsightServer, HindsightClient
|
||||
|
||||
with HindsightServer(
|
||||
llm_provider="openai",
|
||||
llm_model="gpt-5-mini",
|
||||
llm_model="gpt-5-mini",
|
||||
llm_api_key=os.environ["OPENAI_API_KEY"]
|
||||
) as server:
|
||||
client = HindsightClient(base_url=server.url)
|
||||
@@ -158,12 +194,182 @@ with HindsightServer(
|
||||
results = client.recall(bank_id="my-bank", query="Where does Alice work?")
|
||||
```
|
||||
|
||||
A [Node.js equivalent](https://hindsight.vectorize.io/sdks/hindsight-all-npm) and a [daemon CLI](https://hindsight.vectorize.io/sdks/embed) are also available.
|
||||
|
||||
---
|
||||
|
||||
## Adding Hindsight to Your Agent
|
||||
|
||||
### LLM Wrapper (2 lines of code)
|
||||
|
||||
The easiest way to add memory to an existing agent is the LLM Wrapper. Swap your LLM client for a wrapped one — memories are then stored and retrieved automatically on every call, with no other changes to your code.
|
||||
|
||||
```bash
|
||||
pip install hindsight-litellm
|
||||
```
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
from hindsight_litellm import wrap_openai
|
||||
|
||||
# Wrap your existing LLM client and you're done.
|
||||
# Defaults to Hindsight Cloud; pass hindsight_api_url for a self-hosted server.
|
||||
client = wrap_openai(
|
||||
OpenAI(),
|
||||
bank_id="user-123",
|
||||
hindsight_api_url="http://localhost:8888",
|
||||
)
|
||||
|
||||
# Hindsight recalls relevant memories before the call
|
||||
# and retains the conversation after it.
|
||||
response = client.chat.completions.create(
|
||||
model="gpt-5-mini",
|
||||
messages=[{"role": "user", "content": "What do you know about me?"}],
|
||||
)
|
||||
```
|
||||
|
||||
`wrap_anthropic()` does the same for the Anthropic SDK, and every setting — bank, recall budget, fact types, reflect instead of recall — can be overridden per call with `hindsight_*` kwargs. LiteLLM sits underneath, so the same integration covers **100+ models**. See the [LiteLLM integration](https://hindsight.vectorize.io/sdks/integrations/litellm).
|
||||
|
||||
If you need explicit control over *when* memories are stored and recalled, use the [SDKs or REST API](#2-connect-a-client) directly instead.
|
||||
|
||||
### Integrations
|
||||
|
||||
**60+ integrations** — most need no code changes.
|
||||
|
||||
| | |
|
||||
|---|---|
|
||||
| **Coding agents** | [Claude Code](https://hindsight.vectorize.io/sdks/integrations/claude-code) · [Codex](https://hindsight.vectorize.io/sdks/integrations/codex) · [Cursor](https://hindsight.vectorize.io/sdks/integrations/cursor) · [GitHub Copilot](https://hindsight.vectorize.io/sdks/integrations/github-copilot) · [opencode](https://hindsight.vectorize.io/sdks/integrations/opencode) · [Cline](https://hindsight.vectorize.io/sdks/integrations/cline) · [Aider](https://hindsight.vectorize.io/sdks/integrations/aider) · [Zed](https://hindsight.vectorize.io/sdks/integrations/zed) · [Continue](https://hindsight.vectorize.io/sdks/integrations/continue) · [Roo Code](https://hindsight.vectorize.io/sdks/integrations/roo-code) · [OpenHands](https://hindsight.vectorize.io/sdks/integrations/openhands) |
|
||||
| **Agent frameworks** | [LangGraph / LangChain](https://hindsight.vectorize.io/sdks/integrations/langgraph) · [LlamaIndex](https://hindsight.vectorize.io/sdks/integrations/llamaindex) · [CrewAI](https://hindsight.vectorize.io/sdks/integrations/crewai) · [Pydantic AI](https://hindsight.vectorize.io/sdks/integrations/pydantic-ai) · [OpenAI Agents SDK](https://hindsight.vectorize.io/sdks/integrations/openai-agents) · [Google ADK](https://hindsight.vectorize.io/sdks/integrations/google-adk) · [Agno](https://hindsight.vectorize.io/sdks/integrations/agno) · [Strands](https://hindsight.vectorize.io/sdks/integrations/strands) · [AutoGen](https://hindsight.vectorize.io/sdks/integrations/autogen) · [Microsoft Agent Framework](https://hindsight.vectorize.io/sdks/integrations/agent-framework) · [Vercel AI SDK](https://hindsight.vectorize.io/sdks/integrations/ai-sdk) · [Haystack](https://hindsight.vectorize.io/sdks/integrations/haystack) |
|
||||
| **No-code / low-code** | [n8n](https://hindsight.vectorize.io/sdks/integrations/n8n) · [Zapier](https://hindsight.vectorize.io/sdks/integrations/zapier) · [Dify](https://hindsight.vectorize.io/sdks/integrations/dify) · [Flowise](https://hindsight.vectorize.io/sdks/integrations/flowise) |
|
||||
| **Apps & tools** | [ChatGPT](https://hindsight.vectorize.io/sdks/integrations/chatgpt) · [Perplexity](https://hindsight.vectorize.io/sdks/integrations/perplexity) · [Obsidian](https://hindsight.vectorize.io/sdks/integrations/obsidian) · [Pipecat](https://hindsight.vectorize.io/sdks/integrations/pipecat) · [Vapi](https://hindsight.vectorize.io/sdks/integrations/vapi) |
|
||||
|
||||
👉 [**Browse all integrations**](https://hindsight.vectorize.io/integrations)
|
||||
|
||||
### Coding Agents
|
||||
|
||||
One package gives CLI coding agents long-term project memory: a per-repo bank built automatically from git history and past sessions, injected into the agent as it starts working, plus curated knowledge pages covering architecture, conventions and in-flight work.
|
||||
|
||||
```bash
|
||||
npx @vectorize-io/hindsight-coding-agents install all # every detected agent, wired natively
|
||||
npx @vectorize-io/hindsight-coding-agents install claude-code # or just one
|
||||
```
|
||||
|
||||
Supports Claude Code, Codex CLI, Cursor CLI, GitHub Copilot CLI, opencode, Kilo CLI, Cline CLI, Antigravity CLI, Devin CLI, Prime Agent, Grok Build and DeepSeek Harness. Ingestion is automatic — there is no setup command. See the [coding agents integration](https://hindsight.vectorize.io/sdks/integrations/coding-agents).
|
||||
|
||||
### MCP Server
|
||||
|
||||
Every server ships a built-in [Model Context Protocol](https://modelcontextprotocol.io/) endpoint, one per bank, enabled by default:
|
||||
|
||||
```
|
||||
http://localhost:8888/mcp/{bank_id}/
|
||||
```
|
||||
|
||||
Point any MCP client at it to expose retain, recall and reflect as tools. See the [MCP server docs](https://hindsight.vectorize.io/developer/mcp-server).
|
||||
|
||||
---
|
||||
|
||||
## Core Concepts
|
||||
|
||||
|
||||
|
||||
### Memory Types
|
||||
|
||||
Most agent memory implementations rely on basic vector search or sometimes use a knowledge graph. Hindsight uses biomimetic data structures to organize agent memories in a way that is more like how human memory works:
|
||||
|
||||
- **World facts:** facts about the world ("The stove gets hot")
|
||||
- **Experiences:** the agent's own experiences ("I touched the stove and it really hurt")
|
||||
- **Observations:** consolidated, evidence-backed beliefs formed from many memories
|
||||
- **Mental models:** learned understanding of the agent's world, synthesized from observations and facts
|
||||
|
||||
Memories live in **banks**. When memories are added, they are pushed into either the world facts or the experiences pathway, then represented as a combination of entities, relationships, and time series with sparse/dense vector representations to aid in later recall.
|
||||
|
||||
### The Three Operations
|
||||
|
||||
#### Retain
|
||||
|
||||
The `retain` operation is used to push new memories into Hindsight. It tells Hindsight to _retain_ the information you pass in as an input.
|
||||
|
||||
```python
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content="Alice got promoted to senior engineer",
|
||||
context="career update",
|
||||
timestamp="2025-06-15T10:00:00Z",
|
||||
)
|
||||
```
|
||||
|
||||
Behind the scenes, retain uses an LLM to extract key facts, temporal data, entities, and relationships. It passes these through a normalization process to transform extracted data into canonical entities, time series, and search indexes along with metadata. These representations create the pathways for accurate memory retrieval in the recall and reflect operations.
|
||||
|
||||
|
||||
|
||||
[Retain docs →](https://hindsight.vectorize.io/developer/retain)
|
||||
|
||||
#### Recall
|
||||
|
||||
The recall operation is used to retrieve memories. These memories can come from any of the memory types (world, experiences, etc.)
|
||||
|
||||
```python
|
||||
client.recall(bank_id="my-bank", query="What does Alice do?")
|
||||
client.recall(bank_id="my-bank", query="What happened in June?") # temporal
|
||||
```
|
||||
|
||||
Recall performs 4 retrieval strategies in parallel:
|
||||
- Semantic: Vector similarity
|
||||
- Keyword: BM25 exact matching
|
||||
- Graph: Entity/temporal/causal links
|
||||
- Temporal: Time range filtering
|
||||
|
||||
|
||||
|
||||
The individual results are merged, ordered by relevance using reciprocal rank fusion and a cross-encoder reranking model, then trimmed as needed to fit within the token limit.
|
||||
|
||||
[Recall docs →](https://hindsight.vectorize.io/developer/retrieval)
|
||||
|
||||
#### Reflect
|
||||
|
||||
The reflect operation performs a more thorough analysis of existing memories. This allows the agent to form new connections between memories and build a more thorough understanding of its world — or to answer a question that needs deep thinking rather than lookup.
|
||||
|
||||
```python
|
||||
client.reflect(bank_id="my-bank", query="What should I know about Alice?")
|
||||
```
|
||||
|
||||
For example, reflect supports use cases such as:
|
||||
|
||||
- An **AI Project Manager** reflecting on what risks need to be mitigated on a project.
|
||||
- A **Sales Agent** reflecting on why certain outreach messages have gotten responses while others haven't.
|
||||
- A **Support Agent** reflecting on opportunities where customers have questions not answered by current product documentation.
|
||||
|
||||
|
||||
|
||||
[Reflect docs →](https://hindsight.vectorize.io/developer/reflect)
|
||||
|
||||
### Observations
|
||||
|
||||
Retained facts don't stay a flat pile. In the background, Hindsight consolidates related facts into **observations** — deduplicated beliefs the bank has built up over time. Each observation keeps its supporting evidence with exact quotes and a proof count, and is *refined* rather than overwritten when new evidence arrives, so new information strengthens, weakens or extends an existing belief instead of silently replacing it.
|
||||
|
||||
[Observations docs →](https://hindsight.vectorize.io/developer/observations)
|
||||
|
||||
### Mental Models & Knowledge Pages
|
||||
|
||||
A **mental model** is a standing answer to a question about a bank ("What are this user's preferences?"). You define the question once; Hindsight writes the answer, stores it, and rewrites it in the background as the bank learns more. Reading one is a database read — no retrieval, no LLM call — so an agent can boot with a page of settled knowledge instead of rediscovering it every session.
|
||||
|
||||
**Knowledge pages** are mental models with the mechanics hidden: living documents a bank writes about itself, organized in folders like a wiki, searchable, and projectable onto disk as ordinary markdown. Supply a name and a question; every other decision is a default you can override.
|
||||
|
||||
[Mental models →](https://hindsight.vectorize.io/developer/mental-models) · [Knowledge pages →](https://hindsight.vectorize.io/developer/knowledge-pages)
|
||||
|
||||
### Memory Banks
|
||||
|
||||
A **bank** is an isolated memory store — one "brain" for one user, agent, or project. Isolation is strict: no cross-bank leakage. Banks carry background context and **disposition traits** (skepticism, literalism, empathy) that shape how reflect reasons over their memories, and can be created from declarative [bank templates](https://hindsight.vectorize.io/developer/api/bank-templates).
|
||||
|
||||
Two more things worth knowing:
|
||||
|
||||
- **Multilingual by default.** Input language is detected and preserved end to end — facts stay in their original language and entities keep their native script (张伟 stays 张伟, not "Zhang Wei"). [Docs →](https://hindsight.vectorize.io/developer/multilingual)
|
||||
- **Memory Defense.** An opt-in, per-bank policy that scans every retain for secrets and PII against 45 patterns and either redacts the match (`[REDACTED:github_token]`) or blocks the item before it reaches storage. [Docs →](https://hindsight.vectorize.io/developer/memory-defense)
|
||||
|
||||
---
|
||||
|
||||
## Use Cases
|
||||
|
||||
|
||||
Hindsight is built to support conversational AI agents as well as agents that are intended to perform tasks autonomously. The ideal use case for Hindsight are agents that require a blend of these features such as AI employees that need to handle open-ended tasks, change behavior based on user feedback, and learn to perform complex tasks to automate work at a level that approximates a human work. Hindsight can be used with simple AI workflows like those built with n8n and other similar tools, but may be overkill for such applications.
|
||||
|
||||
### Per-User Memories and Chat History
|
||||
@@ -176,141 +382,46 @@ The requirements for this use case usually look something like this:
|
||||
|
||||
<video src="https://github.com/user-attachments/assets/4805e8e1-e7d1-47c6-a4f8-2344a5ec8906" controls></video>
|
||||
|
||||
Satisfying these requirements in Hindsight is straightforward. When new user inputs and tool calls are ingested into Hindsight using the retain operation, custom metadata can be used to enrich the new memories. Metadata provides a convenient way to isolate memories that need to be restricted to a given user. Once these are fed into the retain operation, any raw memories and mental models that get created can be filtered when retrieving relevant memories.
|
||||
Satisfying these requirements in Hindsight is straightforward. When new user inputs and tool calls are ingested into Hindsight using the retain operation, custom metadata can be used to enrich the new memories. Metadata provides a convenient way to isolate memories that need to be restricted to a given user. Once these are fed into the retain operation, any raw memories and mental models that get created can be filtered when retrieving relevant memories.
|
||||
|
||||
|
||||
|
||||
More patterns in the [Cookbook](https://hindsight.vectorize.io/cookbook) and [Best Practices](https://hindsight.vectorize.io/best-practices).
|
||||
|
||||
---
|
||||
|
||||
## Architecture & Operations
|
||||
## Running in Production
|
||||
|
||||
|
||||
|
||||
Most agent memory implementations rely on basic vector search or sometimes use a knowledge graph. Hindsight uses biomimetic data structures to organize agent memories in a way that is more like how human memory works:
|
||||
|
||||
- **World:** Facts about the world ("The stove gets hot")
|
||||
- **Experiences:** Agent's own experiences ("I touched the stove and it really hurt")
|
||||
- **Mental Models:** Learned understanding of the agent's world formed by reflecting on raw memories and experiences.
|
||||
|
||||
Memories in Hindsight are stored in banks (i.e. memory banks). When memories are added to Hindsight, they are pushed into either the world facts or experiences memory pathway. They are then represented as a combination of entities, relationships, and time series with sparse/dense vector representations to aid in later recall.
|
||||
|
||||
Hindsight provides three simple methods to interact with the system:
|
||||
|
||||
- **Retain:** Provide information to Hindsight that you want it to remember
|
||||
- **Recall:** Retrieve memories from Hindsight
|
||||
- **Reflect:** Reflect on memories and experiences to generate new observations and insights from existing memories.
|
||||
|
||||
### Retain
|
||||
|
||||
The `retain` operation is used to push new memories into Hindsight. It tells Hindsight to _retain_ the information you pass in as an input.
|
||||
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
client = Hindsight(base_url="http://localhost:8888")
|
||||
|
||||
# Simple
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content="Alice works at Google as a software engineer"
|
||||
)
|
||||
|
||||
# With context and timestamp
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content="Alice got promoted to senior engineer",
|
||||
context="career update",
|
||||
timestamp="2025-06-15T10:00:00Z"
|
||||
)
|
||||
```
|
||||
|
||||
Behind the scenes, the retain operation uses an LLM to extract key facts, temporal data, entities, and relationships. It passes these through a normalization process to transform extracted data into canonical entities, time series, and search indexes along with metadata. These representations create the pathways for accurate memory retrieval in the recall and reflect operations.
|
||||
|
||||
|
||||
|
||||
### Recall
|
||||
|
||||
The recall operation is used to retrieve memories. These memories can come from any of the memory types (world, experiences, etc.)
|
||||
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
client = Hindsight(base_url="http://localhost:8888")
|
||||
|
||||
# Simple
|
||||
client.recall(bank_id="my-bank", query="What does Alice do?")
|
||||
|
||||
# Temporal
|
||||
client.recall(bank_id="my-bank", query="What happened in June?")
|
||||
```
|
||||
|
||||
Recall performs 4 retrieval strategies in parallel:
|
||||
- Semantic: Vector similarity
|
||||
- Keyword: BM25 exact matching
|
||||
- Graph: Entity/temporal/causal links
|
||||
- Temporal: Time range filtering
|
||||
|
||||
|
||||
|
||||
The individual results from the retrievals are merged, then ordered by relevance using reciprocal rank fusion and a cross-encoder reranking model.
|
||||
|
||||
The final output is trimmed as needed to fit within the token limit.
|
||||
|
||||
### Reflect
|
||||
|
||||
The reflect operation is used to perform a more thorough analysis of existing memories. This allows the agent to form new connections between memories and build a more thorough understanding of its world.
|
||||
|
||||
For example, the `reflect` operation can be used to support use cases such as:
|
||||
|
||||
- An **AI Project Manager** reflecting on what risks need to be mitigated on a project.
|
||||
- A **Sales Agent** reflecting on why certain outreach messages have gotten responses while others haven't.
|
||||
- A **Support Agent** reflecting on opportunities where customers have questions not answered by current product documentation.
|
||||
|
||||
The `reflect` operation can also be used to handle on-demand question answering or analysis which require more deep thinking.
|
||||
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
client = Hindsight(base_url="http://localhost:8888")
|
||||
|
||||
client.reflect(bank_id="my-bank", query="What should I know about Alice?")
|
||||
```
|
||||
|
||||
|
||||
| | |
|
||||
|---|---|
|
||||
| **Storage** | PostgreSQL + pgvector, or Oracle AI Database 23ai with full feature parity — [storage](https://hindsight.vectorize.io/developer/storage) |
|
||||
| **Configuration** | Hierarchical: global env vars → per-tenant → per-bank — [configuration](https://hindsight.vectorize.io/developer/configuration) |
|
||||
| **Monitoring** | Prometheus metrics and dashboards for LLM calls, tokens and latency — [monitoring](https://hindsight.vectorize.io/developer/monitoring) |
|
||||
| **Operations** | Admin CLI for migrations, bank repair and stuck operations — [admin CLI](https://hindsight.vectorize.io/developer/admin-cli) |
|
||||
| **Events** | Webhooks for retain, consolidation and refresh lifecycle events — [webhooks](https://hindsight.vectorize.io/developer/api/webhooks) |
|
||||
| **Extensibility** | Tenant, auth and storage extension points — [extensions](https://hindsight.vectorize.io/developer/extensions) |
|
||||
| **Managed** | Skip all of it with [Hindsight Cloud](https://vectorize.io/pricing) — managed, usage-based, 99.9% uptime SLA |
|
||||
|
||||
---
|
||||
|
||||
## Resources
|
||||
|
||||
**Documentation:**
|
||||
- [https://hindsight.vectorize.io](https://hindsight.vectorize.io)
|
||||
**Documentation:**
|
||||
- [Docs](https://hindsight.vectorize.io) · [FAQ](https://hindsight.vectorize.io/faq) · [Best Practices](https://hindsight.vectorize.io/best-practices) · [Cookbook](https://hindsight.vectorize.io/cookbook) · [Blog](https://hindsight.vectorize.io/blog)
|
||||
- [Paper](https://arxiv.org/abs/2512.12818) · [Benchmarks](https://benchmarks.hindsight.vectorize.io/) · [RAG vs Memory](https://hindsight.vectorize.io/developer/rag-vs-hindsight)
|
||||
|
||||
**Clients:**
|
||||
- [Python](http://hindsight.vectorize.io/sdks/python)
|
||||
- [Node.js](http://hindsight.vectorize.io/sdks/nodejs)
|
||||
- [REST API](https://hindsight.vectorize.io/api-reference)
|
||||
- [CLI](https://hindsight.vectorize.io/sdks/cli)
|
||||
- [Python](https://hindsight.vectorize.io/sdks/python) · [Node.js](https://hindsight.vectorize.io/sdks/nodejs) · [Go](https://hindsight.vectorize.io/sdks/go) · [CLI](https://hindsight.vectorize.io/sdks/cli) · [REST API](https://hindsight.vectorize.io/api-reference)
|
||||
|
||||
**Community:**
|
||||
- [Slack](https://join.slack.com/t/hindsight-space/shared_invite/zt-3nhbm4w29-LeSJ5Ixi6j8PdiYOCPlOgg)
|
||||
- [GitHub Issues](https://github.com/vectorize-io/hindsight/issues)
|
||||
|
||||
---
|
||||
|
||||
## Star History
|
||||
|
||||
[](https://github.com/vectorize-io/hindsight/stargazers)
|
||||
---
|
||||
|
||||
## Supported Platforms
|
||||
|
||||
| Platform | Docker | Bare Metal (pip) | Embedded DB (pg0) |
|
||||
|----------|--------|------------------|--------------------|
|
||||
| **Linux** (x86_64, ARM64) | ✅ | ✅ | ✅ |
|
||||
| **macOS** (Apple Silicon / arm64) | ✅ | ✅ | ✅ |
|
||||
| **macOS** (Intel / x86_64) | ✅ | ⚠️ | ✅ |
|
||||
| **Windows** (x86_64) | ✅ | ✅ | ✅ |
|
||||
|
||||
⚠️ Intel Macs: use `hindsight-all-slim` — see the [installation guide](https://hindsight.vectorize.io/developer/installation#supported-platforms) for details.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -64,15 +64,26 @@ class HindsightEmbedded:
|
||||
- create_directive(), list_directives(), etc.
|
||||
- And all async variants (aretain, arecall, areflect, etc.)
|
||||
|
||||
Only the settings you pass explicitly are forwarded to the daemon. Anything
|
||||
left at its default is resolved by the daemon instead, in this order: the
|
||||
profile's .env file, then the parent process environment, then the daemon's
|
||||
own default. That is what lets a client constructed without credentials run
|
||||
against a profile (or a shell) that already has them configured, rather than
|
||||
overwriting them with placeholders (#3253).
|
||||
|
||||
Args:
|
||||
profile: Profile name for data isolation (default: "default")
|
||||
llm_provider: LLM provider ("groq", "openai", "ollama", "gemini", "anthropic", "lmstudio")
|
||||
llm_api_key: API key for the LLM provider
|
||||
llm_model: Model name to use
|
||||
llm_provider: LLM provider ("groq", "openai", "ollama", "gemini", "anthropic",
|
||||
"lmstudio"). Omit to inherit; the server default is "openai".
|
||||
llm_api_key: API key for the LLM provider. Omit to inherit; pass "" to
|
||||
explicitly run without a key (local services that need no auth).
|
||||
llm_model: Model name to use. Omit to inherit; the server picks a default
|
||||
for the resolved provider.
|
||||
llm_base_url: Optional custom base URL for LLM API
|
||||
database_url: Optional database URL override (default: profile-specific pg0)
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (default: 0, disabled)
|
||||
log_level: Daemon log level (default: "info")
|
||||
idle_timeout: Seconds before daemon auto-exits when idle. Omit to inherit
|
||||
(daemon default: 0, disabled).
|
||||
log_level: Daemon log level. Omit to inherit (daemon default: "info").
|
||||
ui: Whether to start the control plane web UI alongside the daemon (default: False)
|
||||
ui_port: Port for the UI. Defaults to daemon_port + 10000.
|
||||
ui_hostname: Hostname to bind the UI to. Defaults to "0.0.0.0".
|
||||
@@ -81,13 +92,13 @@ class HindsightEmbedded:
|
||||
def __init__(
|
||||
self,
|
||||
profile: str = "default",
|
||||
llm_provider: str = "groq",
|
||||
llm_api_key: str = "",
|
||||
llm_model: str = "openai/gpt-oss-120b",
|
||||
llm_provider: Optional[str] = None,
|
||||
llm_api_key: Optional[str] = None,
|
||||
llm_model: Optional[str] = None,
|
||||
llm_base_url: Optional[str] = None,
|
||||
database_url: Optional[str] = None,
|
||||
idle_timeout: int = 0,
|
||||
log_level: str = "info",
|
||||
idle_timeout: Optional[int] = None,
|
||||
log_level: Optional[str] = None,
|
||||
ui: bool = False,
|
||||
ui_port: Optional[int] = None,
|
||||
ui_hostname: str = "0.0.0.0",
|
||||
@@ -95,29 +106,50 @@ class HindsightEmbedded:
|
||||
"""
|
||||
Initialize the embedded client (daemon starts on first use).
|
||||
|
||||
Every LLM/daemon setting left as None is omitted from the daemon config so
|
||||
the daemon resolves it from the profile .env, then the parent environment,
|
||||
then its own default.
|
||||
|
||||
Args:
|
||||
profile: Profile name for data isolation
|
||||
llm_provider: LLM provider
|
||||
llm_api_key: API key for the LLM provider
|
||||
llm_model: Model name to use
|
||||
llm_provider: LLM provider. Omit to inherit.
|
||||
llm_api_key: API key for the LLM provider. Omit to inherit; pass "" to
|
||||
explicitly run without a key.
|
||||
llm_model: Model name to use. Omit to inherit.
|
||||
llm_base_url: Optional custom base URL for LLM API
|
||||
database_url: Optional database URL override
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (0 = disabled)
|
||||
log_level: Daemon log level
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (0 = disabled).
|
||||
Omit to inherit.
|
||||
log_level: Daemon log level. Omit to inherit.
|
||||
ui: Whether to start the control plane web UI alongside the daemon
|
||||
ui_port: Port for the UI (defaults to daemon_port + 10000)
|
||||
ui_hostname: Hostname to bind the UI to (defaults to "0.0.0.0")
|
||||
"""
|
||||
self.profile = profile
|
||||
|
||||
# Build config dict for daemon (matches CLI format)
|
||||
self.config = {
|
||||
"HINDSIGHT_API_LLM_PROVIDER": llm_provider,
|
||||
"HINDSIGHT_API_LLM_API_KEY": llm_api_key,
|
||||
"HINDSIGHT_API_LLM_MODEL": llm_model,
|
||||
"HINDSIGHT_API_LOG_LEVEL": log_level,
|
||||
"HINDSIGHT_EMBED_DAEMON_IDLE_TIMEOUT": str(idle_timeout),
|
||||
}
|
||||
# Build the config dict for the daemon (matches CLI format), omitting
|
||||
# every setting the caller did not specify. An omitted key is inherited
|
||||
# by the daemon from the profile .env / parent environment; sending a
|
||||
# placeholder instead would overwrite it, and _register_profile would
|
||||
# then persist that placeholder into the profile's .env file (#3253).
|
||||
# An explicit "" is still an override — that is how a local LLM service
|
||||
# with no authentication clears an inherited API key.
|
||||
self.config: dict[str, str] = {}
|
||||
|
||||
if llm_provider is not None:
|
||||
self.config["HINDSIGHT_API_LLM_PROVIDER"] = llm_provider
|
||||
|
||||
if llm_api_key is not None:
|
||||
self.config["HINDSIGHT_API_LLM_API_KEY"] = llm_api_key
|
||||
|
||||
if llm_model is not None:
|
||||
self.config["HINDSIGHT_API_LLM_MODEL"] = llm_model
|
||||
|
||||
if log_level is not None:
|
||||
self.config["HINDSIGHT_API_LOG_LEVEL"] = log_level
|
||||
|
||||
if idle_timeout is not None:
|
||||
self.config["HINDSIGHT_EMBED_DAEMON_IDLE_TIMEOUT"] = str(idle_timeout)
|
||||
|
||||
if llm_base_url:
|
||||
self.config["HINDSIGHT_API_LLM_BASE_URL"] = llm_base_url
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
"""Configuration forwarding rules for HindsightEmbedded.
|
||||
|
||||
Regression coverage for #3253: a setting the caller does not pass must be left
|
||||
out of the daemon config, so the daemon can resolve it from the profile's .env
|
||||
file or the parent environment instead of receiving a client-side placeholder
|
||||
that overwrites it — and that the daemon then persists back into the profile.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight import HindsightEmbedded
|
||||
from hindsight_embed.daemon_embed_manager import DaemonEmbedManager
|
||||
|
||||
LLM_PROVIDER = "HINDSIGHT_API_LLM_PROVIDER"
|
||||
LLM_API_KEY = "HINDSIGHT_API_LLM_API_KEY"
|
||||
LLM_MODEL = "HINDSIGHT_API_LLM_MODEL"
|
||||
LOG_LEVEL = "HINDSIGHT_API_LOG_LEVEL"
|
||||
IDLE_TIMEOUT = "HINDSIGHT_EMBED_DAEMON_IDLE_TIMEOUT"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_home(tmp_path, monkeypatch):
|
||||
"""Isolate HOME so profile .env files never touch the real user profile.
|
||||
|
||||
USERPROFILE is set as well because Path.home() consults it on Windows.
|
||||
"""
|
||||
home = tmp_path / "home"
|
||||
home.mkdir()
|
||||
monkeypatch.setenv("HOME", str(home))
|
||||
monkeypatch.setenv("USERPROFILE", str(home))
|
||||
return home
|
||||
|
||||
|
||||
def _write_profile(home, name, port, env_contents=None):
|
||||
"""Create a registered profile, optionally with a pre-populated .env file."""
|
||||
profile_dir = home / ".hindsight" / "profiles"
|
||||
profile_dir.mkdir(parents=True, exist_ok=True)
|
||||
(profile_dir / "metadata.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"version": 1,
|
||||
"profiles": {
|
||||
name: {
|
||||
"port": port,
|
||||
"created_at": "2024-01-01T00:00:00+00:00",
|
||||
"last_used": "2024-01-01T00:00:00+00:00",
|
||||
}
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
env_path = profile_dir / f"{name}.env"
|
||||
if env_contents is not None:
|
||||
env_path.write_text(env_contents)
|
||||
return env_path
|
||||
|
||||
|
||||
def _daemon_env(client):
|
||||
"""Run the real daemon start path with Popen stubbed, returning the child env.
|
||||
|
||||
Asserting on client.config alone would not catch a regression in how the
|
||||
embed manager merges that config with the profile and the parent
|
||||
environment, which is where the reported bug actually surfaced.
|
||||
"""
|
||||
manager = DaemonEmbedManager()
|
||||
captured: dict[str, dict[str, str]] = {}
|
||||
spawned = [False]
|
||||
|
||||
def fake_popen(cmd, env, **kwargs):
|
||||
captured["env"] = env
|
||||
spawned[0] = True
|
||||
process = MagicMock()
|
||||
process.pid = 12345
|
||||
return process
|
||||
|
||||
with (
|
||||
patch("hindsight_embed.daemon_embed_manager.subprocess.Popen", side_effect=fake_popen),
|
||||
patch("hindsight_embed.daemon_embed_manager.time.sleep"),
|
||||
patch.object(manager, "_clear_port", return_value=True),
|
||||
patch.object(manager, "_find_api_command", return_value=["hindsight-api"]),
|
||||
patch.object(manager, "is_running", side_effect=lambda profile="": spawned[0]),
|
||||
patch("hindsight_embed.daemon_embed_manager.platform.system", return_value="Linux"),
|
||||
):
|
||||
assert manager.ensure_running(client.config, client.profile)
|
||||
|
||||
return captured["env"]
|
||||
|
||||
|
||||
def test_nothing_is_forwarded_when_nothing_is_specified(temp_home):
|
||||
assert HindsightEmbedded(profile="test").config == {}
|
||||
|
||||
|
||||
def test_explicitly_passed_settings_are_forwarded(temp_home):
|
||||
client = HindsightEmbedded(
|
||||
profile="test",
|
||||
llm_provider="openai",
|
||||
llm_api_key="sk-real",
|
||||
llm_model="gpt-4o-mini",
|
||||
log_level="debug",
|
||||
idle_timeout=300,
|
||||
)
|
||||
|
||||
assert client.config == {
|
||||
LLM_PROVIDER: "openai",
|
||||
LLM_API_KEY: "sk-real",
|
||||
LLM_MODEL: "gpt-4o-mini",
|
||||
LOG_LEVEL: "debug",
|
||||
IDLE_TIMEOUT: "300",
|
||||
}
|
||||
|
||||
|
||||
def test_empty_api_key_is_forwarded_as_an_override(temp_home):
|
||||
"""An empty string is an explicit choice, not an omission.
|
||||
|
||||
Local LLM services that need no authentication rely on it to clear a key
|
||||
inherited from the environment.
|
||||
"""
|
||||
assert HindsightEmbedded(profile="test", llm_api_key="").config[LLM_API_KEY] == ""
|
||||
|
||||
|
||||
def test_idle_timeout_zero_is_forwarded(temp_home):
|
||||
"""0 is falsy but meaningful ("never auto-exit"), so it must survive."""
|
||||
assert HindsightEmbedded(profile="test", idle_timeout=0).config[IDLE_TIMEOUT] == "0"
|
||||
|
||||
|
||||
def test_omitted_key_inherits_the_parent_environment(temp_home, monkeypatch):
|
||||
monkeypatch.setenv(LLM_API_KEY, "sk-parent")
|
||||
_write_profile(temp_home, "inherit-env", 9871)
|
||||
|
||||
env = _daemon_env(HindsightEmbedded(profile="inherit-env", llm_provider="openai"))
|
||||
|
||||
assert env[LLM_API_KEY] == "sk-parent"
|
||||
|
||||
|
||||
def test_omitted_settings_inherit_the_profile_env(temp_home, monkeypatch):
|
||||
for var in (LLM_PROVIDER, LLM_API_KEY, LLM_MODEL):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
env_path = _write_profile(
|
||||
temp_home,
|
||||
"prod",
|
||||
9872,
|
||||
"HINDSIGHT_API_LLM_PROVIDER=anthropic\n"
|
||||
"HINDSIGHT_API_LLM_MODEL=claude-sonnet-4-20250514\n"
|
||||
"HINDSIGHT_API_LLM_API_KEY=sk-ant-prod\n",
|
||||
)
|
||||
|
||||
env = _daemon_env(HindsightEmbedded(profile="prod"))
|
||||
|
||||
assert env[LLM_PROVIDER] == "anthropic"
|
||||
assert env[LLM_MODEL] == "claude-sonnet-4-20250514"
|
||||
assert env[LLM_API_KEY] == "sk-ant-prod"
|
||||
|
||||
# A successful start rewrites the profile's .env; it must not come back with
|
||||
# client-side placeholders in place of the configured values.
|
||||
persisted = env_path.read_text()
|
||||
assert "HINDSIGHT_API_LLM_PROVIDER=anthropic" in persisted
|
||||
assert "HINDSIGHT_API_LLM_MODEL=claude-sonnet-4-20250514" in persisted
|
||||
assert "HINDSIGHT_API_LLM_API_KEY=sk-ant-prod" in persisted
|
||||
|
||||
|
||||
def test_explicit_empty_key_overrides_the_parent_environment(temp_home, monkeypatch):
|
||||
monkeypatch.setenv(LLM_API_KEY, "sk-parent")
|
||||
_write_profile(temp_home, "no-auth", 9873)
|
||||
|
||||
env = _daemon_env(
|
||||
HindsightEmbedded(profile="no-auth", llm_provider="lmstudio", llm_api_key="")
|
||||
)
|
||||
|
||||
assert env[LLM_API_KEY] == ""
|
||||
|
||||
|
||||
def test_explicit_settings_still_win_over_the_profile(temp_home, monkeypatch):
|
||||
monkeypatch.delenv(LLM_PROVIDER, raising=False)
|
||||
_write_profile(temp_home, "override", 9874, "HINDSIGHT_API_LLM_PROVIDER=anthropic\n")
|
||||
|
||||
env = _daemon_env(HindsightEmbedded(profile="override", llm_provider="openai"))
|
||||
|
||||
assert env[LLM_PROVIDER] == "openai"
|
||||
@@ -47,13 +47,64 @@ _INDEX_TYPE_KEYWORDS = {
|
||||
"scann": "scann",
|
||||
}
|
||||
|
||||
# Ceiling on how many tuples one resumed ANN scan may visit (hnsw.max_scan_tuples).
|
||||
# Only iterative scans consult it, and it is approximate — the initial round is not
|
||||
# counted. pgvector defaults to 20000; this is deliberately lower.
|
||||
#
|
||||
# The filters that thin a semantic arm (the similarity floor, tags, date ranges) are
|
||||
# applied *after* the index scan, so a selective query resumes repeatedly to fill its
|
||||
# LIMIT. Unbounded, that turns the cheapest queries today into the most expensive:
|
||||
# ~20x the standing batch is enough to fill even a large recall budget on an
|
||||
# unfiltered query, and caps the pathological filtered case at a scan that returns
|
||||
# short — which is exactly what those queries did before iterative scans were on.
|
||||
# The GUCs that make a scan resumable — dropped wholesale when the operator turns the
|
||||
# behaviour off, so a connection is left exactly as it was before it existed (and a
|
||||
# pgvector too old to define them is never sent them either).
|
||||
_ITERATIVE_SCAN_GUCS = frozenset({"hnsw.iterative_scan", "hnsw.max_scan_tuples"})
|
||||
|
||||
|
||||
def iterative_scan_enabled() -> bool:
|
||||
"""Whether ANN scans may resume to satisfy a query's LIMIT.
|
||||
|
||||
Turning it off restores the previous depth exactly: a scan stops when its first
|
||||
candidate list drains, so no recall retrieves more rows than that list holds,
|
||||
whatever its budget.
|
||||
|
||||
Resolved through the config object rather than read from the environment, so a
|
||||
value set any other way — a CLI override applied with dataclasses.replace, a
|
||||
programmatically built config — is honoured, and the parsing and validation live
|
||||
in one place. Imported inside the function because config imports this module.
|
||||
"""
|
||||
from .config import get_config
|
||||
|
||||
return get_config().ann_iterative_scan
|
||||
|
||||
|
||||
def ann_max_scan_tuples() -> int:
|
||||
"""Ceiling on tuples one resumed scan may visit (hnsw.max_scan_tuples).
|
||||
|
||||
This is the knob that governs the cost of the behaviour. It bounds the CPU a
|
||||
selective query can spend resuming, and with it the scan's memory — pgvector
|
||||
otherwise caps that at ``work_mem * hnsw.scan_mem_multiplier``, but at this
|
||||
default the memory ceiling is never approached: squeezing work_mem to 256kB
|
||||
changes neither the rows returned nor the latency.
|
||||
|
||||
Approximate, and the initial scan is not counted, so even 1 leaves intact the
|
||||
depth a query had before scans could resume.
|
||||
"""
|
||||
from .config import get_config
|
||||
|
||||
return get_config().ann_max_scan_tuples
|
||||
|
||||
|
||||
# Per-backend ANN search-time tuning GUCs. Each entry is a tuple of
|
||||
# (guc_name, value) pairs the caller can apply with SET or SET LOCAL.
|
||||
#
|
||||
# - pgvector exposes hnsw.ef_search. The 60 / 200 pair is unchanged from the
|
||||
# pre-dispatcher code (internal benchmarks tuned around our embedding count
|
||||
# and recall floor; see the link_utils / pool init call sites for the
|
||||
# latency-vs-recall framing).
|
||||
# latency-vs-recall framing). With iterative scans on (below) the ef value is a
|
||||
# batch size rather than a ceiling, so a query's own LIMIT decides its depth.
|
||||
# - vchord exposes vchordrq.probes, but its shape must match the index's
|
||||
# build.internal.lists hierarchy. VectorChord 1.1 added per-index fallback
|
||||
# parameters for this reason: a session GUC overrides every vchordrq index,
|
||||
@@ -63,11 +114,30 @@ _INDEX_TYPE_KEYWORDS = {
|
||||
# indexes should attach probes to the index storage parameters instead.
|
||||
# - pgvectorscale / pg_diskann / scann do not expose an equivalent per-statement
|
||||
# knob in the engine today, so the dispatcher returns no statements for them.
|
||||
#
|
||||
# hnsw.iterative_scan is what makes ef_search a *batch* size rather than a ceiling.
|
||||
# With it off (pgvector's default, and what Hindsight ran until now) the ground-layer
|
||||
# search runs once and the scan ends when its list drains, so a query could never get
|
||||
# more rows than ef_search however large its LIMIT — the recall budget moved the SQL
|
||||
# and nothing else. With it on, the scan resumes in ef_search-sized rounds until the
|
||||
# LIMIT is met, so each query gets the depth it asks for with no per-query setting.
|
||||
# strict_order, not relaxed_order: the arms are trimmed in Python on the assumption
|
||||
# that rows arrive ordered by distance.
|
||||
#
|
||||
# Retain-side link probing wants the opposite — it is tuned for latency, not depth,
|
||||
# and resuming past its small candidate list would defeat that — so the low-latency
|
||||
# profile pins it off. Both profiles set it explicitly rather than relying on the
|
||||
# server default, so neither depends on what the other last left on the connection.
|
||||
_ANN_TUNING_LOW_LATENCY: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (("hnsw.ef_search", "60"),),
|
||||
"pgvector": (("hnsw.ef_search", "60"), ("hnsw.iterative_scan", "off")),
|
||||
}
|
||||
_ANN_TUNING_HIGH_RECALL: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (("hnsw.ef_search", "200"),),
|
||||
"pgvector": (
|
||||
("hnsw.ef_search", "200"),
|
||||
("hnsw.iterative_scan", "strict_order"),
|
||||
# Value filled in per call by ann_search_tuning_settings().
|
||||
("hnsw.max_scan_tuples", ""),
|
||||
),
|
||||
}
|
||||
|
||||
_EXTENSION_INSTALL_SQL = {
|
||||
@@ -167,7 +237,12 @@ def ann_search_tuning_settings(ext: str, *, kind: str) -> tuple[tuple[str, str],
|
||||
table = _ANN_TUNING_HIGH_RECALL
|
||||
else:
|
||||
raise ValueError(f"Unknown ANN tuning kind: {kind!r}")
|
||||
return table.get(_normalize_resolved(ext), ())
|
||||
settings = table.get(_normalize_resolved(ext), ())
|
||||
if not iterative_scan_enabled():
|
||||
return tuple(pair for pair in settings if pair[0] not in _ITERATIVE_SCAN_GUCS)
|
||||
return tuple(
|
||||
(name, str(ann_max_scan_tuples()) if name == "hnsw.max_scan_tuples" else value) for name, value in settings
|
||||
)
|
||||
|
||||
|
||||
def uses_per_bank_vector_indexes(ext: str) -> bool:
|
||||
@@ -175,6 +250,83 @@ def uses_per_bank_vector_indexes(ext: str) -> bool:
|
||||
return _normalize_resolved(ext) != "scann"
|
||||
|
||||
|
||||
def per_bank_index_min_rows() -> int:
|
||||
"""Rows a (bank, fact_type) needs before it earns its own partial vector index.
|
||||
|
||||
Distinct from :func:`minimum_rows_for_index`, which is ScaNN's *build*
|
||||
requirement for its single global index (AlloyDB cannot construct one below
|
||||
a floor). This is a cost policy for the per-bank backends: the indexes sit on
|
||||
the shared ``memory_units`` table, so each one is enumerated and locked at
|
||||
plan time by queries belonging to every *other* bank, and opened by every DML
|
||||
statement against the table. A small bank's index cannot repay that — the
|
||||
``(bank_id, fact_type)`` B-tree plus a top-N sort answers the same query
|
||||
exactly and faster. See issue #3485.
|
||||
|
||||
Read from config rather than passed in because the write path's pre-check,
|
||||
the maintenance operation and the admin command must all apply the same
|
||||
number; a threshold that differed between the one deciding to queue work and
|
||||
the one deciding what to do would either oscillate or never converge.
|
||||
"""
|
||||
from .config import get_config
|
||||
|
||||
return get_config().vector_index_min_rows
|
||||
|
||||
|
||||
def per_bank_index_drop_rows() -> int:
|
||||
"""Row count below which an existing per-bank vector index is dropped.
|
||||
|
||||
Strictly below :func:`per_bank_index_min_rows` so the build and drop
|
||||
decisions cannot both be true at one row count. Without the gap, a bank
|
||||
hovering at the threshold — consolidation prunes a few facts, retain adds
|
||||
them back — would rebuild and drop the same ANN index on alternating sweeps.
|
||||
"""
|
||||
from .config import VECTOR_INDEX_DROP_RATIO
|
||||
|
||||
return int(per_bank_index_min_rows() * VECTOR_INDEX_DROP_RATIO)
|
||||
|
||||
|
||||
def should_keep_per_bank_index(row_count: int) -> bool:
|
||||
"""Whether an *existing* index on a partition of ``row_count`` rows is kept.
|
||||
|
||||
The counterpart to :func:`qualifies_for_per_bank_index`, and deliberately a
|
||||
separate, lower bound: keeping starts below building, so a partition
|
||||
hovering at the threshold does not rebuild and drop the same ANN index on
|
||||
alternating writes.
|
||||
|
||||
The ``row_count > 0`` term is not redundant with the ratio. At the default
|
||||
threshold of 0 the drop floor is also 0, so a bare ``row_count >= floor``
|
||||
keeps an index over an *emptied* partition forever — every bank ever written
|
||||
to and then cleared would hold three indexes over nothing, which is the
|
||||
accumulation the threshold exists to prevent. An emptied partition loses its
|
||||
index at every threshold.
|
||||
"""
|
||||
return row_count > 0 and row_count >= per_bank_index_drop_rows()
|
||||
|
||||
|
||||
def qualifies_for_per_bank_index(row_count: int) -> bool:
|
||||
"""Whether a (bank, fact_type) holding ``row_count`` rows should have an index.
|
||||
|
||||
At the default threshold of 0 this is true for every partition that holds
|
||||
any rows at all, which is the behaviour before the threshold existed.
|
||||
|
||||
An empty partition is excluded explicitly rather than by arithmetic: at a
|
||||
threshold of 0, ``row_count >= minimum`` alone is true for zero rows, so
|
||||
every bank in the deployment would be entitled to three indexes over nothing
|
||||
the moment it was created — the exact index explosion the threshold exists
|
||||
to prevent, reintroduced by its own default.
|
||||
|
||||
Only the build side: an existing index is kept until the count falls under
|
||||
:func:`per_bank_index_drop_rows`, so callers reconciling live state must
|
||||
consult both bounds rather than treating this as the full policy.
|
||||
|
||||
Takes no extension: the backend question is settled before any reconcile
|
||||
runs (``uses_per_bank_vector_indexes`` gates the maintenance operation and
|
||||
``_vector_index_clause`` gates the admin command), so re-asking it here
|
||||
would be a second, weaker copy of a decision already made.
|
||||
"""
|
||||
return row_count > 0 and row_count >= per_bank_index_min_rows()
|
||||
|
||||
|
||||
def bootstrap_extension(conn: Connection, ext: str) -> None:
|
||||
"""Install the configured vector extension and any prerequisites if possible."""
|
||||
normalized = validate_extension(ext)
|
||||
|
||||
@@ -23,7 +23,12 @@ from ..engine.memory_engine import _current_schema
|
||||
from ..engine.retain.bank_utils import _vector_index_clause
|
||||
from ..engine.schema import fq_table_explicit as _fq_table
|
||||
from ..engine.transfer import export_bank
|
||||
from ..engine.vector_index_health import SchemaVectorIndexResult, repair_vector_indexes
|
||||
from ..engine.vector_index_health import (
|
||||
BankIndexResult,
|
||||
drop_orphaned_bank_indexes,
|
||||
list_bank_ids,
|
||||
reconcile_bank_vector_indexes,
|
||||
)
|
||||
from ..extensions import TenantExtension, load_extension
|
||||
from ..pg0 import parse_pg0_url, resolve_database_url
|
||||
|
||||
@@ -618,11 +623,14 @@ async def _run_repair_bank(
|
||||
schema: str | None,
|
||||
bank_id: str | None,
|
||||
dry_run: bool,
|
||||
) -> list[SchemaVectorIndexResult]:
|
||||
) -> list[BankIndexResult]:
|
||||
"""Reconcile per-(bank, fact_type) vector index coverage over a raw connection.
|
||||
|
||||
A single autocommit connection is used because ``CREATE INDEX CONCURRENTLY``
|
||||
(used by ``repair_vector_indexes``) cannot run inside a transaction block.
|
||||
cannot run inside a transaction block.
|
||||
|
||||
Deliberately unbudgeted, unlike the background operation: this is an operator
|
||||
asking for convergence now, across as many banks as they named.
|
||||
"""
|
||||
schemas = [schema] if schema else await _resolve_schemas(base_schema)
|
||||
index_clause = _vector_index_clause()
|
||||
@@ -631,13 +639,38 @@ async def _run_repair_bank(
|
||||
assert index_clause is not None
|
||||
|
||||
conn = await _admin_connect(db_url)
|
||||
results: list[BankIndexResult] = []
|
||||
try:
|
||||
results = await repair_vector_indexes(conn, schemas, index_clause, dry_run=dry_run, bank_id=bank_id)
|
||||
for result in results:
|
||||
for target_schema in schemas:
|
||||
try:
|
||||
bank_ids = [bank_id] if bank_id else await list_bank_ids(conn, target_schema)
|
||||
except Exception as exc: # noqa: BLE001 — one bad schema must not abort the sweep
|
||||
typer.echo(f" schema '{target_schema}': skipped ({exc})", err=True)
|
||||
continue
|
||||
schema_results = [
|
||||
await reconcile_bank_vector_indexes(conn, target_schema, bid, index_clause, dry_run=dry_run)
|
||||
for bid in bank_ids
|
||||
]
|
||||
results.extend(schema_results)
|
||||
# Only in --all mode: an index whose bank row is gone is unreachable
|
||||
# from every bank-scoped path, so this is the one place that can
|
||||
# collect it. Normally finds nothing — delete_bank drops a bank's
|
||||
# indexes while it still knows their names — but a deployment that
|
||||
# hit the #3485 wall could not run delete_bank at all.
|
||||
orphans = [] if bank_id else await drop_orphaned_bank_indexes(conn, target_schema, dry_run=dry_run)
|
||||
if orphans:
|
||||
typer.echo(
|
||||
f" schema '{target_schema}': {len(orphans)} orphaned index(es) "
|
||||
f"{'to drop (dry-run)' if dry_run else 'dropped'} (no matching bank)"
|
||||
)
|
||||
typer.echo(
|
||||
f" schema '{result.schema}': {result.banks_scanned} bank(s) scanned, "
|
||||
f"{result.already_present} present, {result.created} created, "
|
||||
f"{result.skipped} to-create (dry-run), {result.failed} failed"
|
||||
f" schema '{target_schema}': {len(bank_ids)} bank(s) scanned, "
|
||||
f"{sum(r.already_present for r in schema_results)} present, "
|
||||
f"{sum(r.created for r in schema_results)} created, "
|
||||
f"{sum(r.dropped for r in schema_results)} dropped, "
|
||||
f"{sum(r.skipped for r in schema_results)} to-create (dry-run), "
|
||||
f"{sum(r.would_drop for r in schema_results)} to-drop (dry-run), "
|
||||
f"{sum(r.failed for r in schema_results)} failed"
|
||||
)
|
||||
return results
|
||||
finally:
|
||||
@@ -669,17 +702,23 @@ def repair_bank(
|
||||
help="Report what would be repaired without creating or dropping any index.",
|
||||
),
|
||||
):
|
||||
"""Verify and repair a bank's per-(bank, fact_type) vector index coverage.
|
||||
"""Reconcile per-(bank, fact_type) vector index coverage against the size threshold.
|
||||
|
||||
Per-bank partial vector indexes are created when a bank is first created
|
||||
(instant on an empty bank). Banks that arrive populated — via logical
|
||||
restore, a cross-version upgrade, or a vector-extension switch — never hit
|
||||
that path, so their recall silently falls back to a global index +
|
||||
post-filter (slower, under-returning). This command detects missing OR
|
||||
invalid coverage (an INVALID leftover or an index whose access method
|
||||
drifted counts as missing) and rebuilds it with CREATE INDEX CONCURRENTLY,
|
||||
so it never blocks the live fleet. Idempotent and safe to re-run — the
|
||||
escape hatch after a restore, upgrade, or backend switch.
|
||||
A (bank, fact_type) earns a partial vector index once it holds
|
||||
HINDSIGHT_API_VECTOR_INDEX_MIN_ROWS rows; below that the planner answers the
|
||||
same query exactly, and faster, from the (bank_id, fact_type) B-tree plus a
|
||||
top-N sort. This command builds what qualifies and drops what no longer does
|
||||
— including indexes orphaned by a deleted bank — detecting invalid coverage
|
||||
too (an INVALID leftover, or an index whose access method drifted after a
|
||||
backend switch, counts as missing). All DDL is CONCURRENTLY, so it never
|
||||
blocks the live fleet.
|
||||
|
||||
Writes keep this converged on their own — every insert that could move a bank
|
||||
across the threshold queues a vector_index_maintenance operation. Reach for
|
||||
the command when you want convergence without waiting for a write: after a
|
||||
restore or upgrade, after a backend switch, or to shed indexes in bulk on a
|
||||
deployment recovering from lock-table exhaustion (#3485). Idempotent and safe
|
||||
to re-run.
|
||||
"""
|
||||
if bool(bank_id) == all_banks:
|
||||
typer.echo("Error: pass exactly one of --bank <id> or --all.", err=True)
|
||||
@@ -713,15 +752,18 @@ def repair_bank(
|
||||
)
|
||||
)
|
||||
|
||||
total_banks = sum(r.banks_scanned for r in results)
|
||||
total_banks = len(results)
|
||||
total_present = sum(r.already_present for r in results)
|
||||
total_created = sum(r.created for r in results)
|
||||
total_dropped = sum(r.dropped for r in results)
|
||||
total_skipped = sum(r.skipped for r in results)
|
||||
total_would_drop = sum(r.would_drop for r in results)
|
||||
total_failed = sum(r.failed for r in results)
|
||||
typer.echo(
|
||||
f"Done: {len(results)} schema(s), {total_banks} bank(s) scanned, "
|
||||
f"{total_present} already present, {total_created} created, "
|
||||
f"{total_skipped} to-create (dry-run), {total_failed} failed"
|
||||
f"{total_present} already present, {total_created} created, {total_dropped} dropped, "
|
||||
f"{total_skipped} to-create (dry-run), {total_would_drop} to-drop (dry-run), "
|
||||
f"{total_failed} failed"
|
||||
)
|
||||
if total_failed:
|
||||
failed_names = [name for r in results for name in r.failed_indexes]
|
||||
|
||||
@@ -654,6 +654,17 @@ class MemoryItem(BaseModel):
|
||||
default=None,
|
||||
description="Optional entities to combine with auto-extracted entities.",
|
||||
)
|
||||
resolve_entities: bool = Field(
|
||||
default=True,
|
||||
description="Whether the names in 'entities' are resolved against the entities already in "
|
||||
"the bank. True (default) matches each name to a similar existing entity when it scores "
|
||||
"above the match threshold, so a name close to one already in the bank may resolve to that "
|
||||
"one instead of the one you wrote. False takes your names literally — an existing entity is "
|
||||
"reused only on a case-insensitive name match, any other name creates a new entity, and "
|
||||
"your names are never merged with each other. This applies only to the entities you supply "
|
||||
"here; auto-extracted entities are always resolved, since they are the extractor's guess at "
|
||||
"a name rather than yours. Ignored when 'entities' is omitted.",
|
||||
)
|
||||
tags: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Optional tags for visibility scoping. Memories with tags can be filtered during recall.",
|
||||
@@ -1282,7 +1293,7 @@ class BankListItem(BaseModel):
|
||||
|
||||
|
||||
class BankListResponse(BaseModel):
|
||||
"""Response model for listing all banks."""
|
||||
"""Response model for listing banks, one page at a time."""
|
||||
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
@@ -1299,12 +1310,18 @@ class BankListResponse(BaseModel):
|
||||
"last_document_at": "2024-01-16T14:20:00Z",
|
||||
"last_write_at": "2024-01-17T09:05:00Z",
|
||||
}
|
||||
]
|
||||
],
|
||||
"total": 50,
|
||||
"limit": 100,
|
||||
"offset": 0,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
banks: list[BankListItem]
|
||||
total: int = Field(description="Total number of banks visible to the caller, ignoring `limit`/`offset`.")
|
||||
limit: int
|
||||
offset: int
|
||||
|
||||
|
||||
class CreateBankRequest(BaseModel):
|
||||
@@ -1793,8 +1810,19 @@ class UpdateMemoryRequest(BaseModel):
|
||||
)
|
||||
entities: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Replace the fact's entities. Names are resolved/find-or-created "
|
||||
"the same way retain does; '[]' detaches all entities. Omit to leave unchanged.",
|
||||
description="Replace the fact's entities. How each name is matched to an entity is "
|
||||
"governed by 'resolve_entities'. '[]' detaches all entities. Omit to leave unchanged.",
|
||||
)
|
||||
resolve_entities: bool = Field(
|
||||
default=True,
|
||||
description="Whether the names in 'entities' are resolved against the entities already in "
|
||||
"the bank. True (default) is what retain does: a similar existing entity is reused when it "
|
||||
"scores above the match threshold, so a name close to one already in the bank may resolve "
|
||||
"to that one instead of the one you wrote. False takes the names literally — an existing "
|
||||
"entity is reused only on a case-insensitive name match, any other name creates a new "
|
||||
"entity, and names in the same request are never merged with each other. Use False for "
|
||||
"hand-authored corrections, where the name you sent is the answer rather than a guess. "
|
||||
"Ignored when 'entities' is omitted.",
|
||||
)
|
||||
state: str | None = Field(
|
||||
default=None,
|
||||
@@ -4524,6 +4552,7 @@ def _register_routes(app: FastAPI):
|
||||
occurred_end=occurred_end,
|
||||
new_fact_type=request.fact_type,
|
||||
entities=request.entities,
|
||||
resolve_entities=request.resolve_entities,
|
||||
state=request.state,
|
||||
reason=request.reason,
|
||||
request_context=request_context,
|
||||
@@ -4937,16 +4966,26 @@ def _register_routes(app: FastAPI):
|
||||
@app.get(
|
||||
"/v1/default/banks",
|
||||
response_model=BankListResponse,
|
||||
summary="List all memory banks",
|
||||
description="Get a list of all agents with their profiles",
|
||||
summary="List memory banks",
|
||||
description=(
|
||||
"List banks with their profiles and summary stats, most recently written first "
|
||||
"(`last_write_at` descending), with pagination and optional search."
|
||||
),
|
||||
operation_id="list_banks",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_list_banks(request_context: RequestContext = Depends(get_request_context)):
|
||||
"""Get list of all banks with their profiles."""
|
||||
async def api_list_banks(
|
||||
q: str | None = Query(None, description="Case-insensitive substring filter on bank ID or name (e.g. 'alice')"),
|
||||
limit: int = Query(default=100, ge=0, description="Maximum number of banks to return"),
|
||||
offset: int = Query(default=0, ge=0, description="Offset for pagination"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Get one page of banks with their profiles."""
|
||||
try:
|
||||
banks = await app.state.memory.list_banks(request_context=request_context)
|
||||
return BankListResponse(banks=banks)
|
||||
data = await app.state.memory.list_banks(
|
||||
search_query=q, limit=limit, offset=offset, request_context=request_context
|
||||
)
|
||||
return BankListResponse(**data)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
@@ -8005,6 +8044,7 @@ def _register_routes(app: FastAPI):
|
||||
content_dict["document_id"] = item.document_id
|
||||
if item.entities:
|
||||
content_dict["entities"] = [{"text": e.text, "type": e.type or "CONCEPT"} for e in item.entities]
|
||||
content_dict["resolve_entities"] = item.resolve_entities
|
||||
if item.tags:
|
||||
content_dict["tags"] = item.tags
|
||||
if item.observation_scopes is not None:
|
||||
|
||||
@@ -154,6 +154,13 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"update_bank",
|
||||
"delete_bank",
|
||||
"clear_memories",
|
||||
"get_knowledge_base_tree",
|
||||
"search_knowledge_base",
|
||||
"get_knowledge_page",
|
||||
"create_knowledge_folder",
|
||||
"create_knowledge_page",
|
||||
"update_knowledge_node",
|
||||
"delete_knowledge_node",
|
||||
}
|
||||
)
|
||||
base_tools: frozenset[str] | None = None if multi_bank else _SINGLE_BANK_TOOLS
|
||||
|
||||
@@ -533,6 +533,8 @@ ENV_RERANKER_GOOGLE_PROJECT_ID = "HINDSIGHT_API_RERANKER_GOOGLE_PROJECT_ID"
|
||||
ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY"
|
||||
|
||||
ENV_VECTOR_EXTENSION = "HINDSIGHT_API_VECTOR_EXTENSION"
|
||||
ENV_ANN_ITERATIVE_SCAN = "HINDSIGHT_API_ANN_ITERATIVE_SCAN"
|
||||
ENV_ANN_MAX_SCAN_TUPLES = "HINDSIGHT_API_ANN_MAX_SCAN_TUPLES"
|
||||
ENV_TEXT_SEARCH_EXTENSION = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION"
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE"
|
||||
ENV_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"
|
||||
@@ -755,6 +757,7 @@ WORKER_SLOT_TYPE_DEFAULTS: dict[str, int] = {
|
||||
"file_convert_retain": 0,
|
||||
"refresh_mental_model": 0,
|
||||
"graph_maintenance": 0,
|
||||
"vector_index_maintenance": 0,
|
||||
"import_documents": 0,
|
||||
"export_documents": 0,
|
||||
}
|
||||
@@ -876,6 +879,7 @@ ENV_MENTAL_MODEL_REFRESH_TICK_SECONDS = "HINDSIGHT_API_MENTAL_MODEL_REFRESH_TICK
|
||||
ENV_RETENTION_SWEEP_INTERVAL_SECONDS = "HINDSIGHT_API_RETENTION_SWEEP_INTERVAL_SECONDS"
|
||||
ENV_OPERATION_CLEANUP_INTERVAL_SECONDS = "HINDSIGHT_API_OPERATION_CLEANUP_INTERVAL_SECONDS"
|
||||
ENV_MAINTENANCE_START_JITTER_SECONDS = "HINDSIGHT_API_MAINTENANCE_START_JITTER_SECONDS"
|
||||
ENV_VECTOR_INDEX_MIN_ROWS = "HINDSIGHT_API_VECTOR_INDEX_MIN_ROWS"
|
||||
|
||||
# Disposition settings
|
||||
ENV_DISPOSITION_SKEPTICISM = "HINDSIGHT_API_DISPOSITION_SKEPTICISM"
|
||||
@@ -1142,6 +1146,22 @@ DEFAULT_RERANKER_GOOGLE_MODEL = "semantic-ranker-default-004"
|
||||
|
||||
# Vector extension (pgvector, vchord, pgvectorscale, or AlloyDB ScaNN)
|
||||
DEFAULT_VECTOR_EXTENSION = "pgvector" # Options: "pgvector", "vchord", "pgvectorscale", "scann"
|
||||
# Let an ANN scan resume until the query's LIMIT is met, instead of stopping when its
|
||||
# first candidate list drains. Off, a recall can never retrieve more rows than the
|
||||
# candidate list holds (pgvector: hnsw.ef_search, 200), so a larger recall budget
|
||||
# widens the SQL and changes nothing. On is the intended behaviour; this exists as an
|
||||
# operational kill switch, because turning it off restores exactly the previous
|
||||
# retrieval depth without a deploy.
|
||||
DEFAULT_ANN_ITERATIVE_SCAN = True
|
||||
# Ceiling on how many tuples one resumed scan may visit. Bounds both the CPU a
|
||||
# selective query can spend resuming (the filters that thin a result are applied after
|
||||
# the index scan, so a selective one resumes repeatedly) and the scan's memory, which
|
||||
# pgvector otherwise caps at work_mem * hnsw.scan_mem_multiplier. Measured at this
|
||||
# value the memory ceiling is never approached — squeezing work_mem to 256kB changes
|
||||
# neither rows nor latency — so this is the knob that governs the cost, not work_mem.
|
||||
# Lower it to trade retrieval depth back for latency; the initial scan is not counted,
|
||||
# so even 1 leaves the pre-existing behaviour intact. pgvector's own default is 20000.
|
||||
DEFAULT_ANN_MAX_SCAN_TUPLES = 4000
|
||||
|
||||
# Text search extension (native PostgreSQL, vchord BM25, Timescale pg_textsearch,
|
||||
# pgroonga, or ParadeDB pg_search)
|
||||
@@ -1484,6 +1504,31 @@ DEFAULT_OPERATION_CLEANUP_INTERVAL_SECONDS = 900
|
||||
# the same instant. 0 disables the jitter (deterministic start).
|
||||
DEFAULT_MAINTENANCE_START_JITTER_SECONDS = 60
|
||||
|
||||
# Rows a (bank, fact_type) partition needs before it gets its own partial vector
|
||||
# index. These indexes live on the *shared* memory_units table: PostgreSQL locks
|
||||
# and builds an IndexOptInfo for every index on a relation at plan time, and
|
||||
# opens every one of them for each DML statement, so one bank's index is a cost
|
||||
# paid by every other bank in the deployment. Three per bank exhausts the lock
|
||||
# table at a few thousand banks (issue #3485).
|
||||
#
|
||||
# 0 is the default and means "no minimum": every partition that holds rows gets
|
||||
# an index, which is the behaviour before the threshold existed. Deployments
|
||||
# holding thousands of banks raise it — above the threshold ANN wins, and below
|
||||
# it PostgreSQL answers the same query from the (bank_id, fact_type) B-tree plus
|
||||
# a top-N sort, which is exact rather than approximate *and* faster, because
|
||||
# sorting a few thousand rows by distance costs less than descending an ANN
|
||||
# graph. 10_000 is a reasonable starting point (it is also ScaNN's own build
|
||||
# floor, SCANN_MIN_ROWS_FOR_AUTO_INDEX).
|
||||
DEFAULT_VECTOR_INDEX_MIN_ROWS = 0
|
||||
|
||||
# A partition that falls back below MIN_ROWS * this ratio loses its index. The
|
||||
# gap between the build and drop thresholds is hysteresis: with a single
|
||||
# boundary, consolidation pruning a bank back and forth across it would rebuild
|
||||
# and drop the same ANN index on alternating writes. At the default threshold of
|
||||
# 0 there is no gap and nothing to flap — a partition either holds rows or does
|
||||
# not.
|
||||
VECTOR_INDEX_DROP_RATIO = 0.5
|
||||
|
||||
# Default MCP tool descriptions (can be customized via env vars)
|
||||
DEFAULT_MCP_RETAIN_DESCRIPTION = """Store important information to long-term memory.
|
||||
|
||||
@@ -2163,6 +2208,8 @@ class HindsightConfig:
|
||||
migration_database_url: str | None
|
||||
database_schema: str
|
||||
vector_extension: str # "pgvector", "vchord", "pgvectorscale", or "scann"
|
||||
ann_iterative_scan: bool
|
||||
ann_max_scan_tuples: int
|
||||
text_search_extension: str # "native", "vchord", "pg_textsearch", "pgroonga", or "pg_search"
|
||||
# PostgreSQL text search dictionary for the "native" backend (ignored by
|
||||
# other backends). Only the "native" backend reads this field; pgroonga
|
||||
@@ -2655,6 +2702,9 @@ class HindsightConfig:
|
||||
# How often the maintenance loop checks for cron-scheduled mental models due for
|
||||
# refresh (the per-model schedule lives in the mental model trigger). 0 = disabled.
|
||||
mental_model_refresh_tick_seconds: int
|
||||
# Rows a (bank, fact_type) needs before it gets its own partial vector index.
|
||||
# 0 (default) = no minimum: every partition holding rows is indexed.
|
||||
vector_index_min_rows: int
|
||||
|
||||
# Webhook configuration (static - server-level only, not per-bank)
|
||||
webhook_url: str | None # Global webhook URL (None = disabled)
|
||||
@@ -2977,6 +3027,12 @@ class HindsightConfig:
|
||||
# Validate vector_extension
|
||||
validate_extension(self.vector_extension)
|
||||
|
||||
if self.ann_iterative_scan and self.ann_max_scan_tuples < 1:
|
||||
raise ValueError(
|
||||
f"Invalid ann_max_scan_tuples: {self.ann_max_scan_tuples}. Must be >= 1 when "
|
||||
f"iterative ANN scans are enabled (set {ENV_ANN_ITERATIVE_SCAN}=false to disable them)"
|
||||
)
|
||||
|
||||
# pg_trgm requires the similarity threshold in (0, 1]. Fail fast here
|
||||
# rather than let an out-of-range value raise on every pool connection's
|
||||
# setup (which would leave the API unable to serve any request).
|
||||
@@ -3145,6 +3201,10 @@ class HindsightConfig:
|
||||
migration_database_url=os.getenv(ENV_MIGRATION_DATABASE_URL) or None,
|
||||
database_schema=os.getenv(ENV_DATABASE_SCHEMA, DEFAULT_DATABASE_SCHEMA),
|
||||
vector_extension=os.getenv(ENV_VECTOR_EXTENSION, DEFAULT_VECTOR_EXTENSION).lower(),
|
||||
ann_iterative_scan=_parse_boolean_env(ENV_ANN_ITERATIVE_SCAN, DEFAULT_ANN_ITERATIVE_SCAN),
|
||||
ann_max_scan_tuples=_parse_non_negative_int(
|
||||
ENV_ANN_MAX_SCAN_TUPLES, os.getenv(ENV_ANN_MAX_SCAN_TUPLES), DEFAULT_ANN_MAX_SCAN_TUPLES
|
||||
),
|
||||
text_search_extension=os.getenv(ENV_TEXT_SEARCH_EXTENSION, DEFAULT_TEXT_SEARCH_EXTENSION).lower(),
|
||||
text_search_extension_native_language=os.getenv(
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
@@ -4017,6 +4077,11 @@ class HindsightConfig:
|
||||
os.getenv(ENV_RETENTION_SWEEP_INTERVAL_SECONDS),
|
||||
DEFAULT_RETENTION_SWEEP_INTERVAL_SECONDS,
|
||||
),
|
||||
vector_index_min_rows=_parse_non_negative_int(
|
||||
ENV_VECTOR_INDEX_MIN_ROWS,
|
||||
os.getenv(ENV_VECTOR_INDEX_MIN_ROWS),
|
||||
DEFAULT_VECTOR_INDEX_MIN_ROWS,
|
||||
),
|
||||
operation_cleanup_interval_seconds=_parse_non_negative_int(
|
||||
ENV_OPERATION_CLEANUP_INTERVAL_SECONDS,
|
||||
os.getenv(ENV_OPERATION_CLEANUP_INTERVAL_SECONDS),
|
||||
|
||||
@@ -55,6 +55,7 @@ if TYPE_CHECKING:
|
||||
from asyncpg import Connection
|
||||
|
||||
from ...api.http import RequestContext
|
||||
from ..memories.base import StoredMemory
|
||||
from ..memory_engine import MemoryEngine
|
||||
from ..response_models import MemoryFact, RecallResult
|
||||
|
||||
@@ -221,6 +222,52 @@ def _dedup_active(config: Any) -> bool:
|
||||
return get_config().database_backend != "oracle"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _TemporalBounds:
|
||||
"""The temporal columns an observation inherits from the facts behind it.
|
||||
|
||||
Merging two observations (or an observation and a fresh set of source facts) must widen
|
||||
these, never replace them: ``event_date``/``occurred_start`` keep the earliest known value
|
||||
and ``occurred_end``/``mentioned_at`` the latest, with a missing value on either side
|
||||
ignored. That is exactly the ``_aggregate_source_fields`` rule, and the Python mirror of the
|
||||
``LEAST``/``GREATEST`` the SQL paths apply.
|
||||
|
||||
The SQL spelling differs by reach, deliberately. The dedup folds only ever run on PostgreSQL
|
||||
(``_dedup_active`` disables dedup on Oracle) and use the plain
|
||||
``LEAST(col, COALESCE(x, col))``, which is enough there because PostgreSQL ignores NULL
|
||||
arguments. ``_execute_update_action`` also runs on Oracle, where LEAST/GREATEST return NULL
|
||||
if any argument is NULL, so it wraps the whole expression in one more COALESCE — see the
|
||||
comment there.
|
||||
"""
|
||||
|
||||
event_date: "datetime | None" = None
|
||||
occurred_start: "datetime | None" = None
|
||||
occurred_end: "datetime | None" = None
|
||||
mentioned_at: "datetime | None" = None
|
||||
|
||||
@classmethod
|
||||
def of(cls, row: "StoredMemory | _SourceAggregation") -> "_TemporalBounds":
|
||||
"""The bounds carried by a stored memory or by an aggregation over source facts.
|
||||
|
||||
Deliberately not a recall ``MemoryFact``: that model has no ``event_date`` at all and
|
||||
keeps the rest as ISO strings, so it has to be read field by field where it is used.
|
||||
"""
|
||||
return cls(
|
||||
event_date=row.event_date,
|
||||
occurred_start=row.occurred_start,
|
||||
occurred_end=row.occurred_end,
|
||||
mentioned_at=row.mentioned_at,
|
||||
)
|
||||
|
||||
def merged_with(self, other: "_TemporalBounds") -> "_TemporalBounds":
|
||||
return _TemporalBounds(
|
||||
event_date=_merge_min(self.event_date, other.event_date),
|
||||
occurred_start=_merge_min(self.occurred_start, other.occurred_start),
|
||||
occurred_end=_merge_max(self.occurred_end, other.occurred_end),
|
||||
mentioned_at=_merge_max(self.mentioned_at, other.mentioned_at),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _DedupOutcome:
|
||||
"""Result of probing one observation against its in-scope neighbours.
|
||||
@@ -322,6 +369,7 @@ async def _dedup_reconcile_create(
|
||||
create_text: str,
|
||||
create_source_ids: list[uuid.UUID],
|
||||
tags: list[str] | None,
|
||||
source_bounds: _TemporalBounds,
|
||||
txn=None,
|
||||
) -> str | None:
|
||||
"""Semantic dedup for a single CREATE (create-time, focused 1-by-1).
|
||||
@@ -330,6 +378,10 @@ async def _dedup_reconcile_create(
|
||||
observation and returns its id (caller skips the CREATE). Returns None when there is
|
||||
no near twin or the LLM keeps them distinct.
|
||||
|
||||
``source_bounds`` are the dates the skipped CREATE would have been stamped with. They are
|
||||
folded into the twin too: this path bypasses the CREATE writer, so without them the twin
|
||||
would cite dated source facts while reporting the dates of its original sources only (#3477).
|
||||
|
||||
The probe/embed/LLM adjudication runs with no connection held; the fold takes a
|
||||
short-lived connection and re-checks source liveness inside the fold transaction.
|
||||
"""
|
||||
@@ -362,6 +414,10 @@ async def _dedup_reconcile_create(
|
||||
SET text = $1,
|
||||
source_memory_ids = (SELECT array_agg(DISTINCT e) FROM unnest(source_memory_ids || $2::uuid[]) e),
|
||||
proof_count = (SELECT count(DISTINCT e) FROM unnest(source_memory_ids || $2::uuid[]) e),
|
||||
event_date = LEAST(event_date, COALESCE($5, event_date)),
|
||||
occurred_start = LEAST(occurred_start, COALESCE($6, occurred_start)),
|
||||
occurred_end = GREATEST(occurred_end, COALESCE($7, occurred_end)),
|
||||
mentioned_at = GREATEST(mentioned_at, COALESCE($8, mentioned_at)),
|
||||
updated_at = now(){search_vector_clause}
|
||||
WHERE id = $3::uuid AND text = $4
|
||||
RETURNING id
|
||||
@@ -370,6 +426,10 @@ async def _dedup_reconcile_create(
|
||||
live_source_ids,
|
||||
uuid.UUID(outcome.best_id),
|
||||
outcome.best_text,
|
||||
source_bounds.event_date,
|
||||
source_bounds.occurred_start,
|
||||
source_bounds.occurred_end,
|
||||
source_bounds.mentioned_at,
|
||||
)
|
||||
if folded is None:
|
||||
# The twin vanished (or was rewritten) during the connection-free LLM window.
|
||||
@@ -382,7 +442,15 @@ async def _dedup_reconcile_create(
|
||||
return None
|
||||
else:
|
||||
await _reconcile_merge_via_store(
|
||||
store, conn, memory_engine, bank_id, outcome.best_id, outcome.merged_text, live_source_ids, txn=txn
|
||||
store,
|
||||
conn,
|
||||
memory_engine,
|
||||
bank_id,
|
||||
outcome.best_id,
|
||||
outcome.merged_text,
|
||||
live_source_ids,
|
||||
source_bounds,
|
||||
txn=txn,
|
||||
)
|
||||
return outcome.best_id
|
||||
|
||||
@@ -426,7 +494,8 @@ async def _dedup_reconcile_update(
|
||||
# Fold the updated observation's live sources into the twin (keeping the twin's embedding, as
|
||||
# in the create path) then delete the now-redundant updated row. The all_strict/any tag match
|
||||
# guarantees twin and updated share scope, so dropping the updated row's tags loses no
|
||||
# visibility. Temporal fields follow the surviving twin (minimal scope; matches create).
|
||||
# visibility. Temporal fields are the UNION of both rows' bounds: the updated row is about to
|
||||
# be deleted, so anything only it knew about would otherwise be lost with it (#3477).
|
||||
# The fold + delete share one short transaction so the twin gains the sources exactly as the
|
||||
# redundant row is removed; the slow adjudication above already ran connection-free.
|
||||
store = get_memories()
|
||||
@@ -469,6 +538,10 @@ async def _dedup_reconcile_update(
|
||||
proof_count = (
|
||||
SELECT count(DISTINCT e) FROM unnest(t.source_memory_ids || $6::uuid[]) e
|
||||
),
|
||||
event_date = LEAST(t.event_date, COALESCE(u.event_date, t.event_date)),
|
||||
occurred_start = LEAST(t.occurred_start, COALESCE(u.occurred_start, t.occurred_start)),
|
||||
occurred_end = GREATEST(t.occurred_end, COALESCE(u.occurred_end, t.occurred_end)),
|
||||
mentioned_at = GREATEST(t.mentioned_at, COALESCE(u.mentioned_at, t.mentioned_at)),
|
||||
updated_at = now(){search_vector_clause}
|
||||
FROM {fq_table("memory_units")} u
|
||||
WHERE t.id = $2::uuid AND u.id = $3::uuid AND t.text = $4 AND u.text = $5
|
||||
@@ -494,7 +567,15 @@ async def _dedup_reconcile_update(
|
||||
if not live_u_sources:
|
||||
return
|
||||
await _reconcile_merge_via_store(
|
||||
store, conn, memory_engine, bank_id, outcome.best_id, outcome.merged_text, live_u_sources, txn=txn
|
||||
store,
|
||||
conn,
|
||||
memory_engine,
|
||||
bank_id,
|
||||
outcome.best_id,
|
||||
outcome.merged_text,
|
||||
live_u_sources,
|
||||
_TemporalBounds.of(updated_obs[0]),
|
||||
txn=txn,
|
||||
)
|
||||
await _execute_delete_action(conn, bank_id, updated_id, txn=txn)
|
||||
logger.info(
|
||||
@@ -1001,17 +1082,22 @@ async def _reconcile_merge_via_store(
|
||||
observation_id: str,
|
||||
merged_text: str,
|
||||
add_source_ids: list,
|
||||
add_bounds: _TemporalBounds,
|
||||
txn=None,
|
||||
) -> None:
|
||||
"""Dedup merge for a store that owns its rows: fold the extra source facts and the merged text
|
||||
into the twin observation and re-upsert it, preserving its other fields. Re-embeds the merged
|
||||
text because ``get_memories`` does not return the stored vector (the SQL path reuses it in
|
||||
place instead)."""
|
||||
place instead).
|
||||
|
||||
``add_bounds`` are the folded-in side's dates, widened onto the twin exactly as the SQL
|
||||
path's LEAST/GREATEST does."""
|
||||
current = await store.get_memories(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[observation_id])
|
||||
cur = current[0] if current else None
|
||||
if cur is None:
|
||||
return
|
||||
merged_sources = list(dict.fromkeys([*(cur.source_memory_ids or []), *(str(s) for s in add_source_ids)]))
|
||||
merged_bounds = _TemporalBounds.of(cur).merged_with(add_bounds)
|
||||
embeddings = await embedding_utils.generate_embeddings_batch(memory_engine.embeddings, [merged_text])
|
||||
await store.upsert_observation(
|
||||
conn=conn,
|
||||
@@ -1025,10 +1111,10 @@ async def _reconcile_merge_via_store(
|
||||
tags=list(cur.tags or []),
|
||||
proof_count=len(merged_sources),
|
||||
source_memory_ids=merged_sources,
|
||||
event_date=cur.event_date,
|
||||
occurred_start=cur.occurred_start,
|
||||
occurred_end=cur.occurred_end,
|
||||
mentioned_at=cur.mentioned_at,
|
||||
event_date=merged_bounds.event_date,
|
||||
occurred_start=merged_bounds.occurred_start,
|
||||
occurred_end=merged_bounds.occurred_end,
|
||||
mentioned_at=merged_bounds.mentioned_at,
|
||||
created_at=cur.created_at,
|
||||
),
|
||||
)
|
||||
@@ -2134,9 +2220,7 @@ async def _process_memory_batch(
|
||||
new_text=update.text,
|
||||
observations=union_observations,
|
||||
source_fact_tags=agg.tags,
|
||||
source_occurred_start=agg.occurred_start,
|
||||
source_occurred_end=agg.occurred_end,
|
||||
source_mentioned_at=agg.mentioned_at,
|
||||
source_bounds=_TemporalBounds.of(agg),
|
||||
perf=perf,
|
||||
txn=txn,
|
||||
)
|
||||
@@ -2204,6 +2288,7 @@ async def _process_memory_batch(
|
||||
create.text,
|
||||
create_source_ids,
|
||||
agg.tags,
|
||||
_TemporalBounds.of(agg),
|
||||
txn=txn,
|
||||
)
|
||||
if merged_into is not None:
|
||||
@@ -2338,17 +2423,15 @@ async def _execute_update_action(
|
||||
new_text: str,
|
||||
observations: list["MemoryFact"],
|
||||
source_fact_tags: list[str] | None = None,
|
||||
source_occurred_start: datetime | None = None,
|
||||
source_occurred_end: datetime | None = None,
|
||||
source_mentioned_at: datetime | None = None,
|
||||
source_bounds: _TemporalBounds = _TemporalBounds(),
|
||||
perf: ConsolidationPerfLog | None = None,
|
||||
txn=None,
|
||||
) -> str | None:
|
||||
"""
|
||||
Update an existing observation.
|
||||
|
||||
Extends source_memory_ids with all contributing memories, updates temporal fields
|
||||
(LEAST for occurred_start, GREATEST for occurred_end / mentioned_at), and merges tags.
|
||||
Extends source_memory_ids with all contributing memories, widens the observation's temporal
|
||||
bounds by ``source_bounds`` (see :class:`_TemporalBounds`), and merges tags.
|
||||
|
||||
The embedding is computed off-connection (a slow embedder must never pin a pooled
|
||||
connection); the liveness check + UPDATE + history + observation_sources sync then run
|
||||
@@ -2418,6 +2501,15 @@ async def _execute_update_action(
|
||||
|
||||
t0 = time.time()
|
||||
if store.writes_memory_rows_in_sql_for(bank_id):
|
||||
# Unlike the dedup folds this statement also runs on Oracle, where LEAST/GREATEST
|
||||
# return NULL as soon as ANY argument is NULL (PostgreSQL ignores NULL arguments).
|
||||
# The inner COALESCE covers a NULL *parameter*; the outer one covers a NULL
|
||||
# *column* — an observation with no occurred interval yet, which is precisely the
|
||||
# #3477 case. Without it Oracle would compute LEAST(NULL, <source date>) = NULL and
|
||||
# silently drop the date it was told to inherit. Keep the inner
|
||||
# ``COALESCE($n, col)`` spelled exactly like this: the Oracle driver shim keys its
|
||||
# TIMESTAMP-TZ input-size hint off that pattern (db/oracle.py::_apply_clob_input_sizes),
|
||||
# and a NULL parameter binds as VARCHAR2 (ORA-00932) without it.
|
||||
updated_rows = await conn.execute_rows_affected(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
@@ -2425,11 +2517,12 @@ async def _execute_update_action(
|
||||
embedding = $2::vector,
|
||||
source_memory_ids = $3,
|
||||
proof_count = $4,
|
||||
tags = $9,
|
||||
tags = $10,
|
||||
updated_at = now(),
|
||||
occurred_start = LEAST(occurred_start, COALESCE($6, occurred_start)),
|
||||
occurred_end = GREATEST(occurred_end, COALESCE($7, occurred_end)),
|
||||
mentioned_at = GREATEST(mentioned_at, COALESCE($8, mentioned_at)){search_vector_clause}
|
||||
event_date = COALESCE(LEAST(event_date, COALESCE($6, event_date)), $6),
|
||||
occurred_start = COALESCE(LEAST(occurred_start, COALESCE($7, occurred_start)), $7),
|
||||
occurred_end = COALESCE(GREATEST(occurred_end, COALESCE($8, occurred_end)), $8),
|
||||
mentioned_at = COALESCE(GREATEST(mentioned_at, COALESCE($9, mentioned_at)), $9){search_vector_clause}
|
||||
WHERE id = $5
|
||||
""",
|
||||
new_text,
|
||||
@@ -2437,9 +2530,10 @@ async def _execute_update_action(
|
||||
source_ids,
|
||||
len(source_ids),
|
||||
uuid.UUID(observation_id),
|
||||
source_occurred_start,
|
||||
source_occurred_end,
|
||||
source_mentioned_at,
|
||||
source_bounds.event_date,
|
||||
source_bounds.occurred_start,
|
||||
source_bounds.occurred_end,
|
||||
source_bounds.mentioned_at,
|
||||
merged_tags,
|
||||
)
|
||||
# The source-liveness checks above guard the *source* memories; the
|
||||
@@ -2457,12 +2551,24 @@ async def _execute_update_action(
|
||||
return None
|
||||
else:
|
||||
# Upsert overwrites the whole observation, so start from its current state (fetched
|
||||
# from the store) and apply the same merge the SQL does — LEAST/GREATEST on the times
|
||||
# — while preserving fields the update never touches (event_date, created_at).
|
||||
# from the store) and apply the same merge the SQL does — LEAST/GREATEST on the
|
||||
# times — while preserving fields the update never touches (created_at).
|
||||
current = await store.get_memories(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[observation_id]
|
||||
)
|
||||
cur = current[0] if current else None
|
||||
# Widen the row the store still holds. If it has vanished, fall back to the
|
||||
# pre-update recall snapshot — ISO strings, and no event_date on that model.
|
||||
current_bounds = (
|
||||
_TemporalBounds.of(cur)
|
||||
if cur
|
||||
else _TemporalBounds(
|
||||
occurred_start=_as_dt(model.occurred_start),
|
||||
occurred_end=_as_dt(model.occurred_end),
|
||||
mentioned_at=_as_dt(model.mentioned_at),
|
||||
)
|
||||
)
|
||||
merged_bounds = current_bounds.merged_with(source_bounds)
|
||||
await store.upsert_observation(
|
||||
conn=conn,
|
||||
bank_id=bank_id,
|
||||
@@ -2475,10 +2581,10 @@ async def _execute_update_action(
|
||||
tags=merged_tags,
|
||||
proof_count=len(source_ids),
|
||||
source_memory_ids=[str(s) for s in source_ids],
|
||||
event_date=cur.event_date if cur else None,
|
||||
occurred_start=_merge_min(model.occurred_start, source_occurred_start),
|
||||
occurred_end=_merge_max(model.occurred_end, source_occurred_end),
|
||||
mentioned_at=_merge_max(model.mentioned_at, source_mentioned_at),
|
||||
event_date=merged_bounds.event_date,
|
||||
occurred_start=merged_bounds.occurred_start,
|
||||
occurred_end=merged_bounds.occurred_end,
|
||||
mentioned_at=merged_bounds.mentioned_at,
|
||||
created_at=cur.created_at if cur else None,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -481,23 +481,11 @@ class DataAccessOps(ABC):
|
||||
|
||||
# -- Bank index management -------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def create_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
internal_id: str,
|
||||
index_clause: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
"""Create per-bank partial vector indexes.
|
||||
|
||||
PG creates per-(bank, fact_type) partial indexes.
|
||||
Non-PG is a no-op (uses global index).
|
||||
"""
|
||||
...
|
||||
|
||||
# No create counterpart: per-(bank, fact_type) partial vector indexes are
|
||||
# earned by size and built by the maintenance sweep over its own autocommit
|
||||
# connection (see engine/vector_index_health.py), never on a request path.
|
||||
# The drop stays here because bank deletion must remove a large bank's
|
||||
# indexes while it still knows the internal_id they are named after.
|
||||
@abstractmethod
|
||||
async def drop_bank_vector_indexes(
|
||||
self,
|
||||
|
||||
@@ -696,6 +696,21 @@ class OracleOps(DataAccessOps):
|
||||
# Entity expansion via observation_sources junction table.
|
||||
# Previously used JSON_TABLE to explode source_memory_ids CLOB. The junction
|
||||
# table approach uses standard SQL joins, identical to the PG backend.
|
||||
#
|
||||
# Two PostgreSQL fixes are deliberately NOT mirrored here, because neither
|
||||
# was measured against Oracle and both are tuned to PostgreSQL's planner:
|
||||
# - #3085 made PG score set-wise; the scoring below is still the
|
||||
# correlated per-observation COUNT(*). On Oracle that counts rows of
|
||||
# the indexed observation_sources junction table rather than scanning
|
||||
# an unpruned array, so it is a much weaker version of that problem.
|
||||
# - #3510 replaced PG's `DISTINCT` over a `LATERAL ... LIMIT` with a
|
||||
# row_number() window, because PostgreSQL cannot estimate the row count
|
||||
# of that shape and mis-planned the scoring join into a nested loop.
|
||||
# `connected_sources` below has the same shape, so the same collapse is
|
||||
# structurally possible, but Oracle's cardinality estimation differs and
|
||||
# no Oracle instance was available to measure it.
|
||||
# If observation recall is reported slow on Oracle, start by capturing the
|
||||
# plan for connected_sources and checking its estimated vs actual rows.
|
||||
from ..schema import fq_table
|
||||
|
||||
obs_sources_table = fq_table("observation_sources")
|
||||
@@ -829,23 +844,6 @@ class OracleOps(DataAccessOps):
|
||||
bank_prefix="mu.",
|
||||
)
|
||||
|
||||
async def create_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
internal_id: str,
|
||||
index_clause: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
# Oracle 23ai supports HNSW vector indexes but does NOT support partial
|
||||
# indexes (WHERE clause on CREATE INDEX for vector indexes). Uses a single
|
||||
# global HNSW index with ORGANIZATION NEIGHBOR PARTITIONS created during
|
||||
# migrations. memory_units is partitioned by LIST (bank_id) AUTOMATIC,
|
||||
# so Oracle creates partitions per bank on INSERT and the optimizer can
|
||||
# prune partitions on bank_id-scoped queries.
|
||||
return
|
||||
|
||||
async def drop_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
@@ -853,7 +851,12 @@ class OracleOps(DataAccessOps):
|
||||
internal_id: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
# Oracle uses a single global vector index (no per-bank indexes to drop).
|
||||
# Oracle uses a single global vector index — it does not support partial
|
||||
# (WHERE-clause) vector indexes, so there are no per-bank ones to drop.
|
||||
# Bank scoping comes from the table itself instead: memory_units is
|
||||
# partitioned LIST (bank_id) AUTOMATIC, so Oracle creates a partition per
|
||||
# bank on INSERT and the optimizer prunes on bank_id. That is why the
|
||||
# size threshold and its sweep are PostgreSQL-only concerns.
|
||||
return
|
||||
|
||||
def get_entity_resolution_strategy(self) -> str:
|
||||
|
||||
@@ -914,6 +914,27 @@ class PostgreSQLOps(DataAccessOps):
|
||||
# ~1.7B element comparisons, 2.6s of one saturated backend (issue #3085).
|
||||
# Unnesting once and hash-joining connected_sources makes the work linear in
|
||||
# the number of source ids instead.
|
||||
#
|
||||
# `connected_sources` caps each entity with row_number() rather than the
|
||||
# LATERAL + LIMIT that reads more naturally. Do not "simplify" it back
|
||||
# (issue #3510). The scoring join above is O(C + U) when planned as a hash
|
||||
# join and O(U x C) when planned as a nested loop — 15s and ~15M rejected
|
||||
# rows on a realistically-shaped bank — and PostgreSQL picks between them
|
||||
# from its row estimate for this CTE. Out of a LATERAL + LIMIT subquery the
|
||||
# capped column carries no n_distinct statistic, so DISTINCT over it was
|
||||
# estimated at 2 and the NOT EXISTS took that to 1 against an actual ~3,700;
|
||||
# a 1-row inner side makes the nested loop look free, so it won on cost and
|
||||
# lost by four orders of magnitude at runtime. Ranking with a window keeps
|
||||
# the column traceable to unit_entities.unit_id, so the estimate comes from
|
||||
# real statistics (207-3,449 against 2,242-4,193 actual) and the nested loop
|
||||
# is priced honestly.
|
||||
#
|
||||
# The trade is that this reads every unit_entities row of a matched entity
|
||||
# to rank it, where the LATERAL stopped at per_entity_limit off the index:
|
||||
# O(sum of degree) rather than O(entities x per_entity_limit). Measured at
|
||||
# parity up to ~12k-degree hubs and +50% traversal cost at 38k. If banks
|
||||
# grow hubs far past that, re-measure before assuming this is still the
|
||||
# right shape.
|
||||
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -930,17 +951,20 @@ class PostgreSQLOps(DataAccessOps):
|
||||
),
|
||||
connected_sources AS (
|
||||
SELECT DISTINCT t.unit_id AS source_id
|
||||
FROM source_entities se
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue_target.unit_id
|
||||
FROM (
|
||||
SELECT
|
||||
ue_target.unit_id,
|
||||
row_number() OVER (
|
||||
PARTITION BY ue_target.entity_id
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
) AS rn
|
||||
FROM {ue_table} ue_target
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
LIMIT {per_entity_limit}
|
||||
JOIN source_entities se ON se.entity_id = ue_target.entity_id
|
||||
) t
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM seed_sources ss WHERE ss.source_id = t.unit_id
|
||||
)
|
||||
WHERE t.rn <= {per_entity_limit}
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM seed_sources ss WHERE ss.source_id = t.unit_id
|
||||
)
|
||||
),
|
||||
connected_array AS (
|
||||
SELECT array_agg(source_id) AS source_ids FROM connected_sources
|
||||
@@ -1047,26 +1071,6 @@ class PostgreSQLOps(DataAccessOps):
|
||||
bank_prefix="",
|
||||
)
|
||||
|
||||
async def create_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
internal_id: str,
|
||||
index_clause: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
escaped = bank_id.replace("'", "''")
|
||||
async with self._index_ddl_lock(table):
|
||||
for ft, suffix in fact_types.items():
|
||||
uid = str(internal_id).replace("-", "")[:16]
|
||||
idx = f"idx_mu_emb_{suffix}_{uid}"
|
||||
await conn.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS {idx} "
|
||||
f"ON {table} {index_clause} "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped}'"
|
||||
)
|
||||
|
||||
async def drop_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
@@ -1080,8 +1084,14 @@ class PostgreSQLOps(DataAccessOps):
|
||||
# table; CONCURRENTLY does not conflict with DML. The caller
|
||||
# (delete_bank) runs this on an autocommit connection after its delete
|
||||
# transaction has committed — CONCURRENTLY cannot run inside a tx.
|
||||
# The lock key must match create_bank_vector_indexes', whose `table`
|
||||
# is the fq name this reconstructs from `schema`.
|
||||
#
|
||||
# The in-process lock serializes concurrent bank deletes against each
|
||||
# other. It does not cover the maintenance sweep, which reconciles the
|
||||
# same indexes over its own raw connection: an in-process lock could not
|
||||
# help there anyway, since the sweep runs in every process and the real
|
||||
# contention is cross-process. Both paths retry the transient deadlock
|
||||
# (40P01) instead, which is the only lock-free option available — the
|
||||
# project forbids advisory locks (unreliable behind poolers, #2817).
|
||||
async with self._index_ddl_lock(f"{schema}.memory_units"):
|
||||
for ft, suffix in fact_types.items():
|
||||
uid = str(internal_id).replace("-", "")[:16]
|
||||
|
||||
@@ -20,6 +20,24 @@ from .pool_instrumentation import PoolStats, instrument_acquire
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# GUC names this server rejected as unknown. Process-wide and never cleared: the
|
||||
# server's extension set does not change under a running process, and re-probing
|
||||
# would reintroduce the per-acquire cost this exists to avoid.
|
||||
_unsupported_settings: set[str] = set()
|
||||
|
||||
|
||||
def setting_rejected_by_server(name: str) -> bool:
|
||||
"""Whether this server has already rejected ``name`` as an unknown GUC.
|
||||
|
||||
For callers that apply a setting outside this helper — notably retain's link
|
||||
probing, which uses SET LOCAL inside its own transaction so the value cannot leak
|
||||
onto a pooled backend. Such a caller cannot simply let the statement fail: an error
|
||||
inside a transaction poisons it, so an unknown GUC would abort its work rather than
|
||||
merely fail to apply. The pool's setup runs on acquire and names the same GUCs, so
|
||||
by the time one of those callers runs, an unknown one is already recorded here.
|
||||
"""
|
||||
return name in _unsupported_settings
|
||||
|
||||
|
||||
async def apply_session_settings(conn: asyncpg.Connection, settings: list[tuple[str, str]]) -> None:
|
||||
"""Apply session-scoped GUCs to ``conn`` in a single round trip.
|
||||
@@ -37,6 +55,7 @@ async def apply_session_settings(conn: asyncpg.Connection, settings: list[tuple[
|
||||
statement fails as a whole, so on error fall back to applying them one by
|
||||
one, skipping only the ones the server rejects.
|
||||
"""
|
||||
settings = [pair for pair in settings if pair[0] not in _unsupported_settings]
|
||||
if not settings:
|
||||
return
|
||||
|
||||
@@ -53,8 +72,19 @@ async def apply_session_settings(conn: asyncpg.Connection, settings: list[tuple[
|
||||
for name, value in settings:
|
||||
try:
|
||||
await conn.execute("SELECT set_config($1, $2, false)", name, value)
|
||||
except asyncpg.exceptions.UndefinedObjectError:
|
||||
# The server does not define this GUC — an extension we tune for is absent
|
||||
# or predates it (hnsw.iterative_scan needs pgvector 0.8+, and pgvector
|
||||
# reserves the "hnsw." prefix, so an older one rejects it rather than
|
||||
# accepting a placeholder). Remember it: otherwise every acquire from here
|
||||
# on re-pays a failed batch plus one statement per setting, which behind a
|
||||
# transaction-mode pooler is a server-side transaction each — the burn
|
||||
# #3499 removed. Narrow to UndefinedObjectError so a transient failure
|
||||
# does not disable a setting the server does support.
|
||||
logger.info("Server does not know %s — not sending it again on this process", name)
|
||||
_unsupported_settings.add(name)
|
||||
except asyncpg.exceptions.PostgresError:
|
||||
logger.debug("Could not set %s — the server may not know this setting", name)
|
||||
logger.debug("Could not set %s — retrying it on the next acquire", name)
|
||||
|
||||
|
||||
class PostgresConnection(DatabaseConnection):
|
||||
@@ -164,6 +194,12 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
self._pool: asyncpg.Pool | None = None
|
||||
self._acquire_warn_threshold_s: float = 1.0
|
||||
self._acquire_timeout_s: float | None = None
|
||||
self._dsn: str | None = None
|
||||
|
||||
@property
|
||||
def dsn(self) -> str | None:
|
||||
"""The DSN this backend's pool was opened with, if it has been initialized."""
|
||||
return self._dsn
|
||||
|
||||
async def initialize(
|
||||
self,
|
||||
@@ -179,6 +215,13 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
from ...config import get_config
|
||||
|
||||
config = get_config()
|
||||
# Kept so code that needs its *own* connection — CREATE/DROP INDEX
|
||||
# CONCURRENTLY cannot run on a pooled one inside a transaction — can
|
||||
# reach the database this engine is actually attached to. Re-deriving it
|
||||
# from HINDSIGHT_API_DATABASE_URL is wrong whenever the engine was handed
|
||||
# a DSN directly (embedders, and the test suite, which resolves pg0 in a
|
||||
# fixture and never sets the env var).
|
||||
self._dsn = dsn
|
||||
self._acquire_warn_threshold_s = config.db_acquire_warn_threshold_ms / 1000.0
|
||||
# Kept for acquire() below: asyncpg's ``timeout`` create_pool kwarg is a
|
||||
# *connect* kwarg (how long establishing a new connection may take), and
|
||||
|
||||
@@ -45,6 +45,9 @@ class _EntityToCreate:
|
||||
# Also stored on the row as entities.entity_kind so label rows stay out of the
|
||||
# partial trigram index (#3208).
|
||||
is_label: bool = False
|
||||
# False when the caller wrote this name literally: it is created as spelled and never
|
||||
# merged with a same-batch near-duplicate (#3479).
|
||||
resolve: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -428,6 +431,18 @@ class EntityResolver:
|
||||
unit_event_date: When this unit was created
|
||||
conn: Optional connection to use (if None, acquires from pool)
|
||||
|
||||
Each mention may carry ``"resolve": False`` to opt out of resolution. The
|
||||
default, True, treats a name as a *guess* at which entity is meant, so
|
||||
similar existing entities are scored on name similarity + co-occurrence +
|
||||
recency and the best above threshold is reused. False takes the name
|
||||
literally: an existing entity is reused only when its canonical name
|
||||
matches case-insensitively, any other name creates its own entity, and it
|
||||
is never merged with a same-batch near-duplicate. Callers who authored the
|
||||
names deliberately want False (#3479) — resolution would otherwise let
|
||||
what the graph already believes outscore, and silently discard, their
|
||||
correction. It is per mention because retain resolves caller-supplied and
|
||||
extracted names in one batch, and only the caller's half is authoritative.
|
||||
|
||||
Returns:
|
||||
Resolved entity identities (id + stored canonical name) in the same
|
||||
order as input.
|
||||
@@ -457,6 +472,25 @@ class EntityResolver:
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[ResolvedEntity]:
|
||||
# `entities_data and` matters: an empty batch must fall through to the normal strategy
|
||||
# dispatch (which the pg_trgm auto-detection hangs off), not take the shortcut vacuously.
|
||||
if entities_data and not any(e.get("resolve", True) for e in entities_data):
|
||||
# Nothing in this batch resolves, so the trigram/UTL_MATCH probe and the
|
||||
# co-occurrence fetch would both be dead work. _resolve_from_candidates routes every
|
||||
# mention straight to its find-or-create path, which matches on LOWER(canonical_name)
|
||||
# equality. A *mixed* batch still probes — the per-mention check below skips the
|
||||
# literal names when scoring, which costs a little wasted lookup but keeps the
|
||||
# common all-resolving case on one code path.
|
||||
return await self._resolve_from_candidates(
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates={},
|
||||
cooccurrence_map={},
|
||||
taxonomy_lookup=taxonomy_lookup,
|
||||
labels_cfg=labels_cfg,
|
||||
)
|
||||
if self.entity_lookup == "trigram":
|
||||
# Route to backend-specific fuzzy strategy.
|
||||
# Non-PG backends (Oracle) use UTL_MATCH instead of pg_trgm.
|
||||
@@ -872,7 +906,7 @@ class EntityResolver:
|
||||
rep_by_lower: dict[str, str] = {}
|
||||
count_by_lower: dict[str, int] = {}
|
||||
for e in entities_to_create:
|
||||
if e.is_label:
|
||||
if e.is_label or not e.resolve:
|
||||
continue
|
||||
name_lower = e.name.lower()
|
||||
rep_by_lower.setdefault(name_lower, e.name)
|
||||
@@ -904,7 +938,12 @@ class EntityResolver:
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[ResolvedEntity]:
|
||||
"""Shared scoring + upsert logic used by both lookup strategies."""
|
||||
"""Shared scoring + upsert logic used by every lookup strategy.
|
||||
|
||||
A mention carrying ``"resolve": False`` skips the scoring entirely and takes the
|
||||
find-or-create path below, which matches an existing row on ``LOWER(canonical_name)``
|
||||
equality and inserts one otherwise.
|
||||
"""
|
||||
|
||||
# Resolve each entity using pre-fetched candidates. A slot stays None
|
||||
# only if find-or-create fails to produce a row for a mention (a DB
|
||||
@@ -923,6 +962,9 @@ class EntityResolver:
|
||||
# Use per-entity date if available, otherwise fall back to batch-level date
|
||||
entity_event_date = entity_data.get("event_date", unit_event_date)
|
||||
|
||||
# Per mention, not per batch: retain resolves the caller's entities and the
|
||||
# extractor's in one pass, and only the caller's are meant literally (#3479).
|
||||
resolve = entity_data.get("resolve", True)
|
||||
candidates = all_candidates.get(entity_text, [])
|
||||
|
||||
# Backstop truncation for candidate sets that were not capped at the
|
||||
@@ -951,10 +993,14 @@ class EntityResolver:
|
||||
# classify by key prefix (see _label_texts).
|
||||
is_label = bool(labels_cfg and _is_label_entity(entity_text, labels_cfg, taxonomy_lookup or set()))
|
||||
|
||||
if not candidates:
|
||||
# Will create new entity
|
||||
if not resolve or not candidates:
|
||||
# Nothing to score against — or the caller named the entity literally, so
|
||||
# similarity must not get a vote. Either way the find-or-create pass below
|
||||
# reuses an identically-named row and otherwise inserts this exact name.
|
||||
entities_to_create.append(
|
||||
_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date, is_label=is_label)
|
||||
_EntityToCreate(
|
||||
idx=idx, name=entity_text, event_date=entity_event_date, is_label=is_label, resolve=resolve
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
@@ -1057,7 +1103,9 @@ class EntityResolver:
|
||||
# variants (case/emoji/suffix/typo of one name) collapse to a single entity. Without
|
||||
# this, resolution only compares against already-persisted rows, so the first sighting
|
||||
# of each variant in a batch always creates a distinct entity (issue #3107). Labels are
|
||||
# excluded and keep exact grouping.
|
||||
# excluded and keep exact grouping, and so are names the caller wrote literally:
|
||||
# "Alice" and "Alice Smith" listed side by side are two entities because they were
|
||||
# written as two (#3479).
|
||||
canonical_by_member = self._intrabatch_canonical_map(entities_to_create)
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -160,16 +160,22 @@ class MemoryEngineInterface(ABC):
|
||||
async def list_banks(
|
||||
self,
|
||||
*,
|
||||
search_query: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: "RequestContext",
|
||||
) -> list[dict[str, Any]]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
List all memory banks.
|
||||
List memory banks, one page at a time.
|
||||
|
||||
Args:
|
||||
search_query: Case-insensitive substring matched against bank ID and name.
|
||||
limit: Maximum number of banks to return (0 returns none).
|
||||
offset: Number of banks to skip.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
List of bank info dicts.
|
||||
Dict with ``banks`` (the page), ``total``, ``limit`` and ``offset``.
|
||||
"""
|
||||
...
|
||||
|
||||
@@ -322,7 +328,7 @@ class MemoryEngineInterface(ABC):
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
fact_type: str | None = None,
|
||||
fact_type: str | list[str] | None = None,
|
||||
search_query: str | None = None,
|
||||
entity_id: str | None = None,
|
||||
created_before: datetime | None = None,
|
||||
@@ -335,7 +341,8 @@ class MemoryEngineInterface(ABC):
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
fact_type: Filter by fact type.
|
||||
fact_type: Filter by fact type. A list matches any of them; an empty
|
||||
list is treated as no filter.
|
||||
search_query: Full-text search query.
|
||||
entity_id: Filter to memory units linked to this entity ID.
|
||||
created_before: Keep units with ``created_at`` before this instant.
|
||||
|
||||
@@ -1043,7 +1043,7 @@ class MemoriesExtension(Extension, ABC):
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | None = None,
|
||||
fact_type: str | list[str] | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
|
||||
@@ -84,7 +84,7 @@ async def list_memory_units(
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | None = None,
|
||||
fact_type: str | list[str] | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
@@ -104,7 +104,8 @@ async def list_memory_units(
|
||||
ops: Dialect ops. Unused by this query; part of the interface signature.
|
||||
fq_table: Table-name resolver.
|
||||
bank_id: Filter by bank ID
|
||||
fact_type: Filter by fact type (world, experience)
|
||||
fact_type: Filter by fact type (world, experience). A list matches any of
|
||||
them; an empty list is treated as no filter.
|
||||
search_query: Full-text search query (searches text and context fields)
|
||||
document_id: Optional filter to a single source document.
|
||||
tags: Optional list of tag names to filter by. When omitted, no tag
|
||||
@@ -154,8 +155,14 @@ async def list_memory_units(
|
||||
|
||||
if fact_type:
|
||||
param_count += 1
|
||||
query_conditions.append(f"fact_type = ${param_count}")
|
||||
query_params.append(fact_type)
|
||||
if isinstance(fact_type, str):
|
||||
query_conditions.append(f"fact_type = ${param_count}")
|
||||
query_params.append(fact_type)
|
||||
else:
|
||||
# A list is "any of these" — one array parameter rather than an IN list
|
||||
# whose placeholder count varies with the caller's argument.
|
||||
query_conditions.append(f"fact_type = ANY(${param_count}::text[])")
|
||||
query_params.append(list(fact_type))
|
||||
|
||||
if document_id:
|
||||
param_count += 1
|
||||
@@ -240,7 +247,8 @@ async def list_memory_units(
|
||||
f"""
|
||||
SELECT id, text, event_date, context, fact_type, document_id,
|
||||
mentioned_at, occurred_start, occurred_end, chunk_id, proof_count,
|
||||
tags, metadata, consolidated_at, consolidation_failed_at, edited_at, {curation_cols}
|
||||
tags, metadata, consolidated_at, consolidation_failed_at, edited_at,
|
||||
updated_at, source_memory_ids, {curation_cols}
|
||||
FROM {source_table}
|
||||
{where_clause}
|
||||
ORDER BY mentioned_at DESC NULLS LAST, created_at DESC
|
||||
@@ -304,6 +312,12 @@ async def list_memory_units(
|
||||
"invalidation_reason": row["invalidation_reason"],
|
||||
"invalidated_at": row["invalidated_at"].isoformat() if row["invalidated_at"] else None,
|
||||
"edited_at": row["edited_at"].isoformat() if row["edited_at"] else None,
|
||||
# Both come off the row already selected above, so neither adds a
|
||||
# query: updated_at is the write watermark curation and freshness
|
||||
# checks compare against, and source_memory_ids is an observation's
|
||||
# lineage (empty for a source fact).
|
||||
"updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
|
||||
"source_memory_ids": [str(sid) for sid in row["source_memory_ids"] or []],
|
||||
}
|
||||
)
|
||||
|
||||
@@ -463,7 +477,7 @@ async def list_entities(
|
||||
# Get paginated entities
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
|
||||
SELECT id, canonical_name, entity_kind, mention_count, first_seen, last_seen, metadata
|
||||
FROM {fq_table("entities")}
|
||||
WHERE {where_clause}
|
||||
ORDER BY mention_count DESC, last_seen DESC, id ASC
|
||||
@@ -490,6 +504,9 @@ async def list_entities(
|
||||
{
|
||||
"id": str(row["id"]),
|
||||
"canonical_name": row["canonical_name"],
|
||||
# How the entity was classified (label vs free-form, etc.); same row,
|
||||
# so listing it costs nothing extra.
|
||||
"entity_kind": row["entity_kind"],
|
||||
"mention_count": row["mention_count"],
|
||||
"first_seen": row["first_seen"].isoformat() if row["first_seen"] else None,
|
||||
"last_seen": row["last_seen"].isoformat() if row["last_seen"] else None,
|
||||
|
||||
@@ -107,8 +107,8 @@ class PostgresMemories(MemoriesExtension):
|
||||
The per-arm split is Postgres's own business, kept off the interface: this reproduces the
|
||||
exact orchestration recall used before it was unified — one dense+BM25 UNION query and the
|
||||
temporal query share a single connection, then the graph retriever runs per fact_type on the
|
||||
pool in parallel, seeded by the same dense over-fetch. Result is byte-identical to running
|
||||
the arms separately; fusion/rerank still happen downstream.
|
||||
pool in parallel, seeded by the same dense results. Result is byte-identical to running the
|
||||
arms separately; fusion/rerank still happen downstream.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
@@ -172,7 +172,7 @@ class PostgresMemories(MemoriesExtension):
|
||||
)
|
||||
|
||||
# Graph per fact_type in parallel, on the pool, after the dense connection is released —
|
||||
# seeded by the dense over-fetch (preselected_semantic_seeds), matching the prior path.
|
||||
# seeded by the dense results (preselected_semantic_seeds), matching the prior path.
|
||||
graph_by_ft: dict[str, list] = {ft: [] for ft in fact_types}
|
||||
if enable_graph:
|
||||
assert retriever is not None # only resolved when the arm is on
|
||||
@@ -228,6 +228,14 @@ class PostgresMemories(MemoriesExtension):
|
||||
min_keyword: float | None = None,
|
||||
graph_seed_min_similarity: float | None = None,
|
||||
) -> "dict[str, SemanticBm25Result]":
|
||||
"""The dense + keyword arms, as one UNION query.
|
||||
|
||||
How deep the ANN scan goes is not decided here: the connection carries
|
||||
``hnsw.iterative_scan``, which lets the scan resume until this query's own LIMIT
|
||||
is met (see ``_ANN_TUNING_HIGH_RECALL``). Before that was enabled the scan
|
||||
stopped at ``hnsw.ef_search`` rows — a fixed 200 — so a larger recall budget
|
||||
widened the SQL and changed nothing.
|
||||
"""
|
||||
# Imported here: retrieval imports this package, so a module-level import
|
||||
# would close the cycle.
|
||||
from ..search.retrieval import retrieve_semantic_bm25_combined_sql
|
||||
@@ -459,7 +467,7 @@ class PostgresMemories(MemoriesExtension):
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | None = None,
|
||||
fact_type: str | list[str] | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
|
||||
@@ -2513,6 +2513,13 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
)
|
||||
|
||||
logger.info(f"[CONSOLIDATION] bank={bank_id} completed: {result.get('memories_processed', 0)} processed")
|
||||
|
||||
# Consolidation is the other writer of memory_units rows: it mints
|
||||
# observations, which are their own fact_type and so their own indexed
|
||||
# partition. Retain's post-insert hook cannot see them — a bank whose
|
||||
# observations crossed the threshold here would otherwise wait for an
|
||||
# unrelated retain to notice (issue #3485).
|
||||
await self._submit_vector_index_maintenance_quietly(bank_id, internal_context, after="consolidation")
|
||||
return result
|
||||
|
||||
async def _handle_graph_maintenance(self, task_dict: dict[str, Any]):
|
||||
@@ -2690,6 +2697,8 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
consolidation_result = await self._handle_consolidation(task_dict)
|
||||
elif task_type == "graph_maintenance":
|
||||
await self._handle_graph_maintenance(task_dict)
|
||||
elif task_type == "vector_index_maintenance":
|
||||
await self._handle_vector_index_maintenance(task_dict)
|
||||
elif task_type == "refresh_mental_model":
|
||||
await self._handle_refresh_mental_model(task_dict)
|
||||
elif task_type == "webhook_delivery":
|
||||
@@ -4621,11 +4630,16 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
* auto-consolidation (when observations + auto-consolidation are enabled
|
||||
for the bank) so freshly inserted facts get observations;
|
||||
* graph maintenance, which short-circuits when no cleanup work was
|
||||
enqueued, so a plain insert pays a single cheap indexed SELECT here.
|
||||
enqueued, so a plain insert pays a single cheap indexed SELECT here;
|
||||
* per-bank vector index coverage, which short-circuits when the bank's
|
||||
indexes already match its size. Inserts are what move a bank across
|
||||
the size threshold, so this is where coverage is decided — there is
|
||||
nothing a periodic sweep could discover that the writer does not
|
||||
already know (issue #3485).
|
||||
|
||||
Both are non-critical: failures are logged, never raised, so they can't
|
||||
fail the operation that produced the facts. Pass ``config`` when the caller
|
||||
already resolved it to avoid a redundant lookup.
|
||||
All three are non-critical: failures are logged, never raised, so they
|
||||
can't fail the operation that produced the facts. Pass ``config`` when the
|
||||
caller already resolved it to avoid a redundant lookup.
|
||||
"""
|
||||
if config is None:
|
||||
config = await self._config_resolver.resolve_full_config(bank_id, request_context)
|
||||
@@ -4638,6 +4652,32 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
await self.submit_async_graph_maintenance(bank_id=bank_id, request_context=request_context)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to submit graph maintenance task for bank {bank_id}: {e}")
|
||||
await self._submit_vector_index_maintenance_quietly(bank_id, request_context, after="retain")
|
||||
|
||||
async def _submit_vector_index_maintenance_quietly(
|
||||
self,
|
||||
bank_id: str,
|
||||
request_context: "RequestContext",
|
||||
*,
|
||||
after: str,
|
||||
) -> None:
|
||||
"""Queue a vector-index reconcile for ``bank_id``, swallowing failures.
|
||||
|
||||
Called from every path that changes how many memory_units rows a bank
|
||||
holds — inserts (retain, import), consolidation (which mints
|
||||
observations) and deletes. Deletes matter as much as inserts: a bank
|
||||
pruned back under the threshold keeps indexes it no longer earns, and
|
||||
with nothing else scanning for that, an emptied bank that is never
|
||||
written to again would carry them forever.
|
||||
|
||||
Never raises. Index coverage is an optimisation — a bank without it
|
||||
falls back to exact search — so it must not be able to fail the delete or
|
||||
retain that produced the change.
|
||||
"""
|
||||
try:
|
||||
await self.submit_async_vector_index_maintenance(bank_id=bank_id, request_context=request_context)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to submit vector index maintenance after {after} for bank {bank_id}: {e}")
|
||||
|
||||
async def _resolve_retain_config(
|
||||
self,
|
||||
@@ -7350,6 +7390,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to submit graph maintenance after document deletion for bank {bank_id}: {e}")
|
||||
await self._submit_vector_index_maintenance_quietly(bank_id, request_context, after="document deletion")
|
||||
|
||||
return result
|
||||
|
||||
@@ -7708,6 +7749,9 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
f"Failed to submit graph maintenance after memory deletion "
|
||||
f"for bank {bank_id_for_graph_maintenance}: {e}"
|
||||
)
|
||||
await self._submit_vector_index_maintenance_quietly(
|
||||
bank_id_for_graph_maintenance, request_context, after="memory deletion"
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
@@ -7892,6 +7936,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to submit graph maintenance after bulk memory deletion for bank {bank_id}: {e}")
|
||||
await self._submit_vector_index_maintenance_quietly(bank_id, request_context, after="bulk memory deletion")
|
||||
|
||||
return {
|
||||
"requested": len(unit_ids),
|
||||
@@ -8125,6 +8170,19 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to submit consolidation after bank deletion for bank {bank_id}: {e}")
|
||||
|
||||
# A bank that survives this call (clear-memories, or a fact_type-scoped
|
||||
# delete) has lost rows and may no longer earn the indexes it has —
|
||||
# frequently all of them, since clearing a bank empties every partition.
|
||||
# The full-delete path above already dropped them by name while it still
|
||||
# knew the internal_id; this is the path where the bank stays, so the
|
||||
# reconcile has to be asked. Without it an emptied-but-kept bank holds
|
||||
# three ANN indexes over nothing until someone writes to it again, and
|
||||
# an emptied bank is exactly the one nobody writes to again.
|
||||
if not delete_bank_profile:
|
||||
await self._submit_vector_index_maintenance_quietly(
|
||||
bank_id, request_context, after="clearing bank memories"
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
async def clear_observations(
|
||||
@@ -8439,6 +8497,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
occurred_end: str | None = None,
|
||||
new_fact_type: str | None = None,
|
||||
entities: list[str] | None = None,
|
||||
resolve_entities: bool = True,
|
||||
state: str | None = None,
|
||||
reason: str | None = None,
|
||||
request_context: "RequestContext",
|
||||
@@ -8456,9 +8515,18 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
observations + temporal/semantic links, and re-consolidates. For date/context fields,
|
||||
``""`` clears to NULL and ``None`` leaves unchanged; ``new_fact_type``
|
||||
must be world/experience. ``entities`` (when not None) replaces the
|
||||
unit's entity set: names are resolved/find-or-created via the same
|
||||
resolver retain uses, ``unit_entities`` + cooccurrence are rebuilt, and
|
||||
``[]`` detaches all entities. Entities orphaned by the swap, and any
|
||||
unit's entity set, ``unit_entities`` + cooccurrence are rebuilt, and
|
||||
``[]`` detaches all entities. ``resolve_entities`` decides how those
|
||||
names find their entities. When True (the default, and what retain
|
||||
does) each name is resolved against the bank: a similar existing entity
|
||||
that scores above the match threshold is reused. When False the names
|
||||
are taken literally — an existing entity is reused only on a
|
||||
case-insensitive name match, any other name creates its own entity, and
|
||||
same-request names are never merged with each other. Hand-authored
|
||||
corrections want False: with resolution on, a similar-but-wrong entity
|
||||
that is well connected to the other names in the same edit outscores
|
||||
the one the caller named, and the correction lands on it silently
|
||||
(#3479). Entities orphaned by the swap, and any
|
||||
now-stale cooccurrence rows, are reclaimed by the graph-maintenance
|
||||
sweep that this edit submits (entity edges live in ``unit_entities``,
|
||||
not ``memory_links``, so there is nothing to relink directly).
|
||||
@@ -8607,6 +8675,14 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
# resolve_entities_only find-or-creates the corrected entities (idempotent) and
|
||||
# autocommits them on this short connection; the Phase-2 relink writes exactly
|
||||
# this resolved set, keeping the stored embedding consistent with the links.
|
||||
#
|
||||
# resolve_entities decides whether these names are a correction or another
|
||||
# guess (#3479). It defaults to True — retain's behaviour, kept as the default
|
||||
# so existing callers are unaffected — under which a similar-but-wrong entity
|
||||
# that is well-connected to the other names in this same list outscores the one
|
||||
# the caller actually named, and the edit lands on it with a 200 and no warning.
|
||||
# Callers correcting a fact by hand should pass False, which reuses an existing
|
||||
# entity only on a case-insensitive name match.
|
||||
entities_resolved = True
|
||||
entity_resolution = await resolve_entities_only(
|
||||
self.entity_resolver,
|
||||
@@ -8616,7 +8692,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
[new_text],
|
||||
new_context or "",
|
||||
[entity_date],
|
||||
[[{"text": name, "type": "CONCEPT"} for name in new_entities]],
|
||||
[[{"text": name, "type": "CONCEPT", "resolve": resolve_entities} for name in new_entities]],
|
||||
entity_labels=entity_labels,
|
||||
)
|
||||
resolved_for_unit = entity_resolution.unit_to_entity_ids.get(str(memory_uuid), [])
|
||||
@@ -9383,7 +9459,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
fact_type: str | None = None,
|
||||
fact_type: str | list[str] | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
@@ -9401,7 +9477,9 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
Args:
|
||||
bank_id: Filter by bank ID
|
||||
fact_type: Filter by fact type (world, experience)
|
||||
fact_type: Filter by fact type (world, experience). A list matches any
|
||||
of them (e.g. ``['world', 'experience']`` for source facts); an
|
||||
empty list is treated as no filter.
|
||||
search_query: Full-text search query (searches text and context fields)
|
||||
document_id: Optional filter to a single source document.
|
||||
entity_id: Optional filter to memory units linked to this entity ID
|
||||
@@ -11080,20 +11158,28 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
async def list_banks(
|
||||
self,
|
||||
*,
|
||||
search_query: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: "RequestContext",
|
||||
) -> list[dict[str, Any]]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
List all agents in the system.
|
||||
List memory banks, most recently written first.
|
||||
|
||||
Args:
|
||||
search_query: Case-insensitive substring matched against bank ID and name.
|
||||
limit: Maximum number of banks to return (0 returns none).
|
||||
offset: Number of banks to skip.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
List of dicts with bank_id, name, disposition, mission, created_at, updated_at
|
||||
Dict with ``banks`` (one page of bank_id, name, disposition, mission,
|
||||
created_at, updated_at and stats), ``total`` (banks matching the search
|
||||
that are visible to the caller, before paging), ``limit`` and ``offset``.
|
||||
"""
|
||||
await self._authenticate_tenant(request_context)
|
||||
await self._get_backend()
|
||||
banks = await bank_utils.list_banks(self._backend)
|
||||
banks = await bank_utils.list_banks(self._backend, search_query=search_query)
|
||||
if self._operation_validator:
|
||||
from hindsight_api.extensions import BankListContext
|
||||
|
||||
@@ -11101,18 +11187,31 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
BankListContext(banks=banks, request_context=request_context)
|
||||
)
|
||||
banks = result.banks
|
||||
# Paging happens here rather than in SQL because filter_bank_list may drop any
|
||||
# bank: a SQL page would hand back short (or empty) pages and a total counting
|
||||
# banks the caller isn't allowed to see.
|
||||
total = len(banks)
|
||||
# Clamped because the page is a Python slice, not a SQL LIMIT: a negative value
|
||||
# from a caller the HTTP layer doesn't validate (the MCP tool) would silently
|
||||
# trim from the end instead of raising.
|
||||
limit = max(limit, 0)
|
||||
offset = max(offset, 0)
|
||||
page = banks[offset : offset + limit]
|
||||
# Per-bank work below is done for the returned page only — a live store count
|
||||
# for banks whose memories live outside SQL, plus config resolution.
|
||||
await bank_utils.apply_store_fact_counts(self._backend, page)
|
||||
# Overlay resolved bank config (reflect_mission + disposition_*) on top of the
|
||||
# legacy banks.disposition / banks.mission columns, mirroring get_bank_profile so
|
||||
# the list and get paths return identical disposition + mission for a bank.
|
||||
# Resolve every bank's config in one batch (single config-column query + a single
|
||||
# Resolve the page's config in one batch (single config-column query + a single
|
||||
# tenant-config resolve) rather than one round-trip per bank.
|
||||
configs = await self._config_resolver.get_bank_configs([bank["bank_id"] for bank in banks], request_context)
|
||||
for bank in banks:
|
||||
configs = await self._config_resolver.get_bank_configs([bank["bank_id"] for bank in page], request_context)
|
||||
for bank in page:
|
||||
resolved = _overlay_bank_config_disposition_mission(
|
||||
bank["disposition"], bank["mission"], configs.get(bank["bank_id"], {})
|
||||
)
|
||||
bank["disposition"], bank["mission"] = resolved.disposition, resolved.mission
|
||||
return banks
|
||||
return {"banks": page, "total": total, "limit": limit, "offset": offset}
|
||||
|
||||
# ==================== Reflect Methods ====================
|
||||
|
||||
@@ -16933,6 +17032,190 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
dedupe_excludes_operation_id=dedupe_excludes_operation_id,
|
||||
)
|
||||
|
||||
async def submit_async_vector_index_maintenance(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
dedupe_excludes_operation_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Bring a bank's per-(bank, fact_type) vector indexes back in line with its size.
|
||||
|
||||
Called after a write that could have changed the bank's coverage —
|
||||
retain, import, consolidation, curation. Only writes move a bank across
|
||||
the threshold, so there is nothing for a periodic sweep to discover that
|
||||
the writer did not already know; this replaces one.
|
||||
|
||||
Idempotent and self-limiting: it plans first and short-circuits with
|
||||
``no_work=True`` when the bank's coverage already matches, so callers
|
||||
can invoke it unconditionally without paying for an async_operations
|
||||
row. At the default threshold of 0 that means one operation per
|
||||
(bank, fact_type) at first write and silence forever after.
|
||||
|
||||
Deduplicates by bank against a job that is pending *or* already running:
|
||||
the job re-plans from live row counts when it starts, so a job in flight
|
||||
already covers a write that landed after it was queued.
|
||||
|
||||
A no-op on backends without per-bank indexes (ScaNN keeps one global
|
||||
index) and on Oracle (partitioned by bank, no partial vector indexes).
|
||||
"""
|
||||
await self._authenticate_tenant(request_context)
|
||||
|
||||
index_clause = bank_utils._vector_index_clause()
|
||||
if index_clause is None:
|
||||
return {"operation_id": None, "no_work": True}
|
||||
|
||||
# Cheap pre-check: two bank-scoped index-only queries plus a catalog
|
||||
# lookup. Mirrors submit_async_graph_maintenance — an unconditional
|
||||
# caller must not create an empty worker task on every write.
|
||||
from .vector_index_health import plan_bank_vector_indexes
|
||||
|
||||
backend = await self._get_backend()
|
||||
try:
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
plan = await plan_bank_vector_indexes(conn, get_current_schema(), bank_id)
|
||||
except Exception as e:
|
||||
# Planning is advisory: a bank whose coverage we could not read is
|
||||
# picked up by the next write, or by `hindsight-admin repair-bank`.
|
||||
logger.warning(f"Vector index planning failed for bank {bank_id}: {e}")
|
||||
return {"operation_id": None, "no_work": True}
|
||||
if plan.is_empty:
|
||||
return {"operation_id": None, "no_work": True}
|
||||
|
||||
if self._operation_validator:
|
||||
from hindsight_api.extensions import BankWriteContext, BankWriteOperation
|
||||
|
||||
ctx = BankWriteContext(
|
||||
bank_id=bank_id,
|
||||
operation=BankWriteOperation.SUBMIT_ASYNC_VECTOR_INDEX_MAINTENANCE,
|
||||
request_context=request_context,
|
||||
)
|
||||
await self._validate_operation(self._operation_validator.validate_bank_write(ctx))
|
||||
|
||||
task_payload: dict[str, Any] = {}
|
||||
if request_context.tenant_id:
|
||||
task_payload["_tenant_id"] = request_context.tenant_id
|
||||
if request_context.api_key_id:
|
||||
task_payload["_api_key_id"] = request_context.api_key_id
|
||||
|
||||
return await self._submit_async_operation(
|
||||
bank_id=bank_id,
|
||||
operation_type="vector_index_maintenance",
|
||||
task_type="vector_index_maintenance",
|
||||
task_payload=task_payload,
|
||||
dedupe_by_bank=True,
|
||||
# Safe (unlike consolidation, which carries a watermark): the job
|
||||
# re-plans from live row counts at start, so a running job already
|
||||
# covers writes that landed after it was queued.
|
||||
dedupe_by_bank_includes_processing=True,
|
||||
# Set by the job's own hand-off so it does not match its own
|
||||
# still-'processing' row and suppress its successor.
|
||||
dedupe_excludes_operation_id=dedupe_excludes_operation_id,
|
||||
)
|
||||
|
||||
async def _handle_vector_index_maintenance(self, task_dict: dict) -> None:
|
||||
"""Reconcile one bank's vector indexes against the size threshold.
|
||||
|
||||
Runs on its own raw autocommit connection rather than a pooled one:
|
||||
CREATE INDEX CONCURRENTLY cannot run inside a transaction block, and
|
||||
both it and DROP INDEX CONCURRENTLY need a real backend session for the
|
||||
whole statement. HINDSIGHT_API_MIGRATION_DATABASE_URL is preferred when
|
||||
set, for the same reason migrations prefer it — a transaction-pooled URL
|
||||
cannot hold a session across the statement.
|
||||
"""
|
||||
import asyncpg
|
||||
|
||||
from ..pg0 import resolve_database_url
|
||||
from .vector_index_health import reconcile_bank_vector_indexes
|
||||
|
||||
bank_id = task_dict.get("bank_id")
|
||||
if not bank_id:
|
||||
return
|
||||
index_clause = bank_utils._vector_index_clause()
|
||||
if index_clause is None:
|
||||
return
|
||||
|
||||
# HINDSIGHT_API_MIGRATION_DATABASE_URL first when set — CREATE/DROP INDEX
|
||||
# CONCURRENTLY needs a real backend session for the whole statement, which
|
||||
# a transaction-pooled URL cannot give, and that env var is the documented
|
||||
# direct-connection escape hatch (migrations use it for the same reason).
|
||||
# Otherwise the DSN this engine is actually attached to, NOT
|
||||
# config.database_url: the two differ whenever the engine was handed a DSN
|
||||
# directly rather than reading the env var — embedders, and the test suite,
|
||||
# which resolves pg0 in a fixture. Reading config there connected to a
|
||||
# different database entirely and every reconcile died on
|
||||
# `relation "public.banks" does not exist`.
|
||||
backend = await self._get_backend()
|
||||
url = get_config().migration_database_url or getattr(backend, "dsn", None)
|
||||
if not url:
|
||||
logger.debug("Vector index maintenance skipped: no database URL available")
|
||||
return
|
||||
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
request_context = RequestContext(
|
||||
internal=True,
|
||||
tenant_id=task_dict.get("_tenant_id"),
|
||||
api_key_id=task_dict.get("_api_key_id"),
|
||||
)
|
||||
schema = get_current_schema()
|
||||
conn = await asyncpg.connect(await resolve_database_url(url))
|
||||
try:
|
||||
result = await reconcile_bank_vector_indexes(conn, schema, bank_id, index_clause)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
if result.created or result.dropped or result.failed:
|
||||
logger.info(
|
||||
f"Vector index maintenance for bank {bank_id}: "
|
||||
f"{result.created} built, {result.dropped} dropped, {result.failed} failed"
|
||||
)
|
||||
if result.failed:
|
||||
# Logged, never raised. A failed build leaves the bank on the exact
|
||||
# (bank_id, fact_type) B-tree path — slower on a large bank, but
|
||||
# correct — so it is not worth failing the operation, running it back
|
||||
# through the worker's retry/backoff, and surfacing a broken async op
|
||||
# to the user. The usual cause is a transient deadlock against
|
||||
# another session's concurrent index DDL on the shared memory_units
|
||||
# table, and the next write to this bank re-queues the work anyway.
|
||||
# `hindsight-admin repair-bank` is the path that treats a failed
|
||||
# build as an error, because there a human is waiting on the answer.
|
||||
logger.warning(
|
||||
f"Vector index maintenance left {result.failed} index(es) unbuilt for bank {bank_id} "
|
||||
f"({', '.join(result.failed_indexes)}); the next write to this bank retries"
|
||||
)
|
||||
return
|
||||
|
||||
# Hand off if the bank moved under us. The plan is a snapshot, and a
|
||||
# multi-statement delete that is still committing when this job planned
|
||||
# leaves it acting on a stale count — two jobs racing one delete can
|
||||
# rebuild what the other just dropped. Nothing else is looking: with no
|
||||
# periodic sweep, a bank that is never written again keeps whatever the
|
||||
# last racing job decided. Same gap, and same fix, as graph maintenance's
|
||||
# re-submit when work lands between its final claim and completion.
|
||||
#
|
||||
# Bounded two ways. The successor's own pre-check short-circuits once
|
||||
# coverage matches, so a converged bank stops the chain; and this is
|
||||
# skipped entirely when a build failed (returned above), so a permanently
|
||||
# failing index cannot spin submits forever.
|
||||
from .task_backend import SyncTaskBackend
|
||||
|
||||
# A synchronous task backend (tests, embedded) runs the successor inline
|
||||
# and would recurse inside this handler; there the caller is serial
|
||||
# anyway, so the next write reconciles.
|
||||
if isinstance(self._task_backend, SyncTaskBackend):
|
||||
return
|
||||
try:
|
||||
await self.submit_async_vector_index_maintenance(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
dedupe_excludes_operation_id=task_dict.get("operation_id"),
|
||||
)
|
||||
except Exception:
|
||||
# Never fail a completed reconcile over the hand-off; the next write
|
||||
# picks it up. Logged loudly so a persistent failure is visible.
|
||||
logger.exception(f"Vector index maintenance follow-up submit failed for bank {bank_id}")
|
||||
|
||||
async def submit_async_refresh_mental_model(
|
||||
self,
|
||||
bank_id: str,
|
||||
|
||||
@@ -45,36 +45,6 @@ def _vector_index_clause() -> str | None:
|
||||
return index_using_clause(ext)
|
||||
|
||||
|
||||
async def create_bank_vector_indexes(conn, bank_id: str, internal_id: str, ops=None) -> None:
|
||||
"""Create per-(bank, fact_type) partial vector indexes for a newly created bank.
|
||||
|
||||
Respects the HINDSIGHT_API_VECTOR_EXTENSION config to use the appropriate
|
||||
index type (HNSW for pgvector, DiskANN for pgvectorscale, vchordrq for vchord).
|
||||
|
||||
AlloyDB ScaNN uses global vector indexes with filtered vector search; it
|
||||
cannot safely create per-bank indexes at bank-creation time because new
|
||||
banks have no embedding rows.
|
||||
bank_id is escaped for SQL literal safety (apostrophes doubled).
|
||||
|
||||
On Oracle 23ai, this is a no-op — Oracle uses a single global vector index
|
||||
created during migrations. Partial indexes (WHERE clause) are not supported
|
||||
for Oracle vector indexes.
|
||||
"""
|
||||
index_clause = _vector_index_clause()
|
||||
if index_clause is None:
|
||||
logger.debug("Skipping per-bank vector indexes for configured backend")
|
||||
return
|
||||
|
||||
await ops.create_bank_vector_indexes(
|
||||
conn,
|
||||
fq_table("memory_units"),
|
||||
bank_id,
|
||||
internal_id,
|
||||
index_clause,
|
||||
_BANK_INDEX_FACT_TYPES,
|
||||
)
|
||||
|
||||
|
||||
async def drop_bank_vector_indexes(conn, internal_id: str, ops=None) -> None:
|
||||
"""Drop per-(bank, fact_type) partial vector indexes for a bank being deleted.
|
||||
|
||||
@@ -190,12 +160,12 @@ async def get_or_create_bank_profile(pool, bank_id: str) -> BankProfileResult:
|
||||
``get_or_create_bank_profile_on_conn`` instead.
|
||||
"""
|
||||
|
||||
# A fresh bank builds its per-(bank, fact_type) partial vector indexes with
|
||||
# a plain CREATE INDEX (it must — this runs inside the bank-create tx, and
|
||||
# CONCURRENTLY cannot). That CREATE takes a ShareLock on the shared
|
||||
# memory_units table, which can deadlock with concurrent writers. The build
|
||||
# is idempotent (INSERT ... ON CONFLICT + CREATE INDEX IF NOT EXISTS), so a
|
||||
# transient deadlock (40P01 / ORA-00060) is safe to retry as a whole tx.
|
||||
# Retried as a whole transaction. This used to guard the per-bank CREATE
|
||||
# INDEX that ran inline here and took a ShareLock on the shared memory_units
|
||||
# table; that DDL is gone (#3485), but the lazy create can still lose a
|
||||
# deadlock (40P01 / ORA-00060) to a concurrent writer touching the same
|
||||
# bank row, and the body is idempotent (INSERT ... ON CONFLICT DO NOTHING),
|
||||
# so retrying stays correct and cheap.
|
||||
async def _create() -> BankProfileResult:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
@@ -242,10 +212,16 @@ async def get_or_create_bank_profile_on_conn(conn, bank_id: str, *, ops) -> Bank
|
||||
created=False,
|
||||
)
|
||||
|
||||
# Bank doesn't exist, create with defaults.
|
||||
# Generate internal_id here so we control the value and can use it
|
||||
# immediately for vector index creation without a RETURNING round-trip.
|
||||
internal_id = uuid.uuid4()
|
||||
# Bank doesn't exist, create with defaults. internal_id is minted here rather
|
||||
# than defaulted server-side so its value is known without a RETURNING
|
||||
# round-trip; the vector-index sweep derives index names from it.
|
||||
#
|
||||
# No vector-index DDL here. A fresh bank holds no rows, so it cannot meet
|
||||
# the size threshold that earns a per-(bank, fact_type) partial index; the
|
||||
# maintenance sweep builds one if and when the bank grows into it. Keeping
|
||||
# DDL out of this path also takes CREATE INDEX's ShareLock on the shared
|
||||
# memory_units table off the retain hot path, where it deadlocked against
|
||||
# concurrent writers. See issue #3485.
|
||||
inserted = await conn.fetchval(
|
||||
f"""
|
||||
INSERT INTO {fq_table("banks")} (bank_id, name, disposition, mission, internal_id)
|
||||
@@ -257,14 +233,10 @@ async def get_or_create_bank_profile_on_conn(conn, bank_id: str, *, ops) -> Bank
|
||||
bank_id, # Default name is the bank_id
|
||||
json.dumps(DEFAULT_DISPOSITION),
|
||||
"",
|
||||
internal_id,
|
||||
uuid.uuid4(),
|
||||
)
|
||||
|
||||
created = inserted is not None
|
||||
if created:
|
||||
# Fresh insert — create per-bank vector indexes (instant on empty bank)
|
||||
await create_bank_vector_indexes(conn, bank_id, str(internal_id), ops=ops)
|
||||
|
||||
return BankProfileResult(
|
||||
profile=BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission=""),
|
||||
created=created,
|
||||
@@ -424,9 +396,9 @@ def _as_utc(ts: datetime | None) -> datetime | None:
|
||||
return ts if ts.tzinfo is not None else ts.replace(tzinfo=UTC)
|
||||
|
||||
|
||||
async def list_banks(pool) -> list:
|
||||
async def list_banks(pool, *, search_query: str | None = None) -> list:
|
||||
"""
|
||||
List all banks in the system with summary stats.
|
||||
List banks with summary stats, optionally narrowed by a search string.
|
||||
|
||||
``last_document_at`` is document *ingestion* time (when a document first
|
||||
landed), while ``last_write_at`` is the last time anything was written to
|
||||
@@ -434,8 +406,14 @@ async def list_banks(pool) -> list:
|
||||
to a long-lived document does not move ``last_document_at``, which is why
|
||||
the two differ and why UIs showing "last write" must use ``last_write_at``.
|
||||
|
||||
``fact_count`` comes from the ``memory_units`` join, which is empty for a bank
|
||||
whose memories live outside SQL. Those banks need :func:`apply_store_fact_counts`
|
||||
to get a real count; callers run it on the page they actually return so the live
|
||||
per-bank count query doesn't fire for every bank in the system.
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
search_query: Case-insensitive substring matched against bank ID and name
|
||||
|
||||
Returns:
|
||||
List of dicts with bank info and stats (fact_count, last_document_at, last_write_at),
|
||||
@@ -445,6 +423,15 @@ async def list_banks(pool) -> list:
|
||||
docs_table = fq_table("documents")
|
||||
mu_table = fq_table("memory_units")
|
||||
|
||||
# Spelled out as UPPER(...) LIKE UPPER(...) rather than ILIKE: the Oracle
|
||||
# rewriter only recognizes ILIKE on an unqualified column, and these are
|
||||
# alias-qualified.
|
||||
where_clause = ""
|
||||
params: list[str] = []
|
||||
if search_query:
|
||||
where_clause = "WHERE (UPPER(b.bank_id) LIKE UPPER($1) OR UPPER(COALESCE(b.name, '')) LIKE UPPER($2))"
|
||||
params = [f"%{search_query}%", f"%{search_query}%"]
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -470,19 +457,16 @@ async def list_banks(pool) -> list:
|
||||
FROM {mu_table}
|
||||
GROUP BY bank_id
|
||||
) m ON m.bank_id = b.bank_id
|
||||
{where_clause}
|
||||
ORDER BY b.bank_id
|
||||
"""
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
result = []
|
||||
# Banks are ordered by last write in Python rather than SQL: GREATEST() has
|
||||
# different NULL semantics on PostgreSQL vs Oracle, and the bank list is small.
|
||||
sort_keys: dict[str, datetime] = {}
|
||||
# A store that keeps memories outside SQL leaves the memory_units join empty, so its
|
||||
# per-bank fact_count comes from the store instead (one live count per bank).
|
||||
from ..memories import get_memories
|
||||
|
||||
_store = get_memories()
|
||||
|
||||
for row in rows:
|
||||
disposition_data = row["disposition"]
|
||||
@@ -498,12 +482,6 @@ async def list_banks(pool) -> list:
|
||||
write_times = [t for t in (_as_utc(row["last_document_write_at"]), _as_utc(row["last_fact_at"])) if t]
|
||||
last_write = max(write_times) if write_times else None
|
||||
|
||||
fact_count = row["fact_count"]
|
||||
if not _store.writes_memory_rows_in_sql_for(row["bank_id"]):
|
||||
fact_count = sum(
|
||||
(await _store.count_memories(conn=conn, fq_table=fq_table, bank_id=row["bank_id"])).values()
|
||||
)
|
||||
|
||||
sort_keys[row["bank_id"]] = last_write or created_at or _UNIX_EPOCH
|
||||
result.append(
|
||||
{
|
||||
@@ -513,7 +491,7 @@ async def list_banks(pool) -> list:
|
||||
"mission": row["mission"] or "",
|
||||
"created_at": created_at.isoformat() if created_at else None,
|
||||
"updated_at": updated_at.isoformat() if updated_at else None,
|
||||
"fact_count": fact_count,
|
||||
"fact_count": row["fact_count"],
|
||||
"last_document_at": last_doc.isoformat() if last_doc else None,
|
||||
"last_write_at": last_write.isoformat() if last_write else None,
|
||||
}
|
||||
@@ -521,3 +499,23 @@ async def list_banks(pool) -> list:
|
||||
|
||||
result.sort(key=lambda bank: sort_keys[bank["bank_id"]], reverse=True)
|
||||
return result
|
||||
|
||||
|
||||
async def apply_store_fact_counts(pool, banks: list[dict]) -> None:
|
||||
"""Replace ``fact_count`` in-place for banks that keep their memories outside SQL.
|
||||
|
||||
Those banks leave the ``memory_units`` join empty, so the count has to come
|
||||
from the store — one live count per bank, which is why this runs on a single
|
||||
page of :func:`list_banks` rather than on every bank in the system.
|
||||
"""
|
||||
from ..memories import get_memories
|
||||
|
||||
store = get_memories()
|
||||
external = [bank for bank in banks if not store.writes_memory_rows_in_sql_for(bank["bank_id"])]
|
||||
if not external:
|
||||
return
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
for bank in external:
|
||||
counts = await store.count_memories(conn=conn, fq_table=fq_table, bank_id=bank["bank_id"])
|
||||
bank["fact_count"] = sum(counts.values())
|
||||
|
||||
@@ -7,18 +7,23 @@ Handles entity extraction and resolution for stored facts.
|
||||
import logging
|
||||
|
||||
from . import link_utils
|
||||
from .types import EntityResolutionResult, ProcessedFact
|
||||
from .types import EntityResolutionResult, ProcessedFact, UserEntities
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _prepare_facts_for_entity_processing(
|
||||
facts: list[ProcessedFact],
|
||||
user_entities_per_content: dict[int, list[dict]] | None = None,
|
||||
user_entities_per_content: dict[int, UserEntities] | None = None,
|
||||
) -> tuple[list[str], list, list[list[dict]]]:
|
||||
"""
|
||||
Extract fact texts, dates, and merged entity lists from ProcessedFact objects.
|
||||
|
||||
Extracted names always carry ``resolve=True`` — they are the extractor's guess at a name, so
|
||||
matching them onto the bank's existing entities is the point. Caller-supplied names carry the
|
||||
content item's ``resolve_entities`` flag, so a caller can have their own names taken literally
|
||||
without turning off resolution for the extractor's (#3479).
|
||||
|
||||
Returns:
|
||||
Tuple of (fact_texts, fact_dates, entities_per_fact)
|
||||
"""
|
||||
@@ -29,20 +34,29 @@ def _prepare_facts_for_entity_processing(
|
||||
|
||||
entities_per_fact = []
|
||||
for fact in facts:
|
||||
llm_entities = [{"text": entity.name, "type": "CONCEPT"} for entity in (fact.entities or [])]
|
||||
llm_entities = [{"text": entity.name, "type": "CONCEPT", "resolve": True} for entity in (fact.entities or [])]
|
||||
|
||||
user_entities = user_entities_per_content.get(fact.content_index, [])
|
||||
supplied = user_entities_per_content.get(fact.content_index)
|
||||
user_entities = supplied.entities if supplied else []
|
||||
user_resolve = supplied.resolve if supplied else True
|
||||
|
||||
seen_texts = {e["text"].lower() for e in llm_entities}
|
||||
by_text = {e["text"].lower(): e for e in llm_entities}
|
||||
for user_entity in user_entities:
|
||||
if user_entity["text"].lower() not in seen_texts:
|
||||
llm_entities.append(
|
||||
{
|
||||
"text": user_entity["text"],
|
||||
"type": user_entity.get("type", "CONCEPT"),
|
||||
}
|
||||
)
|
||||
seen_texts.add(user_entity["text"].lower())
|
||||
text_lower = user_entity["text"].lower()
|
||||
existing = by_text.get(text_lower)
|
||||
if existing is None:
|
||||
entity = {
|
||||
"text": user_entity["text"],
|
||||
"type": user_entity.get("type", "CONCEPT"),
|
||||
"resolve": user_resolve,
|
||||
}
|
||||
llm_entities.append(entity)
|
||||
by_text[text_lower] = entity
|
||||
else:
|
||||
# The extractor produced this name too. The caller still authored it, so their
|
||||
# intent wins: a literal name must not become resolvable just because extraction
|
||||
# happened to agree on the spelling.
|
||||
existing["resolve"] = existing["resolve"] and user_resolve
|
||||
|
||||
entities_per_fact.append(llm_entities)
|
||||
|
||||
@@ -56,7 +70,7 @@ async def resolve_entities(
|
||||
unit_ids: list[str],
|
||||
facts: list[ProcessedFact],
|
||||
log_buffer: list[str] = None,
|
||||
user_entities_per_content: dict[int, list[dict]] = None,
|
||||
user_entities_per_content: dict[int, UserEntities] | None = None,
|
||||
entity_labels: list | None = None,
|
||||
) -> EntityResolutionResult:
|
||||
"""
|
||||
@@ -72,7 +86,8 @@ async def resolve_entities(
|
||||
unit_ids: Placeholder unit IDs (used only for grouping)
|
||||
facts: List of ProcessedFact objects
|
||||
log_buffer: Optional buffer for detailed logging
|
||||
user_entities_per_content: Dict mapping content_index to user-provided entities
|
||||
user_entities_per_content: Dict mapping content_index to the caller-supplied
|
||||
entities for that content item and whether to resolve them
|
||||
entity_labels: Optional entity label taxonomy
|
||||
|
||||
Returns:
|
||||
|
||||
@@ -13,7 +13,7 @@ from typing import Any
|
||||
from ...config import _get_raw_config
|
||||
from ..memory_engine import fq_table
|
||||
from ..metadata_utils import drop_null_values
|
||||
from .bank_utils import DEFAULT_DISPOSITION, create_bank_vector_indexes
|
||||
from .bank_utils import DEFAULT_DISPOSITION
|
||||
from .fact_extraction import _sanitize_text
|
||||
from .types import ProcessedFact
|
||||
|
||||
@@ -132,25 +132,26 @@ async def ensure_bank_exists(conn, bank_id: str, ops=None) -> None:
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
"""
|
||||
# Generate internal_id here so we control the value and can use it
|
||||
# immediately for HNSW index creation without a RETURNING round-trip.
|
||||
internal_id = uuid.uuid4()
|
||||
inserted = await conn.fetchval(
|
||||
# internal_id is generated here rather than defaulted server-side so the
|
||||
# value is known without a RETURNING round-trip; the vector-index sweep
|
||||
# derives index names from it.
|
||||
#
|
||||
# No vector-index DDL on this path. A fresh bank holds no rows, so it cannot
|
||||
# meet the size threshold that earns a per-(bank, fact_type) partial index;
|
||||
# the maintenance sweep builds one if the bank later grows into it. See
|
||||
# issue #3485.
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("banks")} (bank_id, name, disposition, mission, internal_id)
|
||||
VALUES ($1, $2, $3::jsonb, $4, $5)
|
||||
ON CONFLICT (bank_id) DO NOTHING
|
||||
RETURNING bank_id
|
||||
""",
|
||||
bank_id,
|
||||
bank_id, # Default name is the bank_id (matches get_or_create_bank_profile)
|
||||
json.dumps(DEFAULT_DISPOSITION),
|
||||
"",
|
||||
internal_id,
|
||||
uuid.uuid4(),
|
||||
)
|
||||
if inserted:
|
||||
# Fresh insert — create per-bank vector indexes
|
||||
await create_bank_vector_indexes(conn, bank_id, str(internal_id), ops=ops)
|
||||
|
||||
|
||||
async def delete_stale_observations_for_memories(
|
||||
|
||||
@@ -17,6 +17,7 @@ from ..causal_links import (
|
||||
)
|
||||
from ..db.base import DatabaseConnection
|
||||
from ..db.ops import DataAccessOps
|
||||
from ..db.postgresql import setting_rejected_by_server
|
||||
from ..memory_engine import fq_table
|
||||
from .types import CausalRelation, EntityResolutionResult
|
||||
|
||||
@@ -43,6 +44,16 @@ def _normalize_entity_name(name: str) -> str:
|
||||
return _WHITESPACE_RUN_RE.sub(" ", name).strip()
|
||||
|
||||
|
||||
def _entity_resolve_flag(ent) -> bool:
|
||||
"""Whether this candidate name should be resolved against existing entities.
|
||||
|
||||
Defaults to True (extraction's behaviour). Only dict candidates can opt out, which is how
|
||||
retain marks the entities its *caller* supplied: those are authoritative names, not guesses
|
||||
at which entity is meant (#3479).
|
||||
"""
|
||||
return bool(ent.get("resolve", True)) if isinstance(ent, dict) else True
|
||||
|
||||
|
||||
# Maximum number of temporal links to keep per unit (from_unit_id).
|
||||
# Retrieval only reads top 10-20 per unit via LATERAL join, so keeping
|
||||
# more is wasted storage and write amplification.
|
||||
@@ -213,7 +224,7 @@ def _prepare_entities_for_resolution(
|
||||
# own entity list) identical, and the upstream dedup in
|
||||
# entity_processing runs on the raw text. Without this, the same entity
|
||||
# would be resolved twice for one fact and its mention_count bumped twice.
|
||||
seen_in_fact: set[str] = set()
|
||||
seen_in_fact: dict[str, dict] = {}
|
||||
for ent in entity_list:
|
||||
if hasattr(ent, "text"):
|
||||
raw_text, entity_type = ent.text, "CONCEPT"
|
||||
@@ -230,11 +241,19 @@ def _prepare_entities_for_resolution(
|
||||
dropped_empty += 1
|
||||
continue
|
||||
|
||||
if normalized_text.lower() in seen_in_fact:
|
||||
resolve = _entity_resolve_flag(ent)
|
||||
kept = seen_in_fact.get(normalized_text.lower())
|
||||
if kept is not None:
|
||||
# Same name after normalization. Keep the first spelling but carry the stricter
|
||||
# flag: entity_processing dedups on the RAW text, so a caller's literal
|
||||
# "Acme Corp" and the extractor's "Acme\nCorp" both reach here, and dropping the
|
||||
# caller's outright would let the name be resolved away after all (#3479).
|
||||
kept["resolve"] = kept["resolve"] and resolve
|
||||
continue
|
||||
seen_in_fact.add(normalized_text.lower())
|
||||
|
||||
formatted_entities.append({"text": normalized_text, "type": entity_type})
|
||||
entity = {"text": normalized_text, "type": entity_type, "resolve": resolve}
|
||||
seen_in_fact[normalized_text.lower()] = entity
|
||||
formatted_entities.append(entity)
|
||||
all_entities.append(formatted_entities)
|
||||
|
||||
if dropped_empty:
|
||||
@@ -263,6 +282,7 @@ def _prepare_entities_for_resolution(
|
||||
{
|
||||
"text": entity["text"],
|
||||
"type": entity["type"],
|
||||
"resolve": entity["resolve"],
|
||||
"nearby_entities": entities,
|
||||
}
|
||||
)
|
||||
@@ -574,7 +594,14 @@ async def compute_semantic_links_ann(
|
||||
# are safe to apply at session/transaction scope for the configured
|
||||
# backend. VectorChord probe values are index-shaped, so vchordrq uses
|
||||
# index storage fallback parameters instead of a blanket SET LOCAL.
|
||||
#
|
||||
# A GUC the server has already rejected is skipped rather than attempted:
|
||||
# hnsw.iterative_scan needs pgvector 0.8+, and pgvector reserves the "hnsw."
|
||||
# prefix, so an older server errors on it — which inside this transaction would
|
||||
# abort the whole link computation rather than merely fail to apply.
|
||||
for guc, value in ann_search_tuning_settings(configured_vector_extension(), kind="low_latency"):
|
||||
if setting_rejected_by_server(guc):
|
||||
continue
|
||||
await conn.execute(f"SET LOCAL {guc} = {value}")
|
||||
|
||||
t_setup = time_mod.time()
|
||||
|
||||
@@ -282,6 +282,7 @@ from .types import (
|
||||
ResolvedEntity,
|
||||
RetainContent,
|
||||
RetainContentDict,
|
||||
UserEntities,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -403,7 +404,11 @@ async def _pre_resolve_phase1(
|
||||
set_stage("retain.phase1.resolve")
|
||||
from .link_utils import compute_semantic_links_ann
|
||||
|
||||
user_entities_per_content = {idx: content.entities for idx, content in enumerate(contents) if content.entities}
|
||||
user_entities_per_content = {
|
||||
idx: UserEntities(entities=content.entities, resolve=content.resolve_entities)
|
||||
for idx, content in enumerate(contents)
|
||||
if content.entities
|
||||
}
|
||||
|
||||
# Use placeholder unit_ids for grouping during resolution. The actual
|
||||
# unit_ids are created later by insert_facts_batch inside the transaction,
|
||||
@@ -1911,6 +1916,7 @@ async def _streaming_retain_batch(
|
||||
event_date=source.event_date,
|
||||
metadata=source.metadata,
|
||||
entities=source.entities,
|
||||
resolve_entities=source.resolve_entities,
|
||||
tags=source.tags,
|
||||
observation_scopes=source.observation_scopes,
|
||||
)
|
||||
@@ -3383,6 +3389,7 @@ def _build_contents(contents_dicts: list[RetainContentDict], document_tags: list
|
||||
event_date=event_date_value,
|
||||
metadata=item.get("metadata", {}),
|
||||
entities=item.get("entities", []),
|
||||
resolve_entities=item.get("resolve_entities", True),
|
||||
tags=merged_tags,
|
||||
observation_scopes=item.get("observation_scopes"),
|
||||
)
|
||||
@@ -3444,6 +3451,7 @@ def _build_delta_contents(
|
||||
event_date=template_content.event_date,
|
||||
metadata=template_content.metadata,
|
||||
entities=template_content.entities,
|
||||
resolve_entities=template_content.resolve_entities,
|
||||
tags=template_content.tags,
|
||||
observation_scopes=template_content.observation_scopes,
|
||||
)
|
||||
|
||||
@@ -26,6 +26,8 @@ class RetainContentDict(TypedDict, total=False):
|
||||
metadata: Custom key-value metadata (optional)
|
||||
document_id: Document ID for this content item (optional)
|
||||
entities: User-provided entities to merge with extracted entities (optional)
|
||||
resolve_entities: Whether the supplied `entities` are resolved against the bank's
|
||||
existing entities (optional, default True). False takes them literally.
|
||||
tags: Visibility scope tags for this content item (optional)
|
||||
observation_scopes: How to scope observations for consolidation (optional).
|
||||
"per_tag" runs one pass per individual tag; "combined" (default) runs a
|
||||
@@ -43,6 +45,7 @@ class RetainContentDict(TypedDict, total=False):
|
||||
metadata: dict[str, str]
|
||||
document_id: str
|
||||
entities: list[dict[str, str]] # [{"text": "...", "type": "..."}]
|
||||
resolve_entities: bool
|
||||
tags: list[str] # Visibility scope tags
|
||||
observation_scopes: (
|
||||
Literal["per_tag", "combined", "all_combinations", "shared"] | list[list[str]]
|
||||
@@ -50,6 +53,18 @@ class RetainContentDict(TypedDict, total=False):
|
||||
update_mode: Literal["replace", "append"]
|
||||
|
||||
|
||||
@dataclass
|
||||
class UserEntities:
|
||||
"""The entities a caller supplied for one retain content item, and how to match them.
|
||||
|
||||
Kept together so the resolution choice travels with the names it applies to: retain merges
|
||||
these with the extractor's own entities into one batch, and only these are authoritative.
|
||||
"""
|
||||
|
||||
entities: list[dict[str, str]]
|
||||
resolve: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainContent:
|
||||
"""
|
||||
@@ -63,6 +78,9 @@ class RetainContent:
|
||||
event_date: datetime | None = None
|
||||
metadata: dict[str, str] = field(default_factory=dict)
|
||||
entities: list[dict[str, str]] = field(default_factory=list) # User-provided entities
|
||||
# Whether the supplied `entities` are matched against the bank's existing entities. False
|
||||
# takes them literally; extracted entities are always resolved either way (#3479).
|
||||
resolve_entities: bool = True
|
||||
tags: list[str] = field(default_factory=list) # Visibility scope tags
|
||||
observation_scopes: Literal["per_tag", "combined", "all_combinations", "shared"] | list[list[str]] | None = (
|
||||
None # Observation scopes
|
||||
|
||||
@@ -147,10 +147,14 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
idx_mu_emb_observation, idx_mu_emb_experience), created automatically by
|
||||
Alembic migration a3b4c5d6e7f8_add_partial_hnsw_indexes.py.
|
||||
|
||||
HNSW is approximate — semantic arms over-fetch by 5x (min 100) and trim to
|
||||
limit in Python to compensate. ef_search=200 is set globally on pool
|
||||
connections at init time (see memory_engine.py) to improve recall on sparse
|
||||
graphs.
|
||||
Each semantic arm asks for exactly ``limit`` rows. It used to ask for ``limit * 5``
|
||||
and trim back to ``limit`` in Python "to compensate for HNSW approximation", but that
|
||||
could never work: the rows arrive already ordered by distance within their arm, so
|
||||
keeping the first ``limit`` of ``limit * 5`` returns precisely what ``LIMIT limit``
|
||||
would have — the extra rows were fetched, decoded and dropped, unread. What actually
|
||||
governs ANN quality is the size of the candidate list the scan explores, which is a
|
||||
connection setting, not a row count; the caller sizes it for this query (see
|
||||
``PostgresMemories.search``) rather than over-fetching rows here.
|
||||
|
||||
fact_type values are inlined as literals (safe: they come from a controlled
|
||||
internal enum, never from user input).
|
||||
@@ -180,8 +184,17 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
sem_min = min_semantic if min_semantic is not None else config.semantic_min_similarity
|
||||
bm25_min = min_keyword if min_keyword is not None else config.bm25_min_score
|
||||
|
||||
# Over-fetch for HNSW approximation; semantic results trimmed to limit in Python.
|
||||
hnsw_fetch = max(limit * 5, 100)
|
||||
# How many semantic rows each arm must return. Two consumers read them: the semantic
|
||||
# list itself (``limit``), and — when the dense rows also clear the graph arm's
|
||||
# threshold — its entry points (``GRAPH_SEED_LIMIT``), derived from the same ordered
|
||||
# rows instead of a duplicate ANN query per fact type. A budget below GRAPH_SEED_LIMIT
|
||||
# would otherwise starve the graph arm of seeds.
|
||||
graph_seed_threshold = (
|
||||
graph_seed_min_similarity
|
||||
if graph_seed_min_similarity is not None and sem_min <= graph_seed_min_similarity
|
||||
else None
|
||||
)
|
||||
semantic_fetch = max(limit, GRAPH_SEED_LIMIT if graph_seed_threshold is not None else 0)
|
||||
|
||||
cols = (
|
||||
"id, text, context, event_date, occurred_start, occurred_end, mentioned_at, "
|
||||
@@ -199,7 +212,7 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
# $1 = query_emb_str (semantic arms)
|
||||
# $2 = bank_id
|
||||
# When tokens present:
|
||||
# $3 = limit (BM25 LIMIT; semantic uses inlined hnsw_fetch literal)
|
||||
# $3 = limit (BM25 LIMIT; semantic inlines the same limit as a literal)
|
||||
# $4 = bm25_text
|
||||
# $5 = tags (if present)
|
||||
# $6+ = tag_groups params (one per leaf)
|
||||
@@ -242,7 +255,7 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
fact_type=ft,
|
||||
embedding_param="$1",
|
||||
bank_id_param="$2",
|
||||
fetch_limit=hnsw_fetch,
|
||||
fetch_limit=semantic_fetch,
|
||||
min_similarity=sem_min,
|
||||
tags_clause=tags_clause,
|
||||
groups_clause=groups_clause,
|
||||
@@ -346,7 +359,7 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
fact_type=ft,
|
||||
embedding_param="$1",
|
||||
bank_id_param="$2",
|
||||
fetch_limit=hnsw_fetch,
|
||||
fetch_limit=semantic_fetch,
|
||||
min_similarity=sem_min,
|
||||
tags_clause=fb_tags_clause,
|
||||
groups_clause=fb_groups_clause,
|
||||
@@ -364,17 +377,7 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
else:
|
||||
raise
|
||||
|
||||
# Group results. The semantic SQL deliberately over-fetches for HNSW recall;
|
||||
# when that pool also covers the graph threshold, derive graph entry points
|
||||
# from the same ordered rows instead of issuing one duplicate ANN query per
|
||||
# fact type. Convert only the prefix either consumer can observe, not the
|
||||
# entire HNSW over-fetch pool.
|
||||
graph_seed_threshold = (
|
||||
graph_seed_min_similarity
|
||||
if graph_seed_min_similarity is not None and sem_min <= graph_seed_min_similarity
|
||||
else None
|
||||
)
|
||||
semantic_candidate_limit = max(limit, GRAPH_SEED_LIMIT if graph_seed_threshold is not None else 0)
|
||||
# Group results, converting only the prefix either consumer can observe.
|
||||
semantic_candidates: dict[str, list[RetrievalResult]] = {ft: [] for ft in fact_types}
|
||||
for r in rows:
|
||||
row = dict(r)
|
||||
@@ -383,7 +386,7 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
if ft not in result_dict:
|
||||
continue
|
||||
if source == "semantic":
|
||||
if len(semantic_candidates[ft]) < semantic_candidate_limit:
|
||||
if len(semantic_candidates[ft]) < semantic_fetch:
|
||||
semantic_candidates[ft].append(RetrievalResult.from_db_row(row))
|
||||
else:
|
||||
result_dict[ft].bm25.append(RetrievalResult.from_db_row(row))
|
||||
|
||||
@@ -374,7 +374,8 @@ class SQLDialect(ABC):
|
||||
fact_type: Fact type literal (inlined, not parameterized).
|
||||
embedding_param: Parameter placeholder for query embedding.
|
||||
bank_id_param: Parameter placeholder for bank_id.
|
||||
fetch_limit: Max rows to fetch (over-fetched for HNSW approximation).
|
||||
fetch_limit: Max rows the arm returns. Its ANN candidate list must be at
|
||||
least this wide or the scan cannot fill it — see PostgresMemories.search.
|
||||
min_similarity: Minimum cosine similarity to include.
|
||||
tags_clause: Optional WHERE clause fragment for tag filtering.
|
||||
groups_clause: Optional WHERE clause fragment for tag group filtering.
|
||||
|
||||
@@ -22,7 +22,7 @@ from typing import Any, Literal
|
||||
from ..causal_links import CANONICAL_CAUSAL_LINK_TYPE, LEGACY_CAUSAL_LINK_TYPES
|
||||
from ..db.ops_postgresql import pg_search_vector_expr
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..retain import bank_utils, chunk_storage, embedding_processing, fact_storage, link_utils, orchestrator
|
||||
from ..retain import chunk_storage, embedding_processing, fact_storage, link_utils, orchestrator
|
||||
from ..retain.types import (
|
||||
CausalRelation,
|
||||
ChunkMetadata,
|
||||
@@ -559,16 +559,15 @@ async def import_bank(
|
||||
parsed.bank_rows.get("banks", []),
|
||||
bank_rows_json_encoding=bank_rows_json_encoding,
|
||||
)
|
||||
# The restored banks row bypasses the fresh-INSERT gate that normally
|
||||
# creates per-bank vector indexes, so create them explicitly here while
|
||||
# the bank is still empty (facts are imported below, so the build is
|
||||
# instant). get_or_create_bank_profile would NOT do this: the row now
|
||||
# exists, so it takes the SELECT branch and skips index creation —
|
||||
# leaving the restored bank falling back to the global index +
|
||||
# post-filter (slower, under-returning recall). See #2645.
|
||||
internal_id = await conn.fetchval(f"SELECT internal_id FROM {fq_table('banks')} WHERE bank_id = $1", bank_id)
|
||||
if internal_id is not None:
|
||||
await bank_utils.create_bank_vector_indexes(conn, bank_id, str(internal_id), ops=ops)
|
||||
# No vector-index DDL here. #2645 needed it because every bank was
|
||||
# entitled to indexes and a restored bank bypassed the fresh-INSERT gate
|
||||
# that created them; now entitlement is by size, and a restored bank's
|
||||
# rows land through the normal import path where the maintenance sweep
|
||||
# picks them up. Building inline would also be wrong twice over: the
|
||||
# bank is empty at this point (the facts arrive below), and CREATE INDEX
|
||||
# inside the import transaction takes a ShareLock on the shared
|
||||
# memory_units table. An import large enough to deserve an index gets
|
||||
# one on the next sweep. See #3485.
|
||||
|
||||
# Only now does the bank row — and with it the archive's own config — exist, so
|
||||
# this is where the config the documents are replayed with has to come from.
|
||||
|
||||
@@ -1,16 +1,26 @@
|
||||
"""Per-bank vector index coverage checks and repair.
|
||||
"""Per-bank vector index coverage: what a bank should have, and making it so.
|
||||
|
||||
Per-(bank, fact_type) partial vector indexes are created only when a bank is
|
||||
first created (instant on an empty bank). A bank that becomes *populated*
|
||||
outside that fresh-INSERT path — via a logical restore, a cross-version upgrade,
|
||||
or a vector-extension switch (e.g. ScaNN→pgvector) — never gets them, so its
|
||||
bank-scoped recall silently falls back to the global index + post-filter, which
|
||||
is both slower and under-returns results. See issue #2645.
|
||||
A (bank, fact_type) partition gets its own partial vector index once it holds
|
||||
``HINDSIGHT_API_VECTOR_INDEX_MIN_ROWS`` rows. At the default of 0 that is every
|
||||
partition holding any rows — the behaviour before the threshold existed — and a
|
||||
deployment with thousands of banks raises it, because these indexes live on the
|
||||
shared ``memory_units`` table: PostgreSQL locks and plans against every index on
|
||||
a relation, and opens every one for each DML statement, so one bank's index is
|
||||
charged to every other bank's queries. Three per bank exhausts the lock table at
|
||||
a few thousand banks (issue #3485). Below the threshold the planner answers the
|
||||
same query from the ``(bank_id, fact_type)`` B-tree plus a top-N sort, which is
|
||||
exact rather than approximate and faster.
|
||||
|
||||
This module is the shared engine for detecting and repairing that gap. It is
|
||||
driven by the ``repair-bank`` admin command; the build always uses
|
||||
``CREATE INDEX CONCURRENTLY`` on a raw autocommit connection so it never takes
|
||||
``ACCESS EXCLUSIVE`` on the shared ``memory_units`` table.
|
||||
Nothing here runs on a request path. Index DDL is issued only by the
|
||||
``vector_index_maintenance`` async operation (submitted after a write that could
|
||||
have changed a bank's coverage) and by the ``repair-bank`` admin command, both
|
||||
of which reconcile a bank against the plan this module computes.
|
||||
|
||||
All DDL is ``CREATE/DROP INDEX CONCURRENTLY`` on a raw autocommit connection, so
|
||||
it never takes ``ACCESS EXCLUSIVE`` on the shared table. That is also what keeps
|
||||
the drop path usable on an instance that has already hit the #3485 wall:
|
||||
``DROP INDEX`` is a utility statement that locks its own index plus the table,
|
||||
rather than planning against all of the table's indexes the way any DML must.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -19,6 +29,7 @@ import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from .._vector_index import qualifies_for_per_bank_index, should_keep_per_bank_index
|
||||
from .db_utils import retry_with_backoff
|
||||
from .retain.bank_utils import _BANK_INDEX_FACT_TYPES, _bank_index_name
|
||||
|
||||
@@ -46,14 +57,38 @@ _SUPPORTED_INDEX_AM: tuple[str, ...] = (
|
||||
|
||||
|
||||
@dataclass
|
||||
class SchemaVectorIndexResult:
|
||||
"""Per-schema outcome of a vector-index repair pass."""
|
||||
class BankIndexPlan:
|
||||
"""What one bank's vector-index coverage should become.
|
||||
|
||||
schema: str
|
||||
banks_scanned: int = 0
|
||||
Computed without issuing any DDL so the same plan can answer two questions:
|
||||
"is there anything to do?" (the cheap pre-check that keeps every write from
|
||||
queueing an empty operation) and "what exactly?" (the operation itself).
|
||||
"""
|
||||
|
||||
bank_id: str
|
||||
# fact_types at or above the build threshold whose index is missing or unhealthy.
|
||||
to_build: list[str] = field(default_factory=list)
|
||||
# Index names present in the catalog that this bank should no longer carry.
|
||||
to_drop: list[str] = field(default_factory=list)
|
||||
# Indexes already present and healthy — reported, never touched.
|
||||
already_present: int = 0
|
||||
|
||||
@property
|
||||
def is_empty(self) -> bool:
|
||||
return not self.to_build and not self.to_drop
|
||||
|
||||
|
||||
@dataclass
|
||||
class BankIndexResult:
|
||||
"""Outcome of applying a :class:`BankIndexPlan`."""
|
||||
|
||||
bank_id: str
|
||||
created: int = 0
|
||||
skipped: int = 0 # would-create, reported under --dry-run
|
||||
dropped: int = 0
|
||||
already_present: int = 0
|
||||
# Would-create / would-drop, reported under dry_run.
|
||||
skipped: int = 0
|
||||
would_drop: int = 0
|
||||
failed: int = 0
|
||||
failed_indexes: list[str] = field(default_factory=list)
|
||||
|
||||
@@ -97,138 +132,238 @@ async def _index_health(conn: Any, schema: str, index_names: list[str]) -> dict[
|
||||
return {row["index_name"]: bool(row["healthy"]) for row in rows}
|
||||
|
||||
|
||||
async def _repair_schema(
|
||||
async def plan_bank_vector_indexes(conn: Any, schema: str, bank_id: str) -> BankIndexPlan:
|
||||
"""Work out what ``bank_id``'s vector-index coverage should become.
|
||||
|
||||
Two cheap, bank-scoped queries plus one catalog lookup, so the write path
|
||||
can call this on every write to decide whether an operation is worth
|
||||
queueing at all. The row count is an index-only scan of
|
||||
``idx_memory_units_bank_fact_type``; deliberately unfiltered by
|
||||
``embedding IS NOT NULL``, since that predicate is not in the index and
|
||||
would turn the scan into a heap read without changing a threshold decision
|
||||
by enough to matter.
|
||||
|
||||
A bank whose row is gone yields an empty plan: its indexes are dropped by
|
||||
``delete_bank`` while the internal_id they are named after is still known,
|
||||
and a bank-scoped reconcile has no way to name them afterwards.
|
||||
"""
|
||||
plan = BankIndexPlan(bank_id=bank_id)
|
||||
qschema = _quote_identifier(schema)
|
||||
|
||||
internal_id = await conn.fetchval(
|
||||
f"SELECT internal_id FROM {qschema}.banks WHERE bank_id = $1", # noqa: S608 — schema is a quoted identifier
|
||||
bank_id,
|
||||
)
|
||||
if internal_id is None:
|
||||
return plan
|
||||
|
||||
counts = {
|
||||
row["fact_type"]: int(row["row_count"])
|
||||
for row in await conn.fetch(
|
||||
f"""
|
||||
SELECT fact_type, COUNT(*) AS row_count
|
||||
FROM {qschema}.memory_units
|
||||
WHERE bank_id = $1 AND fact_type = ANY($2::text[])
|
||||
GROUP BY fact_type
|
||||
""", # noqa: S608 — schema is a quoted identifier
|
||||
bank_id,
|
||||
list(_BANK_INDEX_FACT_TYPES),
|
||||
)
|
||||
}
|
||||
|
||||
names = {ft: _bank_index_name(ft, str(internal_id)) for ft in _BANK_INDEX_FACT_TYPES}
|
||||
health = await _index_health(conn, schema, list(names.values()))
|
||||
|
||||
for fact_type, index_name in names.items():
|
||||
row_count = counts.get(fact_type, 0)
|
||||
healthy = health.get(index_name)
|
||||
if qualifies_for_per_bank_index(row_count):
|
||||
if healthy is True:
|
||||
plan.already_present += 1
|
||||
else:
|
||||
plan.to_build.append(fact_type)
|
||||
elif healthy is not None and not should_keep_per_bank_index(row_count):
|
||||
# Present but no longer earned. Keeping has its own, lower bound than
|
||||
# building (see should_keep_per_bank_index) so a partition hovering
|
||||
# at the threshold does not rebuild and drop the same ANN index on
|
||||
# alternating writes.
|
||||
plan.to_drop.append(index_name)
|
||||
|
||||
return plan
|
||||
|
||||
|
||||
async def apply_bank_index_plan(
|
||||
conn: Any,
|
||||
schema: str,
|
||||
index_clause: str,
|
||||
plan: BankIndexPlan,
|
||||
*,
|
||||
dry_run: bool,
|
||||
bank_id: str | None,
|
||||
) -> SchemaVectorIndexResult:
|
||||
result = SchemaVectorIndexResult(schema=schema)
|
||||
dry_run: bool = False,
|
||||
) -> BankIndexResult:
|
||||
"""Build and drop what ``plan`` calls for, on a raw autocommit connection.
|
||||
|
||||
``conn`` must not be inside a transaction: ``CREATE INDEX CONCURRENTLY``
|
||||
cannot run in one, and both it and ``DROP INDEX CONCURRENTLY`` need a real
|
||||
backend session for the whole statement (a transaction-pooled URL will not
|
||||
do — that is what ``HINDSIGHT_API_MIGRATION_DATABASE_URL`` is for).
|
||||
|
||||
Concurrency is handled by idempotency, not a lock: the project forbids
|
||||
advisory locks, which are unreliable behind connection poolers, and leaning
|
||||
on one is why #2803's version of this was rejected. Every build is
|
||||
``CREATE INDEX CONCURRENTLY IF NOT EXISTS`` guarded by a valid/ready health
|
||||
check and every drop is ``DROP INDEX CONCURRENTLY IF EXISTS``, so a second
|
||||
concurrent run is a no-op on work the first already did.
|
||||
"""
|
||||
result = BankIndexResult(bank_id=plan.bank_id, already_present=plan.already_present)
|
||||
qschema = _quote_identifier(schema)
|
||||
|
||||
if bank_id is not None:
|
||||
banks = await conn.fetch(
|
||||
f"SELECT bank_id, internal_id FROM {qschema}.banks WHERE bank_id = $1", # noqa: S608 — schema is a quoted identifier
|
||||
bank_id,
|
||||
)
|
||||
else:
|
||||
banks = await conn.fetch(f"SELECT bank_id, internal_id FROM {qschema}.banks ORDER BY bank_id") # noqa: S608
|
||||
result.banks_scanned = len(banks)
|
||||
for index_name in plan.to_drop:
|
||||
if dry_run:
|
||||
result.would_drop += 1
|
||||
continue
|
||||
qualified = f"{qschema}.{_quote_identifier(index_name)}"
|
||||
try:
|
||||
await retry_with_backoff(
|
||||
lambda qualified=qualified: conn.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {qualified}")
|
||||
)
|
||||
result.dropped += 1
|
||||
except Exception as exc: # noqa: BLE001 — one failed drop must not abort the rest
|
||||
result.failed += 1
|
||||
result.failed_indexes.append(qualified)
|
||||
logger.warning("Failed to drop stale vector index %s: %s", qualified, exc)
|
||||
|
||||
# Resolve expected index names for every bank, then check them all in one
|
||||
# catalog query rather than one round-trip per index.
|
||||
expected_by_bank: list[tuple[str, dict[str, str]]] = []
|
||||
all_index_names: list[str] = []
|
||||
for bank in banks:
|
||||
expected = {ft: _bank_index_name(ft, str(bank["internal_id"])) for ft in _BANK_INDEX_FACT_TYPES}
|
||||
expected_by_bank.append((bank["bank_id"], expected))
|
||||
all_index_names.extend(expected.values())
|
||||
health = await _index_health(conn, schema, all_index_names)
|
||||
if not plan.to_build:
|
||||
return result
|
||||
|
||||
for bid, expected in expected_by_bank:
|
||||
# Render the bank_id literal server-side so escaping does not depend on
|
||||
# standard_conforming_strings (the predicate is inlined into the DDL).
|
||||
bank_id_literal = await conn.fetchval("SELECT quote_literal($1::text)", bid)
|
||||
for ft in _BANK_INDEX_FACT_TYPES:
|
||||
index_name = expected[ft]
|
||||
healthy = health.get(index_name)
|
||||
if healthy is True:
|
||||
result.already_present += 1
|
||||
continue
|
||||
if dry_run:
|
||||
result.skipped += 1
|
||||
continue
|
||||
# Render the bank_id literal server-side so escaping does not depend on
|
||||
# standard_conforming_strings (the predicate is inlined into the DDL).
|
||||
bank_id_literal = await conn.fetchval("SELECT quote_literal($1::text)", plan.bank_id)
|
||||
internal_id = await conn.fetchval(
|
||||
f"SELECT internal_id FROM {qschema}.banks WHERE bank_id = $1", # noqa: S608 — quoted identifier
|
||||
plan.bank_id,
|
||||
)
|
||||
if internal_id is None:
|
||||
# The bank was deleted between planning and applying; delete_bank has
|
||||
# already dropped its indexes and there is nothing left to name.
|
||||
return result
|
||||
|
||||
qindex = _quote_identifier(index_name)
|
||||
qualified = f"{qschema}.{qindex}"
|
||||
for fact_type in plan.to_build:
|
||||
if dry_run:
|
||||
result.skipped += 1
|
||||
continue
|
||||
qindex = _quote_identifier(_bank_index_name(fact_type, str(internal_id)))
|
||||
qualified = f"{qschema}.{qindex}"
|
||||
|
||||
async def _rebuild(
|
||||
qindex: str = qindex,
|
||||
qualified: str = qualified,
|
||||
ft: str = ft,
|
||||
bank_id_literal: str = bank_id_literal,
|
||||
) -> None:
|
||||
# Always drop first. An unhealthy-but-present index (INVALID
|
||||
# leftover, wrong access method) can't be repaired by
|
||||
# IF NOT EXISTS, and a prior deadlocked CONCURRENTLY build leaves
|
||||
# an INVALID stub that IF NOT EXISTS would likewise skip — so a
|
||||
# retry must clear it. DROP ... IF EXISTS is a no-op when the
|
||||
# index is simply absent (healthy is None).
|
||||
await conn.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {qualified}")
|
||||
await conn.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {qindex} "
|
||||
f"ON {qschema}.memory_units {index_clause} "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = {bank_id_literal}"
|
||||
)
|
||||
async def _rebuild(qindex: str = qindex, qualified: str = qualified, fact_type: str = fact_type) -> None:
|
||||
# Always drop first. An unhealthy-but-present index (INVALID
|
||||
# leftover, wrong access method) can't be repaired by IF NOT EXISTS,
|
||||
# and a prior deadlocked CONCURRENTLY build leaves an INVALID stub
|
||||
# that IF NOT EXISTS would likewise skip — so a retry must clear it.
|
||||
# DROP ... IF EXISTS is a no-op when the index is simply absent.
|
||||
await conn.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {qualified}")
|
||||
await conn.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {qindex} "
|
||||
f"ON {qschema}.memory_units {index_clause} "
|
||||
f"WHERE fact_type = '{fact_type}' AND bank_id = {bank_id_literal}"
|
||||
)
|
||||
|
||||
try:
|
||||
# CREATE INDEX CONCURRENTLY on the live, concurrently-written
|
||||
# memory_units table can be chosen as a deadlock victim (40P01).
|
||||
# That is transient — Postgres aborts one side to break the cycle —
|
||||
# so retry the drop+build before recording a permanent failure.
|
||||
await retry_with_backoff(_rebuild)
|
||||
result.created += 1
|
||||
logger.info("Built vector index %s (bank=%s, fact_type=%s)", qualified, plan.bank_id, fact_type)
|
||||
except Exception as exc: # noqa: BLE001 — one failed index must not abort the rest
|
||||
result.failed += 1
|
||||
result.failed_indexes.append(qualified)
|
||||
logger.warning(
|
||||
"Failed to build vector index %s (bank=%s, fact_type=%s): %s — "
|
||||
"dropping the invalid leftover so a re-run can retry.",
|
||||
qualified,
|
||||
plan.bank_id,
|
||||
fact_type,
|
||||
exc,
|
||||
)
|
||||
# A failed concurrent build leaves an INVALID index behind that
|
||||
# would shadow the good one; drop it so a re-run retries cleanly.
|
||||
try:
|
||||
# CREATE INDEX CONCURRENTLY on the live, concurrently-written
|
||||
# memory_units table can be chosen as a deadlock victim
|
||||
# (sqlstate 40P01 / ORA-00060). That is transient — Postgres
|
||||
# aborts one side to break the cycle — so retry the drop+build a
|
||||
# few times before recording a permanent failure.
|
||||
await retry_with_backoff(_rebuild)
|
||||
result.created += 1
|
||||
except Exception as exc: # noqa: BLE001 — one failed index must not abort the rest
|
||||
result.failed += 1
|
||||
result.failed_indexes.append(qualified)
|
||||
logger.warning(
|
||||
"Failed to repair vector index %s (bank=%s, fact_type=%s): %s — "
|
||||
"dropping the invalid leftover so a re-run can retry.",
|
||||
qualified,
|
||||
bid,
|
||||
ft,
|
||||
exc,
|
||||
)
|
||||
# A failed concurrent build leaves an INVALID index behind that
|
||||
# would shadow the good one; drop it so a re-run retries cleanly.
|
||||
try:
|
||||
await conn.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {qualified}")
|
||||
except Exception as cleanup_exc: # noqa: BLE001
|
||||
logger.warning("Cleanup DROP INDEX for %s also failed: %s", qualified, cleanup_exc)
|
||||
await conn.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {qualified}")
|
||||
except Exception as cleanup_exc: # noqa: BLE001
|
||||
logger.warning("Cleanup DROP INDEX for %s also failed: %s", qualified, cleanup_exc)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
async def _safe_repair_schema(
|
||||
async def reconcile_bank_vector_indexes(
|
||||
conn: Any,
|
||||
schema: str,
|
||||
index_clause: str,
|
||||
*,
|
||||
dry_run: bool,
|
||||
bank_id: str | None,
|
||||
) -> SchemaVectorIndexResult:
|
||||
try:
|
||||
return await _repair_schema(conn, schema, index_clause, dry_run=dry_run, bank_id=bank_id)
|
||||
except Exception as exc: # noqa: BLE001 — one bad schema must not abort the whole sweep
|
||||
logger.warning("Vector index repair aborted for schema %s: %s", schema, exc)
|
||||
return SchemaVectorIndexResult(schema=schema, failed=1, failed_indexes=[f"{schema}.<schema-error>"])
|
||||
|
||||
|
||||
async def repair_vector_indexes(
|
||||
conn: Any,
|
||||
schemas: list[str],
|
||||
bank_id: str,
|
||||
index_clause: str,
|
||||
*,
|
||||
dry_run: bool = False,
|
||||
bank_id: str | None = None,
|
||||
) -> list[SchemaVectorIndexResult]:
|
||||
"""Rebuild missing or invalid per-bank vector indexes across ``schemas``.
|
||||
) -> BankIndexResult:
|
||||
"""Plan and apply one bank's vector-index coverage."""
|
||||
plan = await plan_bank_vector_indexes(conn, schema, bank_id)
|
||||
return await apply_bank_index_plan(conn, schema, index_clause, plan, dry_run=dry_run)
|
||||
|
||||
``conn`` must be a raw autocommit PostgreSQL connection: ``CREATE INDEX
|
||||
CONCURRENTLY`` cannot run inside a transaction block. When ``bank_id`` is
|
||||
given, only that bank is reconciled (in each schema); otherwise every bank
|
||||
is scanned.
|
||||
|
||||
Concurrency is handled by idempotency, not a lock (project rule: no advisory
|
||||
locks — they are unreliable behind connection poolers). Every build is
|
||||
``CREATE INDEX CONCURRENTLY IF NOT EXISTS`` guarded by a valid/ready health
|
||||
check, so a second concurrent run is a no-op on already-built indexes; if two
|
||||
runs race the *same* missing index, Postgres rejects one build and the
|
||||
per-index handler drops the leftover so a re-run converges cleanly.
|
||||
async def list_bank_ids(conn: Any, schema: str) -> list[str]:
|
||||
"""Every bank in ``schema``, for the admin command's ``--all`` mode."""
|
||||
rows = await conn.fetch(
|
||||
f"SELECT bank_id FROM {_quote_identifier(schema)}.banks ORDER BY bank_id" # noqa: S608 — quoted identifier
|
||||
)
|
||||
return [row["bank_id"] for row in rows]
|
||||
|
||||
|
||||
async def drop_orphaned_bank_indexes(conn: Any, schema: str, *, dry_run: bool = False) -> list[str]:
|
||||
"""Drop per-bank vector indexes whose bank no longer exists.
|
||||
|
||||
``delete_bank`` drops a bank's indexes while the ``internal_id`` they are
|
||||
named after is still known, so this should find nothing. It exists for when
|
||||
that did not happen: a deployment that hit the #3485 wall could not run
|
||||
``delete_bank`` at all (the delete DML could not plan), so operators dropped
|
||||
banks by other means and left the indexes behind — and an orphan is
|
||||
unreachable by every bank-scoped path, because there is no bank row to plan
|
||||
from.
|
||||
|
||||
Catalog-only, matching each index's name suffix against the live
|
||||
``internal_id`` set, so it answers even on an instance whose lock table is
|
||||
exhausted. Only the admin command calls it; the write path has no reason to.
|
||||
"""
|
||||
return [
|
||||
await _safe_repair_schema(conn, schema, index_clause, dry_run=dry_run, bank_id=bank_id) for schema in schemas
|
||||
]
|
||||
qschema = _quote_identifier(schema)
|
||||
live = {
|
||||
str(row["internal_id"]).replace("-", "")[:16]
|
||||
for row in await conn.fetch(f"SELECT internal_id FROM {qschema}.banks") # noqa: S608 — quoted identifier
|
||||
}
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT c.relname AS index_name
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
JOIN pg_index i ON i.indexrelid = c.oid
|
||||
JOIN pg_class t ON t.oid = i.indrelid
|
||||
WHERE n.nspname = $1
|
||||
AND t.relname = 'memory_units'
|
||||
AND c.relname LIKE 'idx\\_mu\\_emb\\_%'
|
||||
""",
|
||||
schema,
|
||||
)
|
||||
|
||||
orphans = [row["index_name"] for row in rows if row["index_name"].rsplit("_", 1)[-1] not in live]
|
||||
if dry_run:
|
||||
return orphans
|
||||
|
||||
dropped = []
|
||||
for index_name in orphans:
|
||||
qualified = f"{qschema}.{_quote_identifier(index_name)}"
|
||||
try:
|
||||
await retry_with_backoff(
|
||||
lambda qualified=qualified: conn.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {qualified}")
|
||||
)
|
||||
dropped.append(index_name)
|
||||
logger.info("Dropped orphaned vector index %s (no matching bank)", qualified)
|
||||
except Exception as exc: # noqa: BLE001 — one failure must not abort the rest
|
||||
logger.warning("Failed to drop orphaned vector index %s: %s", qualified, exc)
|
||||
return dropped
|
||||
|
||||
@@ -390,6 +390,7 @@ class BankWriteOperation(StrEnum):
|
||||
SET_BANK_MISSION = "set_bank_mission"
|
||||
SUBMIT_ASYNC_CONSOLIDATION = "submit_async_consolidation"
|
||||
SUBMIT_ASYNC_GRAPH_MAINTENANCE = "submit_async_graph_maintenance"
|
||||
SUBMIT_ASYNC_VECTOR_INDEX_MAINTENANCE = "submit_async_vector_index_maintenance"
|
||||
UPDATE_BANK = "update_bank"
|
||||
UPDATE_BANK_CONFIG = "update_bank_config"
|
||||
UPDATE_BANK_DISPOSITION = "update_bank_disposition"
|
||||
|
||||
@@ -16,6 +16,7 @@ from mcp.types import ToolAnnotations
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.api import page_markdown
|
||||
from hindsight_api.config import (
|
||||
DEFAULT_MCP_RECALL_DESCRIPTION,
|
||||
DEFAULT_MCP_RETAIN_DESCRIPTION,
|
||||
@@ -65,6 +66,13 @@ _ALL_TOOLS: frozenset[str] = frozenset(
|
||||
"update_bank",
|
||||
"delete_bank",
|
||||
"clear_memories",
|
||||
"get_knowledge_base_tree",
|
||||
"search_knowledge_base",
|
||||
"get_knowledge_page",
|
||||
"create_knowledge_folder",
|
||||
"create_knowledge_page",
|
||||
"update_knowledge_node",
|
||||
"delete_knowledge_node",
|
||||
}
|
||||
)
|
||||
|
||||
@@ -228,6 +236,9 @@ _READ_ONLY_TOOLS = {
|
||||
"list_operations",
|
||||
"get_operation",
|
||||
"list_tags",
|
||||
"get_knowledge_base_tree",
|
||||
"search_knowledge_base",
|
||||
"get_knowledge_page",
|
||||
}
|
||||
_DESTRUCTIVE_TOOLS = {
|
||||
"delete_bank",
|
||||
@@ -237,6 +248,7 @@ _DESTRUCTIVE_TOOLS = {
|
||||
"delete_directive",
|
||||
"delete_document",
|
||||
"invalidate_memory",
|
||||
"delete_knowledge_node",
|
||||
}
|
||||
|
||||
|
||||
@@ -295,6 +307,13 @@ def register_mcp_tools(
|
||||
"update_bank",
|
||||
"delete_bank",
|
||||
"clear_memories",
|
||||
"get_knowledge_base_tree",
|
||||
"search_knowledge_base",
|
||||
"get_knowledge_page",
|
||||
"create_knowledge_folder",
|
||||
"create_knowledge_page",
|
||||
"update_knowledge_node",
|
||||
"delete_knowledge_node",
|
||||
}
|
||||
|
||||
if "retain" in tools_to_register:
|
||||
@@ -399,6 +418,28 @@ def register_mcp_tools(
|
||||
if "clear_memories" in tools_to_register:
|
||||
_register_clear_memories(mcp, memory, config)
|
||||
|
||||
# Knowledge base tools
|
||||
if "get_knowledge_base_tree" in tools_to_register:
|
||||
_register_get_knowledge_base_tree(mcp, memory, config)
|
||||
|
||||
if "search_knowledge_base" in tools_to_register:
|
||||
_register_search_knowledge_base(mcp, memory, config)
|
||||
|
||||
if "get_knowledge_page" in tools_to_register:
|
||||
_register_get_knowledge_page(mcp, memory, config)
|
||||
|
||||
if "create_knowledge_folder" in tools_to_register:
|
||||
_register_create_knowledge_folder(mcp, memory, config)
|
||||
|
||||
if "create_knowledge_page" in tools_to_register:
|
||||
_register_create_knowledge_page(mcp, memory, config)
|
||||
|
||||
if "update_knowledge_node" in tools_to_register:
|
||||
_register_update_knowledge_node(mcp, memory, config)
|
||||
|
||||
if "delete_knowledge_node" in tools_to_register:
|
||||
_register_delete_knowledge_node(mcp, memory, config)
|
||||
|
||||
_apply_bank_tool_filtering(mcp, memory, config)
|
||||
_apply_audit_logging(mcp, memory, config)
|
||||
|
||||
@@ -509,6 +550,10 @@ _AUDITABLE_MCP_TOOLS: frozenset[str] = frozenset(
|
||||
"delete_directive",
|
||||
"delete_document",
|
||||
"cancel_operation",
|
||||
"create_knowledge_folder",
|
||||
"create_knowledge_page",
|
||||
"update_knowledge_node",
|
||||
"delete_knowledge_node",
|
||||
}
|
||||
)
|
||||
|
||||
@@ -1211,19 +1256,30 @@ def _register_list_banks(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
"""Register the list_banks tool."""
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("list_banks"))
|
||||
async def list_banks() -> str:
|
||||
async def list_banks(query: str | None = None, limit: int = 100, offset: int = 0) -> str:
|
||||
"""
|
||||
List all available memory banks.
|
||||
List available memory banks, most recently written first.
|
||||
|
||||
Use this tool to discover what memory banks exist in the system.
|
||||
Each bank is an isolated memory store (like a separate "brain").
|
||||
|
||||
Args:
|
||||
query: Optional case-insensitive substring to match against bank ID and name.
|
||||
limit: Maximum number of banks to return (default 100).
|
||||
offset: Number of banks to skip, for paging through `total`.
|
||||
|
||||
Returns:
|
||||
JSON list of banks with their IDs, names, dispositions, and missions.
|
||||
JSON with the page of banks (IDs, names, dispositions, missions) plus
|
||||
the total number of matching banks and the limit/offset used.
|
||||
"""
|
||||
try:
|
||||
banks = await memory.list_banks(request_context=_get_request_context(config))
|
||||
return json.dumps({"banks": banks}, indent=2)
|
||||
data = await memory.list_banks(
|
||||
search_query=query,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps(data, indent=2)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e), "banks": []})
|
||||
@@ -2034,6 +2090,823 @@ def _register_clear_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# KNOWLEDGE BASE TOOLS
|
||||
# =========================================================================
|
||||
# A tree of folders and pages over mental models (see the HTTP
|
||||
# /knowledge-base endpoints). Pages read as markdown documents; folders are
|
||||
# containers. ``export_knowledge_base`` is deliberately NOT exposed here — it
|
||||
# returns the whole bank as one markdown bundle, which belongs on the HTTP/CLI
|
||||
# path rather than in an agent's context window.
|
||||
|
||||
# MCP tool arguments cannot express an explicit null: an omitted argument and a
|
||||
# null one both arrive as ``None``, so update_knowledge_node reads this literal
|
||||
# as "move to the top level". Node ids are prefixed (``kf-``/``kp-``), so it
|
||||
# cannot collide with a real folder id.
|
||||
KNOWLEDGE_ROOT_PARENT = "root"
|
||||
|
||||
|
||||
def _knowledge_node_json(node: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Project an engine node dict into the compact JSON an MCP client sees.
|
||||
|
||||
Mirrors the HTTP ``KnowledgeNode`` projection, minus the fields an agent has
|
||||
no use for (bank_id, sort_order): page metadata comes from the backing
|
||||
mental model, folders carry structure only.
|
||||
"""
|
||||
is_page = node.get("kind") == "page"
|
||||
projected: dict[str, Any] = {
|
||||
"id": node["id"],
|
||||
"kind": node["kind"],
|
||||
"name": node["name"],
|
||||
"parent_id": node.get("parent_id"),
|
||||
"managed": bool(node.get("managed")),
|
||||
}
|
||||
if is_page:
|
||||
projected["mental_model_id"] = node.get("mental_model_id")
|
||||
projected["description"] = node.get("source_query")
|
||||
projected["tags"] = list(node.get("tags") or [])
|
||||
projected["timestamp"] = node.get("last_refreshed_at")
|
||||
if node.get("trigger") is not None:
|
||||
projected["trigger"] = node["trigger"]
|
||||
if node.get("is_stale") is not None:
|
||||
projected["is_stale"] = node["is_stale"]
|
||||
else:
|
||||
projected["timestamp"] = node.get("updated_at")
|
||||
projected["children"] = []
|
||||
return projected
|
||||
|
||||
|
||||
def _knowledge_tree_json(nodes: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Assemble the flat node list into nested roots (mirrors the HTTP tree)."""
|
||||
projected = {n["id"]: _knowledge_node_json(n) for n in nodes}
|
||||
roots: list[dict[str, Any]] = []
|
||||
for node in nodes:
|
||||
parent_id = node.get("parent_id")
|
||||
# Only folders can be parents (enforced on write), so the parent normally
|
||||
# carries a children list; setdefault keeps a malformed row from raising.
|
||||
if parent_id and parent_id in projected:
|
||||
projected[parent_id].setdefault("children", []).append(projected[node["id"]])
|
||||
else:
|
||||
roots.append(projected[node["id"]])
|
||||
return roots
|
||||
|
||||
|
||||
def _page_trigger_patch(refresh_after_consolidation: bool | None) -> dict[str, Any] | None:
|
||||
"""Build the trigger patch for the one refresh setting exposed over MCP.
|
||||
|
||||
The engine merges a patch over the page's defaults (create) or its current
|
||||
trigger (update), so sending only this field leaves the rest — delta mode,
|
||||
observation-only facts — as they were.
|
||||
"""
|
||||
if refresh_after_consolidation is None:
|
||||
return None
|
||||
return {"refresh_after_consolidation": refresh_after_consolidation}
|
||||
|
||||
|
||||
async def _do_get_knowledge_base_tree(
|
||||
memory: MemoryEngine, target_bank: str, request_context: RequestContext
|
||||
) -> dict[str, Any]:
|
||||
"""Shared implementation for the get_knowledge_base_tree MCP tool variants."""
|
||||
nodes = await memory.list_knowledge_nodes(bank_id=target_bank, with_staleness=True, request_context=request_context)
|
||||
return {"roots": _knowledge_tree_json(nodes)}
|
||||
|
||||
|
||||
async def _do_search_knowledge_base(
|
||||
memory: MemoryEngine, target_bank: str, request_context: RequestContext, *, query: str, limit: int
|
||||
) -> dict[str, Any]:
|
||||
"""Shared implementation for the search_knowledge_base MCP tool variants.
|
||||
|
||||
``limit`` is clamped rather than rejected: the HTTP route answers an
|
||||
out-of-range limit with a 422, but an agent that asked for 500 pages wants
|
||||
results, not a validation round trip.
|
||||
"""
|
||||
results = await memory.search_knowledge_pages(
|
||||
bank_id=target_bank, query=query, limit=max(1, min(limit, 50)), request_context=request_context
|
||||
)
|
||||
return {"results": results, "total": len(results)}
|
||||
|
||||
|
||||
async def _do_get_knowledge_page(
|
||||
memory: MemoryEngine, target_bank: str, request_context: RequestContext, *, page_id: str
|
||||
) -> dict[str, Any]:
|
||||
"""Shared implementation for the get_knowledge_page MCP tool variants."""
|
||||
node = await memory.get_knowledge_page(bank_id=target_bank, page_id=page_id, request_context=request_context)
|
||||
if node is None:
|
||||
return {"error": f"Knowledge page '{page_id}' not found in bank '{target_bank}'"}
|
||||
page = page_markdown.page_type(node.get("tags"))
|
||||
# The rendered document already carries the body under a frontmatter block,
|
||||
# so it is returned once rather than alongside a duplicate `body` field.
|
||||
return {
|
||||
"id": node["id"],
|
||||
"name": node["name"],
|
||||
"type": page.type,
|
||||
"description": node.get("source_query"),
|
||||
"tags": page.display_tags,
|
||||
"timestamp": node.get("last_refreshed_at") or node.get("created_at"),
|
||||
"markdown": page_markdown.render_document(node),
|
||||
}
|
||||
|
||||
|
||||
async def _do_create_knowledge_folder(
|
||||
memory: MemoryEngine, target_bank: str, request_context: RequestContext, *, name: str, parent_id: str | None
|
||||
) -> dict[str, Any]:
|
||||
"""Shared implementation for the create_knowledge_folder MCP tool variants."""
|
||||
node = await memory.create_knowledge_folder(
|
||||
bank_id=target_bank, name=name, parent_id=parent_id, request_context=request_context
|
||||
)
|
||||
return _knowledge_node_json(node)
|
||||
|
||||
|
||||
async def _do_create_knowledge_page(
|
||||
memory: MemoryEngine,
|
||||
target_bank: str,
|
||||
request_context: RequestContext,
|
||||
*,
|
||||
name: str,
|
||||
source_query: str,
|
||||
parent_id: str | None,
|
||||
tags: list[str] | None,
|
||||
max_tokens: int | None,
|
||||
refresh_after_consolidation: bool | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Shared implementation for the create_knowledge_page MCP tool variants."""
|
||||
node = await memory.create_knowledge_page(
|
||||
bank_id=target_bank,
|
||||
name=name,
|
||||
source_query=source_query,
|
||||
content="Generating content...",
|
||||
parent_id=parent_id,
|
||||
tags=tags or None,
|
||||
max_tokens=max_tokens,
|
||||
trigger=_page_trigger_patch(refresh_after_consolidation),
|
||||
request_context=request_context,
|
||||
)
|
||||
if node is None:
|
||||
return {"error": f"A page named '{name}' already exists in this folder"}
|
||||
result = await memory.submit_async_refresh_mental_model(
|
||||
bank_id=target_bank, mental_model_id=node["mental_model_id"], request_context=request_context
|
||||
)
|
||||
return {
|
||||
"page_id": node["id"],
|
||||
"mental_model_id": node["mental_model_id"],
|
||||
"operation_id": result["operation_id"],
|
||||
"status": "created",
|
||||
"message": f"Page '{name}' created. Content is being generated asynchronously.",
|
||||
}
|
||||
|
||||
|
||||
async def _do_update_knowledge_node(
|
||||
memory: MemoryEngine,
|
||||
target_bank: str,
|
||||
request_context: RequestContext,
|
||||
*,
|
||||
node_id: str,
|
||||
name: str | None,
|
||||
parent_id: str | None,
|
||||
source_query: str | None,
|
||||
tags: list[str] | None,
|
||||
max_tokens: int | None,
|
||||
refresh_after_consolidation: bool | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Shared implementation for the update_knowledge_node MCP tool variants.
|
||||
|
||||
Each field is applied only when provided, so a rename never resets a page's
|
||||
query and moving a page never drops its tags.
|
||||
"""
|
||||
trigger = _page_trigger_patch(refresh_after_consolidation)
|
||||
page_update = source_query is not None or tags is not None or max_tokens is not None or trigger is not None
|
||||
if name is None and parent_id is None and not page_update:
|
||||
return {
|
||||
"error": "Provide name, parent_id, source_query, tags, max_tokens, "
|
||||
"and/or refresh_after_consolidation to update"
|
||||
}
|
||||
|
||||
updated: dict[str, Any] | None = None
|
||||
if name is not None:
|
||||
updated = await memory.rename_knowledge_node(
|
||||
bank_id=target_bank, node_id=node_id, name=name, request_context=request_context
|
||||
)
|
||||
if parent_id is not None:
|
||||
updated = await memory.move_knowledge_node(
|
||||
bank_id=target_bank,
|
||||
node_id=node_id,
|
||||
new_parent_id=None if parent_id == KNOWLEDGE_ROOT_PARENT else parent_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
if page_update:
|
||||
updated = await memory.update_knowledge_page(
|
||||
bank_id=target_bank,
|
||||
page_id=node_id,
|
||||
source_query=source_query,
|
||||
tags=tags,
|
||||
max_tokens=max_tokens,
|
||||
trigger=trigger,
|
||||
request_context=request_context,
|
||||
)
|
||||
# A new source query means the page's content no longer answers it — rebuild.
|
||||
if updated is not None and source_query is not None and updated.get("mental_model_id"):
|
||||
await memory.submit_async_refresh_mental_model(
|
||||
bank_id=target_bank, mental_model_id=updated["mental_model_id"], request_context=request_context
|
||||
)
|
||||
if updated is None:
|
||||
return {"error": f"Knowledge node '{node_id}' not found in bank '{target_bank}'"}
|
||||
return _knowledge_node_json(updated)
|
||||
|
||||
|
||||
async def _do_delete_knowledge_node(
|
||||
memory: MemoryEngine, target_bank: str, request_context: RequestContext, *, node_id: str
|
||||
) -> dict[str, Any]:
|
||||
"""Shared implementation for the delete_knowledge_node MCP tool variants."""
|
||||
deleted = await memory.delete_knowledge_node(bank_id=target_bank, node_id=node_id, request_context=request_context)
|
||||
if not deleted:
|
||||
return {"error": f"Knowledge node '{node_id}' not found in bank '{target_bank}'"}
|
||||
return {"status": "deleted", "node_id": node_id}
|
||||
|
||||
|
||||
def _register_get_knowledge_base_tree(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the get_knowledge_base_tree tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("get_knowledge_base_tree"))
|
||||
async def get_knowledge_base_tree(
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Browse the knowledge base as a nested tree of folders and pages.
|
||||
|
||||
Start here to discover what the bank documents: each page is a living
|
||||
markdown document synthesized from the bank's memories, and folders
|
||||
group them. Use get_knowledge_page to read a page's content, or
|
||||
search_knowledge_base when you know what you are looking for.
|
||||
|
||||
Pages report `is_stale`: false means the page is provably up to date;
|
||||
true means something was written since its last refresh, so it MAY be
|
||||
out of date.
|
||||
|
||||
Args:
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
tree = await _do_get_knowledge_base_tree(memory, target_bank, _get_request_context(config))
|
||||
return json.dumps(tree, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting knowledge base tree: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("get_knowledge_base_tree"))
|
||||
async def get_knowledge_base_tree() -> dict:
|
||||
"""
|
||||
Browse the knowledge base as a nested tree of folders and pages.
|
||||
|
||||
Start here to discover what the bank documents: each page is a living
|
||||
markdown document synthesized from the bank's memories, and folders
|
||||
group them. Use get_knowledge_page to read a page's content, or
|
||||
search_knowledge_base when you know what you are looking for.
|
||||
|
||||
Pages report `is_stale`: false means the page is provably up to date;
|
||||
true means something was written since its last refresh, so it MAY be
|
||||
out of date.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
return await _do_get_knowledge_base_tree(memory, target_bank, _get_request_context(config))
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting knowledge base tree: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_search_knowledge_base(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the search_knowledge_base tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("search_knowledge_base"))
|
||||
async def search_knowledge_base(
|
||||
query: str,
|
||||
limit: int = 10,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Find knowledge pages by relevance (hybrid keyword + semantic search).
|
||||
|
||||
Searches page names and content, returning ranked pages with a short
|
||||
snippet each. Read a hit in full with get_knowledge_page. This searches
|
||||
the curated knowledge base only — use recall to search raw memories.
|
||||
|
||||
Args:
|
||||
query: What to search for
|
||||
limit: Maximum pages to return (1-50, default: 10)
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
results = await _do_search_knowledge_base(
|
||||
memory, target_bank, _get_request_context(config), query=query, limit=limit
|
||||
)
|
||||
return json.dumps(results, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error searching knowledge base: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("search_knowledge_base"))
|
||||
async def search_knowledge_base(
|
||||
query: str,
|
||||
limit: int = 10,
|
||||
) -> dict:
|
||||
"""
|
||||
Find knowledge pages by relevance (hybrid keyword + semantic search).
|
||||
|
||||
Searches page names and content, returning ranked pages with a short
|
||||
snippet each. Read a hit in full with get_knowledge_page. This searches
|
||||
the curated knowledge base only — use recall to search raw memories.
|
||||
|
||||
Args:
|
||||
query: What to search for
|
||||
limit: Maximum pages to return (1-50, default: 10)
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
return await _do_search_knowledge_base(
|
||||
memory, target_bank, _get_request_context(config), query=query, limit=limit
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error searching knowledge base: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_get_knowledge_page(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the get_knowledge_page tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("get_knowledge_page"))
|
||||
async def get_knowledge_page(
|
||||
page_id: str,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Read a knowledge page as a markdown document.
|
||||
|
||||
Returns the page's YAML frontmatter (id, type, title, description,
|
||||
tags, timestamp) followed by its synthesized markdown body. Discover
|
||||
page ids with get_knowledge_base_tree or search_knowledge_base.
|
||||
|
||||
Args:
|
||||
page_id: The ID of the page to read (a `kp-...` node id)
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
page = await _do_get_knowledge_page(memory, target_bank, _get_request_context(config), page_id=page_id)
|
||||
return json.dumps(page, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting knowledge page: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("get_knowledge_page"))
|
||||
async def get_knowledge_page(
|
||||
page_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Read a knowledge page as a markdown document.
|
||||
|
||||
Returns the page's YAML frontmatter (id, type, title, description,
|
||||
tags, timestamp) followed by its synthesized markdown body. Discover
|
||||
page ids with get_knowledge_base_tree or search_knowledge_base.
|
||||
|
||||
Args:
|
||||
page_id: The ID of the page to read (a `kp-...` node id)
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
return await _do_get_knowledge_page(memory, target_bank, _get_request_context(config), page_id=page_id)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting knowledge page: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_create_knowledge_folder(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the create_knowledge_folder tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("create_knowledge_folder"))
|
||||
async def create_knowledge_folder(
|
||||
name: str,
|
||||
parent_id: str | None = None,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Create a folder in the knowledge base.
|
||||
|
||||
Folders group pages; they hold no content of their own.
|
||||
|
||||
Args:
|
||||
name: Folder name
|
||||
parent_id: Optional parent folder id (a `kf-...` node id). Omit to create at the top level.
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
node = await _do_create_knowledge_folder(
|
||||
memory, target_bank, _get_request_context(config), name=name, parent_id=parent_id
|
||||
)
|
||||
return json.dumps(node, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except ValueError as e:
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating knowledge folder: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("create_knowledge_folder"))
|
||||
async def create_knowledge_folder(
|
||||
name: str,
|
||||
parent_id: str | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Create a folder in the knowledge base.
|
||||
|
||||
Folders group pages; they hold no content of their own.
|
||||
|
||||
Args:
|
||||
name: Folder name
|
||||
parent_id: Optional parent folder id (a `kf-...` node id). Omit to create at the top level.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
return await _do_create_knowledge_folder(
|
||||
memory, target_bank, _get_request_context(config), name=name, parent_id=parent_id
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except ValueError as e:
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating knowledge folder: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_create_knowledge_page(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the create_knowledge_page tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("create_knowledge_page"))
|
||||
async def create_knowledge_page(
|
||||
name: str,
|
||||
source_query: str,
|
||||
parent_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
max_tokens: int | None = None,
|
||||
refresh_after_consolidation: bool | None = None,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Create a knowledge page — a living document answering a question.
|
||||
|
||||
The page's content is synthesized from the bank's memories by running
|
||||
source_query, asynchronously: use the returned operation_id to track
|
||||
completion, then read it with get_knowledge_page. By default the page
|
||||
keeps itself current, rebuilding after each consolidation.
|
||||
|
||||
EXAMPLES:
|
||||
- name="Deployment Runbook", source_query="How is this service deployed and rolled back?"
|
||||
- name="Team Preferences", source_query="What tools and conventions does the team prefer?"
|
||||
|
||||
Args:
|
||||
name: Page name (must be unique within its folder)
|
||||
source_query: The question this page answers and rebuilds itself from
|
||||
parent_id: Optional parent folder id (a `kf-...` node id). Omit to create at the top level.
|
||||
tags: Optional tags scoping which memories the page is built from
|
||||
max_tokens: Maximum tokens for the generated content (default: 4096)
|
||||
refresh_after_consolidation: Whether the page rebuilds itself after each memory
|
||||
consolidation. Omit to keep the knowledge-page default (True).
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
result = await _do_create_knowledge_page(
|
||||
memory,
|
||||
target_bank,
|
||||
_get_request_context(config),
|
||||
name=name,
|
||||
source_query=source_query,
|
||||
parent_id=parent_id,
|
||||
tags=tags,
|
||||
max_tokens=max_tokens,
|
||||
refresh_after_consolidation=refresh_after_consolidation,
|
||||
)
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except ValueError as e:
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating knowledge page: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("create_knowledge_page"))
|
||||
async def create_knowledge_page(
|
||||
name: str,
|
||||
source_query: str,
|
||||
parent_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
max_tokens: int | None = None,
|
||||
refresh_after_consolidation: bool | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Create a knowledge page — a living document answering a question.
|
||||
|
||||
The page's content is synthesized from the bank's memories by running
|
||||
source_query, asynchronously: use the returned operation_id to track
|
||||
completion, then read it with get_knowledge_page. By default the page
|
||||
keeps itself current, rebuilding after each consolidation.
|
||||
|
||||
EXAMPLES:
|
||||
- name="Deployment Runbook", source_query="How is this service deployed and rolled back?"
|
||||
- name="Team Preferences", source_query="What tools and conventions does the team prefer?"
|
||||
|
||||
Args:
|
||||
name: Page name (must be unique within its folder)
|
||||
source_query: The question this page answers and rebuilds itself from
|
||||
parent_id: Optional parent folder id (a `kf-...` node id). Omit to create at the top level.
|
||||
tags: Optional tags scoping which memories the page is built from
|
||||
max_tokens: Maximum tokens for the generated content (default: 4096)
|
||||
refresh_after_consolidation: Whether the page rebuilds itself after each memory
|
||||
consolidation. Omit to keep the knowledge-page default (True).
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
return await _do_create_knowledge_page(
|
||||
memory,
|
||||
target_bank,
|
||||
_get_request_context(config),
|
||||
name=name,
|
||||
source_query=source_query,
|
||||
parent_id=parent_id,
|
||||
tags=tags,
|
||||
max_tokens=max_tokens,
|
||||
refresh_after_consolidation=refresh_after_consolidation,
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except ValueError as e:
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating knowledge page: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_update_knowledge_node(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the update_knowledge_node tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("update_knowledge_node"))
|
||||
async def update_knowledge_node(
|
||||
node_id: str,
|
||||
name: str | None = None,
|
||||
parent_id: str | None = None,
|
||||
source_query: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
max_tokens: int | None = None,
|
||||
refresh_after_consolidation: bool | None = None,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Rename or move a folder/page, and/or update a page's options.
|
||||
|
||||
Only the arguments you pass are changed; everything else keeps its
|
||||
current value. Changing source_query schedules an async refresh so the
|
||||
page rebuilds against the new question.
|
||||
|
||||
Args:
|
||||
node_id: The ID of the folder (`kf-...`) or page (`kp-...`) to update
|
||||
name: New name for the node
|
||||
parent_id: Folder id to move the node into, or "root" to move it to the top level
|
||||
source_query: Pages only — the new question the page answers
|
||||
tags: Pages only — replacement tag list (pass [] to clear)
|
||||
max_tokens: Pages only — new maximum tokens for the generated content
|
||||
refresh_after_consolidation: Pages only — whether the page rebuilds itself
|
||||
after each memory consolidation
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
result = await _do_update_knowledge_node(
|
||||
memory,
|
||||
target_bank,
|
||||
_get_request_context(config),
|
||||
node_id=node_id,
|
||||
name=name,
|
||||
parent_id=parent_id,
|
||||
source_query=source_query,
|
||||
tags=tags,
|
||||
max_tokens=max_tokens,
|
||||
refresh_after_consolidation=refresh_after_consolidation,
|
||||
)
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except ValueError as e:
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating knowledge node: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("update_knowledge_node"))
|
||||
async def update_knowledge_node(
|
||||
node_id: str,
|
||||
name: str | None = None,
|
||||
parent_id: str | None = None,
|
||||
source_query: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
max_tokens: int | None = None,
|
||||
refresh_after_consolidation: bool | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Rename or move a folder/page, and/or update a page's options.
|
||||
|
||||
Only the arguments you pass are changed; everything else keeps its
|
||||
current value. Changing source_query schedules an async refresh so the
|
||||
page rebuilds against the new question.
|
||||
|
||||
Args:
|
||||
node_id: The ID of the folder (`kf-...`) or page (`kp-...`) to update
|
||||
name: New name for the node
|
||||
parent_id: Folder id to move the node into, or "root" to move it to the top level
|
||||
source_query: Pages only — the new question the page answers
|
||||
tags: Pages only — replacement tag list (pass [] to clear)
|
||||
max_tokens: Pages only — new maximum tokens for the generated content
|
||||
refresh_after_consolidation: Pages only — whether the page rebuilds itself
|
||||
after each memory consolidation
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
return await _do_update_knowledge_node(
|
||||
memory,
|
||||
target_bank,
|
||||
_get_request_context(config),
|
||||
node_id=node_id,
|
||||
name=name,
|
||||
parent_id=parent_id,
|
||||
source_query=source_query,
|
||||
tags=tags,
|
||||
max_tokens=max_tokens,
|
||||
refresh_after_consolidation=refresh_after_consolidation,
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except ValueError as e:
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating knowledge node: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_delete_knowledge_node(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the delete_knowledge_node tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("delete_knowledge_node"))
|
||||
async def delete_knowledge_node(
|
||||
node_id: str,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Delete a knowledge-base folder or page and everything under it.
|
||||
|
||||
Deleting a folder also deletes its whole subtree, and each deleted page
|
||||
takes its backing mental model with it. This cannot be undone.
|
||||
|
||||
Args:
|
||||
node_id: The ID of the folder (`kf-...`) or page (`kp-...`) to delete
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
result = await _do_delete_knowledge_node(
|
||||
memory, target_bank, _get_request_context(config), node_id=node_id
|
||||
)
|
||||
return json.dumps(result)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting knowledge node: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(annotations=_tool_annotations("delete_knowledge_node"))
|
||||
async def delete_knowledge_node(
|
||||
node_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Delete a knowledge-base folder or page and everything under it.
|
||||
|
||||
Deleting a folder also deletes its whole subtree, and each deleted page
|
||||
takes its backing mental model with it. This cannot be undone.
|
||||
|
||||
Args:
|
||||
node_id: The ID of the folder (`kf-...`) or page (`kp-...`) to delete
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
return await _do_delete_knowledge_node(
|
||||
memory, target_bank, _get_request_context(config), node_id=node_id
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting knowledge node: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# DIRECTIVE TOOLS
|
||||
# =========================================================================
|
||||
@@ -2494,6 +3367,13 @@ def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
all). The memory is re-embedded and its derived observations, links, and
|
||||
graph are recomputed automatically.
|
||||
|
||||
resolve_entities controls how the names in entities are matched. The
|
||||
default True behaves like retain and may resolve a name onto a similar
|
||||
entity that already exists, which silently discards a correction when the
|
||||
bank holds a near-duplicate name. Pass False whenever you are correcting a
|
||||
fact deliberately: an existing entity is then reused only on a
|
||||
case-insensitive name match, and any other name becomes its own entity.
|
||||
|
||||
Only raw world/experience facts can be edited; observations are derived.
|
||||
To retire or restore a fact, use invalidate_memory instead.
|
||||
"""
|
||||
@@ -2509,6 +3389,7 @@ def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
occurred_end: str | None = None,
|
||||
fact_type: str | None = None,
|
||||
entities: list[str] | None = None,
|
||||
resolve_entities: bool = True,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -2530,6 +3411,7 @@ def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
occurred_end=occurred_end,
|
||||
new_fact_type=fact_type,
|
||||
entities=entities,
|
||||
resolve_entities=resolve_entities,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if result is None:
|
||||
@@ -2555,6 +3437,7 @@ def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
occurred_end: str | None = None,
|
||||
fact_type: str | None = None,
|
||||
entities: list[str] | None = None,
|
||||
resolve_entities: bool = True,
|
||||
) -> dict:
|
||||
"""
|
||||
Args:
|
||||
@@ -2574,6 +3457,7 @@ def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
occurred_end=occurred_end,
|
||||
new_fact_type=fact_type,
|
||||
entities=entities,
|
||||
resolve_entities=resolve_entities,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if result is None:
|
||||
|
||||
@@ -638,11 +638,15 @@ def ensure_vector_extension(
|
||||
|
||||
if table_name == "memory_units" and uses_per_bank_vector_indexes(target_ext):
|
||||
# Per-bank backends never use a GLOBAL memory_units vector index.
|
||||
# Every vector search is bank + fact_type scoped and served by the
|
||||
# per-(bank, fact_type) partial indexes created at bank-creation time
|
||||
# (bank_utils.create_bank_vector_indexes); the planner never picks a
|
||||
# global index when bank_id is in the WHERE clause, which is exactly
|
||||
# why migration d5e6f7a8b9c0 drops it for these backends.
|
||||
# Every vector search is bank + fact_type scoped, and is served
|
||||
# either by the bank's own partial index or — for a bank below
|
||||
# the size threshold, which is most of them — by an exact
|
||||
# (bank_id, fact_type) B-tree scan plus a top-N sort. The planner
|
||||
# never picks a global index when bank_id is in the WHERE clause,
|
||||
# which is exactly why migration d5e6f7a8b9c0 drops it for these
|
||||
# backends. The partial indexes themselves are owned by the
|
||||
# maintenance sweep (engine/vector_index_health.py), not by this
|
||||
# reconcile and not by bank creation.
|
||||
#
|
||||
# The reconcile is strictly hands-off here, in BOTH directions:
|
||||
# never create the index (dead weight — an older version of this
|
||||
|
||||
@@ -189,6 +189,7 @@ markers = [
|
||||
"hs_llm_core: Core pipeline tests that need a real LLM but only one provider",
|
||||
"integration: Live external-API integration tests (require provider credentials; skipped without)",
|
||||
"slow: Slow tests (minutes); not run in fast CI",
|
||||
"memory_backend_incompatible: asserts Postgres-internal state a non-SQL memories backend cannot reproduce — raw memory_links row counts (the graph read path dedupes bidirectional edges), or internal columns like embedding / search_vector that are not part of the public read model. Deselect when running against an alternative memories backend with -m 'not memory_backend_incompatible'.",
|
||||
]
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
|
||||
@@ -116,10 +116,31 @@ DEFAULT_PG0_PORT = int(os.environ.get("HINDSIGHT_TEST_PG_PORT", "5556"))
|
||||
# no job enabled, so the loop never starts. Tests that exercise it call
|
||||
# MaintenanceLoop methods (_run_reconcile / _run_scheduled_mm_refresh /
|
||||
# _purge_expired) directly.
|
||||
#
|
||||
# Every job added to the loop must be switched off here too: one job left on is
|
||||
# enough to start the loop for the whole suite, which reintroduces exactly the
|
||||
# races the others are disabled to avoid.
|
||||
os.environ.setdefault("HINDSIGHT_API_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS", "0")
|
||||
os.environ.setdefault("HINDSIGHT_API_MENTAL_MODEL_REFRESH_TICK_SECONDS", "0")
|
||||
os.environ.setdefault("HINDSIGHT_API_LLM_TRACE_RETENTION_DAYS", "-1")
|
||||
|
||||
# Keep incidental per-bank vector-index DDL out of the suite. The shipped default
|
||||
# is 0 — every bank holding rows earns its three indexes — which is right for a
|
||||
# deployment whose banks are long-lived and get built once, but wrong here: the
|
||||
# suite creates thousands of throwaway banks and writes one or two facts to each,
|
||||
# so every one of them would queue a build. Eight xdist workers issuing
|
||||
# CREATE INDEX CONCURRENTLY against a single shared memory_units deadlock each
|
||||
# other by design (CONCURRENTLY holds ShareUpdateExclusive while it waits out
|
||||
# every session whose snapshot could still see the index, including other
|
||||
# sessions' index DDL), and that lands on whatever unrelated test happens to be
|
||||
# writing at the time.
|
||||
#
|
||||
# A threshold no test bank can reach means the coverage machinery is inert unless
|
||||
# a test asks for it: tests that exercise it patch the threshold themselves (see
|
||||
# the low_threshold / default_threshold fixtures in
|
||||
# test_repair_bank_vector_indexes.py).
|
||||
os.environ.setdefault("HINDSIGHT_API_VECTOR_INDEX_MIN_ROWS", "1000000")
|
||||
|
||||
|
||||
# Load environment variables from .env at the start of test session
|
||||
def pytest_configure(config):
|
||||
@@ -359,9 +380,17 @@ def oracle_db_url(_oracle_admin_dsn):
|
||||
f'CREATE USER {test_user} IDENTIFIED BY "{test_pass}" DEFAULT TABLESPACE USERS QUOTA UNLIMITED ON USERS'
|
||||
)
|
||||
except oracledb.DatabaseError as e:
|
||||
if hasattr(e.args[0], "code") and e.args[0].code == 1920:
|
||||
code = getattr(e.args[0], "code", None)
|
||||
if code == 1920:
|
||||
# ORA-01920: user name conflicts with another user or role name
|
||||
pass
|
||||
elif code == 1031:
|
||||
# ORA-01031: we are not an admin. CI provisions the user with a
|
||||
# privileged account before pytest runs and then points
|
||||
# ORACLE_TEST_DSN at that same unprivileged user, so this bootstrap
|
||||
# cannot (and need not) create it. Assume it exists — if it does
|
||||
# not, run_migrations below fails with a plain login error.
|
||||
pass
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
@@ -233,6 +233,7 @@ async def test_backup_restore_roundtrip(backup_test_schema):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_backup_restore_preserves_all_column_types(backup_test_schema):
|
||||
"""Test that all column types are preserved: vectors, UUIDs, timestamps, JSONB."""
|
||||
db_url, schema_name, _fq, embeddings = backup_test_schema
|
||||
|
||||
@@ -90,7 +90,8 @@ class TestAgentProfile:
|
||||
await memory.get_bank_profile(agent_id_2, request_context=request_context)
|
||||
await memory.get_bank_profile(agent_id_3, request_context=request_context)
|
||||
|
||||
agents = await memory.list_banks(request_context=request_context)
|
||||
page = await memory.list_banks(search_query="test_list", limit=1000, request_context=request_context)
|
||||
agents = page["banks"]
|
||||
|
||||
agent_ids = [a["bank_id"] for a in agents]
|
||||
assert agent_id_1 in agent_ids
|
||||
|
||||
@@ -0,0 +1,319 @@
|
||||
"""ANN scan depth follows each query's LIMIT, instead of a fixed candidate list.
|
||||
|
||||
An ANN scan explores a bounded candidate list and returns what it found, so that list
|
||||
— not the SQL LIMIT — decided how many rows a recall arm could come back with. On
|
||||
pgvector it is ``hnsw.ef_search``, pinned at 200 for the connection's lifetime by the
|
||||
pool's init callback, which silently capped every recall at ~200 dense candidates
|
||||
however large the budget: the budget moved the SQL and nothing else.
|
||||
|
||||
``hnsw.iterative_scan`` (pgvector 0.8+) resolves that without a per-query statement.
|
||||
With it on, a drained candidate list is refilled in ``ef_search``-sized rounds until
|
||||
the query's LIMIT is met, so depth follows the budget on the connection settings the
|
||||
pool already applies — which matters behind a transaction-mode pooler, where a
|
||||
session GUC issued between statements can land on a different backend.
|
||||
|
||||
Covers:
|
||||
- Both tuning profiles: recall resumes, retain-side link probing explicitly does not.
|
||||
- That the arms fetch what both their consumers read (the semantic list, and the graph
|
||||
arm's seeds) and that ``search`` issues no session statement of its own.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import uuid
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api._vector_index import ann_max_scan_tuples, ann_search_tuning_settings
|
||||
from hindsight_api.engine.memories.postgres import PostgresMemories
|
||||
from hindsight_api.engine.search import retrieval as retrieval_mod
|
||||
from hindsight_api.engine.search.link_expansion_retrieval import GRAPH_SEED_LIMIT
|
||||
|
||||
BUDGET_MID = 300
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tuning profiles
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_recall_connections_resume_the_scan():
|
||||
"""Without this the scan stops at ef_search rows and the budget cannot reach the index."""
|
||||
settings = dict(ann_search_tuning_settings("pgvector", kind="high_recall"))
|
||||
|
||||
assert settings["hnsw.iterative_scan"] == "strict_order"
|
||||
# relaxed_order would return rows out of distance order, which the Python-side
|
||||
# trim in retrieve_semantic_bm25_combined_sql assumes it can rely on.
|
||||
assert settings["hnsw.ef_search"] == "200"
|
||||
# Bounded, so a heavily filtered query cannot resume its way into a huge scan.
|
||||
assert settings["hnsw.max_scan_tuples"] == str(ann_max_scan_tuples())
|
||||
assert ann_max_scan_tuples() < 20000 # pgvector's default
|
||||
|
||||
|
||||
def test_retain_link_probing_does_not_resume():
|
||||
"""Link probing is tuned for latency; resuming past its small list would defeat that."""
|
||||
settings = dict(ann_search_tuning_settings("pgvector", kind="low_latency"))
|
||||
|
||||
assert settings["hnsw.iterative_scan"] == "off"
|
||||
assert settings["hnsw.ef_search"] == "60"
|
||||
|
||||
|
||||
def test_backends_without_the_knobs_get_no_settings():
|
||||
for ext in ("vchord", "pgvectorscale", "pg_diskann", "scann"):
|
||||
assert ann_search_tuning_settings(ext, kind="high_recall") == ()
|
||||
assert ann_search_tuning_settings(ext, kind="low_latency") == ()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# What the arms ask for
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class FakeDialect:
|
||||
"""Captures what each semantic arm asks the index for."""
|
||||
|
||||
def __init__(self):
|
||||
self.fetch_limits: list[int] = []
|
||||
|
||||
def build_semantic_arm(self, *, fetch_limit, **kwargs):
|
||||
self.fetch_limits.append(fetch_limit)
|
||||
return "SELECT 'semantic' AS source"
|
||||
|
||||
def build_bm25_arm(self, **kwargs):
|
||||
return "SELECT 'bm25' AS source"
|
||||
|
||||
def prepare_bm25_text(self, tokens, query_text, **kwargs):
|
||||
return " | ".join(tokens)
|
||||
|
||||
|
||||
class FakeConn:
|
||||
"""Fails the test if recall issues a session setting or opens a transaction."""
|
||||
|
||||
backend_type = "postgresql"
|
||||
|
||||
def transaction(self):
|
||||
raise AssertionError("recall must not open a transaction to tune the scan")
|
||||
|
||||
async def execute(self, sql, *params):
|
||||
raise AssertionError(f"recall must not issue session settings per query: {sql!r}")
|
||||
|
||||
async def fetch(self, query, *params):
|
||||
return []
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def search_path(monkeypatch):
|
||||
dialect = FakeDialect()
|
||||
config = SimpleNamespace(
|
||||
semantic_min_similarity=0.0,
|
||||
bm25_min_score=0.0,
|
||||
text_search_extension="native",
|
||||
text_search_extension_native_language="english",
|
||||
)
|
||||
monkeypatch.setattr(retrieval_mod, "create_sql_dialect", lambda backend: dialect)
|
||||
monkeypatch.setattr(retrieval_mod, "get_config", lambda: config)
|
||||
monkeypatch.setattr(retrieval_mod, "fq_table", lambda name: name)
|
||||
monkeypatch.setattr(retrieval_mod, "get_current_schema", lambda: None)
|
||||
return dialect
|
||||
|
||||
|
||||
async def _search(conn, limit: int, **kwargs):
|
||||
await PostgresMemories({}).search(
|
||||
conn=conn,
|
||||
bank_id="bank-1",
|
||||
fact_types=["world", "experience"],
|
||||
query_embedding="[0.0]",
|
||||
query_text="alpha beta",
|
||||
limit=limit,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
async def test_arms_ask_for_exactly_the_rows_they_keep(search_path):
|
||||
"""No row over-fetch: the arms are ordered by distance, so trimming N of 5N in
|
||||
Python returned precisely what LIMIT N would have — the surplus was fetched,
|
||||
decoded and dropped unread."""
|
||||
await _search(FakeConn(), BUDGET_MID)
|
||||
|
||||
assert search_path.fetch_limits == [BUDGET_MID, BUDGET_MID] # one arm per fact_type
|
||||
|
||||
|
||||
async def test_small_budget_still_covers_the_graph_arms_seeds(search_path):
|
||||
"""The graph arm reads its entry points from these same rows, so a budget below
|
||||
GRAPH_SEED_LIMIT must not starve it."""
|
||||
await _search(FakeConn(), 5, graph_seed_min_similarity=0.3)
|
||||
|
||||
assert search_path.fetch_limits == [GRAPH_SEED_LIMIT, GRAPH_SEED_LIMIT]
|
||||
|
||||
|
||||
async def test_no_seed_threshold_means_no_seed_floor(search_path):
|
||||
"""With the graph arm off, nothing reads past the semantic list itself."""
|
||||
await _search(FakeConn(), 5)
|
||||
|
||||
assert search_path.fetch_limits == [5, 5]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# The property the change delivers, against a real index
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
EMBED_DIM = 384
|
||||
# Enough to exceed the 200-row candidate list at a budget of 400, and no more: this
|
||||
# runs alongside timing-sensitive tests, and a bulk load large enough to saturate the
|
||||
# database starves them.
|
||||
_ROWS = 600
|
||||
|
||||
|
||||
def _near_query_vector(seed: int) -> str:
|
||||
"""A distinct vector from one tight cluster.
|
||||
|
||||
Clustered rather than uniformly random on purpose: an HNSW graph over scattered
|
||||
vectors is sparsely connected, so a resumed scan exhausts the reachable set before
|
||||
it reaches the requested budget and the test measures graph connectivity instead of
|
||||
the setting under test.
|
||||
"""
|
||||
rng = random.Random(seed)
|
||||
values = [1.0] + [rng.uniform(-0.05, 0.05) for _ in range(EMBED_DIM - 1)]
|
||||
norm = sum(v * v for v in values) ** 0.5
|
||||
return "[" + ",".join(f"{v / norm:.5f}" for v in values) + "]"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_kill_switch_flips_real_retrieval_depth(memory, request_context, ann_config):
|
||||
"""End to end, through the pool: on, the budget reaches the index; off, it does not.
|
||||
|
||||
Both halves matter. On is the fix — with iterative scans off the ground-layer search
|
||||
runs once and the scan ends when its ef_search-sized list drains, so the arm cannot
|
||||
return more than ~200 rows however large the recall budget. Off is the operational
|
||||
revert, and it has to land on exactly that pre-existing behaviour rather than some
|
||||
third state nobody runs.
|
||||
|
||||
Driven by the environment variable through the pool's own session setup, not by
|
||||
setting the GUCs by hand, so it covers the path production actually takes. Rows and
|
||||
index are built directly: the property belongs to the index scan, and going through
|
||||
retain would drag in extraction and consolidation.
|
||||
"""
|
||||
from hindsight_api.engine.search.retrieval import retrieve_semantic_bm25_combined_sql
|
||||
from hindsight_api.engine.retain.bank_utils import get_or_create_bank_profile
|
||||
from hindsight_api.engine.task_backend import fq_table
|
||||
|
||||
bank_id = f"test_iter_scan_{uuid.uuid4().hex[:8]}"
|
||||
budget = 400 # deliberately above the standing ef_search of 200
|
||||
# Creating the bank also builds its per-(bank, fact_type) partial vector index —
|
||||
# the same one recall uses — so this exercises the production index, not a stand-in.
|
||||
await get_or_create_bank_profile(memory._backend, bank_id)
|
||||
pool = await memory._get_pool()
|
||||
probe = _near_query_vector(0)
|
||||
table = fq_table("memory_units")
|
||||
try:
|
||||
async with pool.acquire() as conn:
|
||||
await conn.executemany(
|
||||
f"INSERT INTO {table} (bank_id, text, fact_type, embedding) VALUES ($1, $2, 'world', $3::vector)",
|
||||
[(bank_id, f"filler fact {i}", _near_query_vector(i)) for i in range(_ROWS)],
|
||||
)
|
||||
await conn.execute(f"ANALYZE {table}")
|
||||
|
||||
async def semantic_rows(iterative: bool) -> int:
|
||||
ann_config("ann_iterative_scan", iterative)
|
||||
# The pool re-applies its session settings on every acquire, so a fresh
|
||||
# connection resolves the flag again rather than inheriting the value the
|
||||
# process started with.
|
||||
async with pool.acquire() as conn:
|
||||
# The property under test belongs to the ANN scan, not to the planner's
|
||||
# choice: on a table this size a full scan plus a sort is genuinely
|
||||
# cheaper, and inflating the fixture until ANN wins would only make the
|
||||
# test slow. Discourage both alternatives so the ordered path is taken.
|
||||
await conn.execute("SET enable_seqscan = off")
|
||||
await conn.execute("SET enable_sort = off")
|
||||
plan = "\n".join(
|
||||
r[0]
|
||||
for r in await conn.fetch(
|
||||
f"EXPLAIN SELECT id FROM {table} WHERE bank_id = $1 AND fact_type = 'world' "
|
||||
f"AND embedding IS NOT NULL ORDER BY embedding <=> $2::vector LIMIT {budget}",
|
||||
bank_id,
|
||||
probe,
|
||||
)
|
||||
)
|
||||
# "Index Scan" alone is not enough — a btree scan plus a Sort also
|
||||
# matches, and it returns every row regardless of the candidate list,
|
||||
# which would make this quietly measure nothing. An ANN scan emits rows
|
||||
# already ordered, so the giveaway is the absence of a Sort node.
|
||||
assert "Index Scan" in plan and "Sort" not in plan, f"expected an ANN scan, got:\n{plan}"
|
||||
result = await retrieve_semantic_bm25_combined_sql(
|
||||
conn, probe, "", bank_id, ["world"], budget, min_semantic=0.0
|
||||
)
|
||||
return len(result["world"].semantic)
|
||||
|
||||
with_resume = await semantic_rows(True)
|
||||
without_resume = await semantic_rows(False)
|
||||
|
||||
# On: the budget reaches the index.
|
||||
assert with_resume == budget, f"expected the full budget, got {with_resume}"
|
||||
# Off: capped by the candidate list, exactly as before the fix existed.
|
||||
assert without_resume <= 250, f"expected the scan to stop at ~ef_search, got {without_resume}"
|
||||
assert without_resume < with_resume
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Operational controls
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ann_config(monkeypatch):
|
||||
"""Override an ANN config field for one test, without disturbing anything else.
|
||||
|
||||
Set on the cached config instance rather than by setting the env var and clearing
|
||||
the cache: clearing it is process-wide, so every engine built earlier in the
|
||||
session would silently start resolving a config rebuilt from the current
|
||||
environment. monkeypatch restores the attribute at teardown.
|
||||
"""
|
||||
from hindsight_api.config import _get_raw_config
|
||||
|
||||
def _set(field: str, value) -> None:
|
||||
monkeypatch.setattr(_get_raw_config(), field, value)
|
||||
|
||||
return _set
|
||||
|
||||
|
||||
def test_the_kill_switch_removes_the_resume_settings(ann_config):
|
||||
"""Turning it off must leave a connection exactly as it was before the feature.
|
||||
|
||||
Dropping the GUCs rather than sending iterative_scan=off matters for two reasons:
|
||||
a pgvector older than 0.8 rejects them outright (it reserves the "hnsw." prefix),
|
||||
and an operator who pinned values server-side keeps them.
|
||||
"""
|
||||
ann_config("ann_iterative_scan", False)
|
||||
settings = ann_search_tuning_settings("pgvector", kind="high_recall")
|
||||
|
||||
assert settings == (("hnsw.ef_search", "200"),)
|
||||
|
||||
|
||||
def test_the_scan_ceiling_is_tunable(ann_config):
|
||||
"""The dial between the previous behaviour and full budget depth."""
|
||||
ann_config("ann_max_scan_tuples", 1500)
|
||||
settings = dict(ann_search_tuning_settings("pgvector", kind="high_recall"))
|
||||
|
||||
assert settings["hnsw.max_scan_tuples"] == "1500"
|
||||
assert settings["hnsw.iterative_scan"] == "strict_order"
|
||||
|
||||
|
||||
def test_an_unreadable_ceiling_is_rejected_at_config_load(monkeypatch):
|
||||
"""Parsing and validation belong to HindsightConfig, not to this module."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_ANN_MAX_SCAN_TUPLES", "not-a-number")
|
||||
with pytest.raises(ValueError):
|
||||
HindsightConfig.from_env()
|
||||
|
||||
|
||||
def test_retain_probing_is_unaffected_by_the_switch(ann_config):
|
||||
"""Link probing never resumed; the switch has nothing to take from it."""
|
||||
ann_config("ann_iterative_scan", False)
|
||||
settings = dict(ann_search_tuning_settings("pgvector", kind="low_latency"))
|
||||
|
||||
assert settings == {"hnsw.ef_search": "60"}
|
||||
@@ -664,6 +664,7 @@ async def test_all_degenerate_facts_still_persist_document_chunks(memory, reques
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_streaming_offsets_chunk_local_causal_fact_indices(memory, request_context, monkeypatch):
|
||||
"""Causal targets from independently extracted chunks must stay within their source chunk."""
|
||||
from hindsight_api.engine.response_models import TokenUsage
|
||||
@@ -765,17 +766,24 @@ async def test_degenerate_fact_preserves_later_chunk_provenance(memory, request_
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
rows = await pool.fetch(
|
||||
"""
|
||||
SELECT units.text AS fact_text, chunks.chunk_index AS chunk_index
|
||||
FROM memory_units units
|
||||
JOIN chunks ON chunks.chunk_id = units.chunk_id
|
||||
WHERE units.bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
assert {(row["fact_text"], row["chunk_index"]) for row in rows} == {
|
||||
# Each fact carries the chunk it came from, and the document's chunks carry
|
||||
# their index, so the join the assertion needs is over two API reads.
|
||||
chunk_index_by_id = {
|
||||
chunk["chunk_id"]: chunk["chunk_index"]
|
||||
for chunk in (
|
||||
await memory.list_document_chunks(
|
||||
bank_id, "degen-provenance-document", limit=500, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
}
|
||||
units = (await memory.list_memory_units(bank_id, limit=500, request_context=request_context))["items"]
|
||||
# Chunkless units (an observation, say) were dropped by the inner join before
|
||||
# and are dropped here for the same reason: they have no provenance to check.
|
||||
assert {
|
||||
(unit["text"], chunk_index_by_id[unit["chunk_id"]])
|
||||
for unit in units
|
||||
if unit["chunk_id"] in chunk_index_by_id
|
||||
} == {
|
||||
("chunk zero real fact", 0),
|
||||
("chunk one real fact", 1),
|
||||
}
|
||||
|
||||
@@ -22,8 +22,8 @@ async def _retain(memory, bank_id, document_id, content, request_context):
|
||||
|
||||
|
||||
async def _bank_entry(memory, bank_id, request_context):
|
||||
banks = await memory.list_banks(request_context=request_context)
|
||||
return next(b for b in banks if b["bank_id"] == bank_id)
|
||||
page = await memory.list_banks(search_query=bank_id, request_context=request_context)
|
||||
return next(b for b in page["banks"] if b["bank_id"] == bank_id)
|
||||
|
||||
|
||||
def _ts(value: str | None) -> datetime:
|
||||
@@ -84,8 +84,8 @@ async def test_banks_are_ordered_by_last_write(memory, request_context):
|
||||
# the most recently written bank.
|
||||
await _retain(memory, older_bank, "doc-a", "xyzabc123 !@# alpha revised", request_context)
|
||||
|
||||
banks = await memory.list_banks(request_context=request_context)
|
||||
ordered = [b["bank_id"] for b in banks if b["bank_id"] in (older_bank, newer_bank)]
|
||||
page = await memory.list_banks(search_query="test_last_write_order_", request_context=request_context)
|
||||
ordered = [b["bank_id"] for b in page["banks"] if b["bank_id"] in (older_bank, newer_bank)]
|
||||
assert ordered == [older_bank, newer_bank]
|
||||
|
||||
finally:
|
||||
|
||||
@@ -47,6 +47,7 @@ async def _insert_memory(memory, bank_id: str, text: str, *, failed: bool = Fals
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_bank_stats_exposes_memory_write_watermark(api_client, memory, test_bank_id):
|
||||
"""/stats must carry the bank's newest memory write time.
|
||||
|
||||
@@ -217,6 +218,7 @@ async def test_memories_timeseries_reflects_retained_memories(api_client, test_b
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_bank_stats_reports_failed_consolidation(api_client, memory, test_bank_id):
|
||||
"""/stats must surface the count of memories with consolidation_failed_at set."""
|
||||
try:
|
||||
@@ -237,6 +239,7 @@ async def test_bank_stats_reports_failed_consolidation(api_client, memory, test_
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_pending_consolidation_matches_pending_memory_list(api_client, memory, test_bank_id):
|
||||
"""pending_consolidation must agree with ?consolidation_state=pending.
|
||||
|
||||
@@ -266,6 +269,7 @@ async def test_pending_consolidation_matches_pending_memory_list(api_client, mem
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_list_memories_filter_by_consolidation_state_failed(api_client, memory, test_bank_id):
|
||||
"""?consolidation_state=failed returns only memories with consolidation_failed_at set."""
|
||||
try:
|
||||
@@ -332,6 +336,7 @@ async def test_bank_stats_link_counts_have_no_join(api_client, test_bank_id):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_get_bank_freshness_returns_only_consolidation_fields(memory, test_bank_id):
|
||||
"""get_bank_freshness must return just the freshness keys, no link aggregation."""
|
||||
from hindsight_api.extensions import RequestContext
|
||||
|
||||
@@ -55,6 +55,7 @@ class TestDistributedBankStatsCache:
|
||||
assert isinstance(memory._bank_stats_cache, DistributedBankStatsCache)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_result_is_written_and_served_from_table(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
if memory._database_backend_type != "postgresql":
|
||||
pytest.skip("distributed cache is PostgreSQL-only")
|
||||
@@ -93,6 +94,7 @@ class TestDistributedBankStatsCache:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_force_refresh_bypasses_and_updates_cache(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
@@ -127,6 +129,7 @@ class TestDistributedBankStatsCache:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_degrades_when_cache_table_unreachable(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
if memory._database_backend_type != "postgresql":
|
||||
pytest.skip("distributed cache is PostgreSQL-only")
|
||||
|
||||
@@ -7,21 +7,8 @@ from asyncpg.exceptions import DeadlockDetectedError
|
||||
from hindsight_api.engine.retain import bank_utils
|
||||
|
||||
|
||||
class _FailingIndexOps:
|
||||
async def create_bank_vector_indexes(self, *args, **kwargs) -> None:
|
||||
raise RuntimeError("simulated per-bank vector index DDL failure")
|
||||
|
||||
|
||||
class _DeadlockOnceIndexOps:
|
||||
"""Raises a deadlock on the first per-bank index DDL, then succeeds."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
async def create_bank_vector_indexes(self, *args, **kwargs) -> None:
|
||||
self.calls += 1
|
||||
if self.calls == 1:
|
||||
raise DeadlockDetectedError("deadlock detected")
|
||||
class _BankOps:
|
||||
"""Dialect ops stub. Bank creation issues no index DDL (#3485), so this is bare."""
|
||||
|
||||
|
||||
class _FakeTransaction:
|
||||
@@ -39,10 +26,13 @@ class _FakeTransaction:
|
||||
|
||||
|
||||
class _FakeConnection:
|
||||
def __init__(self) -> None:
|
||||
def __init__(self, raise_on_insert: BaseException | None = None) -> None:
|
||||
self.committed_bank: str | None = None
|
||||
self.pending_bank: str | None = None
|
||||
self.in_transaction = False
|
||||
self.insert_calls = 0
|
||||
# Raised by the first INSERT only, then cleared, so a retry succeeds.
|
||||
self._raise_on_insert = raise_on_insert
|
||||
|
||||
def transaction(self) -> _FakeTransaction:
|
||||
return _FakeTransaction(self)
|
||||
@@ -58,6 +48,10 @@ class _FakeConnection:
|
||||
}
|
||||
|
||||
async def fetchval(self, query: str, bank_id: str, *args):
|
||||
self.insert_calls += 1
|
||||
if self._raise_on_insert is not None:
|
||||
error, self._raise_on_insert = self._raise_on_insert, None
|
||||
raise error
|
||||
if self.in_transaction:
|
||||
self.pending_bank = bank_id
|
||||
else:
|
||||
@@ -68,13 +62,13 @@ class _FakeConnection:
|
||||
class _FakePool:
|
||||
def __init__(self, conn: _FakeConnection, ops=None) -> None:
|
||||
self.conn = conn
|
||||
self.ops = ops if ops is not None else _FailingIndexOps()
|
||||
self.ops = ops if ops is not None else _BankOps()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lazy_bank_create_rolls_back_on_vector_index_failure(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""A failed per-bank index DDL must not leave an orphaned bank row."""
|
||||
conn = _FakeConnection()
|
||||
async def test_lazy_bank_create_rolls_back_on_failure(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""A failure inside the bank-create transaction must not leave an orphaned bank row."""
|
||||
conn = _FakeConnection(raise_on_insert=RuntimeError("simulated bank insert failure"))
|
||||
pool = _FakePool(conn)
|
||||
|
||||
@asynccontextmanager
|
||||
@@ -83,24 +77,25 @@ async def test_lazy_bank_create_rolls_back_on_vector_index_failure(monkeypatch:
|
||||
|
||||
monkeypatch.setattr(bank_utils, "acquire_with_retry", acquire_without_transaction)
|
||||
|
||||
with pytest.raises(RuntimeError, match="simulated per-bank vector index DDL failure"):
|
||||
with pytest.raises(RuntimeError, match="simulated bank insert failure"):
|
||||
await bank_utils.get_or_create_bank_profile(pool, "atomicity-test-bank")
|
||||
|
||||
profile = await bank_utils.get_bank_profile_if_exists(pool, "atomicity-test-bank")
|
||||
assert profile is None, "bank row should roll back when per-bank vector index creation fails"
|
||||
assert profile is None, "bank row should roll back when the create transaction fails"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_or_create_bank_profile_retries_on_deadlock(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""A deadlock during per-bank index DDL is retried on a fresh transaction.
|
||||
"""A deadlock inside the create transaction is retried on a fresh one.
|
||||
|
||||
Concurrent bank creation issues CREATE INDEX against the shared memory_units
|
||||
table, which can deadlock (asyncpg DeadlockDetectedError). The pool-owning
|
||||
get_or_create_bank_profile must retry rather than surface the deadlock.
|
||||
The retry originally guarded the per-bank CREATE INDEX that ran inline here;
|
||||
that DDL is gone (#3485), but the lazy create can still lose a deadlock to a
|
||||
concurrent writer touching the same bank row, and the body is idempotent
|
||||
(INSERT ... ON CONFLICT DO NOTHING), so it must still retry rather than
|
||||
surface the deadlock to the caller.
|
||||
"""
|
||||
conn = _FakeConnection()
|
||||
ops = _DeadlockOnceIndexOps()
|
||||
pool = _FakePool(conn, ops=ops)
|
||||
conn = _FakeConnection(raise_on_insert=DeadlockDetectedError("deadlock detected"))
|
||||
pool = _FakePool(conn)
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire(*args, **kwargs):
|
||||
@@ -117,6 +112,6 @@ async def test_get_or_create_bank_profile_retries_on_deadlock(monkeypatch: pytes
|
||||
result = await bank_utils.get_or_create_bank_profile(pool, "deadlock-retry-bank")
|
||||
|
||||
assert result.created is True
|
||||
assert ops.calls == 2, "index DDL should be attempted twice (deadlock, then success)"
|
||||
assert conn.insert_calls == 2, "the insert should be attempted twice (deadlock, then success)"
|
||||
profile = await bank_utils.get_bank_profile_if_exists(pool, "deadlock-retry-bank")
|
||||
assert profile is not None, "bank must exist after the retry succeeds"
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Per-bank vector-index DDL is serialized per table within a process.
|
||||
"""Per-bank vector-index drop DDL is serialized per table within a process.
|
||||
|
||||
Concurrent index DDL on the shared ``memory_units`` table deadlocks by design:
|
||||
DROP INDEX CONCURRENTLY holds ShareUpdateExclusive while waiting out every
|
||||
@@ -10,11 +10,15 @@ test-api 3/3). Advisory locks are banned in this codebase (poolers), so the
|
||||
fix is an in-process asyncio lock on ``PostgreSQLOps`` plus a much larger
|
||||
jittered retry budget for the cross-process residue.
|
||||
|
||||
Bank deletion is now the only request path that issues vector-index DDL —
|
||||
indexes are earned by size and built by the maintenance sweep (#3485) — but the
|
||||
teardown storm this guards against is unchanged, because a bank large enough to
|
||||
have indexes still drops three of them when it goes.
|
||||
|
||||
Two layers are proven here:
|
||||
|
||||
* unit (fake conn): create and drop DDL for one table never interleave, the
|
||||
create and drop sides contend on the same lock, and the lock is released
|
||||
when a statement raises;
|
||||
* unit (fake conn): concurrent drops for one table never interleave, different
|
||||
tables do not contend, and the lock is released when a statement raises;
|
||||
* integration: a many-bank concurrent ``delete_bank`` storm — the CI failure
|
||||
shape — completes without ``DeadlockDetectedError``.
|
||||
"""
|
||||
@@ -25,8 +29,14 @@ import uuid
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.db.ops_postgresql import PostgreSQLOps
|
||||
from hindsight_api.engine.db_utils import retry_with_backoff
|
||||
from hindsight_api.engine.retain.bank_utils import _BANK_INDEX_FACT_TYPES, _bank_index_name, _vector_index_clause
|
||||
|
||||
# Shares the reconcile suite's xdist worker: both do heavy CREATE/DROP INDEX
|
||||
# CONCURRENTLY against the single shared public.memory_units, and concurrent
|
||||
# index DDL on one relation deadlocks by design.
|
||||
pytestmark = pytest.mark.xdist_group("vector_index_reconcile")
|
||||
|
||||
_SCHEMA = "public"
|
||||
_INDEX_CLAUSE = "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
@@ -57,7 +67,7 @@ class _OverlapTrackingConn:
|
||||
return "OK"
|
||||
|
||||
|
||||
async def test_concurrent_create_and_drop_on_one_table_serialize():
|
||||
async def test_concurrent_drops_on_one_table_serialize():
|
||||
ops = PostgreSQLOps()
|
||||
conn = _OverlapTrackingConn()
|
||||
table = f"{_SCHEMA}.memory_units"
|
||||
@@ -65,10 +75,9 @@ async def test_concurrent_create_and_drop_on_one_table_serialize():
|
||||
await asyncio.gather(
|
||||
ops.drop_bank_vector_indexes(conn, _SCHEMA, uuid.uuid4().hex, _BANK_INDEX_FACT_TYPES),
|
||||
ops.drop_bank_vector_indexes(conn, _SCHEMA, uuid.uuid4().hex, _BANK_INDEX_FACT_TYPES),
|
||||
# The drop side reconstructs the create side's fq-table key from
|
||||
# `schema`, so a create must queue behind the drops too.
|
||||
ops.create_bank_vector_indexes(conn, table, "bank-1", uuid.uuid4().hex, _INDEX_CLAUSE, _BANK_INDEX_FACT_TYPES),
|
||||
ops.drop_bank_vector_indexes(conn, _SCHEMA, uuid.uuid4().hex, _BANK_INDEX_FACT_TYPES),
|
||||
)
|
||||
assert ops._index_ddl_lock(table) is ops._index_ddl_lock(f"{_SCHEMA}.memory_units")
|
||||
|
||||
assert len(conn.calls) == 3 * len(_BANK_INDEX_FACT_TYPES)
|
||||
assert conn.max_in_flight == 1, "vector-index DDL statements overlapped despite the per-table lock"
|
||||
@@ -86,11 +95,10 @@ async def test_lock_released_when_ddl_raises():
|
||||
with pytest.raises(RuntimeError):
|
||||
await ops.drop_bank_vector_indexes(conn, _SCHEMA, uuid.uuid4().hex, _BANK_INDEX_FACT_TYPES)
|
||||
|
||||
# A subsequent create must not hang on a lock the failed drop never released.
|
||||
# A subsequent drop must not hang on a lock the failed one never released.
|
||||
ok = _OverlapTrackingConn()
|
||||
await asyncio.wait_for(
|
||||
ops.create_bank_vector_indexes(
|
||||
conn, f"{_SCHEMA}.memory_units", "bank-1", uuid.uuid4().hex, _INDEX_CLAUSE, _BANK_INDEX_FACT_TYPES
|
||||
),
|
||||
ops.drop_bank_vector_indexes(ok, _SCHEMA, uuid.uuid4().hex, _BANK_INDEX_FACT_TYPES),
|
||||
timeout=2.0,
|
||||
)
|
||||
|
||||
@@ -115,6 +123,25 @@ async def test_concurrent_bank_delete_storm_does_not_deadlock(memory, request_co
|
||||
bank_id: str(await conn.fetchval("SELECT internal_id FROM banks WHERE bank_id = $1", bank_id))
|
||||
for bank_id in bank_ids
|
||||
}
|
||||
# Bank creation no longer builds these, so the storm has to be staged:
|
||||
# give every bank its three indexes up front, as a bank past the size
|
||||
# threshold would have.
|
||||
for bank_id, internal_id in internal_ids.items():
|
||||
literal = await conn.fetchval("SELECT quote_literal($1::text)", bank_id)
|
||||
for ft in _BANK_INDEX_FACT_TYPES:
|
||||
name = _bank_index_name(ft, internal_id)
|
||||
# CONCURRENTLY, and retried: a plain CREATE INDEX takes ShareLock
|
||||
# on the shared memory_units table, which closes a deadlock cycle
|
||||
# with another xdist worker's DROP INDEX CONCURRENTLY
|
||||
# (ShareUpdateExclusive). Staging the storm must not itself be the
|
||||
# storm.
|
||||
await retry_with_backoff(
|
||||
lambda name=name, ft=ft: conn.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {name} "
|
||||
f"ON {_SCHEMA}.memory_units {_INDEX_CLAUSE} "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = {literal}"
|
||||
)
|
||||
)
|
||||
|
||||
await asyncio.gather(*(memory.delete_bank(bank_id, request_context=request_context) for bank_id in bank_ids))
|
||||
|
||||
|
||||
@@ -43,6 +43,7 @@ async def _seed_bank_and_document(conn, bank_id: str, document_id: str) -> None:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_store_chunks_batch_is_idempotent_for_same_chunk_id(memory):
|
||||
"""
|
||||
Regression for #977.
|
||||
|
||||
@@ -28,6 +28,20 @@ from hindsight_api.engine.reflect.tools import (
|
||||
from tests.llm_judge import assert_meets_criteria
|
||||
|
||||
|
||||
async def _unconsolidated(memory: MemoryEngine, bank_id: str, request_context) -> int:
|
||||
"""How many source facts are still waiting to be consolidated.
|
||||
|
||||
consolidation_state='pending' is the read API's name for the predicate these
|
||||
tests used to spell out in SQL: consolidated_at IS NULL and a source fact type.
|
||||
It additionally excludes facts whose consolidation permanently failed — the
|
||||
stricter reading, and the one the assertions here actually mean.
|
||||
"""
|
||||
page = await memory.list_memory_units(
|
||||
bank_id=bank_id, consolidation_state="pending", limit=1, request_context=request_context
|
||||
)
|
||||
return page["total"]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def enable_observations():
|
||||
"""Enable observations for all tests in this module."""
|
||||
@@ -63,22 +77,18 @@ class TestConsolidationIntegration:
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify observation exists in memory_units
|
||||
# Verify observation exists via the list API
|
||||
# (consolidation already ran as part of retain via SyncTaskBackend)
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, proof_count, fact_type
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
""",
|
||||
bank_id,
|
||||
observations = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
# With the deterministic mock, consolidation always produces observations
|
||||
assert len(observations) >= 1, "Consolidation must create at least one observation"
|
||||
obs = observations[0]
|
||||
assert obs["proof_count"] >= 1
|
||||
assert obs["fact_type"] == "observation"
|
||||
)["items"]
|
||||
# With the deterministic mock, consolidation always produces observations
|
||||
assert len(observations) >= 1, "Consolidation must create at least one observation"
|
||||
obs = observations[0]
|
||||
assert obs["proof_count"] >= 1
|
||||
assert obs["fact_type"] == "observation"
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -106,21 +116,16 @@ class TestConsolidationIntegration:
|
||||
)
|
||||
|
||||
# Check observations after both retains
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, proof_count
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY proof_count DESC
|
||||
""",
|
||||
bank_id,
|
||||
observations = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
|
||||
# Must have at least one observation from consolidation
|
||||
assert len(observations) >= 1, "Consolidation must create observations from retained memories"
|
||||
assert all(obs["text"] for obs in observations)
|
||||
assert all(obs["proof_count"] >= 1 for obs in observations)
|
||||
# Must have at least one observation from consolidation
|
||||
assert len(observations) >= 1, "Consolidation must create observations from retained memories"
|
||||
assert all(obs["text"] for obs in observations)
|
||||
assert all(obs["proof_count"] >= 1 for obs in observations)
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -208,31 +213,17 @@ class TestConsolidationIntegration:
|
||||
)
|
||||
|
||||
# Check observation and its entity links
|
||||
async with memory._pool.acquire() as conn:
|
||||
observation = await conn.fetchrow(
|
||||
"""
|
||||
SELECT id
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
LIMIT 1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
observations = await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1, request_context=request_context
|
||||
)
|
||||
|
||||
# Consolidation must create an observation
|
||||
assert observation is not None, "Consolidation must create an observation"
|
||||
# Consolidation must create an observation
|
||||
assert observations["items"], "Consolidation must create an observation"
|
||||
|
||||
# Check if entity links were copied
|
||||
entity_links = await conn.fetch(
|
||||
"""
|
||||
SELECT entity_id
|
||||
FROM unit_entities
|
||||
WHERE unit_id = $1
|
||||
""",
|
||||
observation["id"],
|
||||
)
|
||||
# Observation should have inherited entity links from source memory
|
||||
assert entity_links is not None
|
||||
# An observation carries the entities of the facts it was drawn from, so the
|
||||
# detail view is where the copied links show up.
|
||||
observation = await memory.get_memory_unit(bank_id, observations["items"][0]["id"], request_context)
|
||||
assert observation["entities"] is not None
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -267,6 +258,7 @@ class TestConsolidationIntegration:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_consolidation_uses_source_memory_ids(self, memory: MemoryEngine, request_context):
|
||||
"""Test that observations use source_memory_ids (not memory_links) to track source facts.
|
||||
|
||||
@@ -366,31 +358,25 @@ class TestConsolidationIntegration:
|
||||
)
|
||||
|
||||
# Check observations - should have separate observations for each person
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, source_memory_ids
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
""",
|
||||
bank_id,
|
||||
observations = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
|
||||
# Should have multiple observations (one per person/fact)
|
||||
# Not everything merged into one
|
||||
assert len(observations) >= 2, (
|
||||
f"Expected multiple observations for different people, got {len(observations)}"
|
||||
)
|
||||
# Should have multiple observations (one per person/fact)
|
||||
# Not everything merged into one
|
||||
assert len(observations) >= 2, f"Expected multiple observations for different people, got {len(observations)}"
|
||||
|
||||
# Fast structural check first: no single observation should name more
|
||||
# than one of {John, Mary, Bob}. This catches the obvious failure mode
|
||||
# cheaply without paying for a judge call per observation.
|
||||
for obs in observations:
|
||||
text = obs["text"].lower()
|
||||
people_mentioned = sum(1 for name in ["john", "mary", "bob"] if name in text)
|
||||
assert people_mentioned <= 1, f"Observation should not merge different people: {obs['text']}"
|
||||
# Fast structural check first: no single observation should name more
|
||||
# than one of {John, Mary, Bob}. This catches the obvious failure mode
|
||||
# cheaply without paying for a judge call per observation.
|
||||
for obs in observations:
|
||||
text = obs["text"].lower()
|
||||
people_mentioned = sum(1 for name in ["john", "mary", "bob"] if name in text)
|
||||
assert people_mentioned <= 1, f"Observation should not merge different people: {obs['text']}"
|
||||
|
||||
obs_listing = "\n".join(f"Observation {i + 1}: {obs['text']}" for i, obs in enumerate(observations))
|
||||
obs_listing = "\n".join(f"Observation {i + 1}: {obs['text']}" for i, obs in enumerate(observations))
|
||||
|
||||
# Semantic backup: catch the case where the LLM merges facts about different
|
||||
# people using pronouns or referent shifts that bypass the proper-noun check
|
||||
@@ -443,15 +429,12 @@ class TestConsolidationIntegration:
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check we have one observation
|
||||
async with memory._pool.acquire() as conn:
|
||||
obs_before = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
""",
|
||||
bank_id,
|
||||
obs_before = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
count_before = len(obs_before)
|
||||
)["items"]
|
||||
count_before = len(obs_before)
|
||||
|
||||
# Add contradicting fact (same person, same topic, opposite sentiment)
|
||||
await memory.retain_async(
|
||||
@@ -462,29 +445,23 @@ class TestConsolidationIntegration:
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check observations after consolidation
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, source_memory_ids
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
observations = (
|
||||
await memory.list_memory_units(bank_id, fact_type="observation", limit=500, request_context=request_context)
|
||||
)["items"]
|
||||
|
||||
# The contradiction should be reflected in observations — either:
|
||||
# 1. Merged into one observation with temporal context (e.g., "used to love, now hates")
|
||||
# 2. The original observation updated to reflect the new state
|
||||
# 3. Two separate observations capturing each state
|
||||
# The key is that the contradiction is tracked, not ignored.
|
||||
assert len(observations) >= 1, "Should have at least one observation after contradiction"
|
||||
# The contradiction should be reflected in observations — either:
|
||||
# 1. Merged into one observation with temporal context (e.g., "used to love, now hates")
|
||||
# 2. The original observation updated to reflect the new state
|
||||
# 3. Two separate observations capturing each state
|
||||
# The key is that the contradiction is tracked, not ignored.
|
||||
assert len(observations) >= 1, "Should have at least one observation after contradiction"
|
||||
|
||||
# Format as numbered list rather than pipe-separated — weaker judge
|
||||
# models read pipe-joins as a single conflated statement.
|
||||
obs_listing = "\n".join(f"Observation {i + 1}: {obs['text']}" for i, obs in enumerate(observations))
|
||||
all_source_ids = []
|
||||
for obs in observations:
|
||||
all_source_ids.extend(obs["source_memory_ids"] or [])
|
||||
# Format as numbered list rather than pipe-separated — weaker judge
|
||||
# models read pipe-joins as a single conflated statement.
|
||||
obs_listing = "\n".join(f"Observation {i + 1}: {obs['text']}" for i, obs in enumerate(observations))
|
||||
all_source_ids = []
|
||||
for obs in observations:
|
||||
all_source_ids.extend(obs["source_memory_ids"])
|
||||
|
||||
# Either the observations reference both sentiments (via text content) or the
|
||||
# consolidation linked both source memories together. The judge evaluates the
|
||||
@@ -543,11 +520,11 @@ class TestConsolidationIntegration:
|
||||
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"SELECT id, text FROM memory_units WHERE bank_id = $1 AND fact_type = 'observation' ORDER BY created_at",
|
||||
bank_id,
|
||||
observations = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
|
||||
obs_count = len(observations)
|
||||
obs_texts = [o["text"] for o in observations]
|
||||
@@ -760,19 +737,16 @@ class TestConsolidationTagRouting:
|
||||
await self._retain_with_tags(memory, bank_id, "Alice likes coffee.", ["alice"], request_context)
|
||||
|
||||
# Check observation has correct tags
|
||||
async with memory._pool.acquire() as conn:
|
||||
obs_before = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
""",
|
||||
bank_id,
|
||||
obs_before = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
count_before = len(obs_before)
|
||||
if obs_before:
|
||||
assert "alice" in (obs_before[0]["tags"] or []), (
|
||||
f"Expected observation to have 'alice' tag, got: {obs_before[0]['tags']}"
|
||||
)
|
||||
count_before = len(obs_before)
|
||||
if obs_before:
|
||||
assert "alice" in (obs_before[0]["tags"] or []), (
|
||||
f"Expected observation to have 'alice' tag, got: {obs_before[0]['tags']}"
|
||||
)
|
||||
|
||||
# Retain related memory with same tags
|
||||
await self._retain_with_tags(
|
||||
@@ -780,25 +754,22 @@ class TestConsolidationTagRouting:
|
||||
)
|
||||
|
||||
# Check observations - should NOT have increased (same scope update)
|
||||
async with memory._pool.acquire() as conn:
|
||||
obs_after = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags, source_memory_ids FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
""",
|
||||
bank_id,
|
||||
obs_after = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
|
||||
# Count of observations should stay same or decrease (merge)
|
||||
assert len(obs_after) <= count_before + 1, (
|
||||
f"Same scope fact should update existing observation, not create new. "
|
||||
f"Before: {count_before}, After: {len(obs_after)}"
|
||||
)
|
||||
# Count of observations should stay same or decrease (merge)
|
||||
assert len(obs_after) <= count_before + 1, (
|
||||
f"Same scope fact should update existing observation, not create new. "
|
||||
f"Before: {count_before}, After: {len(obs_after)}"
|
||||
)
|
||||
|
||||
# The observation(s) should still have alice tag
|
||||
for obs in obs_after:
|
||||
if "coffee" in obs["text"].lower() or "espresso" in obs["text"].lower():
|
||||
assert "alice" in (obs["tags"] or []), f"Updated observation should keep 'alice' tag: {obs['text']}"
|
||||
# The observation(s) should still have alice tag
|
||||
for obs in obs_after:
|
||||
if "coffee" in obs["text"].lower() or "espresso" in obs["text"].lower():
|
||||
assert "alice" in (obs["tags"] or []), f"Updated observation should keep 'alice' tag: {obs['text']}"
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -828,48 +799,41 @@ class TestConsolidationTagRouting:
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check untagged observation exists
|
||||
async with memory._pool.acquire() as conn:
|
||||
obs_before = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
""",
|
||||
bank_id,
|
||||
obs_before = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
count_before = len(obs_before)
|
||||
# Should be untagged or have empty tags
|
||||
if obs_before:
|
||||
assert not obs_before[0]["tags"] or len(obs_before[0]["tags"]) == 0, (
|
||||
f"Expected untagged observation, got: {obs_before[0]['tags']}"
|
||||
)
|
||||
count_before = len(obs_before)
|
||||
# Should be untagged or have empty tags
|
||||
if obs_before:
|
||||
assert not obs_before[0]["tags"] or len(obs_before[0]["tags"]) == 0, (
|
||||
f"Expected untagged observation, got: {obs_before[0]['tags']}"
|
||||
)
|
||||
|
||||
# Retain scoped memory that relates to the global topic
|
||||
await self._retain_with_tags(memory, bank_id, "Pizza originated in Naples.", ["history"], request_context)
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check - global observation should be updated OR new scoped observation created
|
||||
async with memory._pool.acquire() as conn:
|
||||
obs_after = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags, source_memory_ids FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
obs_after = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
|
||||
# At least one observation should exist
|
||||
assert len(obs_after) >= 1, "Expected at least one observation"
|
||||
# At least one observation should exist
|
||||
assert len(obs_after) >= 1, "Expected at least one observation"
|
||||
|
||||
# Check that global observation was updated (source_memory_ids increased)
|
||||
# OR new observation was created with appropriate tags
|
||||
global_observations = [o for o in obs_after if not o["tags"] or len(o["tags"]) == 0]
|
||||
scoped_observations = [o for o in obs_after if o["tags"] and len(o["tags"]) > 0]
|
||||
# Check that global observation was updated (source_memory_ids increased)
|
||||
# OR new observation was created with appropriate tags
|
||||
global_observations = [o for o in obs_after if not o["tags"] or len(o["tags"]) == 0]
|
||||
scoped_observations = [o for o in obs_after if o["tags"] and len(o["tags"]) > 0]
|
||||
|
||||
# Either global was updated or scoped was created
|
||||
assert len(global_observations) >= 1 or len(scoped_observations) >= 1, (
|
||||
"Expected either global observation update or scoped observation creation"
|
||||
)
|
||||
# Either global was updated or scoped was created
|
||||
assert len(global_observations) >= 1 or len(scoped_observations) >= 1, (
|
||||
"Expected either global observation update or scoped observation creation"
|
||||
)
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -897,15 +861,12 @@ class TestConsolidationTagRouting:
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check Alice's observation exists with correct tags
|
||||
async with memory._pool.acquire() as conn:
|
||||
obs_alice = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
""",
|
||||
bank_id,
|
||||
obs_alice = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
count_before = len(obs_alice)
|
||||
)["items"]
|
||||
count_before = len(obs_alice)
|
||||
|
||||
# Retain Bob's memory that relates to Alice's topic (cross-scope)
|
||||
await self._retain_with_tags(
|
||||
@@ -914,28 +875,22 @@ class TestConsolidationTagRouting:
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check observations
|
||||
async with memory._pool.acquire() as conn:
|
||||
obs_after = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags, source_memory_ids FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
obs_after = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
|
||||
# Note: some LLMs may or may not consolidate cross-scope facts.
|
||||
# Just verify structural correctness of any observations that exist.
|
||||
# Note: some LLMs may or may not consolidate cross-scope facts.
|
||||
# Just verify structural correctness of any observations that exist.
|
||||
|
||||
# If observations were created, ensure alice and bob are not merged into same observation
|
||||
# (cross-scope merging should not produce an observation with both tags)
|
||||
if obs_after:
|
||||
observations_with_both = [
|
||||
o for o in obs_after if o["tags"] and "alice" in o["tags"] and "bob" in o["tags"]
|
||||
]
|
||||
assert len(observations_with_both) == 0, (
|
||||
"Should not merge different scopes into one observation with both tags"
|
||||
)
|
||||
# If observations were created, ensure alice and bob are not merged into same observation
|
||||
# (cross-scope merging should not produce an observation with both tags)
|
||||
if obs_after:
|
||||
observations_with_both = [o for o in obs_after if o["tags"] and "alice" in o["tags"] and "bob" in o["tags"]]
|
||||
assert len(observations_with_both) == 0, (
|
||||
"Should not merge different scopes into one observation with both tags"
|
||||
)
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -962,21 +917,18 @@ class TestConsolidationTagRouting:
|
||||
)
|
||||
|
||||
# Check observation was created with correct tags
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
""",
|
||||
bank_id,
|
||||
observations = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
|
||||
assert len(observations) >= 1, "Expected observation to be created"
|
||||
assert len(observations) >= 1, "Expected observation to be created"
|
||||
|
||||
# The observation should have the fact's tags
|
||||
obs = observations[0]
|
||||
assert obs["tags"] is not None, "Observation should have tags"
|
||||
assert "project_x" in obs["tags"], f"Observation should have 'project_x' tag, got: {obs['tags']}"
|
||||
# The observation should have the fact's tags
|
||||
obs = observations[0]
|
||||
assert obs["tags"] is not None, "Observation should have tags"
|
||||
assert "project_x" in obs["tags"], f"Observation should have 'project_x' tag, got: {obs['tags']}"
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -1013,22 +965,18 @@ class TestConsolidationTagRouting:
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check observations
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags, source_memory_ids FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
observations = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
|
||||
# Either alice's observation was updated OR a global observation was created
|
||||
# This is valid LLM behavior - just verify no errors and structure is correct.
|
||||
# Note: with some LLMs, a single simple fact may not generate an observation,
|
||||
# so we don't assert a minimum count - just verify structural correctness if any exist.
|
||||
for obs in observations:
|
||||
assert obs["text"], "Observation should have text"
|
||||
# Either alice's observation was updated OR a global observation was created
|
||||
# This is valid LLM behavior - just verify no errors and structure is correct.
|
||||
# Note: with some LLMs, a single simple fact may not generate an observation,
|
||||
# so we don't assert a minimum count - just verify structural correctness if any exist.
|
||||
for obs in observations:
|
||||
assert obs["text"], "Observation should have text"
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -1099,15 +1047,12 @@ class TestConsolidationTagRouting:
|
||||
await self._retain_with_tags(memory, bank_id, "Alice drinks coffee every morning.", ["alice"], request_context)
|
||||
|
||||
# Check observations before
|
||||
async with memory._pool.acquire() as conn:
|
||||
obs_before = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags, source_memory_ids FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
""",
|
||||
bank_id,
|
||||
obs_before = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
count_before = len(obs_before)
|
||||
)["items"]
|
||||
count_before = len(obs_before)
|
||||
|
||||
# Add fact that could relate to both
|
||||
await self._retain_with_tags(
|
||||
@@ -1115,24 +1060,20 @@ class TestConsolidationTagRouting:
|
||||
)
|
||||
|
||||
# Check observations after
|
||||
async with memory._pool.acquire() as conn:
|
||||
obs_after = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags, source_memory_ids, proof_count FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
obs_after = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
|
||||
# Should have processed without errors
|
||||
assert len(obs_after) >= 1, "Expected at least one observation"
|
||||
# Should have processed without errors
|
||||
assert len(obs_after) >= 1, "Expected at least one observation"
|
||||
|
||||
# Check that consolidation worked (either updates or maintains structure)
|
||||
# The key is no errors and proper tag handling
|
||||
for obs in obs_after:
|
||||
assert obs["text"], "Observation should have text"
|
||||
# Tags should be consistent (not mixing alice and bob, etc.)
|
||||
# Check that consolidation worked (either updates or maintains structure)
|
||||
# The key is no errors and proper tag handling
|
||||
for obs in obs_after:
|
||||
assert obs["text"], "Observation should have text"
|
||||
# Tags should be consistent (not mixing alice and bob, etc.)
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -1442,34 +1383,29 @@ class TestObservationDrillDown:
|
||||
)
|
||||
|
||||
# Get the observation with source_memory_ids
|
||||
async with memory._pool.acquire() as conn:
|
||||
obs_rows = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, proof_count, source_memory_ids
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
obs_rows = (
|
||||
await memory.list_memory_units(bank_id, fact_type="observation", limit=500, request_context=request_context)
|
||||
)["items"]
|
||||
|
||||
assert obs_rows, "Consolidation must create observations"
|
||||
|
||||
# Collect all source_memory_ids across all observations
|
||||
all_source_ids = []
|
||||
for obs in obs_rows:
|
||||
all_source_ids.extend(obs["source_memory_ids"] or [])
|
||||
all_source_ids.extend(obs["source_memory_ids"])
|
||||
|
||||
assert all_source_ids, "Observations must have source_memory_ids"
|
||||
|
||||
# Verify source_memory_ids point to actual memories
|
||||
async with memory._pool.acquire() as conn:
|
||||
source_memories = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text FROM memory_units
|
||||
WHERE id = ANY($1) AND fact_type IN ('world', 'experience')
|
||||
""",
|
||||
all_source_ids,
|
||||
# Verify source_memory_ids point to actual memories. The list arm takes both
|
||||
# source fact types at once, so the lineage is checked against the same set
|
||||
# the SQL IN clause used to name.
|
||||
source_facts = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type=["world", "experience"], limit=500, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
wanted = set(all_source_ids)
|
||||
source_memories = [m for m in source_facts if m["id"] in wanted]
|
||||
|
||||
assert len(source_memories) >= 1, (
|
||||
f"source_memory_ids should point to valid memories. IDs: {all_source_ids}, Found: {len(source_memories)}"
|
||||
@@ -1519,14 +1455,11 @@ class TestHierarchicalRetrieval:
|
||||
)
|
||||
|
||||
# Verify observation was created
|
||||
async with memory._pool.acquire() as conn:
|
||||
obs_count = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
""",
|
||||
bank_id,
|
||||
obs_count = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
)["total"]
|
||||
assert obs_count >= 1, "Consolidation should have created an observation"
|
||||
|
||||
# Create a mental model about John (higher quality, user-curated)
|
||||
@@ -1992,11 +1925,11 @@ async def test_consolidation_with_observations_mission(memory: "MemoryEngine", r
|
||||
content="Alice uses Python for data analysis and loves its simplicity.",
|
||||
request_context=request_context,
|
||||
)
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"SELECT id, text, fact_type FROM memory_units WHERE bank_id = $1 AND fact_type = 'observation'",
|
||||
bank_id,
|
||||
observations = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
assert isinstance(observations, list)
|
||||
finally:
|
||||
memory._config_resolver._global_config = original_global_config
|
||||
@@ -2010,6 +1943,7 @@ async def test_consolidation_with_observations_mission(memory: "MemoryEngine", r
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_observation_scopes_explicit_multi_pass(memory: MemoryEngine, request_context):
|
||||
"""Test that observation_scopes with an explicit list triggers separate consolidation passes.
|
||||
|
||||
@@ -2037,16 +1971,9 @@ async def test_observation_scopes_explicit_multi_pass(memory: MemoryEngine, requ
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
observations = (
|
||||
await memory.list_memory_units(bank_id, fact_type="observation", limit=1000, request_context=request_context)
|
||||
)["items"]
|
||||
|
||||
try:
|
||||
# Must have at least 2 observations (one per tag scope)
|
||||
@@ -2073,6 +2000,7 @@ async def test_observation_scopes_explicit_multi_pass(memory: MemoryEngine, requ
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_observation_scopes_per_tag(memory: MemoryEngine, request_context):
|
||||
"""Test that observation_scopes='per_tag' derives one pass per individual tag.
|
||||
|
||||
@@ -2095,16 +2023,9 @@ async def test_observation_scopes_per_tag(memory: MemoryEngine, request_context)
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
observations = (
|
||||
await memory.list_memory_units(bank_id, fact_type="observation", limit=1000, request_context=request_context)
|
||||
)["items"]
|
||||
|
||||
try:
|
||||
assert len(observations) >= 2, (
|
||||
@@ -2147,16 +2068,9 @@ async def test_observation_scopes_combined(memory: MemoryEngine, request_context
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
observations = (
|
||||
await memory.list_memory_units(bank_id, fact_type="observation", limit=1000, request_context=request_context)
|
||||
)["items"]
|
||||
|
||||
try:
|
||||
assert len(observations) >= 1, "Expected at least 1 observation, got 0"
|
||||
@@ -2178,6 +2092,7 @@ async def test_observation_scopes_combined(memory: MemoryEngine, request_context
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_observation_scopes_all_combinations(memory: MemoryEngine, request_context):
|
||||
"""Test that observation_scopes='all_combinations' generates passes for every tag subset.
|
||||
|
||||
@@ -2201,16 +2116,9 @@ async def test_observation_scopes_all_combinations(memory: MemoryEngine, request
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
observations = (
|
||||
await memory.list_memory_units(bank_id, fact_type="observation", limit=1000, request_context=request_context)
|
||||
)["items"]
|
||||
|
||||
try:
|
||||
# With 2 tags there are 3 subsets: {alice}, {ben}, {alice, ben}
|
||||
@@ -2671,6 +2579,7 @@ def test_max_observations_per_scope_default():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_count_observations_for_scope(memory: MemoryEngine, request_context):
|
||||
"""Test _count_observations_for_scope counts observations filtered by tags."""
|
||||
bank_id = f"test-count-obs-scope-{uuid.uuid4().hex[:8]}"
|
||||
@@ -2762,6 +2671,7 @@ def _make_mock_llm_one_obs_per_fact():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_max_observations_per_scope_limits_creates(memory: MemoryEngine, request_context):
|
||||
"""Mock LLM tries to create 1 obs per fact; with limit=2, only 2 should exist after 5 facts."""
|
||||
bank_id = f"test-max-obs-limit-{uuid.uuid4().hex[:8]}"
|
||||
@@ -2818,6 +2728,7 @@ async def test_max_observations_per_scope_limits_creates(memory: MemoryEngine, r
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_max_observations_per_scope_zero_forbids_all_creates(memory: MemoryEngine, request_context):
|
||||
"""limit=0 means "no new observations": consolidation must create none.
|
||||
|
||||
@@ -2871,6 +2782,7 @@ async def test_max_observations_per_scope_zero_forbids_all_creates(memory: Memor
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_max_observations_per_scope_allows_updates_at_capacity(memory: MemoryEngine, request_context):
|
||||
"""At capacity, the LLM can still update existing observations."""
|
||||
from hindsight_api.engine.consolidation.consolidator import (
|
||||
@@ -2971,6 +2883,7 @@ async def test_max_observations_per_scope_allows_updates_at_capacity(memory: Mem
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_max_observations_per_scope_no_tags_skips_limit(memory: MemoryEngine, request_context):
|
||||
"""With limit=1, memories with no tags should bypass the limit and create freely."""
|
||||
bank_id = f"test-max-obs-no-tags-{uuid.uuid4().hex[:8]}"
|
||||
@@ -3006,12 +2919,12 @@ async def test_max_observations_per_scope_no_tags_skips_limit(memory: MemoryEngi
|
||||
await run_consolidation_job(memory_engine=memory, bank_id=bank_id, request_context=request_context)
|
||||
|
||||
# No tag limit should apply — all 3 observations should be created
|
||||
async with memory._pool.acquire() as conn:
|
||||
obs = await conn.fetch(
|
||||
"SELECT id FROM memory_units WHERE bank_id = $1 AND fact_type = 'observation'",
|
||||
bank_id,
|
||||
total = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
assert len(obs) == 3, f"Expected 3 observations (no limit for no-tag), got {len(obs)}"
|
||||
)["total"]
|
||||
assert total == 3, f"Expected 3 observations (no limit for no-tag), got {total}"
|
||||
finally:
|
||||
memory._config_resolver._global_config = original_global_config
|
||||
memory._consolidation_llm_config = original_llm
|
||||
@@ -3020,6 +2933,7 @@ async def test_max_observations_per_scope_no_tags_skips_limit(memory: MemoryEngi
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_max_observations_unlimited_default(memory: MemoryEngine, request_context):
|
||||
"""With default config (-1), all creates go through."""
|
||||
bank_id = f"test-max-obs-unlimited-{uuid.uuid4().hex[:8]}"
|
||||
@@ -3053,6 +2967,7 @@ async def test_max_observations_unlimited_default(memory: MemoryEngine, request_
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_targeted_consolidation_filters_by_scopes(memory: MemoryEngine, request_context):
|
||||
"""Consolidation with observation_scopes only processes memories matching those scopes."""
|
||||
bank_id = f"test-targeted-{uuid.uuid4().hex[:8]}"
|
||||
@@ -3085,17 +3000,10 @@ async def test_targeted_consolidation_filters_by_scopes(memory: MemoryEngine, re
|
||||
alice_obs = await _count_observations_for_scope(conn, bank_id, ["user:alice"])
|
||||
assert alice_obs == 1
|
||||
|
||||
# Bob and Charlie should still be unconsolidated
|
||||
unconsolidated = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1
|
||||
AND consolidated_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
assert unconsolidated == 2
|
||||
# Bob and Charlie should still be unconsolidated. consolidation_state='pending'
|
||||
# is the read API's name for this predicate; it also excludes facts whose
|
||||
# consolidation failed, which is the stricter (and here equivalent) reading.
|
||||
assert await _unconsolidated(memory, bank_id, request_context) == 2
|
||||
|
||||
# Now consolidate bob
|
||||
result = await run_consolidation_job(
|
||||
@@ -3107,23 +3015,14 @@ async def test_targeted_consolidation_filters_by_scopes(memory: MemoryEngine, re
|
||||
assert result["memories_processed"] == 1
|
||||
|
||||
# Charlie still unconsolidated
|
||||
async with memory._pool.acquire() as conn:
|
||||
unconsolidated = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1
|
||||
AND consolidated_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
assert unconsolidated == 1
|
||||
assert await _unconsolidated(memory, bank_id, request_context) == 1
|
||||
finally:
|
||||
memory._consolidation_llm_config = original_llm
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_targeted_consolidation_multiple_scopes(memory: MemoryEngine, request_context):
|
||||
"""Consolidation with multiple observation_scopes matches memories in any scope."""
|
||||
bank_id = f"test-targeted-multi-{uuid.uuid4().hex[:8]}"
|
||||
@@ -3151,23 +3050,14 @@ async def test_targeted_consolidation_multiple_scopes(memory: MemoryEngine, requ
|
||||
assert result["observations_created"] == 2
|
||||
|
||||
# Bob still unconsolidated
|
||||
async with memory._pool.acquire() as conn:
|
||||
unconsolidated = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1
|
||||
AND consolidated_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
assert unconsolidated == 1
|
||||
assert await _unconsolidated(memory, bank_id, request_context) == 1
|
||||
finally:
|
||||
memory._consolidation_llm_config = original_llm
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_targeted_consolidation_no_scopes_processes_all(memory: MemoryEngine, request_context):
|
||||
"""Consolidation without observation_scopes processes all unconsolidated memories (backward compat)."""
|
||||
bank_id = f"test-targeted-all-{uuid.uuid4().hex[:8]}"
|
||||
@@ -3197,6 +3087,7 @@ async def test_targeted_consolidation_no_scopes_processes_all(memory: MemoryEngi
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_targeted_consolidation_contains_semantics(memory: MemoryEngine, request_context):
|
||||
"""Scope ["user:alice"] matches memories tagged ["user:alice", "team:eng"] (contains)."""
|
||||
bank_id = f"test-targeted-contains-{uuid.uuid4().hex[:8]}"
|
||||
@@ -3253,23 +3144,15 @@ async def test_enable_auto_consolidation_flag(memory: MemoryEngine, request_cont
|
||||
)
|
||||
|
||||
# Check that memories are NOT consolidated
|
||||
async with memory._pool.acquire() as conn:
|
||||
unconsolidated = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1
|
||||
AND consolidated_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
assert unconsolidated > 0, "Memories should remain unconsolidated when auto consolidation is disabled"
|
||||
unconsolidated = await _unconsolidated(memory, bank_id, request_context)
|
||||
assert unconsolidated > 0, "Memories should remain unconsolidated when auto consolidation is disabled"
|
||||
|
||||
observations = await conn.fetchval(
|
||||
"SELECT COUNT(*) FROM memory_units WHERE bank_id = $1 AND fact_type = 'observation'",
|
||||
bank_id,
|
||||
observations = (
|
||||
await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
assert observations == 0, "No observations should be created when auto consolidation is disabled"
|
||||
)["total"]
|
||||
assert observations == 0, "No observations should be created when auto consolidation is disabled"
|
||||
finally:
|
||||
memory._config_resolver._global_config = original_global_config
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -3330,6 +3213,7 @@ def test_consolidation_prompt_split_is_cacheable_and_complete():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_create_observation_populates_search_vector_native(memory, request_context):
|
||||
"""Observations created via consolidation must have search_vector populated
|
||||
when text_search_extension == 'native', so BM25 retrieval finds them."""
|
||||
|
||||
@@ -10,6 +10,7 @@ import types
|
||||
import uuid
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import DEFAULT, AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
@@ -23,10 +24,19 @@ from hindsight_api.engine.consolidation.consolidator import (
|
||||
_DedupDecision,
|
||||
_duplicate_create_target,
|
||||
_norm_obs_text,
|
||||
_TemporalBounds,
|
||||
)
|
||||
from hindsight_api.engine.memories import RecallArms
|
||||
from hindsight_api.engine.search.types import RetrievalResult
|
||||
|
||||
#: Dates the skipped CREATE would have been stamped with; the fold must carry them onto the twin.
|
||||
_SOURCE_BOUNDS = _TemporalBounds(
|
||||
event_date=datetime(2024, 1, 2, tzinfo=timezone.utc),
|
||||
occurred_start=datetime(2023, 1, 2, tzinfo=timezone.utc),
|
||||
occurred_end=datetime(2024, 1, 3, tzinfo=timezone.utc),
|
||||
mentioned_at=datetime(2024, 1, 4, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeObs:
|
||||
@@ -195,6 +205,7 @@ def _ctx(threshold: float = 0.97):
|
||||
create_text="YouTube content in Uzbek is very rich.",
|
||||
create_source_ids=[uuid.uuid4()],
|
||||
tags=["t1"],
|
||||
source_bounds=_SOURCE_BOUNDS,
|
||||
)
|
||||
return kwargs, conn, llm
|
||||
|
||||
@@ -335,6 +346,15 @@ async def test_dedup_llm_merge_folds_into_twin() -> None:
|
||||
assert args[1] == "Uzbek content on YouTube is very rich." # merged text persisted
|
||||
assert args[2] == kwargs["create_source_ids"] # new (live) source facts folded in
|
||||
assert args[3] == uuid.UUID(_TWIN_ID) # onto the twin row
|
||||
# ...along with the dates the skipped CREATE carried, so the twin's interval widens (#3477).
|
||||
# What the SQL *does* with them is covered against a real database in
|
||||
# test_consolidation_temporal_merge.py — a mocked connection cannot check that.
|
||||
assert args[5:] == (
|
||||
_SOURCE_BOUNDS.event_date,
|
||||
_SOURCE_BOUNDS.occurred_start,
|
||||
_SOURCE_BOUNDS.occurred_end,
|
||||
_SOURCE_BOUNDS.mentioned_at,
|
||||
)
|
||||
|
||||
|
||||
async def test_dedup_llm_merge_sanitizes_text_before_write() -> None:
|
||||
@@ -401,6 +421,7 @@ def _update_ctx(threshold: float = 0.97):
|
||||
return kwargs, conn, llm
|
||||
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_dedup_update_merge_folds_into_twin_and_deletes_updated() -> None:
|
||||
kwargs, conn, llm = _update_ctx()
|
||||
llm.call.return_value = _DedupDecision(action="merge", text="Uzbek YouTube content is very rich and growing.")
|
||||
|
||||
@@ -167,12 +167,10 @@ async def _insert_memory(conn, bank_id: str, text: str, tags: list[str]) -> uuid
|
||||
return mem_id
|
||||
|
||||
|
||||
async def _count_observations(memory: MemoryEngine, bank_id: str) -> int:
|
||||
async with memory._pool.acquire() as conn:
|
||||
return await conn.fetchval(
|
||||
"SELECT count(*) FROM memory_units WHERE bank_id = $1 AND fact_type = 'observation'",
|
||||
bank_id,
|
||||
)
|
||||
async def _count_observations(memory: MemoryEngine, bank_id: str, request_context) -> int:
|
||||
return (
|
||||
await memory.list_memory_units(bank_id, fact_type="observation", limit=1000, request_context=request_context)
|
||||
)["total"]
|
||||
|
||||
|
||||
class _TxnSpy:
|
||||
@@ -196,6 +194,7 @@ class _TxnSpy:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_recall_failure_cancels_sibling_tag_groups(memory: MemoryEngine, request_context):
|
||||
"""One group's recall times out → the other two groups are cancelled before
|
||||
they write, and the job propagates the original error without waiting for
|
||||
@@ -253,12 +252,13 @@ async def test_recall_failure_cancels_sibling_tag_groups(memory: MemoryEngine, r
|
||||
# Nothing was written: the cancelled groups never reached their commit,
|
||||
# and no orphan lands a write after the operation has already failed.
|
||||
await asyncio.sleep(0.2)
|
||||
assert await _count_observations(memory, bank_id) == 0
|
||||
assert await _count_observations(memory, bank_id, request_context) == 0
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_failed_batch_aborts_its_write_group(memory: MemoryEngine, request_context):
|
||||
"""A batch that raises before its witness commit decides its write-group
|
||||
abort, rather than leaving it pending for the recovery sweep."""
|
||||
@@ -292,6 +292,7 @@ async def test_failed_batch_aborts_its_write_group(memory: MemoryEngine, request
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_successful_batch_still_commits_its_write_group(memory: MemoryEngine, request_context):
|
||||
"""Guard on the abort path: the happy path must still decide commit=True."""
|
||||
bank_id = f"test-commit-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
@@ -122,6 +122,7 @@ class TestAdaptiveBatchSplitting:
|
||||
"""Verify that a failing batch is halved and retried until batch_size=1 succeeds."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_splitting_recovers_all_memories(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""When a batch of 2 fails, both are retried individually and succeed."""
|
||||
bank_id = f"test-split-recovery-{uuid.uuid4().hex[:8]}"
|
||||
@@ -152,15 +153,11 @@ class TestAdaptiveBatchSplitting:
|
||||
assert result["memories_failed"] == 0
|
||||
|
||||
# Both memories must have consolidated_at set and consolidation_failed_at NULL
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT id, consolidated_at, consolidation_failed_at
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'experience'
|
||||
""",
|
||||
bank_id,
|
||||
rows = (
|
||||
await memory_no_llm_verify.list_memory_units(
|
||||
bank_id, fact_type="experience", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
assert len(rows) == 2
|
||||
for row in rows:
|
||||
assert row["consolidated_at"] is not None, f"Memory {row['id']} should have consolidated_at set"
|
||||
@@ -175,6 +172,7 @@ class TestAdaptiveBatchSplitting:
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_splitting_with_larger_batch(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""A batch of 4 that always fails at size>1 resolves to 4 individual calls."""
|
||||
bank_id = f"test-split-large-{uuid.uuid4().hex[:8]}"
|
||||
@@ -206,12 +204,11 @@ class TestAdaptiveBatchSplitting:
|
||||
assert result["memories_processed"] == 4
|
||||
assert result["memories_failed"] == 0
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT consolidated_at, consolidation_failed_at FROM memory_units "
|
||||
"WHERE bank_id = $1 AND fact_type = 'experience'",
|
||||
bank_id,
|
||||
rows = (
|
||||
await memory_no_llm_verify.list_memory_units(
|
||||
bank_id, fact_type="experience", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
assert all(r["consolidated_at"] is not None for r in rows)
|
||||
assert all(r["consolidation_failed_at"] is None for r in rows)
|
||||
|
||||
@@ -222,6 +219,7 @@ class TestConsolidationFailedAt:
|
||||
"""Verify that consolidation_failed_at is set — and consolidated_at is NOT — when all retries fail."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_single_memory_permanent_failure(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""A single memory that exhausts all LLM retries gets consolidation_failed_at, not consolidated_at."""
|
||||
bank_id = f"test-perm-fail-{uuid.uuid4().hex[:8]}"
|
||||
@@ -243,11 +241,12 @@ class TestConsolidationFailedAt:
|
||||
assert result["memories_failed"] == 1
|
||||
assert result["memories_processed"] == 1
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"SELECT consolidated_at, consolidation_failed_at FROM memory_units WHERE id = $1",
|
||||
mem_id,
|
||||
units = (
|
||||
await memory_no_llm_verify.list_memory_units(
|
||||
bank_id, fact_type="experience", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
row = next(u for u in units if str(u["id"]) == str(mem_id))
|
||||
|
||||
assert row["consolidated_at"] is None, "consolidated_at must NOT be set for a permanently failed memory"
|
||||
assert row["consolidation_failed_at"] is not None, "consolidation_failed_at must be set"
|
||||
@@ -255,6 +254,7 @@ class TestConsolidationFailedAt:
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_failed_memory_excluded_from_next_run(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""A memory marked consolidation_failed_at is not re-processed on the next consolidation run."""
|
||||
bank_id = f"test-excluded-{uuid.uuid4().hex[:8]}"
|
||||
@@ -285,17 +285,19 @@ class TestConsolidationFailedAt:
|
||||
assert result["memories_processed"] == 0
|
||||
|
||||
# Memory still has consolidation_failed_at set and consolidated_at NULL
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"SELECT consolidated_at, consolidation_failed_at FROM memory_units WHERE id = $1",
|
||||
mem_id,
|
||||
units = (
|
||||
await memory_no_llm_verify.list_memory_units(
|
||||
bank_id, fact_type="experience", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
row = next(u for u in units if str(u["id"]) == str(mem_id))
|
||||
assert row["consolidated_at"] is None
|
||||
assert row["consolidation_failed_at"] is not None
|
||||
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_partial_batch_failure(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""In a batch of 2, if only the first individual retry fails, the second still succeeds."""
|
||||
bank_id = f"test-partial-fail-{uuid.uuid4().hex[:8]}"
|
||||
@@ -325,15 +327,12 @@ class TestConsolidationFailedAt:
|
||||
assert result["memories_processed"] == 2
|
||||
assert result["memories_failed"] == 1
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
rows = {
|
||||
str(r["id"]): r
|
||||
for r in await conn.fetch(
|
||||
"SELECT id, consolidated_at, consolidation_failed_at FROM memory_units "
|
||||
"WHERE bank_id = $1 AND fact_type = 'experience'",
|
||||
bank_id,
|
||||
)
|
||||
}
|
||||
items = (
|
||||
await memory_no_llm_verify.list_memory_units(
|
||||
bank_id, fact_type="experience", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
rows = {str(r["id"]): r for r in items}
|
||||
|
||||
# One should have failed, one should have succeeded
|
||||
failed = [r for r in rows.values() if r["consolidation_failed_at"] is not None]
|
||||
@@ -350,6 +349,7 @@ class TestRecoverConsolidation:
|
||||
"""Verify the retry_failed_consolidation() method and the /consolidation/recover endpoint."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_recover_resets_failed_memories(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""retry_failed_consolidation resets consolidation_failed_at and consolidated_at."""
|
||||
bank_id = f"test-recover-reset-{uuid.uuid4().hex[:8]}"
|
||||
@@ -375,12 +375,11 @@ class TestRecoverConsolidation:
|
||||
|
||||
assert result["retried_count"] == 2
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT consolidated_at, consolidation_failed_at FROM memory_units "
|
||||
"WHERE bank_id = $1 AND fact_type = 'experience'",
|
||||
bank_id,
|
||||
rows = (
|
||||
await memory_no_llm_verify.list_memory_units(
|
||||
bank_id, fact_type="experience", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
assert all(r["consolidation_failed_at"] is None for r in rows), "consolidation_failed_at must be cleared"
|
||||
assert all(r["consolidated_at"] is None for r in rows), "consolidated_at must also be cleared"
|
||||
|
||||
@@ -399,6 +398,7 @@ class TestRecoverConsolidation:
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_recover_then_consolidate_succeeds(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""After recovery, the memory is picked up by the next consolidation run."""
|
||||
bank_id = f"test-recover-consolidate-{uuid.uuid4().hex[:8]}"
|
||||
@@ -425,17 +425,19 @@ class TestRecoverConsolidation:
|
||||
assert run_result["memories_processed"] == 1
|
||||
assert run_result["memories_failed"] == 0
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"SELECT consolidated_at, consolidation_failed_at FROM memory_units WHERE id = $1",
|
||||
mem_id,
|
||||
units = (
|
||||
await memory_no_llm_verify.list_memory_units(
|
||||
bank_id, fact_type="experience", limit=1000, request_context=request_context
|
||||
)
|
||||
)["items"]
|
||||
row = next(u for u in units if str(u["id"]) == str(mem_id))
|
||||
assert row["consolidated_at"] is not None, "Memory should be consolidated after recovery"
|
||||
assert row["consolidation_failed_at"] is None
|
||||
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_recover_endpoint_via_http(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""The POST /consolidation/recover endpoint returns the correct retried_count."""
|
||||
import httpx
|
||||
|
||||
@@ -79,16 +79,16 @@ async def _insert_entity_mm(conn, bank_id: str, tag: str) -> str:
|
||||
return mm_id
|
||||
|
||||
|
||||
async def _unconsolidated_count(memory, bank_id: str) -> int:
|
||||
async with memory._pool.acquire() as conn:
|
||||
return await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1 AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
async def _unconsolidated_count(memory, bank_id: str, request_context) -> int:
|
||||
# consolidation_state='pending' is the read API's name for exactly this predicate:
|
||||
# not consolidated, not failed, and a source fact type.
|
||||
page = await memory.list_memory_units(
|
||||
bank_id=bank_id,
|
||||
consolidation_state="pending",
|
||||
limit=1,
|
||||
request_context=request_context,
|
||||
)
|
||||
return page["total"]
|
||||
|
||||
|
||||
async def _pending_consolidations(memory, bank_id: str):
|
||||
@@ -105,6 +105,7 @@ async def _pending_consolidations(memory, bank_id: str):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_multi_round_consolidation_refreshes_all_entity_models(
|
||||
memory: MemoryEngine, request_context, monkeypatch
|
||||
):
|
||||
@@ -203,7 +204,7 @@ async def test_multi_round_consolidation_refreshes_all_entity_models(
|
||||
assert consolidation_runs >= 2, (
|
||||
f"expected a multi-round drain (round limit 3, backlog > 3), but only {consolidation_runs} consolidation(s) ran"
|
||||
)
|
||||
remaining = await _unconsolidated_count(memory, bank_id)
|
||||
remaining = await _unconsolidated_count(memory, bank_id, request_context)
|
||||
assert remaining == 0, f"backlog did not fully drain: {remaining} unconsolidated memories remain"
|
||||
|
||||
# 6. KEY ASSERTION — every refresh_after_consolidation model must be refreshed
|
||||
|
||||
@@ -55,16 +55,16 @@ def enable_observations():
|
||||
config.enable_observations = original
|
||||
|
||||
|
||||
async def _count_unconsolidated(memory, bank_id: str) -> int:
|
||||
async with memory._pool.acquire() as conn:
|
||||
return await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1 AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
async def _count_unconsolidated(memory, bank_id: str, request_context) -> int:
|
||||
# consolidation_state='pending' is the read API's name for exactly this predicate:
|
||||
# not consolidated, not failed, and a source fact type.
|
||||
page = await memory.list_memory_units(
|
||||
bank_id=bank_id,
|
||||
consolidation_state="pending",
|
||||
limit=1,
|
||||
request_context=request_context,
|
||||
)
|
||||
return page["total"]
|
||||
|
||||
|
||||
async def _pending_consolidation_ops(memory, bank_id: str) -> list[str]:
|
||||
@@ -100,7 +100,7 @@ async def test_round_limited_consolidation_leaves_followup_pending_op(memory: Me
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
unconsolidated_before = await _count_unconsolidated(memory, bank_id)
|
||||
unconsolidated_before = await _count_unconsolidated(memory, bank_id, request_context)
|
||||
assert unconsolidated_before >= backlog_size, (
|
||||
f"Expected at least {backlog_size} unconsolidated memories, got {unconsolidated_before}"
|
||||
)
|
||||
@@ -148,7 +148,7 @@ async def test_round_limited_consolidation_leaves_followup_pending_op(memory: Me
|
||||
assert row["status"] == "completed", f"first consolidation op should be marked completed, got {row['status']}"
|
||||
|
||||
# 4. Backlog must remain (round limit kept one round under the total)
|
||||
unconsolidated_after = await _count_unconsolidated(memory, bank_id)
|
||||
unconsolidated_after = await _count_unconsolidated(memory, bank_id, request_context)
|
||||
assert 0 < unconsolidated_after < unconsolidated_before, (
|
||||
f"expected backlog to shrink but still remain after one round; "
|
||||
f"before={unconsolidated_before}, after={unconsolidated_after}"
|
||||
|
||||
@@ -48,15 +48,13 @@ async def test_round_limit_caps_processed_memories(memory: MemoryEngine, request
|
||||
)
|
||||
|
||||
# Verify we have unconsolidated memories
|
||||
async with memory._pool.acquire() as conn:
|
||||
unconsolidated = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1 AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
# consolidation_state='pending' is the read API's name for exactly this predicate:
|
||||
# not consolidated, not failed, and a source fact type.
|
||||
unconsolidated = (
|
||||
await memory.list_memory_units(
|
||||
bank_id=bank_id, consolidation_state="pending", limit=1, request_context=request_context
|
||||
)
|
||||
)["total"]
|
||||
assert unconsolidated >= 6, f"Expected at least 6 unconsolidated memories, got {unconsolidated}"
|
||||
|
||||
# Run consolidation with a round limit of 3
|
||||
@@ -91,15 +89,11 @@ async def test_round_limit_caps_processed_memories(memory: MemoryEngine, request
|
||||
assert result.get("mental_models_refreshed", 0) == 0
|
||||
|
||||
# Verify some memories are still unconsolidated
|
||||
async with memory._pool.acquire() as conn:
|
||||
still_unconsolidated = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1 AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
still_unconsolidated = (
|
||||
await memory.list_memory_units(
|
||||
bank_id=bank_id, consolidation_state="pending", limit=1, request_context=request_context
|
||||
)
|
||||
)["total"]
|
||||
assert still_unconsolidated > 0, "Some memories should still be unconsolidated after hitting round limit"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -140,15 +134,11 @@ async def test_unlimited_round_processes_all(memory: MemoryEngine, request_conte
|
||||
mock_requeue.assert_not_called()
|
||||
|
||||
# All memories should be consolidated
|
||||
async with memory._pool.acquire() as conn:
|
||||
still_unconsolidated = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1 AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
still_unconsolidated = (
|
||||
await memory.list_memory_units(
|
||||
bank_id=bank_id, consolidation_state="pending", limit=1, request_context=request_context
|
||||
)
|
||||
)["total"]
|
||||
assert still_unconsolidated == 0
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -122,14 +122,12 @@ def _mock_llm_one_obs_per_fact():
|
||||
return wrapper, mock_llm
|
||||
|
||||
|
||||
async def _fetch_observation_tag_sets(memory: MemoryEngine, bank_id: str) -> list[frozenset[str]]:
|
||||
async def _fetch_observation_tag_sets(memory: MemoryEngine, bank_id: str, request_context) -> list[frozenset[str]]:
|
||||
"""Return the tag set (as a frozenset) of every observation in the bank."""
|
||||
async with memory._pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT tags FROM memory_units WHERE bank_id = $1 AND fact_type = 'observation'",
|
||||
bank_id,
|
||||
)
|
||||
return [frozenset(r["tags"] or []) for r in rows]
|
||||
items = (
|
||||
await memory.list_memory_units(bank_id, fact_type="observation", limit=1000, request_context=request_context)
|
||||
)["items"]
|
||||
return [frozenset(i["tags"] or []) for i in items]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -139,6 +137,7 @@ async def _fetch_observation_tag_sets(memory: MemoryEngine, bank_id: str) -> lis
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_combined_mode_parallel_writes_to_memory_tag_set(memory: MemoryEngine, request_context):
|
||||
"""combined (default) → each memory yields exactly one observation tagged
|
||||
with the memory's full tag set. With three disjoint tag sets, dispatch
|
||||
@@ -166,7 +165,7 @@ async def test_combined_mode_parallel_writes_to_memory_tag_set(memory: MemoryEng
|
||||
memory._consolidation_llm_config = original_llm
|
||||
|
||||
assert result["status"] == "completed"
|
||||
tag_sets = _ag_sorted(await _fetch_observation_tag_sets(memory, bank_id))
|
||||
tag_sets = _ag_sorted(await _fetch_observation_tag_sets(memory, bank_id, request_context))
|
||||
assert tag_sets == _ag_sorted(
|
||||
[
|
||||
frozenset({"user:alice"}),
|
||||
@@ -179,6 +178,7 @@ async def test_combined_mode_parallel_writes_to_memory_tag_set(memory: MemoryEng
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_shared_mode_parallel_writes_only_untagged_scope(memory: MemoryEngine, request_context):
|
||||
"""shared → every memory writes to the single untagged scope, ignoring its
|
||||
own tags. Three memories with disjoint tags therefore all consolidate into
|
||||
@@ -207,7 +207,7 @@ async def test_shared_mode_parallel_writes_only_untagged_scope(memory: MemoryEng
|
||||
memory._consolidation_llm_config = original_llm
|
||||
|
||||
assert result["status"] == "completed"
|
||||
tag_sets = await _fetch_observation_tag_sets(memory, bank_id)
|
||||
tag_sets = await _fetch_observation_tag_sets(memory, bank_id, request_context)
|
||||
# Every observation lands at the untagged scope — none carries a session tag.
|
||||
assert tag_sets and all(t == frozenset() for t in tag_sets), tag_sets
|
||||
finally:
|
||||
@@ -215,6 +215,7 @@ async def test_shared_mode_parallel_writes_only_untagged_scope(memory: MemoryEng
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_per_tag_mode_parallel_writes_one_observation_per_tag(memory: MemoryEngine, request_context):
|
||||
"""per_tag with tags [a, b] → two observations, tagged [a] and [b] respectively.
|
||||
|
||||
@@ -246,7 +247,7 @@ async def test_per_tag_mode_parallel_writes_one_observation_per_tag(memory: Memo
|
||||
|
||||
assert result["status"] == "completed"
|
||||
|
||||
tag_sets = await _fetch_observation_tag_sets(memory, bank_id)
|
||||
tag_sets = await _fetch_observation_tag_sets(memory, bank_id, request_context)
|
||||
# M1 writes to [alice]; M2 writes to [alice] and [session]. The mock LLM
|
||||
# creates one observation per fact per pass, so we expect:
|
||||
# - one [alice] obs from M1
|
||||
@@ -263,6 +264,7 @@ async def test_per_tag_mode_parallel_writes_one_observation_per_tag(memory: Memo
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_all_combinations_mode_parallel_writes_every_subset(memory: MemoryEngine, request_context):
|
||||
"""all_combinations with tags [a, b] → three observations at [a], [b], [a, b]."""
|
||||
bank_id = f"test-allcombo-{uuid.uuid4().hex[:8]}"
|
||||
@@ -286,7 +288,7 @@ async def test_all_combinations_mode_parallel_writes_every_subset(memory: Memory
|
||||
memory._consolidation_llm_config = original_llm
|
||||
|
||||
assert result["status"] == "completed"
|
||||
tag_sets = set(await _fetch_observation_tag_sets(memory, bank_id))
|
||||
tag_sets = set(await _fetch_observation_tag_sets(memory, bank_id, request_context))
|
||||
assert tag_sets == {
|
||||
frozenset({"alice"}),
|
||||
frozenset({"session"}),
|
||||
@@ -297,6 +299,7 @@ async def test_all_combinations_mode_parallel_writes_every_subset(memory: Memory
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_explicit_scope_list_parallel_writes_declared_scopes(memory: MemoryEngine, request_context):
|
||||
"""Explicit list[list[str]] → observations land at exactly those scopes,
|
||||
regardless of the memory's own tag set."""
|
||||
@@ -327,7 +330,7 @@ async def test_explicit_scope_list_parallel_writes_declared_scopes(memory: Memor
|
||||
memory._consolidation_llm_config = original_llm
|
||||
|
||||
assert result["status"] == "completed"
|
||||
tag_sets = set(await _fetch_observation_tag_sets(memory, bank_id))
|
||||
tag_sets = set(await _fetch_observation_tag_sets(memory, bank_id, request_context))
|
||||
assert tag_sets == {frozenset({"scope_a"}), frozenset({"scope_b", "scope_c"})}
|
||||
# And NOT the memory's own tag.
|
||||
assert frozenset({"tag_ignored"}) not in tag_sets
|
||||
@@ -342,6 +345,7 @@ async def test_explicit_scope_list_parallel_writes_declared_scopes(memory: Memor
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_overlapping_scopes_serialise_under_parallelism(memory: MemoryEngine, request_context):
|
||||
"""Two groups whose write-scope sets intersect on scope S must not have
|
||||
overlapping in-flight LLM-recall windows for S.
|
||||
@@ -424,6 +428,7 @@ async def test_overlapping_scopes_serialise_under_parallelism(memory: MemoryEngi
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_per_batch_log_line_attributes_only_own_work(memory: MemoryEngine, request_context, caplog):
|
||||
"""Per-batch log timings / llm_calls / tokens / processed must reflect only
|
||||
that batch's own work — not totals leaking in from other in-flight batches
|
||||
@@ -502,6 +507,7 @@ async def test_per_batch_log_line_attributes_only_own_work(memory: MemoryEngine,
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_disjoint_scopes_run_concurrently(memory: MemoryEngine, request_context):
|
||||
"""When write-scope sets are pairwise disjoint, the dispatcher must let
|
||||
groups run in parallel — we should observe simultaneous in-flight recalls
|
||||
|
||||
@@ -58,16 +58,16 @@ def enable_observations():
|
||||
config.enable_observations = original
|
||||
|
||||
|
||||
async def _count_unconsolidated(memory, bank_id: str) -> int:
|
||||
async with memory._pool.acquire() as conn:
|
||||
return await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1 AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
async def _count_unconsolidated(memory, bank_id: str, request_context) -> int:
|
||||
# consolidation_state='pending' is the read API's name for exactly this predicate:
|
||||
# not consolidated, not failed, and a source fact type.
|
||||
page = await memory.list_memory_units(
|
||||
bank_id=bank_id,
|
||||
consolidation_state="pending",
|
||||
limit=1,
|
||||
request_context=request_context,
|
||||
)
|
||||
return page["total"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -95,7 +95,7 @@ async def test_requeue_failure_propagates_to_worker_retry(memory: MemoryEngine,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
unconsolidated_before = await _count_unconsolidated(memory, bank_id)
|
||||
unconsolidated_before = await _count_unconsolidated(memory, bank_id, request_context)
|
||||
assert unconsolidated_before >= backlog_size
|
||||
|
||||
op_id = uuid.uuid4()
|
||||
@@ -147,7 +147,7 @@ async def test_requeue_failure_propagates_to_worker_retry(memory: MemoryEngine,
|
||||
# consolidated stay consolidated; the consolidator's per-batch
|
||||
# `UPDATE ... SET consolidated_at = NOW()` commits in its own
|
||||
# transaction (consolidator.py:524-534), not inside the op-level state.
|
||||
unconsolidated_after = await _count_unconsolidated(memory, bank_id)
|
||||
unconsolidated_after = await _count_unconsolidated(memory, bank_id, request_context)
|
||||
assert unconsolidated_after < unconsolidated_before, (
|
||||
f"completed-round work must be durable across the re-queue failure; "
|
||||
f"before={unconsolidated_before}, after={unconsolidated_after}"
|
||||
|
||||
@@ -0,0 +1,363 @@
|
||||
"""Observation merges must never narrow an observation's temporal bounds (#3477).
|
||||
|
||||
Every path that folds evidence into an existing observation — the ordinary UPDATE, the
|
||||
CREATE-time semantic dedup fold, the UPDATE-time dedup fold — has to widen
|
||||
``event_date``/``occurred_start`` to the earliest known value and
|
||||
``occurred_end``/``mentioned_at`` to the latest. Before the fix the dedup folds rewrote only
|
||||
text/sources, so an observation could cite dated source facts while reporting no event
|
||||
interval at all.
|
||||
|
||||
The SQL-executing tests here run against a real database on purpose: the merge lives in SQL,
|
||||
and asserting on a mocked connection cannot tell a working statement from one PostgreSQL
|
||||
refuses to plan.
|
||||
"""
|
||||
|
||||
import types
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.consolidation import consolidator
|
||||
from hindsight_api.engine.consolidation.consolidator import (
|
||||
_aggregate_source_fields,
|
||||
_DedupDecision,
|
||||
_DedupOutcome,
|
||||
_execute_create_action,
|
||||
_execute_update_action,
|
||||
_TemporalBounds,
|
||||
run_consolidation_job,
|
||||
)
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
from hindsight_api.engine.response_models import MemoryFact
|
||||
|
||||
EARLY = datetime(2020, 3, 1, tzinfo=timezone.utc)
|
||||
LATE = datetime(2021, 7, 4, 12, 30, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def observations_enabled():
|
||||
"""Observations on, and a loose dedup gate so the near-duplicate CREATE reaches the fold."""
|
||||
config = _get_raw_config()
|
||||
previous_observations = config.enable_observations
|
||||
previous_threshold = config.consolidation_dedup_threshold
|
||||
config.enable_observations = True
|
||||
config.consolidation_dedup_threshold = 0.5
|
||||
yield config
|
||||
config.enable_observations = previous_observations
|
||||
config.consolidation_dedup_threshold = previous_threshold
|
||||
|
||||
|
||||
async def _retain_fact(
|
||||
memory: MemoryEngine,
|
||||
request_context,
|
||||
config,
|
||||
bank_id: str,
|
||||
content: str,
|
||||
marker: str,
|
||||
bounds: _TemporalBounds = _TemporalBounds(),
|
||||
) -> uuid.UUID:
|
||||
"""Retain ``content`` and return the id of the resulting fact containing ``marker``.
|
||||
|
||||
The mock extractor also emits facts for the extraction prompt's own boilerplate; those are
|
||||
marked consolidated so only the fact under test can join a later consolidation job. Retain
|
||||
runs with observations off so the caller controls when consolidation happens.
|
||||
"""
|
||||
previous = config.enable_observations
|
||||
config.enable_observations = False
|
||||
try:
|
||||
await memory.retain_async(bank_id=bank_id, content=content, request_context=request_context)
|
||||
finally:
|
||||
config.enable_observations = previous
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE memory_units SET consolidated_at = now()
|
||||
WHERE bank_id = $1 AND fact_type <> 'observation' AND consolidated_at IS NULL
|
||||
AND text NOT LIKE $2
|
||||
""",
|
||||
bank_id,
|
||||
f"%{marker}%",
|
||||
)
|
||||
# Repeat the marker predicate rather than leaning on the sweep above: if the extractor
|
||||
# ever emits two facts for this content, stamping both and returning an arbitrary one
|
||||
# would make the caller's source-id assertions flaky rather than fail here.
|
||||
stamped = await conn.fetch(
|
||||
"""
|
||||
UPDATE memory_units
|
||||
SET event_date = $2, occurred_start = $3, occurred_end = $4, mentioned_at = $5
|
||||
WHERE bank_id = $1 AND fact_type <> 'observation' AND consolidated_at IS NULL
|
||||
AND text LIKE $6
|
||||
RETURNING id
|
||||
""",
|
||||
bank_id,
|
||||
bounds.event_date,
|
||||
bounds.occurred_start,
|
||||
bounds.occurred_end,
|
||||
bounds.mentioned_at,
|
||||
f"%{marker}%",
|
||||
)
|
||||
assert len(stamped) == 1, f"expected exactly one unconsolidated fact matching {marker!r}, got {len(stamped)}"
|
||||
return stamped[0]["id"]
|
||||
|
||||
|
||||
async def _seed_undated_observation(
|
||||
memory: MemoryEngine, bank_id: str, source_fact_id: uuid.UUID, text: str
|
||||
) -> uuid.UUID:
|
||||
"""Create an observation the way undated source facts do: stamped with a consolidation-time
|
||||
``event_date``/``mentioned_at`` and no occurred interval at all."""
|
||||
action = await _execute_create_action(
|
||||
pool=await memory._get_backend(),
|
||||
memory_engine=memory,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[source_fact_id],
|
||||
text=text,
|
||||
)
|
||||
assert action == "created"
|
||||
async with memory._pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"SELECT id, occurred_start FROM memory_units WHERE bank_id = $1 AND fact_type = 'observation'",
|
||||
bank_id,
|
||||
)
|
||||
assert row["occurred_start"] is None
|
||||
return row["id"]
|
||||
|
||||
|
||||
async def _observations(memory: MemoryEngine, bank_id: str) -> list:
|
||||
async with memory._pool.acquire() as conn:
|
||||
return await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, event_date, occurred_start, occurred_end, mentioned_at, source_memory_ids
|
||||
FROM memory_units WHERE bank_id = $1 AND fact_type = 'observation' ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_dedup_create_fold_widens_bounds_of_the_twin(memory: MemoryEngine, request_context, observations_enabled):
|
||||
"""#3477: a CREATE folded into a near-twin must hand over its source facts' dates.
|
||||
|
||||
Without this the twin ends up citing a dated source fact while still reporting
|
||||
``occurred_start``/``occurred_end`` as NULL — the exact shape reported in the issue.
|
||||
"""
|
||||
bank_id = f"test-temporal-fold-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
undated_fact = await _retain_fact(
|
||||
memory, request_context, observations_enabled, bank_id, "Alice moved to Berlin for work.", "Alice"
|
||||
)
|
||||
observation_id = await _seed_undated_observation(memory, bank_id, undated_fact, "Alice moved to Berlin for work.")
|
||||
|
||||
dated_fact = await _retain_fact(
|
||||
memory,
|
||||
request_context,
|
||||
observations_enabled,
|
||||
bank_id,
|
||||
"Alice relocated to Berlin for a new job.",
|
||||
"relocated",
|
||||
_TemporalBounds(event_date=EARLY, occurred_start=EARLY, occurred_end=LATE, mentioned_at=LATE),
|
||||
)
|
||||
|
||||
# The adjudicating LLM votes "merge", so this CREATE folds into the twin instead of inserting.
|
||||
with patch.object(
|
||||
consolidator,
|
||||
"_dedup_decision_from_response",
|
||||
lambda _response: _DedupDecision(action="merge", text="Alice moved to Berlin for a new job.", reason="test"),
|
||||
):
|
||||
await run_consolidation_job(memory_engine=memory, bank_id=bank_id, request_context=request_context)
|
||||
|
||||
rows = await _observations(memory, bank_id)
|
||||
assert len(rows) == 1, f"the CREATE should have folded into the twin, not inserted: {[dict(r) for r in rows]}"
|
||||
folded = rows[0]
|
||||
assert folded["id"] == observation_id
|
||||
assert dated_fact in folded["source_memory_ids"], "the dated fact must be cited by the survivor"
|
||||
assert folded["occurred_start"] == EARLY
|
||||
assert folded["occurred_end"] == LATE
|
||||
assert folded["event_date"] == EARLY, "the source's earlier event_date must win over the stamped one"
|
||||
assert folded["mentioned_at"] > LATE, "mentioned_at keeps the later of the two (the stamped one)"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_update_widens_bounds_from_its_source_facts(memory: MemoryEngine, request_context, observations_enabled):
|
||||
"""The ordinary UPDATE path inherits every temporal field from its sources, event_date included."""
|
||||
bank_id = f"test-temporal-update-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
undated_fact = await _retain_fact(
|
||||
memory, request_context, observations_enabled, bank_id, "Bob plays the cello.", "cello"
|
||||
)
|
||||
observation_id = await _seed_undated_observation(memory, bank_id, undated_fact, "Bob plays the cello.")
|
||||
stamped = (await _observations(memory, bank_id))[0]
|
||||
|
||||
dated_fact = await _retain_fact(
|
||||
memory,
|
||||
request_context,
|
||||
observations_enabled,
|
||||
bank_id,
|
||||
"Bob joined a string quartet.",
|
||||
"quartet",
|
||||
_TemporalBounds(event_date=EARLY, occurred_start=EARLY, occurred_end=LATE, mentioned_at=LATE),
|
||||
)
|
||||
|
||||
embedding = await _execute_update_action(
|
||||
pool=await memory._get_backend(),
|
||||
memory_engine=memory,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[dated_fact],
|
||||
observation_id=str(observation_id),
|
||||
new_text="Bob plays the cello in a string quartet.",
|
||||
observations=[
|
||||
MemoryFact(
|
||||
id=str(observation_id),
|
||||
text=stamped["text"],
|
||||
fact_type="observation",
|
||||
source_fact_ids=[str(undated_fact)],
|
||||
tags=[],
|
||||
)
|
||||
],
|
||||
source_bounds=_TemporalBounds(event_date=EARLY, occurred_start=EARLY, occurred_end=LATE, mentioned_at=LATE),
|
||||
)
|
||||
assert embedding is not None, "the update must have landed"
|
||||
|
||||
updated = (await _observations(memory, bank_id))[0]
|
||||
assert updated["occurred_start"] == EARLY
|
||||
assert updated["occurred_end"] == LATE
|
||||
assert updated["event_date"] == EARLY, "event_date must follow the sources, not stay at creation time"
|
||||
assert updated["mentioned_at"] == stamped["mentioned_at"], "the stamped mention is later than the source's"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_dedup_update_fold_unions_the_bounds_of_both_rows(
|
||||
memory: MemoryEngine, request_context, observations_enabled
|
||||
):
|
||||
"""The UPDATE-time fold deletes the folded-from row, so the survivor must absorb its dates."""
|
||||
bank_id = f"test-temporal-updfold-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
twin_fact = await _retain_fact(
|
||||
memory, request_context, observations_enabled, bank_id, "Carla ran the Rome marathon.", "marathon"
|
||||
)
|
||||
updated_fact = await _retain_fact(
|
||||
memory, request_context, observations_enabled, bank_id, "Carla trains every morning.", "trains"
|
||||
)
|
||||
|
||||
twin_text = "Carla ran the Rome marathon."
|
||||
updated_text = "Carla trains every morning and ran the Rome marathon."
|
||||
await _execute_create_action(
|
||||
pool=await memory._get_backend(),
|
||||
memory_engine=memory,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[twin_fact],
|
||||
text=twin_text,
|
||||
occurred_start=EARLY,
|
||||
)
|
||||
await _execute_create_action(
|
||||
pool=await memory._get_backend(),
|
||||
memory_engine=memory,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[updated_fact],
|
||||
text=updated_text,
|
||||
occurred_end=LATE,
|
||||
)
|
||||
rows = {row["text"]: row for row in await _observations(memory, bank_id)}
|
||||
twin_id, updated_id = rows[twin_text]["id"], rows[updated_text]["id"]
|
||||
assert rows[twin_text]["occurred_end"] is None and rows[updated_text]["occurred_start"] is None
|
||||
|
||||
# The re-embedded UPDATE drifted onto the twin and the LLM votes "merge".
|
||||
with patch.object(
|
||||
consolidator,
|
||||
"_dedup_adjudicate",
|
||||
AsyncMock(
|
||||
return_value=_DedupOutcome(
|
||||
best_id=str(twin_id), merged_text=twin_text, should_merge=True, best_text=twin_text
|
||||
)
|
||||
),
|
||||
):
|
||||
await consolidator._dedup_reconcile_update(
|
||||
await memory._get_backend(),
|
||||
memory,
|
||||
bank_id,
|
||||
observations_enabled,
|
||||
None,
|
||||
str(updated_id),
|
||||
updated_text,
|
||||
None,
|
||||
[],
|
||||
)
|
||||
|
||||
survivors = await _observations(memory, bank_id)
|
||||
assert [row["id"] for row in survivors] == [twin_id], "the folded-from row must be gone"
|
||||
assert survivors[0]["occurred_start"] == EARLY, "the twin keeps its own start"
|
||||
assert survivors[0]["occurred_end"] == LATE, "and absorbs the deleted row's end"
|
||||
assert updated_fact in survivors[0]["source_memory_ids"]
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
def test_temporal_bounds_merge_keeps_the_widest_known_interval():
|
||||
"""min/max per field, with a missing value on either side ignored."""
|
||||
merged = _TemporalBounds(event_date=LATE, occurred_start=LATE, occurred_end=EARLY).merged_with(
|
||||
_TemporalBounds(event_date=EARLY, occurred_end=LATE, mentioned_at=LATE)
|
||||
)
|
||||
assert merged == _TemporalBounds(event_date=EARLY, occurred_start=LATE, occurred_end=LATE, mentioned_at=LATE)
|
||||
assert _TemporalBounds().merged_with(_TemporalBounds()) == _TemporalBounds()
|
||||
|
||||
|
||||
def test_temporal_bounds_of_reads_a_source_aggregation():
|
||||
"""``_aggregate_source_fields`` is where an observation's dates come from, so its result has
|
||||
to hand over all four fields unchanged."""
|
||||
aggregation = _aggregate_source_fields(
|
||||
[
|
||||
{"event_date": LATE, "occurred_start": LATE, "occurred_end": LATE, "mentioned_at": EARLY},
|
||||
{"event_date": EARLY, "occurred_start": EARLY, "occurred_end": None, "mentioned_at": LATE},
|
||||
]
|
||||
)
|
||||
assert _TemporalBounds.of(aggregation) == _TemporalBounds(
|
||||
event_date=EARLY, occurred_start=EARLY, occurred_end=LATE, mentioned_at=LATE
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_owned_fold_merges_bounds_like_the_sql_path():
|
||||
"""A memories store that owns its rows re-upserts the whole observation, so the fold has to
|
||||
apply the same widening in Python."""
|
||||
stored = types.SimpleNamespace(
|
||||
source_memory_ids=["existing"],
|
||||
tags=["t1"],
|
||||
created_at=EARLY,
|
||||
event_date=LATE,
|
||||
occurred_start=LATE,
|
||||
occurred_end=EARLY,
|
||||
mentioned_at=EARLY,
|
||||
)
|
||||
store = types.SimpleNamespace(get_memories=AsyncMock(return_value=[stored]), upsert_observation=AsyncMock())
|
||||
|
||||
with patch.object(consolidator.embedding_utils, "generate_embeddings_batch", AsyncMock(return_value=[[0.1, 0.2]])):
|
||||
await consolidator._reconcile_merge_via_store(
|
||||
store,
|
||||
conn=object(),
|
||||
memory_engine=types.SimpleNamespace(embeddings=object()),
|
||||
bank_id="bank-1",
|
||||
observation_id=str(uuid.uuid4()),
|
||||
merged_text="merged",
|
||||
add_source_ids=[uuid.uuid4()],
|
||||
add_bounds=_TemporalBounds(event_date=EARLY, occurred_end=LATE, mentioned_at=LATE),
|
||||
)
|
||||
|
||||
record = store.upsert_observation.await_args.kwargs["record"]
|
||||
assert record.event_date == EARLY
|
||||
assert record.occurred_start == LATE
|
||||
assert record.occurred_end == LATE
|
||||
assert record.mentioned_at == LATE
|
||||
assert record.created_at == EARLY, "fields the fold does not own are preserved"
|
||||
@@ -32,6 +32,7 @@ async def _seed(conn, bank_id: str, *, tags: list[str], consolidated: bool = Fal
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_count_dedupes_across_overlapping_scopes(memory: MemoryEngine, request_context: RequestContext):
|
||||
bank_id = f"test-count-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
@@ -7,6 +7,7 @@ test_memory_curation.py.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
@@ -47,6 +48,7 @@ async def _insert_fact(memory: MemoryEngine, bank_id: str, text: str) -> str:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_patch_invalidate_and_revert_over_http(api_client, memory):
|
||||
bank_id = f"curation-http-{uuid.uuid4().hex[:8]}"
|
||||
mem_id = await _insert_fact(memory, bank_id, "srv-04 runs PostgreSQL 14.")
|
||||
@@ -79,6 +81,7 @@ async def test_patch_invalidate_and_revert_over_http(api_client, memory):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_patch_clears_occurred_dates_with_explicit_null(api_client, memory):
|
||||
bank_id = f"curation-http-clear-dates-{uuid.uuid4().hex[:8]}"
|
||||
mem_id = await _insert_fact(memory, bank_id, "Release v1.2 happened on Monday.")
|
||||
@@ -134,3 +137,41 @@ async def test_patch_empty_body_is_rejected(api_client, memory):
|
||||
)
|
||||
assert resp.status_code == 422
|
||||
await memory.delete_bank(bank_id, request_context=RequestContext())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_patch_resolve_entities_reaches_the_engine(api_client, memory):
|
||||
"""resolve_entities must survive the HTTP boundary, and default to True when omitted (#3479)."""
|
||||
bank_id = f"curation-http-resolve-{uuid.uuid4().hex[:8]}"
|
||||
mem_id = await _insert_fact(memory, bank_id, "Dr. Waller referred the patient.")
|
||||
|
||||
seen: list[bool] = []
|
||||
real_update = memory.update_memory_unit
|
||||
|
||||
async def _capture(*args, **kwargs):
|
||||
seen.append(kwargs.get("resolve_entities"))
|
||||
return await real_update(*args, **kwargs)
|
||||
|
||||
with patch.object(memory, "update_memory_unit", new=_capture):
|
||||
resp = await api_client.patch(
|
||||
f"/v1/default/banks/{bank_id}/memories/{mem_id}",
|
||||
json={"entities": ["Dr. Waller"], "resolve_entities": False},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
resp = await api_client.patch(
|
||||
f"/v1/default/banks/{bank_id}/memories/{mem_id}",
|
||||
json={"entities": ["Dr. Waller"]},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
||||
assert seen == [False, True], "an explicit flag is forwarded; omitting it defaults to resolving"
|
||||
|
||||
# A non-boolean is rejected by the request model, not silently coerced.
|
||||
resp = await api_client.patch(
|
||||
f"/v1/default/banks/{bank_id}/memories/{mem_id}",
|
||||
json={"entities": ["Dr. Waller"], "resolve_entities": "exact"},
|
||||
)
|
||||
assert resp.status_code == 422
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=RequestContext())
|
||||
|
||||
@@ -27,6 +27,10 @@ from hindsight_api.engine.task_backend import SyncTaskBackend
|
||||
from hindsight_api.extensions import TenantContext, TenantExtension
|
||||
from hindsight_api.migrations import ensure_embedding_dimension, run_migrations
|
||||
|
||||
# The whole point of this module is the physical embedding column: it counts
|
||||
# memory_units rows with a non-NULL embedding of a configured dimension.
|
||||
pytestmark = pytest.mark.memory_backend_incompatible
|
||||
|
||||
# =============================================================================
|
||||
# Shared Utilities
|
||||
# =============================================================================
|
||||
|
||||
@@ -13,6 +13,7 @@ import pytest
|
||||
|
||||
from hindsight_api.engine.db import DatabaseBackend, DatabaseConnection, create_database_backend
|
||||
from hindsight_api.engine.db.ops import UpdatedWindow
|
||||
from hindsight_api.engine.db import postgresql as pg_backend
|
||||
from hindsight_api.engine.db.postgresql import PostgreSQLBackend, apply_session_settings
|
||||
from hindsight_api.engine.db.result import DictResultRow as ResultRow
|
||||
from hindsight_api.engine.sql import SQLDialect, create_sql_dialect
|
||||
@@ -645,6 +646,13 @@ class _RecordingConnection:
|
||||
class TestApplySessionSettings:
|
||||
"""The pool's setup callback runs on every acquire — it must be one round trip (#3499)."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _forget_rejected_settings(self):
|
||||
"""The rejected-GUC memo is process-wide, so it must not leak between tests."""
|
||||
pg_backend._unsupported_settings.clear()
|
||||
yield
|
||||
pg_backend._unsupported_settings.clear()
|
||||
|
||||
_SETTINGS = [
|
||||
("hnsw.ef_search", "200"),
|
||||
("statement_timeout", "600s"),
|
||||
@@ -684,6 +692,44 @@ class TestApplySessionSettings:
|
||||
# The rejected one raised and was skipped rather than aborting setup.
|
||||
assert len(conn.calls) == 1 + len(self._SETTINGS)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_setting_the_server_rejects_is_not_sent_again(self):
|
||||
"""Otherwise every acquire re-pays a failed batch plus one statement per setting.
|
||||
|
||||
Reached by any GUC the cluster does not define — pg_trgm when the extension is
|
||||
absent, or hnsw.iterative_scan on a pgvector older than 0.8, which reserves the
|
||||
"hnsw." prefix and so rejects it rather than accepting a placeholder.
|
||||
"""
|
||||
conn = _RecordingConnection(fail_batched=True, reject="pg_trgm.similarity_threshold")
|
||||
await apply_session_settings(conn, self._SETTINGS)
|
||||
|
||||
# Next acquire: one batched statement again, carrying only what the server took.
|
||||
conn = _RecordingConnection()
|
||||
await apply_session_settings(conn, self._SETTINGS)
|
||||
|
||||
assert len(conn.calls) == 1
|
||||
_, args = conn.calls[0]
|
||||
assert "pg_trgm.similarity_threshold" not in args
|
||||
assert args == ("hnsw.ef_search", "200", "statement_timeout", "600s")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_transient_failure_does_not_disable_a_setting(self):
|
||||
"""Only "unrecognized configuration parameter" is permanent; anything else retries."""
|
||||
|
||||
class _FlakyConnection(_RecordingConnection):
|
||||
async def execute(self, query: str, *args) -> None:
|
||||
self.calls.append((query, args))
|
||||
if len(self.calls) == 1:
|
||||
raise asyncpg.exceptions.UndefinedObjectError("unrecognized configuration parameter")
|
||||
if "hnsw.ef_search" in args:
|
||||
raise asyncpg.exceptions.DeadlockDetectedError("transient")
|
||||
|
||||
await apply_session_settings(_FlakyConnection(), self._SETTINGS)
|
||||
|
||||
conn = _RecordingConnection()
|
||||
await apply_session_settings(conn, self._SETTINGS)
|
||||
assert "hnsw.ef_search" in conn.calls[0][1]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config integration test
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
Tests for delta retain — upsert optimization that only re-processes changed chunks.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
@@ -237,6 +236,7 @@ async def test_delta_retain_modified_chunk(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_delta_retain_entities_preserved_for_unchanged_chunks(memory, request_context):
|
||||
"""
|
||||
Entities linked to unchanged chunks should be preserved after delta retain.
|
||||
@@ -256,13 +256,8 @@ async def test_delta_retain_entities_preserved_for_unchanged_chunks(memory, requ
|
||||
assert len(v1_units) > 0
|
||||
|
||||
# Check entities exist
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
v1_entities = await conn.fetch(
|
||||
"SELECT canonical_name FROM entities WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
v1_entity_names = {e["canonical_name"].lower() for e in v1_entities}
|
||||
v1_listing = await memory.list_entities(bank_id, request_context=request_context)
|
||||
v1_entity_names = {e["canonical_name"].lower() for e in v1_listing["items"]}
|
||||
assert len(v1_entity_names) > 0, "Should have entities after v1 retain"
|
||||
|
||||
# Upsert with same content — entities should persist
|
||||
@@ -274,12 +269,8 @@ async def test_delta_retain_entities_preserved_for_unchanged_chunks(memory, requ
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
v2_entities = await conn.fetch(
|
||||
"SELECT canonical_name FROM entities WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
v2_entity_names = {e["canonical_name"].lower() for e in v2_entities}
|
||||
v2_listing = await memory.list_entities(bank_id, request_context=request_context)
|
||||
v2_entity_names = {e["canonical_name"].lower() for e in v2_listing["items"]}
|
||||
|
||||
# All v1 entities should still exist
|
||||
assert v1_entity_names.issubset(v2_entity_names), (
|
||||
@@ -308,13 +299,8 @@ async def test_delta_retain_new_entities_created_for_new_chunks(memory, request_
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
v1_entities = await conn.fetch(
|
||||
"SELECT canonical_name FROM entities WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
v1_entity_names = {e["canonical_name"].lower() for e in v1_entities}
|
||||
v1_listing = await memory.list_entities(bank_id, request_context=request_context)
|
||||
v1_entity_names = {e["canonical_name"].lower() for e in v1_listing["items"]}
|
||||
|
||||
# Append content mentioning new entities
|
||||
v2_content = v1_content + "\n\nBob joined Facebook. He works with Charlie on the Reality Labs project."
|
||||
@@ -326,12 +312,8 @@ async def test_delta_retain_new_entities_created_for_new_chunks(memory, request_
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
v2_entities = await conn.fetch(
|
||||
"SELECT canonical_name FROM entities WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
v2_entity_names = {e["canonical_name"].lower() for e in v2_entities}
|
||||
v2_listing = await memory.list_entities(bank_id, request_context=request_context)
|
||||
v2_entity_names = {e["canonical_name"].lower() for e in v2_listing["items"]}
|
||||
|
||||
# Should have more entities after adding content with new people/orgs
|
||||
assert len(v2_entity_names) > len(v1_entity_names), (
|
||||
@@ -343,6 +325,7 @@ async def test_delta_retain_new_entities_created_for_new_chunks(memory, request_
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_delta_retain_links_preserved_for_unchanged_chunks(memory, request_context):
|
||||
"""
|
||||
Memory links (temporal, semantic, entity) for unchanged chunks should be preserved.
|
||||
@@ -444,6 +427,7 @@ async def test_delta_retain_document_metadata_updated(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_delta_retain_metadata_consistent_for_unchanged_units(memory, request_context):
|
||||
"""Metadata updates should reach facts preserved by metadata-only retain."""
|
||||
bank_id = f"test_delta_unit_meta_{_ts()}"
|
||||
@@ -479,25 +463,19 @@ async def test_delta_retain_metadata_consistent_for_unchanged_units(memory, requ
|
||||
assert doc is not None
|
||||
assert doc["document_metadata"] == {"source": "crm"}
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT metadata FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
listing = await memory.list_memory_units(bank_id, document_id=document_id, request_context=request_context)
|
||||
rows = listing["items"]
|
||||
|
||||
assert rows
|
||||
for row in rows:
|
||||
metadata = row["metadata"]
|
||||
if isinstance(metadata, str):
|
||||
metadata = json.loads(metadata)
|
||||
assert metadata == {"source": "crm"}
|
||||
# list_memory_units already parses the JSON metadata into a dict.
|
||||
assert row["metadata"] == {"source": "crm"}
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_delta_retain_drops_null_metadata_values(memory, request_context):
|
||||
"""A null-valued metadata key must never reach memory_units (issue #3209).
|
||||
|
||||
@@ -519,21 +497,12 @@ async def test_delta_retain_drops_null_metadata_values(memory, request_context):
|
||||
async def _unit_metadata() -> dict[str, dict]:
|
||||
"""The document's memory units, keyed by unit id, so a later call can
|
||||
tell units that survived a delta from ones re-extracted from scratch."""
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT id, metadata FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
listing = await memory.list_memory_units(bank_id, document_id=document_id, request_context=request_context)
|
||||
rows = listing["items"]
|
||||
assert rows, "expected memory units for the document"
|
||||
units = {}
|
||||
for row in rows:
|
||||
metadata = row["metadata"]
|
||||
if isinstance(metadata, str):
|
||||
metadata = json.loads(metadata)
|
||||
units[str(row["id"])] = metadata
|
||||
return units
|
||||
# list_memory_units already parses the JSON metadata into a dict and
|
||||
# returns ids as strings.
|
||||
return {row["id"]: row["metadata"] for row in rows}
|
||||
|
||||
async def _retain(content: str, source: str) -> None:
|
||||
await memory.retain_batch_async(
|
||||
@@ -598,13 +567,8 @@ async def test_delta_retain_tags_propagated_to_existing_units(memory, request_co
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
v1_tags = await conn.fetch(
|
||||
"SELECT tags FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
v1_listing = await memory.list_memory_units(bank_id, document_id=document_id, request_context=request_context)
|
||||
v1_tags = v1_listing["items"]
|
||||
assert all("team-a" in row["tags"] for row in v1_tags), "v1 units should have team-a tag"
|
||||
|
||||
# v2 with same content but different tags
|
||||
@@ -620,12 +584,8 @@ async def test_delta_retain_tags_propagated_to_existing_units(memory, request_co
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
v2_tags = await conn.fetch(
|
||||
"SELECT tags FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
v2_listing = await memory.list_memory_units(bank_id, document_id=document_id, request_context=request_context)
|
||||
v2_tags = v2_listing["items"]
|
||||
for row in v2_tags:
|
||||
assert "team-b" in row["tags"], f"v2 units should have team-b tag, got {row['tags']}"
|
||||
assert "important" in row["tags"], f"v2 units should have important tag, got {row['tags']}"
|
||||
@@ -923,13 +883,8 @@ async def test_delta_retain_with_user_entities(memory, request_context):
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
v1_entities = await conn.fetch(
|
||||
"SELECT canonical_name FROM entities WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
v1_names = {e["canonical_name"].lower() for e in v1_entities}
|
||||
v1_listing = await memory.list_entities(bank_id, request_context=request_context)
|
||||
v1_names = {e["canonical_name"].lower() for e in v1_listing["items"]}
|
||||
|
||||
# v2 with additional entity, same content
|
||||
# Note: same content = delta path (no re-extraction)
|
||||
@@ -951,12 +906,8 @@ async def test_delta_retain_with_user_entities(memory, request_context):
|
||||
)
|
||||
|
||||
# Should have entities from both v1 and v2
|
||||
async with pool.acquire() as conn:
|
||||
v2_entities = await conn.fetch(
|
||||
"SELECT canonical_name FROM entities WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
v2_names = {e["canonical_name"].lower() for e in v2_entities}
|
||||
v2_listing = await memory.list_entities(bank_id, request_context=request_context)
|
||||
v2_names = {e["canonical_name"].lower() for e in v2_listing["items"]}
|
||||
|
||||
# v1 entities should be preserved
|
||||
assert v1_names.issubset(v2_names), f"v1 entities should be preserved: {v1_names} not in {v2_names}"
|
||||
|
||||
@@ -114,13 +114,11 @@ async def test_delta_detects_unchanged_after_first_retain(memory, request_contex
|
||||
)
|
||||
assert len(v1_units) > 0
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
v1_count = await conn.fetchval(
|
||||
"SELECT count(*) FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
async def _unit_count() -> int:
|
||||
listing = await memory.list_memory_units(bank_id, document_id=document_id, request_context=request_context)
|
||||
return listing["total"]
|
||||
|
||||
v1_count = await _unit_count()
|
||||
|
||||
# Second retain — same content, should be detected as unchanged by delta
|
||||
v2_units = await memory.retain_async(
|
||||
@@ -135,12 +133,7 @@ async def test_delta_detects_unchanged_after_first_retain(memory, request_contex
|
||||
assert v2_units == [], f"Delta with unchanged content should return empty, got {len(v2_units)} units"
|
||||
|
||||
# Memory unit count should not change
|
||||
async with pool.acquire() as conn:
|
||||
v2_count = await conn.fetchval(
|
||||
"SELECT count(*) FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
v2_count = await _unit_count()
|
||||
assert v2_count == v1_count, f"Memory unit count changed on same-content upsert: {v1_count} -> {v2_count}"
|
||||
|
||||
# Third retain — verify stability
|
||||
@@ -153,12 +146,7 @@ async def test_delta_detects_unchanged_after_first_retain(memory, request_contex
|
||||
)
|
||||
assert v3_units == [], "Third retain should also detect unchanged"
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
v3_count = await conn.fetchval(
|
||||
"SELECT count(*) FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
v3_count = await _unit_count()
|
||||
assert v3_count == v1_count, f"Memory unit count changed on third upsert: {v1_count} -> {v3_count}"
|
||||
|
||||
finally:
|
||||
@@ -189,18 +177,19 @@ async def test_stale_request_skipped_when_newer_retain_completed(memory, request
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
after_newer_count = await conn.fetchval(
|
||||
"SELECT count(*) FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
async def _unit_count() -> int:
|
||||
listing = await memory.list_memory_units(bank_id, document_id=document_id, request_context=request_context)
|
||||
return listing["total"]
|
||||
|
||||
after_newer_count = await _unit_count()
|
||||
assert after_newer_count > 0, "Should have facts from newer content"
|
||||
|
||||
# Simulate the race condition by pushing the document's updated_at into
|
||||
# the future. This makes any new retain appear "stale" (its start_time
|
||||
# is before updated_at), as if another request already completed.
|
||||
# is before updated_at), as if another request already completed. This
|
||||
# forces internal store state (a document's updated_at) that the public
|
||||
# API has no way to set, so it stays a direct write on purpose.
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"UPDATE documents SET updated_at = NOW() + INTERVAL '10 seconds' WHERE id = $1 AND bank_id = $2",
|
||||
@@ -223,12 +212,7 @@ async def test_stale_request_skipped_when_newer_retain_completed(memory, request
|
||||
assert result == [], f"Stale request should return empty, got {result}"
|
||||
|
||||
# Memory units should be unchanged (newer content preserved)
|
||||
async with pool.acquire() as conn:
|
||||
final_count = await conn.fetchval(
|
||||
"SELECT count(*) FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
final_count = await _unit_count()
|
||||
assert final_count == after_newer_count, (
|
||||
f"Stale request should not change memory units: {after_newer_count} -> {final_count}"
|
||||
)
|
||||
@@ -271,6 +255,7 @@ async def memory_no_llm(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.flaky(reruns=2, reruns_delay=2)
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_concurrent_upserts_no_duplicates(memory_no_llm, request_context):
|
||||
"""
|
||||
Stress test: N concurrent retains of the same document with different content.
|
||||
@@ -350,12 +335,10 @@ async def test_concurrent_upserts_no_duplicates(memory_no_llm, request_context):
|
||||
logger.info(f"Winning version: {winning_version} (out of {num_concurrent} concurrent retains)")
|
||||
|
||||
# 2. All memory units should belong to the winning version
|
||||
async with pool.acquire() as conn:
|
||||
units = await conn.fetch(
|
||||
"SELECT text, chunk_id, id::text as unit_id FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
listing = await memory_no_llm.list_memory_units(
|
||||
bank_id, document_id=document_id, limit=1000, request_context=request_context
|
||||
)
|
||||
units = listing["items"]
|
||||
unit_texts = [r["text"] for r in units]
|
||||
assert len(unit_texts) > 0, "Should have at least 1 memory unit"
|
||||
|
||||
@@ -364,9 +347,7 @@ async def test_concurrent_upserts_no_duplicates(memory_no_llm, request_context):
|
||||
# We check for "Person_N" rather than "VERSION_N" because the text
|
||||
# splitter may cut mid-text, so later chunks might not start with the prefix.
|
||||
winning_person = f"Person_{winning_version}"
|
||||
wrong_version_units = [
|
||||
(r["text"], r["chunk_id"], r["unit_id"]) for r in units if winning_person not in r["text"]
|
||||
]
|
||||
wrong_version_units = [(r["text"], r["chunk_id"], r["id"]) for r in units if winning_person not in r["text"]]
|
||||
assert not wrong_version_units, (
|
||||
f"Found {len(wrong_version_units)} memory units NOT from winning version "
|
||||
f"{winning_version} (expected '{winning_person}' in every unit). "
|
||||
|
||||
@@ -115,6 +115,7 @@ async def _export_async(memory, bank_id, request_context, **kwargs):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_import_filters_degenerate_fact_without_shifting_archive_ordinals(memory, request_context):
|
||||
"""A rejected archive fact must not shift chunks, causal links, or observation sources."""
|
||||
dst = _unique_bank("transfer_degenerate_alignment")
|
||||
@@ -393,6 +394,7 @@ async def test_export_bank_contents(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_export_tolerates_legacy_null_and_numeric_fact_metadata(memory, request_context):
|
||||
"""A bank holding legacy metadata must still be exportable (issue #3209).
|
||||
|
||||
@@ -522,6 +524,7 @@ async def _observation_count(memory, bank_id):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_bank_import_preserves_consolidation_lifecycle(memory, request_context):
|
||||
"""Whole-bank import restores each fact's consolidation lifecycle verbatim, so
|
||||
previously-consolidated and previously-failed facts are never re-consolidated
|
||||
@@ -622,6 +625,7 @@ async def test_bank_import_preserves_consolidation_lifecycle(memory, request_con
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_bank_export_import_exact_roundtrip(memory, request_context):
|
||||
"""A whole-bank archive restores EXACT bank content (config, docs, facts,
|
||||
observations, entities, links, webhooks, directives, mental models) with facts
|
||||
@@ -793,6 +797,7 @@ async def test_bank_roundtrip_carries_mental_model_history(memory, request_conte
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_bank_roundtrip_carries_knowledge_pages(memory, request_context):
|
||||
"""A whole-bank archive restores the Knowledge Pages tree — nested folders +
|
||||
pages, parent_id / mental_model_id / managed / sort_order preserved — and
|
||||
@@ -896,6 +901,7 @@ async def test_import_bank_refuses_existing_bank(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_export_import_roundtrip_without_llm(memory, request_context, monkeypatch):
|
||||
"""Export from one bank and import into another without re-running the LLM."""
|
||||
src = _unique_bank("transfer_src")
|
||||
@@ -1019,6 +1025,7 @@ async def _bank_snapshot(memory, bank_id):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_full_roundtrip_integrity(memory, request_context):
|
||||
"""Full export → import must reproduce every persisted artifact (counts + sizes)."""
|
||||
src = _unique_bank("transfer_integ_src")
|
||||
@@ -1072,6 +1079,7 @@ async def test_full_roundtrip_integrity(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_transfer_preserves_legacy_causal_links(memory, request_context):
|
||||
"""Legacy causal edges survive export/import without becoming retain inputs."""
|
||||
src = _unique_bank("transfer_legacy_causal_src")
|
||||
@@ -1123,6 +1131,7 @@ async def test_transfer_preserves_legacy_causal_links(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_export_import_observations(memory, request_context):
|
||||
"""With include_observations, observations transfer and their sources re-link."""
|
||||
src = _unique_bank("transfer_obs_src")
|
||||
@@ -1203,6 +1212,7 @@ async def test_export_import_observations(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_import_triggers_consolidation(memory, request_context):
|
||||
"""Importing (without observations) triggers consolidation in the target bank,
|
||||
so observations get generated there — same as a normal retain."""
|
||||
@@ -1366,6 +1376,7 @@ async def test_import_on_conflict_modes(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_http_export_import_endpoints(api_client, memory, request_context):
|
||||
"""Round trip through the async HTTP export (POST + poll + download) and import endpoints."""
|
||||
src = _unique_bank("transfer_http_src")
|
||||
@@ -1547,6 +1558,7 @@ async def test_import_rejects_invalid_on_conflict(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_bank_import_classifies_label_entities(memory, request_context):
|
||||
"""An imported bank's label entities are stored with entity_kind='label'.
|
||||
|
||||
@@ -1624,6 +1636,7 @@ async def test_bank_import_classifies_label_entities(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_async_export_roundtrip(memory, request_context):
|
||||
"""The async export operation stashes a real archive that re-imports cleanly.
|
||||
|
||||
@@ -1669,6 +1682,7 @@ async def test_async_export_include_observations_subset_rejected(memory, request
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_export_attach_batching_preserves_entities_and_causal_links(memory, request_context, monkeypatch):
|
||||
"""Batched attach queries carry every fact's entities and cross-batch causal edges.
|
||||
|
||||
|
||||
@@ -1916,7 +1916,6 @@ async def test_retain_multivalue_tag_entities_all_stored(memory_real_llm, reques
|
||||
The original bug: tags are added correctly, but unit_entities only stores
|
||||
a subset (typically the first entity).
|
||||
"""
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
|
||||
bank_id = f"test-1558-multivalue-tag-{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
@@ -1964,32 +1963,14 @@ async def test_retain_multivalue_tag_entities_all_stored(memory_real_llm, reques
|
||||
|
||||
assert len(unit_ids) > 0, "Should have extracted at least one fact"
|
||||
|
||||
async with memory_real_llm._pool.acquire() as conn:
|
||||
# Check entities in unit_entities table
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.canonical_name
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("entities")} e ON e.id = ue.entity_id
|
||||
WHERE ue.unit_id = ANY($1::uuid[])
|
||||
""",
|
||||
[u for u in unit_ids],
|
||||
)
|
||||
entity_names = {r["canonical_name"].lower() for r in entity_rows}
|
||||
|
||||
# Check tags on memory_units
|
||||
tag_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, tags
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
""",
|
||||
[u for u in unit_ids],
|
||||
)
|
||||
all_tags = set()
|
||||
for row in tag_rows:
|
||||
if row["tags"]:
|
||||
all_tags.update(t.lower() for t in row["tags"])
|
||||
# Entities and tags both come back on the unit itself, so one read per
|
||||
# unit covers what the two joins used to.
|
||||
entity_names: set[str] = set()
|
||||
all_tags: set[str] = set()
|
||||
for unit_id in unit_ids:
|
||||
unit = await memory_real_llm.get_memory_unit(bank_id, str(unit_id), request_context)
|
||||
entity_names.update(name.lower() for name in unit["entities"])
|
||||
all_tags.update(tag.lower() for tag in unit["tags"])
|
||||
|
||||
# Filter to use:* entities/tags
|
||||
use_entities = {n for n in entity_names if n.startswith("use:")}
|
||||
@@ -2025,7 +2006,6 @@ async def test_retain_multivalue_tag_entities_second_retain(memory_real_llm, req
|
||||
temporal proximity could exceed the 0.6 merge threshold, causing both to
|
||||
resolve to the same entity ID.
|
||||
"""
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
|
||||
bank_id = f"test-1558-second-{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
@@ -2075,30 +2055,14 @@ async def test_retain_multivalue_tag_entities_second_retain(memory_real_llm, req
|
||||
|
||||
assert len(unit_ids_2) > 0
|
||||
|
||||
async with memory_real_llm._pool.acquire() as conn:
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.canonical_name
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("entities")} e ON e.id = ue.entity_id
|
||||
WHERE ue.unit_id = ANY($1::uuid[])
|
||||
""",
|
||||
[u for u in unit_ids_2],
|
||||
)
|
||||
entity_names = {r["canonical_name"].lower() for r in entity_rows}
|
||||
|
||||
tag_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, tags
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
""",
|
||||
[u for u in unit_ids_2],
|
||||
)
|
||||
all_tags = set()
|
||||
for row in tag_rows:
|
||||
if row["tags"]:
|
||||
all_tags.update(t.lower() for t in row["tags"])
|
||||
# Entities and tags both come back on the unit itself, so one read per
|
||||
# unit covers what the two joins used to.
|
||||
entity_names: set[str] = set()
|
||||
all_tags: set[str] = set()
|
||||
for unit_id in unit_ids_2:
|
||||
unit = await memory_real_llm.get_memory_unit(bank_id, str(unit_id), request_context)
|
||||
entity_names.update(name.lower() for name in unit["entities"])
|
||||
all_tags.update(tag.lower() for tag in unit["tags"])
|
||||
|
||||
use_entities = {n for n in entity_names if n.startswith("use:")}
|
||||
use_tags = {t for t in all_tags if t.startswith("use:")}
|
||||
|
||||
@@ -173,3 +173,26 @@ def test_intake_attaches_fact_dates_after_dropping():
|
||||
# candidate must not shift it.
|
||||
assert _texts(all_entities_flat) == ["Alice"]
|
||||
assert all_entities_flat[0]["event_date"] == when
|
||||
|
||||
|
||||
def test_intake_keeps_the_stricter_resolve_flag_when_normalization_collapses_names():
|
||||
"""A caller's literal name must not become resolvable via a normalization collision (#3479).
|
||||
|
||||
entity_processing dedups caller-supplied against extracted names on the RAW text, so
|
||||
"Acme Corp" and "Acme\nCorp" both arrive here and only collapse after normalization.
|
||||
Keeping the first entry verbatim would drop the caller's resolve=False with it.
|
||||
"""
|
||||
all_entities_flat, _all, _map = _prepare(
|
||||
[
|
||||
{"text": "Acme\nCorp", "type": "ORG"}, # extracted: resolvable
|
||||
{"text": "Acme Corp", "type": "ORG", "resolve": False}, # caller: literal
|
||||
]
|
||||
)
|
||||
|
||||
assert _texts(all_entities_flat) == ["Acme Corp"], "still one mention after normalization"
|
||||
assert all_entities_flat[0]["resolve"] is False, "the caller's literal intent survives the merge"
|
||||
|
||||
|
||||
def test_intake_defaults_resolve_to_true():
|
||||
all_entities_flat, _all, _map = _prepare([{"text": "Alice", "type": "PERSON"}, _FakeEntity("Bob")])
|
||||
assert [e["resolve"] for e in all_entities_flat] == [True, True]
|
||||
|
||||
@@ -250,6 +250,127 @@ async def test_fuzzy_scoring_never_merges_regular_text_into_label_row():
|
||||
ops.bulk_insert_entities.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unresolved_names_beat_a_winning_fuzzy_candidate():
|
||||
"""A caller-authored name must not be re-resolved to a similar entity (#3479).
|
||||
|
||||
"Dr Wall" here is the shape that broke curation: a typo entity that outscores the
|
||||
0.6 threshold on name similarity (0.41) plus co-occurrence with the other name in
|
||||
the same request (0.3), so resolution silently swaps it in for the corrected
|
||||
"Dr. Waller". With resolve=False on that mention the scoring is skipped entirely and
|
||||
the literal name is created/reused instead.
|
||||
"""
|
||||
from types import SimpleNamespace
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
candidates = {"Dr. Waller": [("typo-entity-id", "Dr Wall", {}, now, 40)]}
|
||||
cooccurrence = {"typo-entity-id": {"careorg"}}
|
||||
entities_data = [
|
||||
{
|
||||
"text": "Dr. Waller",
|
||||
"nearby_entities": [{"text": "Dr. Waller"}, {"text": "CareOrg"}],
|
||||
"event_date": now,
|
||||
}
|
||||
]
|
||||
|
||||
# Fuzzy matching (retain's behaviour) picks the typo entity — the bug, pinned here so
|
||||
# the assertion below is measuring a real difference.
|
||||
fuzzy_resolver = EntityResolver(
|
||||
pool=SimpleNamespace(
|
||||
ops=SimpleNamespace(bulk_insert_entities=AsyncMock(), fetch_missing_entity_ids=AsyncMock())
|
||||
),
|
||||
entity_lookup="full",
|
||||
)
|
||||
fuzzy = await fuzzy_resolver._resolve_from_candidates(
|
||||
conn=AsyncMock(),
|
||||
bank_id="bank-1",
|
||||
entities_data=entities_data,
|
||||
unit_event_date=now,
|
||||
all_candidates=candidates,
|
||||
cooccurrence_map=cooccurrence,
|
||||
)
|
||||
assert fuzzy[0].canonical_name == "Dr Wall", "precondition: fuzzy scoring picks the typo entity"
|
||||
|
||||
ops = SimpleNamespace(
|
||||
bulk_insert_entities=AsyncMock(return_value={"dr. waller": "correct-entity-id"}),
|
||||
fetch_missing_entity_ids=AsyncMock(return_value=[]),
|
||||
)
|
||||
exact_resolver = EntityResolver(pool=SimpleNamespace(ops=ops), entity_lookup="full")
|
||||
exact = await exact_resolver._resolve_from_candidates(
|
||||
conn=AsyncMock(),
|
||||
bank_id="bank-1",
|
||||
entities_data=[{**entities_data[0], "resolve": False}],
|
||||
unit_event_date=now,
|
||||
all_candidates=candidates,
|
||||
cooccurrence_map=cooccurrence,
|
||||
)
|
||||
assert exact == [ResolvedEntity(entity_id="correct-entity-id", canonical_name="Dr. Waller")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_mention_resolve_flag_only_spares_the_names_that_opt_out():
|
||||
"""Retain mixes caller-supplied and extracted names in one batch (#3479).
|
||||
|
||||
The flag is per mention precisely so a caller can have their own names taken literally
|
||||
without turning off resolution for the extractor's — which would silently fill the bank
|
||||
with near-duplicate entities. Here the caller's "Dr. Waller" must be created as written
|
||||
while the extractor's "Dr Waler" still resolves onto the existing "Dr Wall".
|
||||
"""
|
||||
from types import SimpleNamespace
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
typo = ("typo-entity-id", "Dr Wall", {}, now, 40)
|
||||
ops = SimpleNamespace(
|
||||
bulk_insert_entities=AsyncMock(return_value={"dr. waller": "literal-id"}),
|
||||
fetch_missing_entity_ids=AsyncMock(return_value=[]),
|
||||
)
|
||||
resolver = EntityResolver(pool=SimpleNamespace(ops=ops), entity_lookup="full")
|
||||
|
||||
resolved = await resolver._resolve_from_candidates(
|
||||
conn=AsyncMock(),
|
||||
bank_id="bank-1",
|
||||
entities_data=[
|
||||
{"text": "Dr. Waller", "nearby_entities": [], "event_date": now, "resolve": False},
|
||||
{"text": "Dr Waler", "nearby_entities": [], "event_date": now},
|
||||
],
|
||||
unit_event_date=now,
|
||||
all_candidates={"Dr. Waller": [typo], "Dr Waler": [typo]},
|
||||
cooccurrence_map={},
|
||||
)
|
||||
|
||||
assert resolved[0] == ResolvedEntity(entity_id="literal-id", canonical_name="Dr. Waller"), (
|
||||
"the opted-out name is created as written"
|
||||
)
|
||||
assert resolved[1].canonical_name == "Dr Wall", "the extracted name still resolves onto the existing entity"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_intrabatch_clustering_skips_names_that_opted_out():
|
||||
"""A literal name must not be folded into a same-batch variant, or pull one into itself."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
ops = SimpleNamespace(
|
||||
bulk_insert_entities=AsyncMock(return_value={"alice": "alice-id", "alice smith": "alice-smith-id"}),
|
||||
fetch_missing_entity_ids=AsyncMock(return_value=[]),
|
||||
)
|
||||
resolver = EntityResolver(pool=SimpleNamespace(ops=ops), entity_lookup="full")
|
||||
|
||||
resolved = await resolver._resolve_from_candidates(
|
||||
conn=AsyncMock(),
|
||||
bank_id="bank-1",
|
||||
entities_data=[
|
||||
{"text": "Alice", "nearby_entities": [], "resolve": False},
|
||||
{"text": "Alice Smith", "nearby_entities": [], "resolve": False},
|
||||
],
|
||||
unit_event_date=None,
|
||||
all_candidates={},
|
||||
cooccurrence_map={},
|
||||
)
|
||||
|
||||
assert [e.canonical_name for e in resolved] == ["Alice", "Alice Smith"]
|
||||
assert [e.entity_id for e in resolved] == ["alice-id", "alice-smith-id"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Oracle fuzzy entity resolution — unit tests (mock conn, no live DB)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -685,6 +685,7 @@ async def test_converter_registry():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_file_conversion_creates_separate_retain_operation(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test that file conversion and retain are two separate async operations.
|
||||
|
||||
|
||||
@@ -150,6 +150,7 @@ async def test_graph_q_and_tags_filter_combined(api_client, test_bank_id):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_graph_document_filter_includes_observations_via_source_memories(
|
||||
memory, api_client, test_bank_id, request_context
|
||||
):
|
||||
@@ -297,6 +298,7 @@ async def _seed_scoped_observations(memory, bank_id, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_observation_scopes_enumeration(memory, api_client, test_bank_id, request_context):
|
||||
"""The scopes endpoint enumerates distinct tag sets (order-normalized) with counts."""
|
||||
await _seed_scoped_observations(memory, test_bank_id, request_context)
|
||||
@@ -313,6 +315,7 @@ async def test_observation_scopes_enumeration(memory, api_client, test_bank_id,
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_graph_exact_scope_filter(memory, api_client, test_bank_id, request_context):
|
||||
"""tags_match=exact filters observations to exactly one scope, not supersets."""
|
||||
await _seed_scoped_observations(memory, test_bank_id, request_context)
|
||||
@@ -337,6 +340,7 @@ async def test_graph_exact_scope_filter(memory, api_client, test_bank_id, reques
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_graph_exact_global_scope_filter(memory, api_client, test_bank_id, request_context):
|
||||
"""tags_match=exact with no tags is the global scope: untagged observations only."""
|
||||
await _seed_scoped_observations(memory, test_bank_id, request_context)
|
||||
|
||||
@@ -27,6 +27,11 @@ from hindsight_api.engine.graph_maintenance import (
|
||||
)
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
# Every test here seeds memory_units / memory_links / entities with raw INSERTs and
|
||||
# asserts raw link-row counts, as the module docstring says — none of it round-trips
|
||||
# through the store, so a backend that keeps those rows outside SQL sees an empty graph.
|
||||
pytestmark = pytest.mark.memory_backend_incompatible
|
||||
|
||||
|
||||
async def _ensure_bank(memory: MemoryEngine, bank_id: str, request_context: RequestContext) -> None:
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
@@ -196,6 +196,7 @@ async def test_unordered_concurrent_sweep_and_upsert_deadlocks(memory: MemoryEng
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_graph_maintenance_sweep_retries_on_deadlock(memory: MemoryEngine, request_context: RequestContext):
|
||||
"""The entity-prune batch must survive a deadlock.
|
||||
|
||||
|
||||
@@ -38,8 +38,7 @@ class MockTenantExtension(TenantExtension):
|
||||
|
||||
|
||||
class _FakeBankOps:
|
||||
async def create_bank_vector_indexes(self, *args, **kwargs):
|
||||
return None
|
||||
"""Dialect ops stub. Bank creation issues no index DDL (#3485), so this is bare."""
|
||||
|
||||
|
||||
class FakeBankConfigBackend:
|
||||
|
||||
@@ -3,8 +3,14 @@ Tests for per-bank vector index lifecycle and UNION ALL retrieval.
|
||||
|
||||
Covers:
|
||||
- _bank_index_name deterministic naming
|
||||
- Per-bank vector indexes created on bank creation (retain_async / ensure_bank_exists)
|
||||
- Per-bank vector indexes dropped on bank deletion
|
||||
- At the default threshold (0 = no minimum), a retained bank still ends up with
|
||||
its three per-(bank, fact_type) indexes — the pre-#3485 outcome — but they are
|
||||
built by the queued vector_index_maintenance operation rather than inline on
|
||||
the request path
|
||||
- An untouched bank gets none: bank creation issues no index DDL, and an index
|
||||
over an empty partition serves nothing
|
||||
- Per-bank vector indexes dropped on bank deletion, the one request path that
|
||||
still issues vector-index DDL
|
||||
- retrieve_semantic_bm25_combined_sql groups results correctly by fact_type and source
|
||||
"""
|
||||
|
||||
@@ -13,7 +19,26 @@ from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.retain.bank_utils import _BANK_INDEX_FACT_TYPES, _bank_index_name
|
||||
from hindsight_api.engine import vector_index_health
|
||||
from hindsight_api.engine.db_utils import retry_with_backoff
|
||||
from hindsight_api.engine.retain.bank_utils import (
|
||||
_BANK_INDEX_FACT_TYPES,
|
||||
_bank_index_name,
|
||||
_vector_index_clause,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def default_threshold(monkeypatch):
|
||||
"""Restore the shipped default (0 = no minimum) for this test.
|
||||
|
||||
conftest raises the threshold out of reach suite-wide so thousands of
|
||||
throwaway banks don't each queue an index build; asserting the default
|
||||
behaviour means putting it back.
|
||||
"""
|
||||
monkeypatch.setattr(vector_index_health, "qualifies_for_per_bank_index", lambda rows: rows > 0)
|
||||
monkeypatch.setattr(vector_index_health, "should_keep_per_bank_index", lambda rows: rows > 0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests — no DB required
|
||||
@@ -78,9 +103,46 @@ async def _get_bank_vector_indexes(pool, bank_id: str) -> list[str]:
|
||||
return [row["indexname"] for row in rows]
|
||||
|
||||
|
||||
async def _build_bank_vector_indexes(pool, bank_id: str) -> list[str]:
|
||||
"""Give a bank its three partial indexes, as the maintenance sweep would.
|
||||
|
||||
Used by the delete-path test: retain no longer creates them, so a bank has
|
||||
to be given them before deletion can be asked to take them away.
|
||||
"""
|
||||
index_clause = _vector_index_clause()
|
||||
assert index_clause is not None
|
||||
async with pool.acquire() as conn:
|
||||
internal_id = str(await conn.fetchval("SELECT internal_id FROM banks WHERE bank_id = $1", bank_id))
|
||||
literal = await conn.fetchval("SELECT quote_literal($1::text)", bank_id)
|
||||
names = []
|
||||
for ft in _BANK_INDEX_FACT_TYPES:
|
||||
name = _bank_index_name(ft, internal_id)
|
||||
# CONCURRENTLY, and retried: a plain CREATE INDEX takes ShareLock on
|
||||
# the shared memory_units table, which forms a deadlock cycle with
|
||||
# another xdist worker's DROP INDEX CONCURRENTLY (ShareUpdateExclusive)
|
||||
# — observed as a three-way cycle in CI.
|
||||
await retry_with_backoff(
|
||||
lambda name=name, ft=ft: conn.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {name} ON memory_units {index_clause} "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = {literal}"
|
||||
)
|
||||
)
|
||||
names.append(name)
|
||||
return names
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_creates_per_bank_vector_indexes(memory, request_context):
|
||||
"""retain_async on a new bank must create 3 per-(bank, fact_type) vector indexes."""
|
||||
async def test_retain_still_ends_up_with_per_bank_indexes(memory, request_context, default_threshold):
|
||||
"""At the shipped default a retained bank has the same coverage it always had.
|
||||
|
||||
The threshold defaults to 0 — no minimum — so every partition holding rows is
|
||||
indexed, exactly as before #3485. What changed is *who* builds them: the
|
||||
index DDL used to run inside the retain transaction, taking a ShareLock on
|
||||
the shared memory_units table that deadlocked against concurrent writers.
|
||||
Now retain queues a vector_index_maintenance operation and returns; the tests
|
||||
run a synchronous task backend, so the operation has completed by the time
|
||||
retain_async does.
|
||||
"""
|
||||
bank_id = f"test_hnsw_create_{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
await memory.retain_async(
|
||||
@@ -89,18 +151,39 @@ async def test_retain_creates_per_bank_vector_indexes(memory, request_context):
|
||||
request_context=request_context,
|
||||
)
|
||||
indexes = await _get_bank_vector_indexes(memory._pool, bank_id)
|
||||
assert len(indexes) == 3, f"Expected 3 per-bank vector indexes, got: {indexes}"
|
||||
for ft_short in _BANK_INDEX_FACT_TYPES.values():
|
||||
assert any(ft_short in idx for idx in indexes), (
|
||||
f"Missing index for fact_type short '{ft_short}' in {indexes}"
|
||||
)
|
||||
assert indexes, "a retained bank should end up with per-bank vector indexes at the default threshold"
|
||||
for name in indexes:
|
||||
assert "_" in name
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bank_creation_alone_creates_no_vector_indexes(memory, request_context):
|
||||
"""Creating a bank must issue no index DDL — it is the request path that hurt.
|
||||
|
||||
A fresh bank holds no rows, so its three indexes would cover nothing while
|
||||
still being locked and planned against by every other bank's queries. This is
|
||||
the difference that makes bank count stop being a ceiling (#3485).
|
||||
"""
|
||||
bank_id = f"test_hnsw_empty_{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
indexes = await _get_bank_vector_indexes(memory._pool, bank_id)
|
||||
assert indexes == [], f"bank creation must not create vector indexes, got: {indexes}"
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_bank_drops_vector_indexes(memory, request_context):
|
||||
"""delete_bank must drop all per-bank vector indexes."""
|
||||
"""delete_bank must drop the per-bank vector indexes a large bank had.
|
||||
|
||||
Still the one request path that issues vector-index DDL: an index outliving
|
||||
its bank would be charged to every surviving bank's query planning forever,
|
||||
and nothing else knows the internal_id it is named after.
|
||||
"""
|
||||
bank_id = f"test_hnsw_drop_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
await memory.retain_async(
|
||||
@@ -108,9 +191,9 @@ async def test_delete_bank_drops_vector_indexes(memory, request_context):
|
||||
content="Bob is a data scientist.",
|
||||
request_context=request_context,
|
||||
)
|
||||
# Verify indexes exist before deletion
|
||||
await _build_bank_vector_indexes(memory._pool, bank_id)
|
||||
indexes_before = await _get_bank_vector_indexes(memory._pool, bank_id)
|
||||
assert len(indexes_before) == 3
|
||||
assert len(indexes_before) == 3, "setup: the bank should have indexes to drop"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -119,8 +202,12 @@ async def test_delete_bank_drops_vector_indexes(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_idempotent_bank_creation(memory, request_context):
|
||||
"""Retaining into the same bank twice must not error and still have exactly 3 indexes."""
|
||||
async def test_retain_idempotent_bank_creation(memory, request_context, default_threshold):
|
||||
"""Retaining twice must not error, and must not duplicate or rebuild indexes.
|
||||
|
||||
The second retain queues another maintenance operation; its plan has to come
|
||||
back empty so a busy bank is not rebuilding ANN indexes on every write.
|
||||
"""
|
||||
bank_id = f"test_hnsw_idem_{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
await memory.retain_async(
|
||||
@@ -128,13 +215,17 @@ async def test_retain_idempotent_bank_creation(memory, request_context):
|
||||
content="Carol is a product manager.",
|
||||
request_context=request_context,
|
||||
)
|
||||
after_first = await _get_bank_vector_indexes(memory._pool, bank_id)
|
||||
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Carol joined the company in 2022.",
|
||||
request_context=request_context,
|
||||
)
|
||||
indexes = await _get_bank_vector_indexes(memory._pool, bank_id)
|
||||
assert len(indexes) == 3
|
||||
|
||||
assert await _get_bank_vector_indexes(memory._pool, bank_id) == after_first
|
||||
submitted = await memory.submit_async_vector_index_maintenance(bank_id=bank_id, request_context=request_context)
|
||||
assert submitted["no_work"] is True, "a settled bank must stop queueing maintenance"
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -187,6 +278,7 @@ async def test_retrieve_semantic_bm25_grouped_by_fact_type(memory, request_conte
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_fetch_unit_dates_ignores_noncanonical_uuid_inputs(memory, request_context):
|
||||
"""The indexed UUID lookup preserves the old text-comparison input behavior."""
|
||||
from hindsight_api.engine.db.ops_postgresql import PostgreSQLOps
|
||||
@@ -245,6 +337,11 @@ async def test_recall_reuses_semantic_pool_for_graph_seeds(memory, request_conte
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
# Asserts *how* the graph arm seeds — that recall calls link_expansion_retrieval's
|
||||
# _find_semantic_seeds — rather than what it returns. A store with its own graph
|
||||
# retrieval never goes through that function, so the assertion is specific to the
|
||||
# SQL retrieval path.
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_recall_keeps_graph_seed_query_for_stricter_semantic_floor(memory, request_context, monkeypatch):
|
||||
"""A semantic floor above the graph floor must retain the dedicated seed query."""
|
||||
from hindsight_api.engine.response_models import MinScores
|
||||
|
||||
@@ -101,7 +101,7 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
# ================================================================
|
||||
|
||||
# List banks (should be empty initially or have other test banks)
|
||||
response = await api_client.get("/v1/default/banks")
|
||||
response = await api_client.get("/v1/default/banks", params={"limit": 1000})
|
||||
assert response.status_code == 200
|
||||
initial_banks_data = response.json()["banks"]
|
||||
initial_banks = [a["bank_id"] for a in initial_banks_data]
|
||||
@@ -211,7 +211,7 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
assert fresh_stats["total_nodes"] == stats["total_nodes"]
|
||||
|
||||
# Verify bank list returns stats (fact_count, last_document_at)
|
||||
response = await api_client.get("/v1/default/banks")
|
||||
response = await api_client.get("/v1/default/banks", params={"limit": 1000})
|
||||
assert response.status_code == 200
|
||||
banks_after = response.json()["banks"]
|
||||
our_bank = next(b for b in banks_after if b["bank_id"] == test_bank_id)
|
||||
@@ -354,7 +354,7 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
# 9. List All Banks (should include our test bank)
|
||||
# ================================================================
|
||||
|
||||
response = await api_client.get("/v1/default/banks")
|
||||
response = await api_client.get("/v1/default/banks", params={"limit": 1000})
|
||||
assert response.status_code == 200
|
||||
final_banks_data = response.json()["banks"]
|
||||
final_banks = [a["bank_id"] for a in final_banks_data]
|
||||
@@ -552,6 +552,7 @@ async def test_document_deletion_with_slashes_in_id(api_client):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_delete_bank(api_client):
|
||||
"""Test delete bank endpoint.
|
||||
|
||||
@@ -601,7 +602,7 @@ async def test_delete_bank(api_client):
|
||||
assert len(response.json()["items"]) > 0
|
||||
|
||||
# Check bank is in list
|
||||
response = await api_client.get("/v1/default/banks")
|
||||
response = await api_client.get("/v1/default/banks", params={"limit": 1000})
|
||||
assert response.status_code == 200
|
||||
bank_ids = [b["bank_id"] for b in response.json()["banks"]]
|
||||
assert test_bank_id in bank_ids
|
||||
@@ -616,7 +617,7 @@ async def test_delete_bank(api_client):
|
||||
|
||||
# 4. Verify bank and all data is deleted
|
||||
# Bank should not be in list
|
||||
response = await api_client.get("/v1/default/banks")
|
||||
response = await api_client.get("/v1/default/banks", params={"limit": 1000})
|
||||
assert response.status_code == 200
|
||||
bank_ids = [b["bank_id"] for b in response.json()["banks"]]
|
||||
assert test_bank_id not in bank_ids
|
||||
@@ -674,7 +675,7 @@ async def test_clear_memories_preserves_bank(api_client):
|
||||
assert response.status_code == 200
|
||||
assert response.json()["total_nodes"] > 0
|
||||
|
||||
response = await api_client.get("/v1/default/banks")
|
||||
response = await api_client.get("/v1/default/banks", params={"limit": 1000})
|
||||
assert response.status_code == 200
|
||||
bank_ids = [b["bank_id"] for b in response.json()["banks"]]
|
||||
assert test_bank_id in bank_ids
|
||||
@@ -685,7 +686,7 @@ async def test_clear_memories_preserves_bank(api_client):
|
||||
assert response.json()["success"] is True
|
||||
|
||||
# 3. Bank should still exist in the list
|
||||
response = await api_client.get("/v1/default/banks")
|
||||
response = await api_client.get("/v1/default/banks", params={"limit": 1000})
|
||||
assert response.status_code == 200
|
||||
bank_ids = [b["bank_id"] for b in response.json()["banks"]]
|
||||
assert test_bank_id in bank_ids, "Bank should still exist after clearing memories"
|
||||
@@ -2106,7 +2107,7 @@ async def test_patch_bank_does_not_create_missing_bank(api_client, memory, monke
|
||||
profile = await api_client.get(f"/v1/default/banks/{test_bank_id}/profile")
|
||||
assert profile.status_code == 404, profile.text
|
||||
|
||||
banks = await api_client.get("/v1/default/banks")
|
||||
banks = await api_client.get("/v1/default/banks", params={"limit": 1000})
|
||||
assert banks.status_code == 200, banks.text
|
||||
assert test_bank_id not in {bank["bank_id"] for bank in banks.json()["banks"]}
|
||||
|
||||
|
||||
@@ -227,6 +227,7 @@ class TestTree:
|
||||
assert loose["trigger"]["refresh_after_consolidation"] is False
|
||||
assert loose["trigger"]["mode"] == "delta" # untouched by the patch
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_tree_staleness_follows_the_bank_watermark(self, api_client, memory, kb_bank):
|
||||
"""The tree answers from one bank-wide watermark, not a scan per page.
|
||||
|
||||
|
||||
@@ -151,6 +151,7 @@ async def test_repeated_large_same_id_replacement_is_idempotent(memory, request_
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_append_after_zero_fact_header_slice_skips_unchanged_history(
|
||||
memory,
|
||||
request_context,
|
||||
|
||||
@@ -444,6 +444,35 @@ class TestComputeSemanticLinksAnnPgBouncerSafety:
|
||||
# And there must not be a RESET — SET LOCAL handles it at commit.
|
||||
assert not any(f"RESET {guc}" in s for s in executed_sql)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skips_a_guc_the_server_has_rejected(self, mock_conn, monkeypatch):
|
||||
"""An unknown GUC must not be attempted inside this transaction.
|
||||
|
||||
hnsw.iterative_scan needs pgvector 0.8+, and pgvector reserves the "hnsw."
|
||||
prefix, so an older server errors on it rather than accepting a placeholder —
|
||||
and an error inside an open transaction aborts the whole link computation, not
|
||||
just the setting. The pool's session setup names the same GUCs on acquire, so
|
||||
by the time this runs an unknown one is already recorded.
|
||||
"""
|
||||
from hindsight_api.engine.db import postgresql as pg_backend
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector")
|
||||
monkeypatch.setattr(pg_backend, "_unsupported_settings", {"hnsw.iterative_scan"})
|
||||
|
||||
await compute_semantic_links_ann(
|
||||
conn=mock_conn,
|
||||
bank_id="bank-1",
|
||||
unit_ids=["u1"],
|
||||
embeddings=[[0.1] * 384],
|
||||
fact_types=["world"],
|
||||
threshold=DEFAULT_SEMANTIC_LINK_MIN_SIMILARITY,
|
||||
)
|
||||
|
||||
executed_sql = [call.args[0] for call in mock_conn.execute.call_args_list]
|
||||
assert not any("hnsw.iterative_scan" in s for s in executed_sql)
|
||||
# The supported one is still applied.
|
||||
assert any("hnsw.ef_search" in s for s in executed_sql)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vchord_ann_does_not_set_fixed_probe_count(self, mock_conn, monkeypatch):
|
||||
"""VectorChord probe counts must come from index/default config.
|
||||
|
||||
@@ -55,8 +55,8 @@ async def test_list_banks_overlays_config_disposition_and_mission(memory):
|
||||
assert profile["disposition"] == {"skepticism": 4, "literalism": 5, "empathy": 2}
|
||||
|
||||
# The list path must agree with the get path for this bank.
|
||||
banks = await memory.list_banks(request_context=request_context)
|
||||
entry = next((b for b in banks if b["bank_id"] == bank_id), None)
|
||||
page = await memory.list_banks(search_query=bank_id, request_context=request_context)
|
||||
entry = next((b for b in page["banks"] if b["bank_id"] == bank_id), None)
|
||||
assert entry is not None, f"bank {bank_id!r} not present in list_banks output"
|
||||
|
||||
assert entry["mission"] == profile["mission"], (
|
||||
|
||||
@@ -57,9 +57,9 @@ async def test_list_banks_counts_via_store_for_non_sql_bank(memory, monkeypatch)
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Must not raise NameError; must reach the store's non-SQL count path.
|
||||
banks = await memory.list_banks(request_context=request_context)
|
||||
page = await memory.list_banks(search_query=bank_id, request_context=request_context)
|
||||
|
||||
entry = next((b for b in banks if b["bank_id"] == bank_id), None)
|
||||
entry = next((b for b in page["banks"] if b["bank_id"] == bank_id), None)
|
||||
assert entry is not None, f"bank {bank_id!r} not present in list_banks output"
|
||||
|
||||
# The capability + count were consulted with the row's real bank id.
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
"""Tests for pagination and search on the bank list.
|
||||
|
||||
``GET /v1/default/banks`` used to return every bank in the system in one
|
||||
response: no limit, no offset, and a SQL query with no LIMIT clause. On an
|
||||
instance with many banks that is an unbounded payload, plus per-bank config
|
||||
resolution (and a live store count for non-SQL stores) for every bank rather
|
||||
than the ones actually being shown.
|
||||
|
||||
Runs via: uv run pytest tests/test_list_banks_pagination.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def three_banks(memory, request_context):
|
||||
"""Three banks sharing a unique prefix, so the search is xdist-safe."""
|
||||
prefix = f"pagebank{uuid.uuid4().hex[:8]}"
|
||||
bank_ids = [f"{prefix}_{i}" for i in range(3)]
|
||||
for bank_id in bank_ids:
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
try:
|
||||
yield prefix, bank_ids
|
||||
finally:
|
||||
for bank_id in bank_ids:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pages_are_disjoint_and_cover_every_match(memory, request_context, three_banks):
|
||||
prefix, bank_ids = three_banks
|
||||
|
||||
first = await memory.list_banks(search_query=prefix, limit=2, offset=0, request_context=request_context)
|
||||
second = await memory.list_banks(search_query=prefix, limit=2, offset=2, request_context=request_context)
|
||||
|
||||
assert first["total"] == 3
|
||||
assert first["limit"] == 2
|
||||
assert first["offset"] == 0
|
||||
assert len(first["banks"]) == 2
|
||||
assert second["total"] == 3
|
||||
assert second["offset"] == 2
|
||||
assert len(second["banks"]) == 1
|
||||
|
||||
paged = [bank["bank_id"] for bank in first["banks"] + second["banks"]]
|
||||
assert len(set(paged)) == 3, f"pages overlap: {paged}"
|
||||
assert set(paged) == set(bank_ids)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_offset_past_the_end_returns_no_banks_but_the_real_total(memory, request_context, three_banks):
|
||||
prefix, _ = three_banks
|
||||
|
||||
page = await memory.list_banks(search_query=prefix, limit=10, offset=3, request_context=request_context)
|
||||
|
||||
assert page["banks"] == []
|
||||
assert page["total"] == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_limit_zero_returns_no_banks(memory, request_context, three_banks):
|
||||
prefix, _ = three_banks
|
||||
|
||||
page = await memory.list_banks(search_query=prefix, limit=0, request_context=request_context)
|
||||
|
||||
assert page["banks"] == []
|
||||
assert page["total"] == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_negative_paging_values_are_clamped(memory, request_context, three_banks):
|
||||
"""The MCP tool takes limit/offset straight from a model, and the page is a Python
|
||||
slice — a negative value must not silently trim the tail."""
|
||||
prefix, _ = three_banks
|
||||
|
||||
page = await memory.list_banks(search_query=prefix, limit=-1, offset=-5, request_context=request_context)
|
||||
|
||||
assert page["banks"] == []
|
||||
assert page["total"] == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_matches_bank_name_case_insensitively(memory, request_context):
|
||||
bank_id = f"searchname{uuid.uuid4().hex[:8]}"
|
||||
display_name = f"Zeta {uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
await memory.update_bank(bank_id, name=display_name, request_context=request_context)
|
||||
|
||||
page = await memory.list_banks(search_query=display_name.upper(), request_context=request_context)
|
||||
|
||||
assert [bank["bank_id"] for bank in page["banks"]] == [bank_id]
|
||||
assert page["total"] == 1
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_endpoint_echoes_paging_and_filters(api_client, three_banks):
|
||||
prefix, _ = three_banks
|
||||
|
||||
response = await api_client.get("/v1/default/banks", params={"q": prefix, "limit": 1, "offset": 1})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
body = response.json()
|
||||
assert body["total"] == 3
|
||||
assert body["limit"] == 1
|
||||
assert body["offset"] == 1
|
||||
assert len(body["banks"]) == 1
|
||||
assert body["banks"][0]["bank_id"].startswith(prefix)
|
||||
@@ -53,6 +53,7 @@ def _ids(result: dict) -> set[str]:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_created_before_filter(memory: MemoryEngine, request_context: RequestContext):
|
||||
bank_id = f"test-lmu-created-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
"""What ``list_memory_units`` / ``list_entities`` put in each item.
|
||||
|
||||
Tests that need a unit's write watermark, its lineage, or an entity's kind used to
|
||||
read the columns straight out of ``memory_units`` / ``entities``. Those are part of
|
||||
the read model, so they are on the item — and asserted here through the engine, on
|
||||
units written by retain rather than seeded with SQL, so the coverage holds for any
|
||||
store behind the memories seam.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
|
||||
async def _retain(memory: MemoryEngine, bank_id: str, content: str, request_context: RequestContext) -> list[str]:
|
||||
return await memory.retain_async(bank_id=bank_id, content=content, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_fact_type_accepts_a_list(memory: MemoryEngine, request_context: RequestContext):
|
||||
"""A list of fact types matches any of them — the source-fact selection callers want."""
|
||||
bank_id = f"test-lmu-facttypes-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
try:
|
||||
await _retain(memory, bank_id, "Alice deployed the release on Friday.", request_context)
|
||||
|
||||
async def listed(fact_type):
|
||||
page = await memory.list_memory_units(
|
||||
bank_id, fact_type=fact_type, limit=500, request_context=request_context
|
||||
)
|
||||
return {item["id"] for item in page["items"]}, page["total"]
|
||||
|
||||
world_ids, world_total = await listed("world")
|
||||
exp_ids, exp_total = await listed("experience")
|
||||
both_ids, both_total = await listed(["world", "experience"])
|
||||
|
||||
# The list arm is exactly the union of the single-value arms, so the
|
||||
# assertion holds however the LLM happened to classify the facts.
|
||||
assert both_ids == world_ids | exp_ids
|
||||
assert both_total == world_total + exp_total
|
||||
|
||||
# An empty list filters nothing, matching the "omitted" case.
|
||||
_, empty_total = await listed([])
|
||||
_, unfiltered_total = await listed(None)
|
||||
assert empty_total == unfiltered_total
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_items_carry_updated_at_and_lineage(memory: MemoryEngine, request_context: RequestContext):
|
||||
"""Each item carries its write watermark and (for observations) its sources."""
|
||||
bank_id = f"test-lmu-readmodel-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
try:
|
||||
await _retain(memory, bank_id, "Bob moved to Berlin in March.", request_context)
|
||||
|
||||
page = await memory.list_memory_units(bank_id, limit=500, request_context=request_context)
|
||||
assert page["items"], "retain must have produced at least one fact"
|
||||
|
||||
source_ids = {item["id"] for item in page["items"] if item["fact_type"] != "observation"}
|
||||
assert source_ids, "retain must have produced at least one source fact"
|
||||
|
||||
for item in page["items"]:
|
||||
# Parseable rather than merely present: callers do date arithmetic on it.
|
||||
assert isinstance(datetime.fromisoformat(item["updated_at"]), datetime)
|
||||
if item["fact_type"] == "observation":
|
||||
# An observation's lineage points at the facts it was drawn from,
|
||||
# and those facts are in this same bank.
|
||||
assert item["source_memory_ids"], "an observation must carry its sources"
|
||||
assert set(item["source_memory_ids"]) <= source_ids
|
||||
else:
|
||||
# A source fact has no lineage; the field is always there, never absent.
|
||||
assert item["source_memory_ids"] == []
|
||||
|
||||
# The list item and the detail view agree on the unit.
|
||||
first = page["items"][0]
|
||||
detail = await memory.get_memory_unit(bank_id, first["id"], request_context)
|
||||
assert detail["text"] == first["text"]
|
||||
if first["fact_type"] == "observation":
|
||||
assert detail["source_memory_ids"] == first["source_memory_ids"]
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_entities_carry_their_kind(memory: MemoryEngine, request_context: RequestContext):
|
||||
"""list_entities reports how each entity was classified."""
|
||||
bank_id = f"test-entities-kind-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
try:
|
||||
await _retain(memory, bank_id, "Carol works with Dave at Acme.", request_context)
|
||||
|
||||
page = await memory.list_entities(bank_id, limit=500, request_context=request_context)
|
||||
assert page["items"], "retain must have produced at least one entity"
|
||||
for item in page["items"]:
|
||||
assert "entity_kind" in item, "entity_kind is part of the entity read model"
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -360,6 +360,26 @@ async def test_middleware_handles_both_endpoints(mock_memory):
|
||||
assert "create_bank" not in single_bank_tools
|
||||
|
||||
|
||||
def test_registered_tool_sets_match_the_allowlists(mock_memory):
|
||||
"""Every allowlisted tool name must resolve to a really-registered tool.
|
||||
|
||||
Three hand-maintained lists have to agree: ``_ALL_TOOLS``, the default set in
|
||||
``register_mcp_tools()``, and the single-bank allowlist in
|
||||
``create_mcp_server()``. A name added to one but not the others is silent —
|
||||
the tool simply never appears on the endpoint and nothing fails — so assert
|
||||
over the whole sets rather than spot-checking tool names.
|
||||
"""
|
||||
from hindsight_api.api.mcp import create_mcp_server
|
||||
from hindsight_api.mcp_tools import _ALL_TOOLS
|
||||
|
||||
multi = set(_tools(create_mcp_server(mock_memory, multi_bank=True)))
|
||||
single = set(_tools(create_mcp_server(mock_memory, multi_bank=False)))
|
||||
|
||||
assert multi == set(_ALL_TOOLS)
|
||||
# Single-bank mode drops exactly the tools that operate across banks.
|
||||
assert multi - single == {"list_banks", "create_bank", "get_bank_stats"}
|
||||
|
||||
|
||||
def test_global_mcp_enabled_tools_filter_restricts_registered_tools(mock_memory):
|
||||
"""Test that global mcp_enabled_tools env setting restricts which tools are registered."""
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
@@ -9,7 +9,9 @@ import pytest
|
||||
|
||||
from hindsight_api.engine.memory_engine import DirectivePage, MentalModelPage
|
||||
from hindsight_api.mcp_tools import (
|
||||
KNOWLEDGE_ROOT_PARENT,
|
||||
MCPToolsConfig,
|
||||
_knowledge_tree_json,
|
||||
_validate_mental_model_inputs,
|
||||
build_content_dict,
|
||||
parse_timestamp,
|
||||
@@ -127,6 +129,43 @@ def _apply_detail(model: dict, detail: str) -> dict:
|
||||
return model
|
||||
|
||||
|
||||
# Knowledge-base fixtures: a root folder with one page under it, shaped like the
|
||||
# dicts MemoryEngine._row_to_knowledge_node returns.
|
||||
_KNOWLEDGE_FOLDER: dict[str, Any] = {
|
||||
"id": "kf-1",
|
||||
"kind": "folder",
|
||||
"name": "Runbooks",
|
||||
"parent_id": None,
|
||||
"mental_model_id": None,
|
||||
"managed": False,
|
||||
"updated_at": "2026-01-01T00:00:00+00:00",
|
||||
}
|
||||
|
||||
_KNOWLEDGE_PAGE_NODE: dict[str, Any] = {
|
||||
"id": "kp-1",
|
||||
"kind": "page",
|
||||
"name": "Deploys",
|
||||
"parent_id": "kf-1",
|
||||
"mental_model_id": "mm-page",
|
||||
"managed": False,
|
||||
"tags": ["ops"],
|
||||
"source_query": "How is the service deployed?",
|
||||
"last_refreshed_at": "2026-01-02T00:00:00+00:00",
|
||||
"trigger": {"refresh_after_consolidation": True},
|
||||
}
|
||||
|
||||
_KNOWLEDGE_NODES: list[dict[str, Any]] = [
|
||||
_KNOWLEDGE_FOLDER,
|
||||
{**_KNOWLEDGE_PAGE_NODE, "is_stale": False},
|
||||
]
|
||||
|
||||
_KNOWLEDGE_PAGE: dict[str, Any] = {
|
||||
**_KNOWLEDGE_PAGE_NODE,
|
||||
"tags": ["type:runbook", "ops"],
|
||||
"content": "Run `make deploy`.",
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_memory():
|
||||
"""Create a mock MemoryEngine with all MCP tool methods."""
|
||||
@@ -221,6 +260,28 @@ def mock_memory():
|
||||
memory.update_bank = AsyncMock(side_effect=_update_bank)
|
||||
memory.list_banks = AsyncMock(return_value=[])
|
||||
|
||||
# Knowledge base methods
|
||||
memory.list_knowledge_nodes = AsyncMock(return_value=list(_KNOWLEDGE_NODES))
|
||||
memory.get_knowledge_page = AsyncMock(return_value=dict(_KNOWLEDGE_PAGE))
|
||||
memory.search_knowledge_pages = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"id": "kp-1",
|
||||
"name": "Deploys",
|
||||
"mental_model_id": "mm-page",
|
||||
"snippet": "How the service is deployed",
|
||||
"score": 0.9,
|
||||
"updated_at": "2026-01-02T00:00:00+00:00",
|
||||
}
|
||||
]
|
||||
)
|
||||
memory.create_knowledge_folder = AsyncMock(return_value=dict(_KNOWLEDGE_FOLDER))
|
||||
memory.create_knowledge_page = AsyncMock(return_value=dict(_KNOWLEDGE_PAGE_NODE))
|
||||
memory.rename_knowledge_node = AsyncMock(return_value=dict(_KNOWLEDGE_PAGE_NODE))
|
||||
memory.move_knowledge_node = AsyncMock(return_value=dict(_KNOWLEDGE_PAGE_NODE))
|
||||
memory.update_knowledge_page = AsyncMock(return_value=dict(_KNOWLEDGE_PAGE_NODE))
|
||||
memory.delete_knowledge_node = AsyncMock(return_value=True)
|
||||
|
||||
return memory
|
||||
|
||||
|
||||
@@ -387,7 +448,10 @@ class TestMentalModelToolRegistration:
|
||||
assert "clear_mental_model" in tools
|
||||
assert "update_memory" in tools
|
||||
assert "invalidate_memory" in tools
|
||||
assert len(tools) == 32
|
||||
assert "get_knowledge_base_tree" in tools
|
||||
assert "create_knowledge_page" in tools
|
||||
assert "delete_knowledge_node" in tools
|
||||
assert len(tools) == 39
|
||||
|
||||
def test_all_tools_have_nonempty_descriptions(self):
|
||||
"""Every registered tool must expose a non-empty description.
|
||||
@@ -1892,6 +1956,229 @@ class TestEmptyListReturns:
|
||||
assert '"items": []' in result or "[]" in result
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Knowledge Base Tool Tests
|
||||
# =========================================================================
|
||||
|
||||
_KB_TOOLS = {
|
||||
"get_knowledge_base_tree",
|
||||
"search_knowledge_base",
|
||||
"get_knowledge_page",
|
||||
"create_knowledge_folder",
|
||||
"create_knowledge_page",
|
||||
"update_knowledge_node",
|
||||
"delete_knowledge_node",
|
||||
}
|
||||
|
||||
|
||||
class TestKnowledgeTreeProjection:
|
||||
"""The flat node list nests into roots, with pages projected from their model."""
|
||||
|
||||
def test_nests_pages_under_their_folder(self):
|
||||
roots = _knowledge_tree_json(_KNOWLEDGE_NODES)
|
||||
assert [r["id"] for r in roots] == ["kf-1"]
|
||||
assert [c["id"] for c in roots[0]["children"]] == ["kp-1"]
|
||||
|
||||
def test_page_carries_mental_model_metadata(self):
|
||||
page = _knowledge_tree_json(_KNOWLEDGE_NODES)[0]["children"][0]
|
||||
assert page["description"] == "How is the service deployed?"
|
||||
assert page["tags"] == ["ops"]
|
||||
assert page["timestamp"] == "2026-01-02T00:00:00+00:00"
|
||||
assert page["is_stale"] is False
|
||||
assert page["trigger"] == {"refresh_after_consolidation": True}
|
||||
|
||||
def test_orphaned_node_becomes_a_root(self):
|
||||
# A node whose parent is not in the list (e.g. a partial read) must still
|
||||
# surface rather than vanish from the tree.
|
||||
roots = _knowledge_tree_json([{**_KNOWLEDGE_PAGE_NODE, "parent_id": "kf-missing"}])
|
||||
assert [r["id"] for r in roots] == ["kp-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestKnowledgeBaseTools:
|
||||
async def test_tools_registered(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, _KB_TOOLS, include_bank_id=True)
|
||||
assert _KB_TOOLS == set(_tools(mcp).keys())
|
||||
|
||||
async def test_tools_registered_single_bank(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, _KB_TOOLS, include_bank_id=False)
|
||||
assert _KB_TOOLS == set(_tools(mcp).keys())
|
||||
|
||||
async def test_get_tree(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"get_knowledge_base_tree"}, include_bank_id=True)
|
||||
result = json.loads(await _tools(mcp)["get_knowledge_base_tree"].fn())
|
||||
assert [r["id"] for r in result["roots"]] == ["kf-1"]
|
||||
assert mock_memory.list_knowledge_nodes.call_args.kwargs["with_staleness"] is True
|
||||
|
||||
async def test_get_tree_single_bank(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"get_knowledge_base_tree"}, include_bank_id=False)
|
||||
result = await _tools(mcp)["get_knowledge_base_tree"].fn()
|
||||
assert isinstance(result, dict)
|
||||
assert result["roots"][0]["kind"] == "folder"
|
||||
|
||||
async def test_search(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"search_knowledge_base"}, include_bank_id=True)
|
||||
result = json.loads(await _tools(mcp)["search_knowledge_base"].fn(query="deploy", limit=5))
|
||||
assert result["total"] == 1
|
||||
assert result["results"][0]["id"] == "kp-1"
|
||||
call_kwargs = mock_memory.search_knowledge_pages.call_args.kwargs
|
||||
assert call_kwargs["query"] == "deploy"
|
||||
assert call_kwargs["limit"] == 5
|
||||
|
||||
async def test_search_clamps_limit(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"search_knowledge_base"}, include_bank_id=True)
|
||||
await _tools(mcp)["search_knowledge_base"].fn(query="deploy", limit=500)
|
||||
assert mock_memory.search_knowledge_pages.call_args.kwargs["limit"] == 50
|
||||
|
||||
async def test_get_page_renders_markdown(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"get_knowledge_page"}, include_bank_id=True)
|
||||
result = json.loads(await _tools(mcp)["get_knowledge_page"].fn(page_id="kp-1"))
|
||||
# The `type:` tag becomes the page type and drops out of the displayed tags.
|
||||
assert result["type"] == "runbook"
|
||||
assert result["tags"] == ["ops"]
|
||||
assert result["markdown"].startswith("---\n")
|
||||
assert "Run `make deploy`." in result["markdown"]
|
||||
|
||||
async def test_get_page_not_found(self, mock_memory):
|
||||
mock_memory.get_knowledge_page.return_value = None
|
||||
mcp = _make_mcp_server(mock_memory, {"get_knowledge_page"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["get_knowledge_page"].fn(page_id="kp-missing")
|
||||
assert "not found" in result
|
||||
|
||||
async def test_create_folder(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"create_knowledge_folder"}, include_bank_id=True)
|
||||
result = json.loads(await _tools(mcp)["create_knowledge_folder"].fn(name="Runbooks"))
|
||||
assert result["id"] == "kf-1"
|
||||
assert result["children"] == []
|
||||
assert mock_memory.create_knowledge_folder.call_args.kwargs["parent_id"] is None
|
||||
|
||||
async def test_create_folder_rejects_bad_parent(self, mock_memory):
|
||||
mock_memory.create_knowledge_folder.side_effect = ValueError("Parent folder 'kf-x' not found")
|
||||
mcp = _make_mcp_server(mock_memory, {"create_knowledge_folder"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["create_knowledge_folder"].fn(name="Runbooks", parent_id="kf-x")
|
||||
assert "not found" in result
|
||||
|
||||
async def test_create_page_schedules_refresh(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"create_knowledge_page"}, include_bank_id=True)
|
||||
result = json.loads(
|
||||
await _tools(mcp)["create_knowledge_page"].fn(name="Deploys", source_query="How is it deployed?")
|
||||
)
|
||||
assert result["page_id"] == "kp-1"
|
||||
assert result["operation_id"] == "op-123"
|
||||
create_kwargs = mock_memory.create_knowledge_page.call_args.kwargs
|
||||
# An unstated refresh setting must not overwrite the engine's page defaults.
|
||||
assert create_kwargs["trigger"] is None
|
||||
assert mock_memory.submit_async_refresh_mental_model.call_args.kwargs["mental_model_id"] == "mm-page"
|
||||
|
||||
async def test_create_page_with_refresh_flag(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"create_knowledge_page"}, include_bank_id=True)
|
||||
await _tools(mcp)["create_knowledge_page"].fn(
|
||||
name="Deploys", source_query="q", tags=["ops"], max_tokens=512, refresh_after_consolidation=False
|
||||
)
|
||||
create_kwargs = mock_memory.create_knowledge_page.call_args.kwargs
|
||||
assert create_kwargs["trigger"] == {"refresh_after_consolidation": False}
|
||||
assert create_kwargs["tags"] == ["ops"]
|
||||
assert create_kwargs["max_tokens"] == 512
|
||||
|
||||
async def test_create_page_duplicate_name(self, mock_memory):
|
||||
mock_memory.create_knowledge_page.return_value = None
|
||||
mcp = _make_mcp_server(mock_memory, {"create_knowledge_page"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["create_knowledge_page"].fn(name="Deploys", source_query="q")
|
||||
assert "already exists" in result
|
||||
mock_memory.submit_async_refresh_mental_model.assert_not_called()
|
||||
|
||||
async def test_update_rename(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"update_knowledge_node"}, include_bank_id=True)
|
||||
result = json.loads(await _tools(mcp)["update_knowledge_node"].fn(node_id="kp-1", name="Deployments"))
|
||||
assert result["id"] == "kp-1"
|
||||
assert mock_memory.rename_knowledge_node.call_args.kwargs["name"] == "Deployments"
|
||||
mock_memory.update_knowledge_page.assert_not_called()
|
||||
mock_memory.move_knowledge_node.assert_not_called()
|
||||
|
||||
async def test_update_move_to_folder(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"update_knowledge_node"}, include_bank_id=True)
|
||||
await _tools(mcp)["update_knowledge_node"].fn(node_id="kp-1", parent_id="kf-2")
|
||||
assert mock_memory.move_knowledge_node.call_args.kwargs["new_parent_id"] == "kf-2"
|
||||
|
||||
async def test_update_move_to_root(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"update_knowledge_node"}, include_bank_id=True)
|
||||
await _tools(mcp)["update_knowledge_node"].fn(node_id="kp-1", parent_id=KNOWLEDGE_ROOT_PARENT)
|
||||
assert mock_memory.move_knowledge_node.call_args.kwargs["new_parent_id"] is None
|
||||
|
||||
async def test_update_source_query_triggers_refresh(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"update_knowledge_node"}, include_bank_id=True)
|
||||
await _tools(mcp)["update_knowledge_node"].fn(node_id="kp-1", source_query="new question?")
|
||||
assert mock_memory.update_knowledge_page.call_args.kwargs["source_query"] == "new question?"
|
||||
assert mock_memory.submit_async_refresh_mental_model.call_args.kwargs["mental_model_id"] == "mm-page"
|
||||
|
||||
async def test_update_tags_only_does_not_refresh(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"update_knowledge_node"}, include_bank_id=True)
|
||||
await _tools(mcp)["update_knowledge_node"].fn(node_id="kp-1", tags=[])
|
||||
assert mock_memory.update_knowledge_page.call_args.kwargs["tags"] == []
|
||||
mock_memory.submit_async_refresh_mental_model.assert_not_called()
|
||||
|
||||
async def test_update_requires_a_field(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"update_knowledge_node"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["update_knowledge_node"].fn(node_id="kp-1")
|
||||
assert "Provide name" in result
|
||||
mock_memory.rename_knowledge_node.assert_not_called()
|
||||
mock_memory.update_knowledge_page.assert_not_called()
|
||||
|
||||
async def test_update_not_found(self, mock_memory):
|
||||
mock_memory.rename_knowledge_node.return_value = None
|
||||
mcp = _make_mcp_server(mock_memory, {"update_knowledge_node"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["update_knowledge_node"].fn(node_id="kp-missing", name="x")
|
||||
assert "not found" in result
|
||||
|
||||
async def test_delete_node(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"delete_knowledge_node"}, include_bank_id=True)
|
||||
result = json.loads(await _tools(mcp)["delete_knowledge_node"].fn(node_id="kp-1"))
|
||||
assert result == {"status": "deleted", "node_id": "kp-1"}
|
||||
|
||||
async def test_delete_node_not_found(self, mock_memory):
|
||||
mock_memory.delete_knowledge_node.return_value = False
|
||||
mcp = _make_mcp_server(mock_memory, {"delete_knowledge_node"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["delete_knowledge_node"].fn(node_id="kp-missing")
|
||||
assert "not found" in result
|
||||
|
||||
async def test_no_bank_configured(self, mock_memory):
|
||||
from fastmcp import FastMCP
|
||||
|
||||
mcp = FastMCP("test")
|
||||
register_mcp_tools(
|
||||
mcp,
|
||||
mock_memory,
|
||||
MCPToolsConfig(bank_id_resolver=lambda: None, include_bank_id_param=True, tools=_KB_TOOLS),
|
||||
)
|
||||
result = await _tools(mcp)["get_knowledge_base_tree"].fn()
|
||||
assert "No bank_id configured" in result
|
||||
|
||||
async def test_request_context_is_propagated(self, mock_memory):
|
||||
from fastmcp import FastMCP
|
||||
|
||||
mcp = FastMCP("test")
|
||||
register_mcp_tools(
|
||||
mcp,
|
||||
mock_memory,
|
||||
MCPToolsConfig(
|
||||
bank_id_resolver=lambda: "test-bank",
|
||||
api_key_resolver=lambda: "secret",
|
||||
include_bank_id_param=True,
|
||||
tools={"get_knowledge_base_tree"},
|
||||
),
|
||||
)
|
||||
await _tools(mcp)["get_knowledge_base_tree"].fn()
|
||||
assert mock_memory.list_knowledge_nodes.call_args.kwargs["request_context"].api_key == "secret"
|
||||
|
||||
async def test_read_only_annotations(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, _KB_TOOLS)
|
||||
tools = _tools(mcp)
|
||||
for name in ("get_knowledge_base_tree", "search_knowledge_base", "get_knowledge_page"):
|
||||
assert tools[name].annotations.readOnlyHint is True, name
|
||||
assert tools["delete_knowledge_node"].annotations.destructiveHint is True
|
||||
assert tools["create_knowledge_page"].annotations.destructiveHint is False
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Bank-Level Tool Filtering Tests
|
||||
# =========================================================================
|
||||
|
||||
@@ -21,6 +21,13 @@ from hindsight_api.engine.retain import embedding_processing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# Most of this module seeds its fixtures by INSERTing memory_units / memory_links /
|
||||
# entities directly with the helpers below, then asserts on those rows (including the
|
||||
# embedding and search_vector columns). Those classes and tests carry
|
||||
# ``memory_backend_incompatible``; the handful that go through the engine end to end
|
||||
# — test_not_found_returns_none, test_recall_excludes_invalidated — deliberately do not.
|
||||
|
||||
|
||||
async def _insert_memory(
|
||||
conn,
|
||||
memory: MemoryEngine,
|
||||
@@ -185,12 +192,11 @@ async def _entity_ids_for(conn, unit_id: uuid.UUID) -> list[uuid.UUID]:
|
||||
return [r["entity_id"] for r in rows]
|
||||
|
||||
|
||||
async def _obs_ids(conn, bank_id: str) -> list[str]:
|
||||
rows = await conn.fetch(
|
||||
"SELECT id FROM memory_units WHERE bank_id = $1 AND fact_type = 'observation'",
|
||||
bank_id,
|
||||
async def _obs_ids(memory: MemoryEngine, bank_id: str, request_context: RequestContext) -> list[str]:
|
||||
listing = await memory.list_memory_units(
|
||||
bank_id, fact_type="observation", limit=1000, request_context=request_context
|
||||
)
|
||||
return [str(r["id"]) for r in rows]
|
||||
return [item["id"] for item in listing["items"]]
|
||||
|
||||
|
||||
async def _consolidated_at(conn, mem_id: uuid.UUID):
|
||||
@@ -207,6 +213,8 @@ async def _ensure_bank(memory: MemoryEngine, bank_id: str, request_context: Requ
|
||||
|
||||
|
||||
class TestInvalidate:
|
||||
pytestmark = pytest.mark.memory_backend_incompatible
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalidate_moves_to_archive_and_prunes(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
bank_id = f"test-curation-inv-{uuid.uuid4().hex[:8]}"
|
||||
@@ -250,7 +258,7 @@ class TestInvalidate:
|
||||
"archive is cold storage with no index; the schema drops search_vector (#2503)"
|
||||
)
|
||||
assert await _link_count(conn, m1) == 0, "links cascade-pruned on move"
|
||||
assert str(obs_id) not in await _obs_ids(conn, bank_id), "derived observation removed"
|
||||
assert str(obs_id) not in await _obs_ids(memory, bank_id, request_context), "derived observation removed"
|
||||
assert await _consolidated_at(conn, m2) is None, "surviving source reset for re-consolidation"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -334,6 +342,8 @@ class TestInvalidate:
|
||||
|
||||
|
||||
class TestEdit:
|
||||
pytestmark = pytest.mark.memory_backend_incompatible
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_changes_text_and_rederives(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
bank_id = f"test-curation-edit-{uuid.uuid4().hex[:8]}"
|
||||
@@ -385,7 +395,7 @@ class TestEdit:
|
||||
assert row["consolidated_at"] is None, "edited memory re-consolidates"
|
||||
assert "'assist'" not in row["search_vector"], "old text must not stay in native FTS search_vector"
|
||||
assert "'user'" in row["search_vector"], "new text must refresh native FTS search_vector"
|
||||
assert str(obs_id) not in await _obs_ids(conn, bank_id), "stale observation re-derived"
|
||||
assert str(obs_id) not in await _obs_ids(memory, bank_id, request_context), "stale observation re-derived"
|
||||
queued_ids = await conn.fetch("SELECT unit_id FROM graph_maintenance_queue WHERE bank_id = $1", bank_id)
|
||||
assert {row["unit_id"] for row in queued_ids} == {m1, m2}, "edited memory and incoming victim both queued"
|
||||
|
||||
@@ -468,6 +478,93 @@ class TestEdit:
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def _near_duplicate_entity_bank(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
"""A bank in the shape reported in #3479, returned as (bank_id, unit_id, typo_entity_id).
|
||||
|
||||
Extraction created a typo entity ("Dr Wall") that co-occurs with the other name in the
|
||||
edit, so fuzzy scoring puts it above the 0.6 match threshold for the corrected spelling
|
||||
"Dr. Waller" — name similarity 0.41 plus the full 0.3 co-occurrence bonus.
|
||||
"""
|
||||
bank_id = f"test-curation-entmode-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
unit_id = await _insert_memory(conn, memory, bank_id, "Dr. Waller referred the patient to CareOrg.")
|
||||
typo = await _insert_entity(conn, bank_id, "Dr Wall")
|
||||
careorg = await _insert_entity(conn, bank_id, "CareOrg")
|
||||
await _link_entity(conn, unit_id, typo)
|
||||
await _link_entity(conn, unit_id, careorg)
|
||||
# The co-occurrence edge is what lets the typo entity outscore the corrected name.
|
||||
await conn.execute(
|
||||
"INSERT INTO entity_cooccurrences (entity_id_1, entity_id_2, cooccurrence_count) VALUES ($1, $2, 5)",
|
||||
*sorted([typo, careorg], key=str),
|
||||
)
|
||||
return bank_id, unit_id, typo
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_entities_unresolved_keeps_the_submitted_names(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""resolve_entities=False links the names the caller wrote (#3479)."""
|
||||
bank_id, m1, typo = await self._near_duplicate_entity_bank(memory, request_context)
|
||||
pool = await memory._get_pool()
|
||||
|
||||
with (
|
||||
patch.object(memory, "submit_async_consolidation", new=AsyncMock()),
|
||||
patch.object(memory, "submit_async_graph_maintenance", new=AsyncMock()),
|
||||
):
|
||||
result = await memory.update_memory_unit(
|
||||
bank_id,
|
||||
str(m1),
|
||||
entities=["Dr. Waller", "CareOrg"],
|
||||
resolve_entities=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert set(result["entities"]) == {"Dr. Waller", "CareOrg"}, "the submitted names are stored verbatim"
|
||||
async with pool.acquire() as conn:
|
||||
linked = await conn.fetch(
|
||||
"SELECT e.canonical_name FROM unit_entities ue "
|
||||
"JOIN entities e ON e.id = ue.entity_id WHERE ue.unit_id = $1",
|
||||
m1,
|
||||
)
|
||||
assert {r["canonical_name"] for r in linked} == {"Dr. Waller", "CareOrg"}
|
||||
assert typo not in await _entity_ids_for(conn, m1), "the typo entity is detached, not reused"
|
||||
assert await conn.fetchval("SELECT 1 FROM entities WHERE id = $1", typo), (
|
||||
"the typo entity itself survives for the graph-maintenance sweep to reclaim"
|
||||
)
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_entities_default_still_resolves(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
"""Omitting the flag keeps retain's resolution, so existing callers are unaffected.
|
||||
|
||||
This pins the backwards-compatible default rather than endorsing the outcome: the same
|
||||
edit that resolve_entities=False gets right lands on the near-duplicate here.
|
||||
"""
|
||||
bank_id, m1, typo = await self._near_duplicate_entity_bank(memory, request_context)
|
||||
pool = await memory._get_pool()
|
||||
|
||||
with (
|
||||
patch.object(memory, "submit_async_consolidation", new=AsyncMock()),
|
||||
patch.object(memory, "submit_async_graph_maintenance", new=AsyncMock()),
|
||||
):
|
||||
result = await memory.update_memory_unit(
|
||||
bank_id,
|
||||
str(m1),
|
||||
entities=["Dr. Waller", "CareOrg"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert set(result["entities"]) == {"Dr Wall", "CareOrg"}, "the default still resolves fuzzily"
|
||||
async with pool.acquire() as conn:
|
||||
assert typo in await _entity_ids_for(conn, m1)
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_empty_entities_detaches_all(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
bank_id = f"test-curation-editent0-{uuid.uuid4().hex[:8]}"
|
||||
@@ -529,6 +626,8 @@ class TestCurationRelinking:
|
||||
returns, the links are already rebuilt.
|
||||
"""
|
||||
|
||||
pytestmark = pytest.mark.memory_backend_incompatible
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_with_only_outgoing_links_queues_itself(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
@@ -622,6 +721,7 @@ class TestCurationRelinking:
|
||||
|
||||
class TestGuardsAndListing:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_cannot_curate_observation(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
bank_id = f"test-curation-obs-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
@@ -647,6 +747,7 @@ class TestGuardsAndListing:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_list_filters_by_state(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
bank_id = f"test-curation-list-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
@@ -682,6 +783,7 @@ class TestGuardsAndListing:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_list_and_get_memory_units_include_metadata(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
@@ -723,6 +825,7 @@ class TestGuardsAndListing:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_list_filters_by_document(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
bank_id = f"test-curation-doc-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
@@ -762,6 +865,7 @@ class TestGuardsAndListing:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_list_filters_by_entity(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
bank_id = f"test-curation-entity-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
@@ -865,6 +969,8 @@ class TestCausalLinkPreservation:
|
||||
invalidate/revert round-trip must carry them through the archive.
|
||||
"""
|
||||
|
||||
pytestmark = pytest.mark.memory_backend_incompatible
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_preserves_causal_links_and_drops_derived(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
|
||||
@@ -569,7 +569,7 @@ async def test_retain_allows_clean_content(api_client) -> None:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_stores_redacted_text(api_client, memory) -> None:
|
||||
async def test_retain_stores_redacted_text(api_client, memory, request_context) -> None:
|
||||
await api_client.put("/v1/default/banks/md-retain-2", json={})
|
||||
await _set_policy(api_client, "md-retain-2", _REDACT_POLICY)
|
||||
secret = "ghp_" + "A" * 36
|
||||
@@ -578,8 +578,8 @@ async def test_retain_stores_redacted_text(api_client, memory) -> None:
|
||||
json={"items": [{"content": f"my token is {secret}"}]},
|
||||
)
|
||||
assert r.status_code == 200, r.text
|
||||
async with memory._pool.acquire() as conn:
|
||||
texts = [row["text"] for row in await conn.fetch("SELECT text FROM memory_units WHERE bank_id = 'md-retain-2'")]
|
||||
listing = await memory.list_memory_units("md-retain-2", limit=1000, request_context=request_context)
|
||||
texts = [item["text"] for item in listing["items"]]
|
||||
assert all(secret not in t for t in texts), texts
|
||||
|
||||
|
||||
|
||||
@@ -37,6 +37,10 @@ from hindsight_api.engine.db.ops_postgresql import PostgreSQLOps
|
||||
from hindsight_api.engine.db.postgresql import PostgresConnection
|
||||
from hindsight_api.engine.retain.link_utils import _bulk_insert_links
|
||||
|
||||
# Asserts a raw memory_links row count around a concurrent delete; the graph read
|
||||
# path dedupes bidirectional edges, so the count is not reproducible through it.
|
||||
pytestmark = pytest.mark.memory_backend_incompatible
|
||||
|
||||
|
||||
async def _insert_unit(conn: asyncpg.Connection, bank_id: str) -> str:
|
||||
"""Insert one committed memory_unit (autocommit) and return its id as text."""
|
||||
|
||||
@@ -84,6 +84,7 @@ async def _bank(memory: MemoryEngine, slug: str, request_context: RequestContext
|
||||
|
||||
class TestWritesThatMustStamp:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_document_tag_propagation_stamps_updated_at(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
@@ -109,6 +110,7 @@ class TestWritesThatMustStamp:
|
||||
assert await _updated_at(conn, mem_id) > _BASELINE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_embedding_write_stamps_updated_at(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
"""The stored vector is part of the memory, so the store method stamps on its own.
|
||||
|
||||
@@ -214,6 +216,7 @@ class TestConsolidationBookkeepingIsExempt:
|
||||
assert await _updated_at(conn, mem_id) == _BASELINE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_requeue_after_observation_cleanup_leaves_updated_at_alone(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
|
||||
@@ -97,6 +97,7 @@ async def test_tagged_strict_model_skipped_when_only_untagged_consolidated(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_tagged_non_strict_model_refreshed_when_only_untagged_consolidated(
|
||||
memory: MemoryEngine, request_context, monkeypatch
|
||||
):
|
||||
@@ -116,6 +117,7 @@ async def test_tagged_non_strict_model_refreshed_when_only_untagged_consolidated
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_tag_groups_model_refreshed_when_only_untagged_consolidated(
|
||||
memory: MemoryEngine, request_context, monkeypatch
|
||||
):
|
||||
@@ -157,6 +159,7 @@ async def test_non_strict_model_skipped_when_nothing_changed(memory: MemoryEngin
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_tagged_model_refreshed_when_its_tag_was_consolidated(memory: MemoryEngine, request_context, monkeypatch):
|
||||
"""The overlap path is unchanged: a strict tagged model is refreshed when a memory
|
||||
carrying its tag was consolidated."""
|
||||
@@ -174,6 +177,7 @@ async def test_tagged_model_refreshed_when_its_tag_was_consolidated(memory: Memo
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_non_strict_model_refreshed_on_mixed_run_with_foreign_tags(
|
||||
memory: MemoryEngine, request_context, monkeypatch
|
||||
):
|
||||
|
||||
@@ -275,6 +275,7 @@ class TestDeltaRefreshPlumbing:
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_delta_no_new_facts_advances_watermark_to_newest_processed(
|
||||
self,
|
||||
memory: MemoryEngine,
|
||||
@@ -403,6 +404,7 @@ class TestDeltaRefreshPlumbing:
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_delta_refresh_watermark_survives_straddling_commit(
|
||||
self,
|
||||
memory: MemoryEngine,
|
||||
|
||||
@@ -193,6 +193,7 @@ async def test_routine_returns_cron_models_excludes_plain_and_in_flight(memory:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_due_and_stale_model_is_refreshed(memory: MemoryEngine, request_context, monkeypatch):
|
||||
"""A model whose cron is due and that has new memories in scope is refreshed."""
|
||||
bank = await _make_bank(memory, request_context)
|
||||
|
||||
@@ -1173,6 +1173,7 @@ class TestMentalModelStaleness:
|
||||
assert got["is_stale"] is False
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_untagged_mm_stale_on_any_new_memory(self, memory: MemoryEngine, request_context):
|
||||
bank_id = f"test-mm-stale-untagged-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
@@ -1201,6 +1202,7 @@ class TestMentalModelStaleness:
|
||||
assert got["is_stale"] is False
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_tagged_mm_defaults_to_all_strict(self, memory: MemoryEngine, request_context):
|
||||
bank_id = f"test-mm-stale-overlap-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
@@ -1221,6 +1223,7 @@ class TestMentalModelStaleness:
|
||||
assert got["is_stale"] is True
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_tags_match_any_keeps_overlap_behavior(self, memory: MemoryEngine, request_context):
|
||||
bank_id = f"test-mm-stale-any-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
@@ -1238,6 +1241,7 @@ class TestMentalModelStaleness:
|
||||
assert got["is_stale"] is True
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_tag_groups_define_stale_scope(self, memory: MemoryEngine, request_context):
|
||||
bank_id = f"test-mm-stale-groups-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
@@ -1262,6 +1266,7 @@ class TestMentalModelStaleness:
|
||||
assert got["is_stale"] is True
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_flat_tags_and_fact_types_share_stale_scope(self, memory: MemoryEngine, request_context):
|
||||
bank_id = f"test-mm-stale-flat-fact-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
@@ -1283,6 +1288,7 @@ class TestMentalModelStaleness:
|
||||
assert got["is_stale"] is True
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_tag_groups_and_fact_types_share_stale_scope(self, memory: MemoryEngine, request_context):
|
||||
bank_id = f"test-mm-stale-group-fact-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
@@ -1328,6 +1334,7 @@ class TestMentalModelStaleness:
|
||||
assert got["is_stale"] is True
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_tags_match_all_strict_requires_all_tags(self, memory: MemoryEngine, request_context):
|
||||
"""tags_match='all_strict' → memory must contain ALL MM tags (and be tagged)."""
|
||||
bank_id = f"test-mm-stale-all-{uuid.uuid4().hex[:8]}"
|
||||
@@ -1353,6 +1360,7 @@ class TestMentalModelStaleness:
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_tags_match_any_strict_excludes_untagged(self, memory: MemoryEngine, request_context):
|
||||
"""tags_match='any_strict' → untagged memory does NOT keep MM in scope."""
|
||||
bank_id = f"test-mm-stale-anystrict-{uuid.uuid4().hex[:8]}"
|
||||
@@ -1376,6 +1384,7 @@ class TestMentalModelStaleness:
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_fact_type_filter_narrows_scope(self, memory: MemoryEngine, request_context):
|
||||
bank_id = f"test-mm-stale-fact-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
@@ -1399,6 +1408,7 @@ class TestMentalModelStaleness:
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_tool_search_mental_models_returns_is_stale_per_mm(self, memory: MemoryEngine, request_context):
|
||||
"""Regression: tool_search_mental_models must compute is_stale per-MM via scope,
|
||||
not via a bank-wide pending_consolidation short-circuit."""
|
||||
@@ -1435,6 +1445,7 @@ class TestMentalModelStaleness:
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_tool_search_mental_models_skips_the_scan_below_the_watermark(
|
||||
self, memory: MemoryEngine, request_context
|
||||
):
|
||||
@@ -1621,6 +1632,7 @@ class TestMentalModelRefreshTimestamps:
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_staleness_keys_off_memories_seen_not_refresh_time(self, memory: MemoryEngine, request_context):
|
||||
"""The inverse regression: making ``last_refreshed_at`` a wall clock must not let
|
||||
a recent refresh mask a memory the document has never seen."""
|
||||
@@ -2585,6 +2597,7 @@ class TestMentalModelRefreshFactTypeFilter:
|
||||
memory._reflect_llm_config = wrapper
|
||||
return mock_llm
|
||||
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_refresh_with_fact_types_experience_grounds_on_experience_facts(
|
||||
self, memory: MemoryEngine, request_context
|
||||
):
|
||||
|
||||
@@ -22,6 +22,10 @@ from alembic import command
|
||||
from alembic.config import Config
|
||||
from sqlalchemy import create_engine, text
|
||||
|
||||
# Drives alembic and asserts on the search_vector column itself — an internal FTS
|
||||
# index column, not part of the public read model.
|
||||
pytestmark = pytest.mark.memory_backend_incompatible
|
||||
|
||||
_SCRIPT_LOCATION = str(Path(__file__).parent.parent / "hindsight_api" / "alembic")
|
||||
|
||||
# Revision immediately before the backfill migration.
|
||||
|
||||
@@ -163,3 +163,73 @@ async def test_wide_source_arrays_do_not_change_results(memory, request_context)
|
||||
assert scores[narrow] == scores[wide] == 1.0
|
||||
finally:
|
||||
await memory.delete_bank(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_entity_cap_bounds_hub_traversal(memory, request_context):
|
||||
"""A hub entity contributes at most ``per_entity_limit`` source facts (#3510).
|
||||
|
||||
The cap used to be a LATERAL + LIMIT and is now a row_number() window, because
|
||||
the planner cannot estimate DISTINCT over a LIMIT subquery and mis-planned the
|
||||
scoring join into a nested loop. The two forms have to select the *same* rows:
|
||||
the highest ``per_entity_limit`` unit_ids of each entity. Ranking is by unit_id
|
||||
descending, which is what the LATERAL ordered by, so a candidate built from the
|
||||
lowest ids of an over-cap entity must fall outside the cap and score nothing.
|
||||
"""
|
||||
from hindsight_api.engine.db.ops import UpdatedWindow
|
||||
from hindsight_api.engine.task_backend import fq_table
|
||||
|
||||
bank_id = f"test_obs_cap_{uuid.uuid4().hex[:8]}"
|
||||
per_entity_limit = 3
|
||||
try:
|
||||
pool = await memory._get_pool()
|
||||
backend = await memory._get_backend()
|
||||
mu, ue, ml = fq_table("memory_units"), fq_table("unit_entities"), fq_table("memory_links")
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
await _ensure_bank(conn, bank_id)
|
||||
entity_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
f"INSERT INTO {fq_table('entities')} (id, bank_id, canonical_name) VALUES ($1, $2, $3)",
|
||||
entity_id,
|
||||
bank_id,
|
||||
"Hub",
|
||||
)
|
||||
# Six facts on one entity with ids we control, so "top 3 by unit_id
|
||||
# descending" is a known set rather than an accident of uuid4().
|
||||
facts = sorted(uuid.UUID(int=i) for i in range(1, 7))
|
||||
for fid in facts:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {mu} (id, bank_id, text, fact_type, source_memory_ids, event_date)
|
||||
VALUES ($1, $2, $3, 'world', NULL, $4)
|
||||
""",
|
||||
fid,
|
||||
bank_id,
|
||||
f"hub fact {fid}",
|
||||
datetime.now(timezone.utc),
|
||||
)
|
||||
await conn.execute(f"INSERT INTO {ue} (unit_id, entity_id) VALUES ($1, $2)", fid, entity_id)
|
||||
|
||||
# The seed reaches the hub through its own source fact.
|
||||
seed = await _insert_unit(conn, mu, bank_id, "seed obs", "observation", [facts[0]])
|
||||
# Inside the cap: built from the highest ids. Outside: the lowest.
|
||||
inside = await _insert_unit(conn, mu, bank_id, "inside cap", "observation", facts[-3:])
|
||||
outside = await _insert_unit(conn, mu, bank_id, "outside cap", "observation", [facts[1]])
|
||||
|
||||
rows = await backend.ops.expand_observations(
|
||||
conn,
|
||||
mu,
|
||||
ue,
|
||||
ml,
|
||||
[seed],
|
||||
100,
|
||||
per_entity_limit,
|
||||
UpdatedWindow(after=None, before=None, first_param_index=3),
|
||||
)
|
||||
|
||||
scores = {r["id"] for r in rows.entity}
|
||||
assert inside in scores, "observations built from the top-ranked ids must be reachable"
|
||||
assert outside not in scores, "the per-entity cap must exclude ids ranked below it"
|
||||
finally:
|
||||
await memory.delete_bank(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
@@ -167,17 +167,15 @@ async def test_observation_fact_type_in_database(memory, request_context, disabl
|
||||
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check that NO observations exist in memory_units
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, fact_type, context
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
""",
|
||||
# Check that NO observations exist as memory units
|
||||
observations = (
|
||||
await memory.list_memory_units(
|
||||
bank_id,
|
||||
fact_type="observation",
|
||||
limit=1000,
|
||||
request_context=request_context,
|
||||
)
|
||||
)["items"]
|
||||
|
||||
print(f"\n=== Observation Records in memory_units ===")
|
||||
print(f"Found {len(observations)} observation records (should be 0)")
|
||||
@@ -194,6 +192,7 @@ async def test_observation_fact_type_in_database(memory, request_context, disabl
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_entity_mention_counts(memory, request_context):
|
||||
"""
|
||||
Test that entity mention counts are tracked correctly.
|
||||
@@ -235,17 +234,7 @@ async def test_entity_mention_counts(memory, request_context):
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check entity mention counts
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
entities = await conn.fetch(
|
||||
"""
|
||||
SELECT e.id, e.canonical_name, e.mention_count
|
||||
FROM entities e
|
||||
WHERE e.bank_id = $1
|
||||
ORDER BY e.mention_count DESC
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
entities = (await memory.list_entities(bank_id, limit=1000, request_context=request_context))["items"]
|
||||
|
||||
print(f"\n=== Entity Mention Counts Test ===")
|
||||
print(f"Total entities: {len(entities)}")
|
||||
@@ -283,6 +272,7 @@ async def test_entity_mention_counts(memory, request_context):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(1200)
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_entity_mention_ranking(memory, request_context):
|
||||
"""
|
||||
Test that entity mention counts correctly rank entities.
|
||||
@@ -322,17 +312,7 @@ async def test_entity_mention_ranking(memory, request_context):
|
||||
|
||||
# Phase 3: Verify entities are ranked by mention count
|
||||
print("\n=== Phase 3: Check entity ranking ===")
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
all_entities = await conn.fetch(
|
||||
"""
|
||||
SELECT canonical_name, mention_count
|
||||
FROM entities
|
||||
WHERE bank_id = $1
|
||||
ORDER BY mention_count DESC
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
all_entities = (await memory.list_entities(bank_id, limit=1000, request_context=request_context))["items"]
|
||||
|
||||
print(f"\nAll entities by mention count:")
|
||||
for e in all_entities:
|
||||
|
||||
@@ -207,6 +207,7 @@ class TestMarkOperationGracefulOnMissingRow:
|
||||
|
||||
class TestConsolidationCheckpoint:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_consolidation_stops_early_when_op_cancelled(self, memory: MemoryEngine, request_context):
|
||||
"""Consolidation returns 'cancelled' status after the first batch if _check_op_alive is False."""
|
||||
from hindsight_api.config import _get_raw_config
|
||||
|
||||
@@ -161,6 +161,7 @@ async def test_progress_absent_returns_null(api_client, memory: MemoryEngine):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.memory_backend_incompatible
|
||||
async def test_consolidation_records_advancing_progress(memory: MemoryEngine, request_context, monkeypatch):
|
||||
"""A real consolidation run emits scanning → processing_batch → refreshing_mental_models,
|
||||
with processed advancing and a durable snapshot left on the operation row."""
|
||||
|
||||
@@ -37,6 +37,9 @@ except ImportError:
|
||||
pytestmark = [
|
||||
pytest.mark.skipif(not ORACLEDB_AVAILABLE, reason="oracledb not installed"),
|
||||
pytest.mark.skipif(not os.getenv("ORACLE_TEST_DSN"), reason="ORACLE_TEST_DSN not set"),
|
||||
# Talks to the backend's own tables in SQL, down to VECTOR_DISTANCE over the
|
||||
# embedding column.
|
||||
pytest.mark.memory_backend_incompatible,
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -35,6 +35,18 @@ def _bank_id(prefix: str = "oracle") -> str:
|
||||
return f"test-{prefix}-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
|
||||
def _as_utc(value: datetime | None) -> datetime | None:
|
||||
"""Attach UTC to a timestamp oracledb handed back naive.
|
||||
|
||||
Values are written as UTC-aware, but the driver reads some timestamp columns back
|
||||
without a tzinfo, and a naive/aware comparison is silently False rather than an error.
|
||||
Normalise before comparing instants.
|
||||
"""
|
||||
if value is None or value.tzinfo is not None:
|
||||
return value
|
||||
return value.replace(tzinfo=timezone.utc)
|
||||
|
||||
|
||||
async def _safe_cleanup(memory: MemoryEngine, bank_id: str, request_context: RequestContext) -> None:
|
||||
"""Delete a bank, suppressing deadlock/lock errors in test teardown.
|
||||
|
||||
@@ -1154,6 +1166,97 @@ class TestOracleSpecific:
|
||||
run_migrations(oracle_db_url)
|
||||
run_migrations(oracle_db_url)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_observation_update_widens_null_bounds(
|
||||
self, oracle_memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""#3477 on Oracle: an observation with NO occurred interval must inherit its source's.
|
||||
|
||||
Oracle's LEAST/GREATEST return NULL as soon as any argument is NULL, where PostgreSQL
|
||||
ignores NULL arguments. A plain ``LEAST(occurred_start, COALESCE(:n, occurred_start))``
|
||||
therefore evaluates to NULL for an observation that has no interval yet — silently
|
||||
dropping the very dates it was told to inherit — while the identical statement is
|
||||
correct on PostgreSQL. The PG-side coverage lives in test_consolidation_temporal_merge.py;
|
||||
only this test can catch the dialect difference.
|
||||
"""
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.consolidation.consolidator import (
|
||||
_execute_create_action,
|
||||
_execute_update_action,
|
||||
_TemporalBounds,
|
||||
)
|
||||
from hindsight_api.engine.response_models import MemoryFact
|
||||
|
||||
early = datetime(2020, 3, 1, tzinfo=timezone.utc)
|
||||
late = datetime(2021, 7, 4, 12, 30, tzinfo=timezone.utc)
|
||||
bank_id = _bank_id("obsbounds")
|
||||
config = _get_raw_config()
|
||||
previous_observations = config.enable_observations
|
||||
config.enable_observations = False # the test drives consolidation itself
|
||||
try:
|
||||
await oracle_memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
await oracle_memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Dana learned to sail.",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
backend = await oracle_memory._get_backend()
|
||||
async with backend.acquire() as conn:
|
||||
source_id = await conn.fetchval(
|
||||
"SELECT id FROM memory_units WHERE bank_id = $1 AND fact_type <> 'observation'",
|
||||
bank_id,
|
||||
)
|
||||
assert source_id is not None, "retain produced no fact to build an observation from"
|
||||
|
||||
# An observation built from an undated fact: event_date/mentioned_at are stamped at
|
||||
# creation time, occurred_start/occurred_end stay NULL.
|
||||
action = await _execute_create_action(
|
||||
pool=backend,
|
||||
memory_engine=oracle_memory,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[source_id],
|
||||
text="Dana learned to sail.",
|
||||
)
|
||||
assert action == "created"
|
||||
async with backend.acquire() as conn:
|
||||
seeded = await conn.fetchrow(
|
||||
"SELECT id, occurred_start, occurred_end FROM memory_units "
|
||||
"WHERE bank_id = $1 AND fact_type = 'observation'",
|
||||
bank_id,
|
||||
)
|
||||
assert seeded["occurred_start"] is None and seeded["occurred_end"] is None
|
||||
|
||||
await _execute_update_action(
|
||||
pool=backend,
|
||||
memory_engine=oracle_memory,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[source_id],
|
||||
observation_id=str(seeded["id"]),
|
||||
new_text="Dana learned to sail on the Adriatic.",
|
||||
observations=[
|
||||
MemoryFact(
|
||||
id=str(seeded["id"]),
|
||||
text="Dana learned to sail.",
|
||||
fact_type="observation",
|
||||
source_fact_ids=[str(source_id)],
|
||||
tags=[],
|
||||
)
|
||||
],
|
||||
source_bounds=_TemporalBounds(occurred_start=early, occurred_end=late),
|
||||
)
|
||||
|
||||
async with backend.acquire() as conn:
|
||||
updated = await conn.fetchrow(
|
||||
"SELECT occurred_start, occurred_end FROM memory_units WHERE id = $1",
|
||||
seeded["id"],
|
||||
)
|
||||
assert _as_utc(updated["occurred_start"]) == early, "a NULL occurred_start must take the source's date"
|
||||
assert _as_utc(updated["occurred_end"]) == late, "a NULL occurred_end must take the source's date"
|
||||
finally:
|
||||
config.enable_observations = previous_observations
|
||||
await _safe_cleanup(oracle_memory, bank_id, request_context)
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# Tier 6 — Edge Cases & Robustness
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user