Compare commits
311
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
45395e2b51 | ||
|
|
1a200a9cd6 | ||
|
|
6df79e329d | ||
|
|
8170fe880e | ||
|
|
95d77233bf | ||
|
|
4a0a599473 | ||
|
|
bfdc1c5e65 | ||
|
|
37e28fac09 | ||
|
|
3ab39a8d3c | ||
|
|
dbfe83a2ae | ||
|
|
7c2d1848ec | ||
|
|
644e37ac19 | ||
|
|
8ccddd2406 | ||
|
|
a2de0b0dd6 | ||
|
|
421cde6de1 | ||
|
|
8cadecb3a1 | ||
|
|
c2524473e7 | ||
|
|
6166972023 | ||
|
|
24abf373da | ||
|
|
858095f3ba | ||
|
|
27d5ac2832 | ||
|
|
102416c428 | ||
|
|
6de5024aaa | ||
|
|
8bd44716a1 | ||
|
|
dbb0ada924 | ||
|
|
796a9eff91 | ||
|
|
e774617625 | ||
|
|
aa024a5cde | ||
|
|
227441a302 | ||
|
|
831f0efa10 | ||
|
|
bd60c7575c | ||
|
|
cc6fc94468 | ||
|
|
50b7eda2ab | ||
|
|
78c27bfa74 | ||
|
|
b07392c97c | ||
|
|
727d3214cd | ||
|
|
3346363d2f | ||
|
|
f62500193f | ||
|
|
e23e7ca909 | ||
|
|
e68d325830 | ||
|
|
3c8ca47dda | ||
|
|
454069af4d | ||
|
|
9622747759 | ||
|
|
854d0a6283 | ||
|
|
36fd445003 | ||
|
|
568fcea422 | ||
|
|
c1089698b5 | ||
|
|
1a39843e58 | ||
|
|
0fb376f5e7 | ||
|
|
ef37dd6347 | ||
|
|
7ac2c6516c | ||
|
|
febf3528a6 | ||
|
|
b708302187 | ||
|
|
18c45c9d01 | ||
|
|
06f36b8b25 | ||
|
|
a74e5e6b5a | ||
|
|
c01fc12f7e | ||
|
|
df7f45e698 | ||
|
|
dfe74b1de9 | ||
|
|
a933a417cd | ||
|
|
05602730e8 | ||
|
|
b67e813a83 | ||
|
|
30d8993272 | ||
|
|
6766e23f92 | ||
|
|
2e011d279c | ||
|
|
c94935bfa2 | ||
|
|
a3e1e691ae | ||
|
|
ea3ecdaf02 | ||
|
|
687db27cc9 | ||
|
|
1fb1bf080d | ||
|
|
f9c06113a3 | ||
|
|
2a7c496e28 | ||
|
|
e30f8af148 | ||
|
|
4f50034800 | ||
|
|
3b2830c7d8 | ||
|
|
c255d35525 | ||
|
|
7e1145c08a | ||
|
|
75a7c19d6a | ||
|
|
82800ba864 | ||
|
|
2860c9ae16 | ||
|
|
049901802f | ||
|
|
a3d3d42b39 | ||
|
|
01296d8d52 | ||
|
|
e77931fa22 | ||
|
|
23710f4a8f | ||
|
|
b5a324b77b | ||
|
|
adbad877d5 | ||
|
|
d3eff9fba2 | ||
|
|
ddef3d8c6b | ||
|
|
ab61330698 | ||
|
|
505a013812 | ||
|
|
a3797e2014 | ||
|
|
1615456384 | ||
|
|
06c88e0435 | ||
|
|
8aa31edd4c | ||
|
|
4c33a4e55b | ||
|
|
72985b6153 | ||
|
|
8872b9d9ef | ||
|
|
56ed38c8c5 | ||
|
|
e0704a445e | ||
|
|
40871e231e | ||
|
|
56b4271d9f | ||
|
|
1226fd96ad | ||
|
|
087a729d57 | ||
|
|
c5a61db2b8 | ||
|
|
86ec97183b | ||
|
|
221acc8f66 | ||
|
|
f4a3329cea | ||
|
|
655435ea49 | ||
|
|
0db70bb88a | ||
|
|
61f9bc8c77 | ||
|
|
d826d648d8 | ||
|
|
ed34756cdc | ||
|
|
28044b1782 | ||
|
|
bfdcb366d7 | ||
|
|
7683f29004 | ||
|
|
b383c6edd9 | ||
|
|
0eb40a7c1b | ||
|
|
d7a3aa5269 | ||
|
|
01134047d1 | ||
|
|
6b8fc53d79 | ||
|
|
602c9f55e2 | ||
|
|
e1d5db5c59 | ||
|
|
2535db2745 | ||
|
|
613a699e9f | ||
|
|
2834192800 | ||
|
|
1b3925f22f | ||
|
|
1d6d73bce4 | ||
|
|
d695611ada | ||
|
|
a14ce623c5 | ||
|
|
a809547aa8 | ||
|
|
70d98c7a27 | ||
|
|
b1f6bbb8b4 | ||
|
|
24d6c2a43b | ||
|
|
23168ebf68 | ||
|
|
dd75f0dbc8 | ||
|
|
d8dadc0a95 | ||
|
|
401c3cd3fb | ||
|
|
7dffc0459d | ||
|
|
f950e0c11c | ||
|
|
ffd7f94572 | ||
|
|
201f5d7cda | ||
|
|
8a1f0461cf | ||
|
|
670c2be5e4 | ||
|
|
99c7367fc0 | ||
|
|
e4b50f8054 | ||
|
|
bddd22a852 | ||
|
|
c032a74f17 | ||
|
|
a4650f2da5 | ||
|
|
0a2ee84581 | ||
|
|
fa0be9f8a4 | ||
|
|
74021bb317 | ||
|
|
41ad2b55a0 | ||
|
|
ed3d2d09f5 | ||
|
|
b7f267b0a1 | ||
|
|
867b7b4ab6 | ||
|
|
08ce81762c | ||
|
|
324769ac5e | ||
|
|
364ccf17c1 | ||
|
|
ae67665145 | ||
|
|
32dbbb50df | ||
|
|
4bc7013e48 | ||
|
|
537b28128c | ||
|
|
d1dff0d010 | ||
|
|
8cf0dcbf83 | ||
|
|
4280ac3f25 | ||
|
|
8d9000a83d | ||
|
|
df73c7924e | ||
|
|
2a9589fbca | ||
|
|
9345b46336 | ||
|
|
a7337b3abf | ||
|
|
bb81b696f9 | ||
|
|
84330d0453 | ||
|
|
79f63249f6 | ||
|
|
607fbdafdd | ||
|
|
5047bbc473 | ||
|
|
3e99a3f490 | ||
|
|
8a8d2f7abf | ||
|
|
c29c173441 | ||
|
|
5e547f71b2 | ||
|
|
85f6769e4f | ||
|
|
86b80236e3 | ||
|
|
f49b85c0db | ||
|
|
dee9a7b9dd | ||
|
|
ec62acb30f | ||
|
|
ef065f39eb | ||
|
|
6e734e1afa | ||
|
|
ed82801b93 | ||
|
|
9571a341ff | ||
|
|
32b5da60a0 | ||
|
|
4b0d2658a4 | ||
|
|
f367ca81c8 | ||
|
|
18b9c59667 | ||
|
|
a82a20213a | ||
|
|
0cba7f3fbf | ||
|
|
a1ee94ab3e | ||
|
|
fb554664d0 | ||
|
|
bf6b90263b | ||
|
|
68f4a00e8e | ||
|
|
635cf9dd57 | ||
|
|
d425cfc3c5 | ||
|
|
dde133da00 | ||
|
|
e4686b92f0 | ||
|
|
7738021155 | ||
|
|
082213bd97 | ||
|
|
bcae23d9fe | ||
|
|
7cdacc4bf9 | ||
|
|
9704d9182e | ||
|
|
5ad0bffcd5 | ||
|
|
93232213c2 | ||
|
|
6f0a0f1c23 | ||
|
|
779e3140c8 | ||
|
|
9ec73e8455 | ||
|
|
0d2ba56f41 | ||
|
|
cf637799f2 | ||
|
|
74525cc049 | ||
|
|
d7dc8514ca | ||
|
|
1890d2b721 | ||
|
|
374c013689 | ||
|
|
4d9f4ab9ac | ||
|
|
a510b07a81 | ||
|
|
fdb5f47b23 | ||
|
|
129d88c56c | ||
|
|
4b19a0fb69 | ||
|
|
830d8472ca | ||
|
|
617939d822 | ||
|
|
ffa6fbf2a8 | ||
|
|
7a4400e08d | ||
|
|
2de19e578b | ||
|
|
dc41f6a534 | ||
|
|
5123e2a753 | ||
|
|
7e0afff340 | ||
|
|
09c9cecf56 | ||
|
|
78c35253ee | ||
|
|
eadb510eb3 | ||
|
|
691cb5394b | ||
|
|
a401b97eb7 | ||
|
|
aa4c1bbaf3 | ||
|
|
99525144b2 | ||
|
|
ded52e8de6 | ||
|
|
cc3ba4a37c | ||
|
|
7ef64f14ca | ||
|
|
16f807697d | ||
|
|
486c3a8b3b | ||
|
|
fbbc7a5e4c | ||
|
|
d7d41e76c2 | ||
|
|
262d4894f2 | ||
|
|
ec49175fa3 | ||
|
|
488f428009 | ||
|
|
d1ef9da95e | ||
|
|
30acca6fd9 | ||
|
|
2538708308 | ||
|
|
9e7aff6bd4 | ||
|
|
a908cdc974 | ||
|
|
4cd260b691 | ||
|
|
0c17e9acfd | ||
|
|
beca4b42f3 | ||
|
|
6e9b741b02 | ||
|
|
4a1b2f39c1 | ||
|
|
28ec22c3dc | ||
|
|
d802f91488 | ||
|
|
cb04cb79d9 | ||
|
|
dabbf9ff49 | ||
|
|
6348f42451 | ||
|
|
41a2ccabf8 | ||
|
|
eaf3048f2c | ||
|
|
9d95149852 | ||
|
|
ac3ab2b54c | ||
|
|
cb037290bb | ||
|
|
2582b45a16 | ||
|
|
0be157eeb5 | ||
|
|
90cb145aa6 | ||
|
|
7bd11bedf6 | ||
|
|
c3b2b1543a | ||
|
|
2743d061f7 | ||
|
|
daf2348bcd | ||
|
|
46dd2dfd94 | ||
|
|
878ef957f7 | ||
|
|
00d327a049 | ||
|
|
31d1e1729e | ||
|
|
592f01bba6 | ||
|
|
da05ee7215 | ||
|
|
19d23921fb | ||
|
|
e1e1a5e02b | ||
|
|
44b34c891c | ||
|
|
67ae2a41d4 | ||
|
|
2e5186a6fc | ||
|
|
9a20180415 | ||
|
|
f61ae2a185 | ||
|
|
dfd7cb52d4 | ||
|
|
5300d401b0 | ||
|
|
0b6bf53bef | ||
|
|
203ddfdd6c | ||
|
|
dcf5588e6c | ||
|
|
3d6c2ba8b0 | ||
|
|
80046797f7 | ||
|
|
db7dabcebd | ||
|
|
b83bb87ddd | ||
|
|
15ec55b703 | ||
|
|
f2596e1fe9 | ||
|
|
21c71f7bb8 | ||
|
|
86b686cd72 | ||
|
|
248c40e670 | ||
|
|
d18a9452ad | ||
|
|
806fbcd41c | ||
|
|
a75c3c85ad | ||
|
|
0db9f3da19 | ||
|
|
8940710c72 | ||
|
|
6252643de0 | ||
|
|
3fce309c0d | ||
|
|
3fc361aabd |
@@ -73,6 +73,11 @@ results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
```
|
||||
|
||||
### API Layer & Data Access
|
||||
- **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).
|
||||
|
||||
### Branch Hygiene
|
||||
- **Always start new feature branches from `origin/main`** — rebase to ensure a clean base.
|
||||
- **Only include commits relevant to the PR/branch/feature** — no unrelated changes. If the branch contains commits that don't belong, they must be removed before merging.
|
||||
@@ -135,6 +140,13 @@ For each new or significantly changed function/endpoint/class:
|
||||
|
||||
Flag any new logic that lacks test coverage.
|
||||
|
||||
**LLM-behaviour changes need a real-LLM judge test, not MockLLM.** If the change alters how the model interprets a prompt — fact/observation extraction, `fact_type` (world/experience) classification, speaker attribution, instruction-following, prompt wording — there MUST be a test marked `pytest.mark.hs_llm_core` that runs the real pipeline and asserts via `tests.llm_judge.assert_meets_criteria` (not string/enum matching). Flag these as findings:
|
||||
- A prompt/classification change verified only by MockLLM or string assertions (MockLLM echoes input — such tests pass spuriously). **Should fix.**
|
||||
- A test that hard-asserts `fact_type == "world"/"experience"` (or other model-decided output) instead of judging it — non-deterministic, will flake across providers/runs. **Should fix** (move the classification check into the judge `criteria`; keep only genuinely deterministic structural asserts direct).
|
||||
- Deterministic mechanics (prompt assembly, suppression/branching logic) that are covered *only* by a slow LLM test — these should also have fast non-LLM unit tests. **Note.**
|
||||
|
||||
See CLAUDE.md → Key Conventions → Testing for the full pattern.
|
||||
|
||||
### 7. Check API consistency
|
||||
|
||||
If any files in `hindsight-api-slim/hindsight_api/api/` were changed:
|
||||
@@ -142,6 +154,12 @@ If any files in `hindsight-api-slim/hindsight_api/api/` were changed:
|
||||
- Were the client SDKs regenerated? (`./scripts/generate-clients.sh`)
|
||||
- Were the control plane proxy routes updated? (`hindsight-control-plane/src/app/api/`)
|
||||
|
||||
### 7b. Check API-layer data-access boundary
|
||||
|
||||
For each changed handler in `hindsight-api-slim/hindsight_api/api/` (e.g. `http.py`, `mcp.py`):
|
||||
- **Flag any direct DB access in the handler** — `acquire_with_retry`, `conn.fetch` / `fetchrow` / `execute`, raw SQL strings, or `fq_table(...)`. These are a **must fix**: the query must be moved into a `MemoryEngine` method that returns a typed model, and the handler must call that method.
|
||||
- **Verify authentication is enforced in the engine** — the handler must delegate to an engine method that authenticates via `request_context` (`_authenticate_tenant`, typically through `get_bank_profile`). A handler that reads/writes tenant-scoped data without an engine method enforcing auth is a **must fix** (tenant data could leak across schemas).
|
||||
|
||||
### 8. Check code comments
|
||||
|
||||
For each non-trivial change:
|
||||
@@ -154,7 +172,8 @@ For each non-trivial change:
|
||||
If any files in `hindsight-integrations/` were added or changed, verify:
|
||||
- **Tests exist** — the integration must have tests that simulate/exercise the external framework (not just pure unit tests of helpers). Check for a `tests/` directory with meaningful test files.
|
||||
- **CI job exists** — check `.github/workflows/test.yml` for a corresponding `test-<name>-integration` job. If missing, flag it.
|
||||
- **Release process** — check that the integration name is in the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh`. If missing, flag it.
|
||||
- **Release process** — check that the integration name is in the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh` AND in the `INTEGRATIONS` dict in `hindsight-dev/hindsight_dev/generate_changelog.py` (the changelog generator keeps its own list; a release fails at the changelog step if the name is missing there). If either is missing, flag it.
|
||||
- **Docs gallery + sidebar entry** — the integration must have an entry in `hindsight-docs/src/data/integrations.json`. This file is the **single source of truth** that drives both the integrations gallery and the docs sidebar (the sidebar category is injected from it at render time across all docs versions). The entry needs an internal `/sdks/integrations/<slug>` `link` and a matching page at `hindsight-docs/docs-integrations/<slug>.md(x)`. The `hindsight-docs/scripts/check-integrations.mjs` build step enforces both directions — forward: every internal JSON entry has a doc page; reverse: every released tag (`integrations/<name>/vX.Y.Z`) appears in the JSON (private infra like `cloudflare-oauth-proxy` is in the script's `EXCLUDED` set). Flag any integration that is released (or being released) but missing from `integrations.json`, and any JSON entry without a doc page. Do **not** hand-edit `versioned_sidebars/*.json` to add integration links — they are positional placeholders filled from the JSON.
|
||||
- **Code standards** — the integration code must follow all Python style rules (type hints, no raw dicts, no tuple returns, etc.).
|
||||
|
||||
### 10. Check MCP tool registration completeness
|
||||
@@ -166,7 +185,14 @@ If any new MCP tools were added or existing tools renamed in `hindsight-api-slim
|
||||
- **`MCP_TOOL_GROUPS`** in `hindsight-control-plane/src/components/bank-config-view.tsx` — must include the new tool in the appropriate group for the UI tool selector
|
||||
- **Tool count assertions** in tests (e.g., `test_mcp_tools.py`) — must be updated to reflect the new count
|
||||
|
||||
### 11. Review against other coding standards
|
||||
### 11. Check backup/restore table coverage
|
||||
|
||||
If a migration adds a new PostgreSQL table (look for `CREATE TABLE` / `op.create_table` in `hindsight-api-slim/hindsight_api/alembic/versions/`):
|
||||
- **`BACKUP_TABLES`** in `hindsight-api-slim/hindsight_api/admin/cli.py` — must include the new table, placed after any table it references via foreign key (parents before children). A missing entry is silent data loss: the table is never backed up, and restore's `TRUNCATE banks CASCADE` wipes any FK-to-banks child (e.g. `mental_models`, `directives`) on restore even though it was never saved.
|
||||
- The guard test `test_backup_tables_covers_entire_schema` in `tests/test_admin_backup_restore.py` enforces this — flag it as a **must fix** if a new table is absent from `BACKUP_TABLES`.
|
||||
- Oracle-only tables (e.g. `observation_sources`) are intentionally excluded — admin backup/restore is PostgreSQL-only.
|
||||
|
||||
### 12. Review against other coding standards
|
||||
|
||||
Check the diff for violations of the standards listed above:
|
||||
- Python files at project root (not allowed)
|
||||
@@ -178,7 +204,7 @@ Check the diff for violations of the standards listed above:
|
||||
- Premature abstractions or speculative helpers
|
||||
- Backwards-compatibility hacks (unused vars, re-exports, "removed" comments)
|
||||
|
||||
### 12. Report findings
|
||||
### 13. Report findings
|
||||
|
||||
Present a clear summary organized by severity:
|
||||
|
||||
@@ -189,7 +215,11 @@ Present a clear summary organized by severity:
|
||||
- Raw dict usage for structured data (including internal code)
|
||||
- Multi-item tuple returns (including internal code)
|
||||
- Missing tests for new endpoints
|
||||
- Direct DB access (raw SQL / `acquire_with_retry` / `fq_table`) in an `api/` handler instead of a `MemoryEngine` method
|
||||
- Tenant-scoped data accessed without authentication enforced in the engine (`_authenticate_tenant` / `get_bank_profile`)
|
||||
- New integration missing tests, CI job, or release-integration.sh entry
|
||||
- Released/added integration missing from `hindsight-docs/src/data/integrations.json`, or a JSON entry with no `docs-integrations/<slug>` page (fails the docs build via `check-integrations.mjs`)
|
||||
- New PostgreSQL table missing from `BACKUP_TABLES` in `admin/cli.py` (silent data loss on restore)
|
||||
|
||||
**Should fix** — issues that hurt code quality:
|
||||
- Dead code / unused imports missed by linter
|
||||
|
||||
+42
-2
@@ -7,6 +7,8 @@ HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# Reasoning effort for providers/models that support it. Examples: low, medium, high, xhigh.
|
||||
# HINDSIGHT_API_LLM_REASONING_EFFORT=low
|
||||
|
||||
# Example: Anthropic Claude configuration
|
||||
# HINDSIGHT_API_LLM_PROVIDER=anthropic
|
||||
@@ -23,7 +25,7 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# Example: MiniMax configuration (1M context window)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=minimax
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-minimax-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=MiniMax-M2.7
|
||||
# HINDSIGHT_API_LLM_MODEL=MiniMax-M3 # or MiniMax-M2.7 for the previous generation
|
||||
|
||||
# Example: DeepSeek configuration (https://api.deepseek.com)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=deepseek
|
||||
@@ -64,11 +66,37 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# For Azure PostgreSQL with DiskANN:
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale # Auto-detects pg_diskann on Azure
|
||||
|
||||
# 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
|
||||
# Native backend dictionary (only used by HINDSIGHT_API_TEXT_SEARCH_EXTENSION=native)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE=english
|
||||
# ParadeDB pg_search tokenizer (only used when creating pg_search BM25 indexes).
|
||||
# Empty uses ParadeDB's default tokenizer: unicode_words.
|
||||
# Supported values: unicode_words, simple, whitespace, literal, literal_normalized,
|
||||
# chinese_compatible, icu, jieba, source_code,
|
||||
# chinese_lindera/lindera(chinese), japanese_lindera/lindera(japanese),
|
||||
# korean_lindera/lindera(korean), ngram(min,max), edge_ngram(min,max)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER=
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default), "tei", "openai", "cohere", "google", "openrouter", "litellm", or "litellm-sdk"
|
||||
# Provider: "local" (default), "onnx", "tei", "openai", "cohere", "google", "openrouter", "zeroentropy", "litellm", or "litellm-sdk"
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
# For local provider:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
|
||||
# For ONNX provider (local CPU embeddings without an Ollama/TEI sidecar):
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=onnx
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_ID=intfloat/multilingual-e5-small
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_FILE=onnx/model.onnx
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_DIMENSIONS=384
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_MAX_TOKENS=512
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_POOLING=mean
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_NORMALIZE=true
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_QUERY_PREFIX="query: "
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_PASSAGE_PREFIX="passage: "
|
||||
# Optional for local model paths or pre-downloaded artifacts:
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_PATH=/models/multilingual-e5-small/onnx/model.onnx
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_TOKENIZER_NAME_OR_PATH=/models/multilingual-e5-small
|
||||
# Optional for China network / restricted HF access:
|
||||
# HF_ENDPOINT=https://hf-mirror.com
|
||||
# For TEI provider:
|
||||
@@ -77,6 +105,13 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=sk-xxxx
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL=text-embedding-3-small
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL=https://api.openai.com/v1
|
||||
# For ZeroEntropy zembed-1:
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=zeroentropy
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_API_KEY=ze-xxxx
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_MODEL=zembed-1
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_DIMENSIONS=1280
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT=float
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_LATENCY=fast
|
||||
#
|
||||
# IMPORTANT: Embedding keys require provider-specific names:
|
||||
# HINDSIGHT_API_EMBEDDINGS_{PROVIDER}_{PARAMETER}
|
||||
@@ -115,6 +150,11 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Dataplane API URL - where the CP proxies requests to
|
||||
# HINDSIGHT_CP_DATAPLANE_API_URL=http://localhost:8888
|
||||
|
||||
# Optional: Bearer token the CP sends as `Authorization: Bearer <key>` to the
|
||||
# dataplane API. Required when the API service is auth-protected; omit for a
|
||||
# public/unauthenticated API.
|
||||
# HINDSIGHT_CP_DATAPLANE_API_KEY=your-dataplane-bearer-token
|
||||
|
||||
# Optional: Require a shared access key to view the Control Plane UI.
|
||||
# When set, visitors see a login page and must enter the key before
|
||||
# accessing the dashboard or any /api/* routes (except /api/health).
|
||||
|
||||
@@ -22,6 +22,8 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0 # fetch tags so check-released-integrations can see them
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: 20
|
||||
|
||||
@@ -23,7 +23,9 @@ on:
|
||||
- retain
|
||||
- recall
|
||||
- recall-with-observations
|
||||
- recall-temporal
|
||||
- consolidation
|
||||
- graph-maintenance
|
||||
default: ""
|
||||
locomo_conversations:
|
||||
description: "LoComo conversation IDs (space-separated). Blank = curated set (conv-26 conv-30 conv-43)."
|
||||
@@ -33,6 +35,18 @@ on:
|
||||
description: "Skip LoComo job"
|
||||
type: boolean
|
||||
default: false
|
||||
obs_skip:
|
||||
description: "Skip observation-dedup benchmark job"
|
||||
type: boolean
|
||||
default: false
|
||||
obs_dataset:
|
||||
description: "Obs benchmark dataset substring (blank = English hermes transcript)."
|
||||
type: string
|
||||
default: ""
|
||||
obs_fraction:
|
||||
description: "Obs benchmark fraction (0-1] of each document to run."
|
||||
type: string
|
||||
default: "1.0"
|
||||
ref:
|
||||
description: "Git ref to test (branch, tag, or SHA). Defaults to main."
|
||||
type: string
|
||||
@@ -198,3 +212,86 @@ jobs:
|
||||
PERF_DASHBOARD_TOKEN: ${{ secrets.PERF_DASHBOARD_TOKEN }}
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: ./scripts/benchmarks/publish-locomo-results.sh hindsight-dev/benchmarks/locomo/results/benchmark_results.json
|
||||
|
||||
obs:
|
||||
# Observation-dedup quality benchmark: ingests a transcript, drains consolidation
|
||||
# (serial SyncTaskBackend + embedded pg0 — no external DB / worker), and reports the
|
||||
# near-duplicate observation rate. Real LLM via VertexAI, mirroring the LoComo job.
|
||||
if: inputs.obs_skip != true
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_ENABLE_OBSERVATIONS: "true"
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.ref }}
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv run --frozen --all-extras --index-strategy unsafe-best-match python -c "
|
||||
from sentence_transformers import SentenceTransformer
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Model downloaded successfully')
|
||||
"
|
||||
|
||||
- name: Install hindsight-dev dependencies
|
||||
run: |
|
||||
cd hindsight-dev && uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
|
||||
- name: Run obs benchmark
|
||||
# Default to the English hermes transcript at full fraction — a clean, deterministic
|
||||
# consolidation-dedup signal (the Chinese variant adds a cross-lingual embedding
|
||||
# confound). Override dataset/fraction via workflow_dispatch.
|
||||
run: |
|
||||
DATASET="${{ inputs.obs_dataset }}"
|
||||
if [ -z "$DATASET" ]; then DATASET="hermes_session_2026-05-15_en"; fi
|
||||
FRACTION="${{ inputs.obs_fraction }}"
|
||||
if [ -z "$FRACTION" ]; then FRACTION="1.0"; fi
|
||||
cd hindsight-dev
|
||||
uv run python -m benchmarks.obs.obs_benchmark \
|
||||
--dataset "$DATASET" --fraction "$FRACTION" --wipe-bank --output obs-results.json
|
||||
|
||||
- name: Upload obs results
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: obs-results-${{ github.sha }}
|
||||
path: hindsight-dev/obs-results.json
|
||||
retention-days: 90
|
||||
|
||||
- name: Publish obs to dashboard
|
||||
if: success() && (github.event_name == 'schedule' || github.event_name == 'push' || github.event_name == 'workflow_dispatch')
|
||||
env:
|
||||
PERF_DASHBOARD_TOKEN: ${{ secrets.PERF_DASHBOARD_TOKEN }}
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: ./scripts/benchmarks/publish-obs-results.sh hindsight-dev/obs-results.json
|
||||
|
||||
+641
-11
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,115 @@
|
||||
name: Windows Smoke Test
|
||||
|
||||
# Daily smoke test that installs the API on Windows and runs the Python client
|
||||
# integration tests against a live server. Windows is only exercised by the
|
||||
# hindsight-embed jobs in test.yml on PRs; this catches Windows-specific
|
||||
# regressions in the API server + client path (e.g. process spawning, console
|
||||
# subsystem / ConPTY behaviour, see #1885) that the Linux client jobs miss.
|
||||
on:
|
||||
schedule:
|
||||
# 06:00 UTC daily.
|
||||
- cron: "0 6 * * *"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
windows-client-smoke:
|
||||
# Don't run on forks: the job needs the org's Vertex AI credentials.
|
||||
if: github.repository == 'vectorize-io/hindsight'
|
||||
runs-on: windows-latest
|
||||
timeout-minutes: 45
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: ${{ github.workspace }}/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_URL: http://localhost:8888
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
# Force UTF-8 I/O so the API/CLI's ✓/box-drawing output doesn't crash the
|
||||
# default Windows cp1252 codec (matches test-embed-windows in test.yml).
|
||||
PYTHONIOENCODING: utf-8
|
||||
PYTHONUTF8: "1"
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Setup GCP credentials
|
||||
shell: bash
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Install API dependencies (all extras - local-ml + embedded pg0)
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
|
||||
- name: Install Python client test dependencies
|
||||
working-directory: ./hindsight-clients/python
|
||||
run: uv sync --frozen --extra test --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
# `uv run` re-syncs the project env to its default (no-extras) state before
|
||||
# running, which drops sentence-transformers / pg0. Pass --all-extras on
|
||||
# every `uv run` so the local-ml + embedded-db deps stay installed (this is
|
||||
# the same reason hindsight-embed launches the daemon with `--extra all`).
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv run --all-extras python -c "from sentence_transformers import SentenceTransformer, CrossEncoder; SentenceTransformer('BAAI/bge-small-en-v1.5'); CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); print('Models downloaded')"
|
||||
|
||||
# Start the server and run the client tests in a SINGLE step. On Windows
|
||||
# runners a process backgrounded with `&` in one step is not reliably kept
|
||||
# alive for later steps (unlike Linux, where it reparents to init), so the
|
||||
# server must live in the same shell that runs pytest.
|
||||
- name: Start API server and run Python client tests
|
||||
shell: bash
|
||||
run: |
|
||||
# Config is read straight from the environment (job-level env + the
|
||||
# PROJECT_ID exported to GITHUB_ENV above), so no .env file is needed.
|
||||
# Embedded pg0 is the default when HINDSIGHT_API_DATABASE_URL is unset.
|
||||
( cd hindsight-api-slim && uv run --all-extras hindsight-api --port 8888 ) > "$RUNNER_TEMP/api-server.log" 2>&1 &
|
||||
server_pid=$!
|
||||
echo "Waiting for API server to be ready (pid $server_pid)..."
|
||||
# pg0 unpacks Postgres + runs initdb on first boot, which is slow on a
|
||||
# cold Windows runner — give it a generous budget before failing.
|
||||
ready=false
|
||||
for i in $(seq 1 300); do
|
||||
if curl -sf http://localhost:8888/health > /dev/null 2>&1; then
|
||||
echo "API server is ready after ${i}s"
|
||||
ready=true
|
||||
break
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
if [ "$ready" != true ]; then
|
||||
echo "API server failed to start after 300s"
|
||||
cat "$RUNNER_TEMP/api-server.log"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
cd hindsight-clients/python && uv run --extra test pytest tests -v
|
||||
|
||||
- name: Show API server logs
|
||||
if: always()
|
||||
shell: bash
|
||||
run: cat "$RUNNER_TEMP/api-server.log" || echo "No API server log found"
|
||||
+4
-1
@@ -54,7 +54,10 @@ hindsight-clients/rust/target
|
||||
!.claude/skills/
|
||||
whats-next.md
|
||||
TASK.md
|
||||
# Parked / draft integrations that aren't ready to ship
|
||||
hindsight-integrations/_drafts/
|
||||
# Changelog is now tracked in hindsight-docs/src/pages/changelog.md
|
||||
# CHANGELOG.md
|
||||
|
||||
blog-post*
|
||||
blog-post*
|
||||
.worktrees/
|
||||
|
||||
@@ -220,6 +220,30 @@ migration file dispatches through `run_for_dialect`, which calls either
|
||||
|
||||
**MANDATORY: Run `/code-review` before pushing code or creating a pull request.** Do not push or create a PR until all "must fix" issues are resolved.
|
||||
|
||||
### Testing
|
||||
|
||||
Most tests are deterministic (MockLLM, pure functions) — assert directly.
|
||||
|
||||
**Tests that verify LLM behaviour use a real LLM + an LLM-as-judge.** When the thing under test is *how the model interprets a prompt* (classification, attribution, dimension preservation, instruction-following), MockLLM can't simulate it and exact string/enum asserts flake across providers and runs. Use this pattern instead:
|
||||
|
||||
1. Mark the test module `pytestmark = pytest.mark.hs_llm_core` (single-provider; CI runs it in the core-LLM job). Use `hs_llm_mat` only for provider-matrix acceptance tests.
|
||||
2. Call the real pipeline (`LLMConfig.from_env()`, `_get_raw_config()`), e.g. `extract_facts_from_text(...)`.
|
||||
3. Assert with the judge, not string matching:
|
||||
```python
|
||||
from tests.llm_judge import assert_meets_criteria
|
||||
facts_summary = "\n".join(f"- [{f.fact_type}] {f.fact}" for f in facts)
|
||||
await assert_meets_criteria(
|
||||
response=facts_summary,
|
||||
criteria="The first-person user statements are classified 'world' and attributed to the user, not the agent.",
|
||||
context="What the input said and who was speaking.",
|
||||
)
|
||||
```
|
||||
|
||||
Rules of thumb:
|
||||
- **Judge anything non-deterministic** — including `fact_type` classification and speaker attribution. Do NOT hard-assert `fact_type == "..."`; pass a `[fact_type] fact` summary to the judge instead. Structural facts that ARE deterministic (counts, presence of a field, that a substring was injected into a prompt) stay as direct asserts in fast unit tests.
|
||||
- **Split the test surface**: cover the deterministic mechanics (prompt assembly, suppression logic) with fast non-LLM unit tests, and the model-following behaviour with one `hs_llm_core` judge test. (Example pair: `test_narrator_resolution.py` + `test_narrator_context_override.py`.)
|
||||
- The judge model is independent of the test provider (defaults to Gemini); never judge with the same call you're testing.
|
||||
|
||||
### Memory Banks
|
||||
- Each bank is an isolated memory store (like a "brain" for one user/agent)
|
||||
- Banks have dispositions (skepticism, literalism, empathy traits 1-5) affecting reflect
|
||||
|
||||
+25
-2
@@ -9,13 +9,36 @@ Thanks for your interest in contributing to Hindsight!
|
||||
git clone [email protected]:vectorize-io/hindsight.git
|
||||
cd hindsight
|
||||
```
|
||||
2. Set up your environment:
|
||||
|
||||
2. Bootstrap your dev environment in one shot:
|
||||
```bash
|
||||
./scripts/dev/setup.sh
|
||||
```
|
||||
This is idempotent (safe to re-run) and gets you ready to develop, including
|
||||
offline. It:
|
||||
- installs the required toolchains if missing (uv/Python, Node/npm, Rust/cargo),
|
||||
- creates `.env` from `.env.example` (remember to add your LLM API key),
|
||||
- configures git hooks,
|
||||
- installs all Python and Node workspace dependencies,
|
||||
- pre-downloads the local ML models + tokenizer so the API runs offline,
|
||||
- builds the TypeScript SDK and the Rust CLI.
|
||||
|
||||
Useful flags: `--skip-build` (deps only), `--skip-models` (skip ML model
|
||||
download), `--with-docs` (also build the docs site), `--force` (rebuild
|
||||
artifacts). Docker image builds are out of scope. Run
|
||||
`./scripts/dev/setup.sh --help` for details.
|
||||
|
||||
### Manual setup
|
||||
|
||||
If you'd rather set things up by hand instead of running the script above:
|
||||
|
||||
1. Set up your environment:
|
||||
```bash
|
||||
cp .env.example .env
|
||||
```
|
||||
Edit the .env to add LLM API key and config as required
|
||||
|
||||
3. Install dependencies:
|
||||
2. Install dependencies:
|
||||
```bash
|
||||
# Python dependencies
|
||||
uv sync --directory hindsight-api/
|
||||
|
||||
@@ -62,9 +62,9 @@ If you need more control over how and when your agent stores and recalls memorie
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
|
||||
docker run --rm -it --pull always -p 8888:8888 -p 9999:9999 \
|
||||
docker run -it --pull always --name hindsight --restart unless-stopped -p 8888:8888 -p 9999:9999 \
|
||||
-e HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY \
|
||||
-v $HOME/.hindsight-docker:/home/hindsight/.pg0 \
|
||||
-v hindsight-data:/home/hindsight/.pg0 \
|
||||
ghcr.io/vectorize-io/hindsight:latest
|
||||
```
|
||||
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
# Hindsight with Claude Code (Claude Pro/Max subscription)
|
||||
|
||||
Run Hindsight inside Docker using the `claude-code` LLM provider, backed by
|
||||
your host machine's Claude Pro or Max subscription credentials.
|
||||
|
||||
The standalone Hindsight Docker image ships `claude-agent-sdk` but does **not**
|
||||
bundle the host `claude` CLI binary or any Claude credentials. This Compose
|
||||
file bind-mounts the host's CLI install and credentials into the container so
|
||||
the `claude-code` provider works without an API key.
|
||||
|
||||
## When to use this
|
||||
|
||||
- You have an active Claude Pro or Max subscription and want to use it for
|
||||
Hindsight without paying separate Anthropic API costs.
|
||||
- You want a one-command `docker compose up` instead of a long `docker run`
|
||||
invocation with many flags.
|
||||
- You are running on **Linux/amd64** — macOS Docker Desktop and Windows host
|
||||
paths differ and are not yet covered (please open an issue if you'd like to
|
||||
contribute a verified recipe for either).
|
||||
|
||||
> **Personal-use only.** Anthropic's
|
||||
> [Agent SDK documentation](https://docs.claude.com/en/api/agent-sdk/overview)
|
||||
> states that third-party developers should not offer claude.ai login or rate
|
||||
> limits for their products. Hindsight does **not** perform any login on your
|
||||
> behalf — it uses credentials you've already authenticated via
|
||||
> `claude auth login`. In January 2026, Anthropic
|
||||
> [enforced restrictions](https://paddo.dev/blog/anthropic-walled-garden-crackdown/)
|
||||
> against tools that spoofed the Claude Code client identity; Hindsight uses
|
||||
> the official Claude Agent SDK instead.
|
||||
>
|
||||
> Do not deploy this configuration to shared environments or production. For
|
||||
> that, use the `anthropic` provider with an API key from the
|
||||
> [Anthropic Console](https://console.anthropic.com/). Usage counts against
|
||||
> your Claude Pro/Max subscription limits.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Host has `claude` CLI installed (e.g., `npm install -g @anthropics/claude-code`)
|
||||
and `claude auth login` has been run successfully.
|
||||
- `~/.claude.json` and `~/.claude/.credentials.json` exist on the host.
|
||||
- Host `claude` CLI version is **2.1.128 or newer** — the version bundled with
|
||||
`claude-agent-sdk` 0.5.x has a protocol incompatibility in containers, so
|
||||
the recipe overrides it with the host binary.
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
# Set your host UID/GID (defaults to 1000:1000 if unset)
|
||||
export HOST_UID=$(id -u)
|
||||
export HOST_GID=$(id -g)
|
||||
|
||||
docker compose -f docker/docker-compose/claude-code/docker-compose.yaml up -d
|
||||
```
|
||||
|
||||
- API: http://localhost:8888
|
||||
- Control Plane: http://localhost:9999
|
||||
|
||||
## Post-setup (one-time)
|
||||
|
||||
After the container starts for the first time, run these commands to fix
|
||||
permissions and symlink the host `claude` binary into `$PATH`:
|
||||
|
||||
```bash
|
||||
# Make ~/.claude writable by your UID (the CLI writes session/project state)
|
||||
docker exec --user 0:0 hindsight-claude-code chown $(id -u):$(id -g) /home/hindsight/.claude
|
||||
docker exec --user 0:0 hindsight-claude-code chmod 755 /home/hindsight/.claude
|
||||
|
||||
# Symlink the host claude binary into PATH
|
||||
docker exec --user 0:0 hindsight-claude-code \
|
||||
ln -sf /home/hindsight/.local/share/claude/versions/2.1.128 /usr/local/bin/claude
|
||||
```
|
||||
|
||||
If you set `CLAUDE_CLI_VERSION` to a version other than `2.1.128`, update the
|
||||
symlink path accordingly.
|
||||
|
||||
## Notes on the bind-mount surface (every flag is load-bearing)
|
||||
|
||||
- **Host `claude` binary required** — the image ships only `claude-agent-sdk`,
|
||||
not the CLI itself.
|
||||
- **SDK bundled-binary override** — the override of
|
||||
`claude_agent_sdk/_bundled/claude` works around a protocol issue in the
|
||||
bundled v2.1.121 binary inside containers. Once `claude-agent-sdk` ships
|
||||
with v2.1.128+ this override can be dropped. Set `CLAUDE_CLI_VERSION` to
|
||||
match your installed version.
|
||||
- **Single-file credential mounts** — credentials are mounted as individual
|
||||
`:ro` files rather than a whole-directory `:ro` mount of `~/.claude`,
|
||||
because the CLI writes session/project state at runtime and a read-only
|
||||
directory mount silently breaks it.
|
||||
- **`--user` / `user:`** — the `user: ${HOST_UID}:${HOST_GID}` pattern
|
||||
requires `chmod 755 /home/hindsight`, which is built into the image since
|
||||
v0.6.0 (see [#1481](https://github.com/vectorize-io/hindsight/issues/1481)).
|
||||
- **`~/.hindsight-docker` data directory** — the pg0 data bind mount must be
|
||||
writable by your host UID (see
|
||||
[#1483](https://github.com/vectorize-io/hindsight/issues/1483)).
|
||||
- **Verified** on `linux/amd64` against `ghcr.io/vectorize-io/hindsight:latest`
|
||||
v0.5.6+.
|
||||
|
||||
## Using a different Claude CLI version
|
||||
|
||||
If your host has a `claude` version other than 2.1.128, set
|
||||
`CLAUDE_CLI_VERSION` before starting:
|
||||
|
||||
```bash
|
||||
export CLAUDE_CLI_VERSION=2.2.0
|
||||
docker compose -f docker/docker-compose/claude-code/docker-compose.yaml up -d
|
||||
```
|
||||
|
||||
Then update the post-setup symlink to match:
|
||||
|
||||
```bash
|
||||
docker exec --user 0:0 hindsight-claude-code \
|
||||
ln -sf /home/hindsight/.local/share/claude/versions/2.2.0 /usr/local/bin/claude
|
||||
```
|
||||
@@ -0,0 +1,44 @@
|
||||
name: hindsight-claude-code
|
||||
# Run Hindsight with the claude-code LLM provider, using your host machine's
|
||||
# Claude Pro/Max subscription credentials. Linux/amd64 only for now.
|
||||
#
|
||||
# Quick start:
|
||||
# docker compose -f docker/docker-compose/claude-code/docker-compose.yaml up -d
|
||||
#
|
||||
# See README.md for prerequisites, post-setup steps, and important caveats.
|
||||
|
||||
services:
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:latest
|
||||
container_name: hindsight-claude-code
|
||||
user: "${HOST_UID:-1000}:${HOST_GID:-1000}"
|
||||
ports:
|
||||
- "127.0.0.1:8888:8888"
|
||||
- "127.0.0.1:9999:9999"
|
||||
environment:
|
||||
HOME: /home/hindsight
|
||||
USER: hindsight
|
||||
LOGNAME: hindsight
|
||||
PATH: /usr/local/bin:/usr/bin:/bin:/app/api/.venv/bin
|
||||
HINDSIGHT_API_LLM_PROVIDER: claude-code
|
||||
volumes:
|
||||
# ── Persistent data ────────────────────────────────────────────
|
||||
# Writable pg0 data directory. Must be writable by HOST_UID.
|
||||
- ${HOME:-.}/.hindsight-docker:/home/hindsight/.pg0
|
||||
|
||||
# ── Claude credentials (read-only, single-file mounts) ────────
|
||||
# A whole-directory :ro mount of ~/.claude silently breaks the
|
||||
# CLI, which writes session/project state at runtime — so we
|
||||
# mount only the two credential files.
|
||||
- ${HOME}/.claude/.credentials.json:/home/hindsight/.claude/.credentials.json:ro
|
||||
- ${HOME}/.claude.json:/home/hindsight/.claude.json:ro
|
||||
|
||||
# ── Claude CLI install (read-only) ─────────────────────────────
|
||||
- ${HOME}/.local/share/claude:/home/hindsight/.local/share/claude:ro
|
||||
|
||||
# ── SDK bundled-binary override ────────────────────────────────
|
||||
# The claude-agent-sdk 0.5.x image bundles v2.1.121 which has a
|
||||
# protocol incompatibility in containers. Override it with the
|
||||
# host's v2.1.128+ binary. Drop this mount once claude-agent-sdk
|
||||
# ships with v2.1.128+.
|
||||
- ${HOME}/.local/share/claude/versions/${CLAUDE_CLI_VERSION:-2.1.128}:/app/api/.venv/lib/python3.11/site-packages/claude_agent_sdk/_bundled/claude:ro
|
||||
@@ -0,0 +1,103 @@
|
||||
# Hindsight with a local llama.cpp server sidecar
|
||||
|
||||
Example Docker Compose setup that runs Hindsight against a **local
|
||||
llama.cpp server**, fully offline, with no external API key required.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌────────────┐ HTTP /v1/chat/completions ┌──────────────────────────────┐
|
||||
│ hindsight │ ──────────────────────────▶ │ llama.cpp server (sidecar) │
|
||||
│ (API + CP) │ │ ghcr.io/ggml-org/llama.cpp │
|
||||
└────────────┘ └──────────────────────────────┘
|
||||
```
|
||||
|
||||
`llama.cpp` runs as its own container and exposes an OpenAI-compatible
|
||||
HTTP API. Hindsight talks to it via the standard `openai` LLM provider
|
||||
with `HINDSIGHT_API_LLM_BASE_URL` pointed at the sidecar.
|
||||
|
||||
This pattern follows
|
||||
[*Hosting llama-server with Docker* (ServiceStack)](https://servicestack.net/posts/hosting-llama-server).
|
||||
|
||||
### Why a sidecar and not the in-process `llamacpp` provider?
|
||||
|
||||
Hindsight does ship an in-process `llamacpp` provider that spawns
|
||||
`llama-cpp-python`, but the **published `ghcr.io/vectorize-io/hindsight`
|
||||
image deliberately omits `llama-cpp-python`** to keep the image small and
|
||||
avoid bundling native inference libraries that most users don't need.
|
||||
Trying to set `HINDSIGHT_API_LLM_PROVIDER=llamacpp` against the published
|
||||
image fails with `ModuleNotFoundError: No module named 'llama_cpp'`.
|
||||
|
||||
The sidecar approach side-steps that entirely: the official llama.cpp
|
||||
image is used as-is for inference, Hindsight is used as-is for memory.
|
||||
Clean separation, no derived images.
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
docker compose -f docker/docker-compose/local-llm/docker-compose.yaml up
|
||||
```
|
||||
|
||||
- API: http://localhost:8888
|
||||
- Control Plane: http://localhost:9999
|
||||
|
||||
**First boot downloads ~3.5 GB** (Gemma 4 E2B Q4_K_M GGUF) into the
|
||||
`llama_models` named volume. Subsequent boots reuse it.
|
||||
|
||||
Hindsight only starts after llama.cpp's `/health` endpoint reports
|
||||
healthy, so the API will appear "stuck" for a few minutes on the first
|
||||
run while the model downloads.
|
||||
|
||||
## Using a different model
|
||||
|
||||
Override the HuggingFace repo / file in `docker-compose.yaml`:
|
||||
|
||||
```yaml
|
||||
environment:
|
||||
LLAMA_ARG_HF_REPO: bartowski/Qwen2.5-7B-Instruct-GGUF
|
||||
LLAMA_ARG_HF_FILE: Qwen2.5-7B-Instruct-Q4_K_M.gguf
|
||||
```
|
||||
|
||||
Also update `HINDSIGHT_API_LLM_MODEL` on the `hindsight` service to a
|
||||
matching alias (the value is sent to llama-server as the OpenAI `model`
|
||||
field — llama-server is lenient about this but it shows up in logs).
|
||||
|
||||
## GPU acceleration
|
||||
|
||||
The default compose file targets CPU because not everyone has a GPU. On
|
||||
CPU, Gemma 4 E2B runs at ~2-3 tokens/sec — fine for a smoke test, but the
|
||||
retain pipeline (which makes several multi-hundred-token LLM calls per
|
||||
memory) will time out against Hindsight's default LLM timeout. **For any
|
||||
real use, run on a GPU.**
|
||||
|
||||
### NVIDIA
|
||||
|
||||
1. Switch the `llama` service image from `:server` to `:server-cuda`.
|
||||
2. Uncomment the `LLAMA_ARG_N_GPU_LAYERS: "999"` env var (offload all
|
||||
layers to GPU).
|
||||
3. Uncomment the `deploy.resources.reservations.devices` block.
|
||||
4. Install the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
|
||||
on the host.
|
||||
|
||||
The compose file has all four spots marked with inline comments.
|
||||
|
||||
### Apple Silicon / ROCm / Vulkan
|
||||
|
||||
The official `ghcr.io/ggml-org/llama.cpp` image only ships CPU and CUDA
|
||||
variants. For Metal (Apple Silicon), ROCm (AMD), or Vulkan backends,
|
||||
build llama.cpp yourself with the appropriate flags and reference the
|
||||
image you build instead. Docker Desktop on macOS cannot pass through the
|
||||
host GPU to a Linux container in any case — for Apple Silicon, run
|
||||
llama-server directly on the host and only put Hindsight in Docker.
|
||||
|
||||
## Caveats
|
||||
|
||||
- llama.cpp's HTTP API is OpenAI-compatible but not 100% feature-parity.
|
||||
Function/tool calling support depends on the chat template baked into
|
||||
the GGUF; some retain/reflect flows may behave differently than against
|
||||
a hosted OpenAI model.
|
||||
- Small GGUFs (~3 B params) are useful for smoke testing but will
|
||||
underperform a hosted frontier model on retain quality. Use a larger
|
||||
GGUF (7-13 B params) for production-quality memory.
|
||||
- The `llama_models` named volume persists the GGUF across `docker
|
||||
compose down`/`up` so the model is downloaded once, not every restart.
|
||||
@@ -0,0 +1,74 @@
|
||||
name: hindsight-local-llm
|
||||
# Example: run Hindsight against a local llama.cpp server sidecar — fully
|
||||
# offline, no external API key needed.
|
||||
#
|
||||
# Pattern follows https://servicestack.net/posts/hosting-llama-server :
|
||||
# llama.cpp runs as its own container exposing an OpenAI-compatible HTTP
|
||||
# API, and Hindsight talks to it via the `openai` LLM provider with a
|
||||
# custom `base_url`. This means we can use the published Hindsight image
|
||||
# unchanged — no derived Dockerfile, no `llama-cpp-python` install on top.
|
||||
#
|
||||
# Quick start:
|
||||
# docker compose -f docker/docker-compose/local-llm/docker-compose.yaml up
|
||||
#
|
||||
# First boot downloads the default Gemma 4 E2B GGUF (~3.5 GB) into the
|
||||
# `llama_models` volume; subsequent boots reuse it.
|
||||
|
||||
services:
|
||||
llama:
|
||||
image: ghcr.io/ggml-org/llama.cpp:server
|
||||
container_name: hindsight-local-llm-llama
|
||||
environment:
|
||||
LLAMA_ARG_HOST: 0.0.0.0
|
||||
LLAMA_ARG_PORT: "8080"
|
||||
# Auto-download a small GGUF from HuggingFace on first start.
|
||||
# Override these to use a different model.
|
||||
LLAMA_ARG_HF_REPO: bartowski/google_gemma-4-E2B-it-GGUF
|
||||
LLAMA_ARG_HF_FILE: google_gemma-4-E2B-it-Q4_K_M.gguf
|
||||
LLAMA_ARG_CTX_SIZE: "8192"
|
||||
# Uncomment for NVIDIA GPU (and switch image to :server-cuda):
|
||||
# LLAMA_ARG_N_GPU_LAYERS: "999"
|
||||
volumes:
|
||||
# llama-server stores HuggingFace downloads under ~/.cache/huggingface
|
||||
# (not ~/.cache/llama.cpp), so mount the named volume there to avoid
|
||||
# re-downloading the GGUF on every recreate.
|
||||
- llama_models:/root/.cache/huggingface
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "curl -fsS http://localhost:8080/health || exit 1"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 60
|
||||
start_period: 30s
|
||||
# For NVIDIA GPU acceleration, swap the image above to
|
||||
# `ghcr.io/ggml-org/llama.cpp:server-cuda` and uncomment:
|
||||
# deploy:
|
||||
# resources:
|
||||
# reservations:
|
||||
# devices:
|
||||
# - driver: nvidia
|
||||
# count: all
|
||||
# capabilities: [gpu]
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:latest
|
||||
container_name: hindsight-local-llm
|
||||
depends_on:
|
||||
llama:
|
||||
condition: service_healthy
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# llama-server is OpenAI-compatible, so use the `openai` provider and
|
||||
# point base_url at the sidecar. The API key is unused by llama-server
|
||||
# but Hindsight requires the env var to be set.
|
||||
HINDSIGHT_API_LLM_PROVIDER: openai
|
||||
HINDSIGHT_API_LLM_BASE_URL: http://llama:8080/v1
|
||||
HINDSIGHT_API_LLM_API_KEY: not-needed
|
||||
HINDSIGHT_API_LLM_MODEL: gemma-4-e2b-it
|
||||
volumes:
|
||||
- pg_data:/home/hindsight/.pg0
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
llama_models:
|
||||
@@ -0,0 +1,7 @@
|
||||
# PostgreSQL with pgvector and ParadeDB pg_search extensions.
|
||||
#
|
||||
# The official ParadeDB image ships PostgreSQL with pg_search and pgvector
|
||||
# already installed, so no build steps are required. We pin to the PG17
|
||||
# variant for parity with the other Hindsight docker-compose examples
|
||||
# (vchord, pg_textsearch).
|
||||
FROM paradedb/paradedb:latest-pg17
|
||||
@@ -0,0 +1,96 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and ParadeDB pg_search.
|
||||
#
|
||||
# pg_search is the only BM25 backend supported by Hindsight that works with
|
||||
# Citus, so this is the recommended setup for horizontally scaled deployments.
|
||||
#
|
||||
# Usage:
|
||||
# docker compose -f docker/docker-compose/pg_search/docker-compose.yaml up -d
|
||||
#
|
||||
# Required environment variables:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see the hindsight service)
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER: ParadeDB pg_search
|
||||
# tokenizer for new BM25 indexes (default: empty, uses ParadeDB default)
|
||||
|
||||
services:
|
||||
db:
|
||||
# Use ParadeDB image which bundles pgvector + pg_search
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
ports:
|
||||
- "5437:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
pg-search-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pg_search CASCADE;';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvector
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pg_search
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER: ${HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER:-}
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -0,0 +1,23 @@
|
||||
# PostgreSQL with pgvector and pgroonga extensions.
|
||||
#
|
||||
# pgroonga is a multilingual full-text search extension built on Groonga.
|
||||
# It works out of the box for CJK (Chinese, Japanese, Korean) and other
|
||||
# non-whitespace-segmented languages via the TokenBigram tokenizer.
|
||||
FROM groonga/pgroonga:latest-debian-pg17
|
||||
|
||||
# Install pgvector on top of the pgroonga base image (which already provides
|
||||
# pgroonga and the Groonga library).
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
build-essential \
|
||||
git \
|
||||
postgresql-server-dev-17 \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN cd /tmp && \
|
||||
git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git && \
|
||||
cd pgvector && \
|
||||
make && \
|
||||
make install
|
||||
|
||||
RUN rm -rf /tmp/pgvector && \
|
||||
apt-get purge -y --auto-remove build-essential git postgresql-server-dev-17
|
||||
@@ -0,0 +1,91 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and pgroonga
|
||||
#
|
||||
# pgroonga provides multilingual BM25 indexing that works out of the box for
|
||||
# CJK (Chinese, Japanese, Korean) and other non-whitespace-segmented languages.
|
||||
# Use this recipe if your bank content is not English/European.
|
||||
#
|
||||
# docker compose -f docker/docker-compose/pgroonga/docker-compose.yaml down && \
|
||||
# sleep 2 && \
|
||||
# docker compose -f docker/docker-compose/pgroonga/docker-compose.yaml up -d
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_DB_PASSWORD: PostgreSQL password (default: hindsight_password)
|
||||
|
||||
services:
|
||||
db:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
ports:
|
||||
- "5439:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
pgroonga-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pgroonga CASCADE;';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvector
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pgroonga
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -50,6 +50,8 @@ WORKDIR /app/api
|
||||
# Sync dependencies using appropriate extras based on INCLUDE_LOCAL_MODELS
|
||||
# local-ml: torch, sentence-transformers, transformers, einops, flashrank, mlx (optional)
|
||||
# embedded-db: pg0-embedded (always included for embedded PostgreSQL support)
|
||||
# ONNX Runtime embeddings are intentionally not bundled into the official
|
||||
# standalone image; install the local-onnx extra in custom images when needed.
|
||||
RUN if [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
|
||||
uv sync --extra local-ml --extra embedded-db; \
|
||||
else \
|
||||
|
||||
@@ -10,19 +10,90 @@ set -e
|
||||
# loss scenarios where a container restart caused the data directory to be
|
||||
# wiped despite a volume mount being present.
|
||||
# =============================================================================
|
||||
PG0_DATA_DIR="${HOME}/.pg0"
|
||||
if [ -d "$PG0_DATA_DIR" ]; then
|
||||
pg0_has_pg_version() {
|
||||
local pg0_data_dir="$1"
|
||||
|
||||
# pg0 has used more than one on-disk layout. Newer standalone images keep
|
||||
# PostgreSQL data under instances/<name>/data, while older volumes may have
|
||||
# placed PG_VERSION at or one level below the mount.
|
||||
[ -f "$pg0_data_dir/PG_VERSION" ] && return 0
|
||||
compgen -G "$pg0_data_dir"/*/PG_VERSION > /dev/null 2>&1 && return 0
|
||||
compgen -G "$pg0_data_dir"/instances/*/data/PG_VERSION > /dev/null 2>&1 && return 0
|
||||
|
||||
return 1
|
||||
}
|
||||
|
||||
check_pg0_data_integrity() {
|
||||
local pg0_data_dir="$1"
|
||||
|
||||
if [ ! -d "$pg0_data_dir" ]; then
|
||||
return 0
|
||||
fi
|
||||
|
||||
# Look for actual PostgreSQL data directories (pg0 creates subdirs per instance)
|
||||
if compgen -G "$PG0_DATA_DIR"/*/PG_VERSION > /dev/null 2>&1; then
|
||||
echo "✅ Existing pg0 data directory detected at $PG0_DATA_DIR"
|
||||
elif [ "$(ls -A "$PG0_DATA_DIR" 2>/dev/null)" ]; then
|
||||
echo "⚠️ WARNING: pg0 data directory exists at $PG0_DATA_DIR but no PG_VERSION found."
|
||||
if pg0_has_pg_version "$pg0_data_dir"; then
|
||||
echo "✅ Existing pg0 data directory detected at $pg0_data_dir"
|
||||
elif [ "$(ls -A "$pg0_data_dir" 2>/dev/null)" ]; then
|
||||
echo "⚠️ WARNING: pg0 data directory exists at $pg0_data_dir but no PG_VERSION found."
|
||||
echo " This may indicate data corruption or an incomplete previous shutdown."
|
||||
echo " If you see all migrations running from scratch after this, your data may have been lost."
|
||||
echo " See: https://github.com/vectorize-io/hindsight/issues/675"
|
||||
fi
|
||||
|
||||
return 0
|
||||
}
|
||||
|
||||
# =============================================================================
|
||||
# Embedded pg0 writability pre-check (#1483)
|
||||
#
|
||||
# The container runs as the unprivileged `hindsight` user (UID 1000). When the
|
||||
# pg0 data directory is a host bind mount (e.g. `-v $HOME/dir:/home/hindsight/.pg0`)
|
||||
# that is not owned by UID 1000 — the default on macOS Docker Desktop and most
|
||||
# non-1000 Linux hosts — pg0 fails with the opaque "Permission denied (os error
|
||||
# 13)". We cannot chown it ourselves without root (and the image is deliberately
|
||||
# rootless), so we surface an actionable message up front instead.
|
||||
#
|
||||
# Docker *named* volumes are seeded with the image directory's ownership (UID
|
||||
# 1000) on first use, so they avoid this entirely — hence the named-volume
|
||||
# recommendation below and in the README.
|
||||
# =============================================================================
|
||||
check_pg0_writable() {
|
||||
local pg0_data_dir="$1"
|
||||
|
||||
# Only relevant for embedded pg0; an external database doesn't use this dir.
|
||||
if [ -n "${HINDSIGHT_API_DATABASE_URL:-}" ]; then
|
||||
return 0
|
||||
fi
|
||||
|
||||
mkdir -p "$pg0_data_dir" 2>/dev/null || true
|
||||
if touch "$pg0_data_dir/.hindsight-write-test" 2>/dev/null; then
|
||||
rm -f "$pg0_data_dir/.hindsight-write-test" 2>/dev/null || true
|
||||
return 0
|
||||
fi
|
||||
|
||||
echo "❌ The embedded database directory $pg0_data_dir is not writable by this container (UID $(id -u))."
|
||||
echo ""
|
||||
echo " A host directory was bind-mounted but is not owned by the container user (UID 1000)."
|
||||
echo " Hindsight runs rootless and cannot fix this for you. Choose one:"
|
||||
echo ""
|
||||
echo " • Recommended — use a Docker named volume (auto-owned by the container):"
|
||||
echo " -v hindsight-data:/home/hindsight/.pg0"
|
||||
echo ""
|
||||
echo " • Or keep the host path and run as your host user, chowning it to match:"
|
||||
echo " sudo chown -R \$(id -u):\$(id -g) <host-directory>"
|
||||
echo " docker run --user \$(id -u):\$(id -g) -e HOME=/home/hindsight ..."
|
||||
echo ""
|
||||
echo " See https://github.com/vectorize-io/hindsight/issues/1483"
|
||||
return 1
|
||||
}
|
||||
|
||||
if [ "${HINDSIGHT_START_ALL_SOURCE_ONLY:-false}" = "true" ]; then
|
||||
return 0 2>/dev/null || exit 0
|
||||
fi
|
||||
|
||||
check_pg0_data_integrity "${HOME}/.pg0"
|
||||
check_pg0_writable "${HOME}/.pg0" || exit 1
|
||||
|
||||
# Service flags (default to true if not set)
|
||||
ENABLE_API="${HINDSIGHT_ENABLE_API:-true}"
|
||||
ENABLE_CP="${HINDSIGHT_ENABLE_CP:-true}"
|
||||
@@ -156,7 +227,7 @@ PIDS=()
|
||||
# Start API if enabled
|
||||
if [ "$ENABLE_API" = "true" ]; then
|
||||
cd /app/api
|
||||
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:8888/health}"
|
||||
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:${HINDSIGHT_API_PORT:-8888}/health}"
|
||||
API_STARTUP_WAIT_SECONDS="${HINDSIGHT_API_STARTUP_WAIT_SECONDS:-300}"
|
||||
|
||||
# Run API directly - Python's PYTHONUNBUFFERED=1 handles output buffering
|
||||
|
||||
Executable
+121
@@ -0,0 +1,121 @@
|
||||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
HINDSIGHT_START_ALL_SOURCE_ONLY=true
|
||||
source "$SCRIPT_DIR/start-all.sh"
|
||||
unset HINDSIGHT_START_ALL_SOURCE_ONLY
|
||||
|
||||
TMP_DIR="$(mktemp -d)"
|
||||
trap 'chmod -R u+rwx "$TMP_DIR" 2>/dev/null || true; rm -rf "$TMP_DIR"' EXIT
|
||||
|
||||
assert_contains() {
|
||||
local output="$1"
|
||||
local expected="$2"
|
||||
|
||||
if [[ "$output" != *"$expected"* ]]; then
|
||||
echo "Expected output to contain: $expected"
|
||||
echo "Actual output:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
assert_not_contains() {
|
||||
local output="$1"
|
||||
local unexpected="$2"
|
||||
|
||||
if [[ "$output" == *"$unexpected"* ]]; then
|
||||
echo "Expected output not to contain: $unexpected"
|
||||
echo "Actual output:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
assert_empty() {
|
||||
local output="$1"
|
||||
|
||||
if [ -n "$output" ]; then
|
||||
echo "Expected no output, got:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
mkdir -p "$TMP_DIR/empty"
|
||||
assert_empty "$(check_pg0_data_integrity "$TMP_DIR/empty")"
|
||||
|
||||
mkdir -p "$TMP_DIR/direct"
|
||||
touch "$TMP_DIR/direct/PG_VERSION"
|
||||
direct_output="$(check_pg0_data_integrity "$TMP_DIR/direct")"
|
||||
assert_contains "$direct_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$direct_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/legacy/instance"
|
||||
touch "$TMP_DIR/legacy/instance/PG_VERSION"
|
||||
legacy_output="$(check_pg0_data_integrity "$TMP_DIR/legacy")"
|
||||
assert_contains "$legacy_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$legacy_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/nested/instances/hindsight/data"
|
||||
touch "$TMP_DIR/nested/instances/hindsight/data/PG_VERSION"
|
||||
nested_output="$(check_pg0_data_integrity "$TMP_DIR/nested")"
|
||||
assert_contains "$nested_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$nested_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/nonempty/instances/hindsight"
|
||||
touch "$TMP_DIR/nonempty/instances/hindsight/instance.json"
|
||||
nonempty_output="$(check_pg0_data_integrity "$TMP_DIR/nonempty")"
|
||||
assert_contains "$nonempty_output" "WARNING: pg0 data directory exists"
|
||||
|
||||
echo "start-all pg0 integrity checks passed"
|
||||
|
||||
# =============================================================================
|
||||
# check_pg0_writable (#1483)
|
||||
# These rely on filesystem permissions, which root bypasses; skip under root.
|
||||
# =============================================================================
|
||||
if [ "$(id -u)" != "0" ]; then
|
||||
# Writable directory: returns 0, prints nothing, leaves no artifact behind.
|
||||
mkdir -p "$TMP_DIR/writable"
|
||||
writable_output="$(check_pg0_writable "$TMP_DIR/writable")"
|
||||
assert_empty "$writable_output"
|
||||
if [ -e "$TMP_DIR/writable/.hindsight-write-test" ]; then
|
||||
echo "check_pg0_writable left its write-test file behind"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Non-writable directory: returns 1 with actionable guidance.
|
||||
mkdir -p "$TMP_DIR/readonly"
|
||||
chmod 000 "$TMP_DIR/readonly"
|
||||
set +e
|
||||
readonly_output="$(check_pg0_writable "$TMP_DIR/readonly" 2>&1)"
|
||||
readonly_rc=$?
|
||||
set -e
|
||||
chmod 755 "$TMP_DIR/readonly"
|
||||
if [ "$readonly_rc" -eq 0 ]; then
|
||||
echo "check_pg0_writable should fail on a non-writable directory"
|
||||
exit 1
|
||||
fi
|
||||
assert_contains "$readonly_output" "not writable"
|
||||
assert_contains "$readonly_output" "hindsight-data:/home/hindsight/.pg0"
|
||||
assert_contains "$readonly_output" "--user"
|
||||
|
||||
# External database configured: skip the check regardless of dir perms.
|
||||
mkdir -p "$TMP_DIR/extdb"
|
||||
chmod 000 "$TMP_DIR/extdb"
|
||||
set +e
|
||||
HINDSIGHT_API_DATABASE_URL="postgres://x" check_pg0_writable "$TMP_DIR/extdb" >/dev/null 2>&1
|
||||
extdb_rc=$?
|
||||
set -e
|
||||
chmod 755 "$TMP_DIR/extdb"
|
||||
if [ "$extdb_rc" -ne 0 ]; then
|
||||
echo "check_pg0_writable should skip when an external database is configured"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "start-all pg0 writability checks passed"
|
||||
else
|
||||
echo "⚠️ Running as root; skipping pg0 writability checks (permissions are bypassed)."
|
||||
fi
|
||||
@@ -1,6 +0,0 @@
|
||||
dependencies:
|
||||
- name: postgresql
|
||||
repository: https://charts.bitnami.com/bitnami
|
||||
version: 15.5.38
|
||||
digest: sha256:f67c7612736803ece8a669f8ca6b0555f3b78557bc0ecb732aa2e43f0df7750d
|
||||
generated: "2025-12-10T17:20:57.058794+01:00"
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.6.2
|
||||
appVersion: "0.6.2"
|
||||
version: 0.8.0
|
||||
appVersion: "0.8.0"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.6.2",
|
||||
"version": "0.8.0",
|
||||
"description": "Node.js programmatic lifecycle manager for Hindsight — embeds a local hindsight daemon in a Node application. Pair with @vectorize-io/hindsight-client for memory operations.",
|
||||
"main": "dist/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.6.2"
|
||||
version = "0.8.0"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim==0.6.2",
|
||||
"hindsight-api-slim==0.8.0",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.6.2"
|
||||
version = "0.8.0"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim[all]==0.6.2",
|
||||
"hindsight-api-slim[all]==0.8.0",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
@@ -21,7 +21,7 @@ hindsight-embed = { workspace = true }
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-llm = [
|
||||
"hindsight-api-slim[local-llm]==0.6.2",
|
||||
"hindsight-api-slim[local-llm]==0.8.0",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -386,7 +386,7 @@ def test_embedded_ui_flag(llm_config):
|
||||
|
||||
# Verify UI is reachable and reports connected dataplane
|
||||
ui_url = client.ui_url
|
||||
assert ui_url, "ui_url should be set"
|
||||
assert isinstance(ui_url, str) and ui_url, "ui_url should be a non-empty string"
|
||||
|
||||
health_url = f"{ui_url}/api/health"
|
||||
with urllib.request.urlopen(health_url, timeout=10) as resp:
|
||||
|
||||
@@ -99,7 +99,7 @@ hindsight-api
|
||||
## Docker
|
||||
|
||||
```bash
|
||||
docker run --rm -it -p 8888:8888 \
|
||||
docker run -it --name hindsight --restart unless-stopped -p 8888:8888 \
|
||||
-e HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY \
|
||||
-v $HOME/.hindsight-docker:/home/hindsight/.pg0 \
|
||||
ghcr.io/vectorize-io/hindsight:latest
|
||||
|
||||
@@ -4,6 +4,13 @@ Memory System for AI Agents.
|
||||
Temporal + Semantic Memory Architecture using PostgreSQL with pgvector.
|
||||
"""
|
||||
|
||||
# Cap native ML thread pools (OpenBLAS/OpenMP/MKL) before any import pulls in
|
||||
# numpy/torch/onnxruntime — they read these env vars only at load time. See
|
||||
# hindsight_api/_thread_limits.py for the rationale.
|
||||
from ._thread_limits import apply_default_thread_limits
|
||||
|
||||
apply_default_thread_limits()
|
||||
|
||||
from .config import HindsightConfig, get_config
|
||||
from .engine.cross_encoder import CrossEncoderModel, LocalSTCrossEncoder, RemoteTEICrossEncoder
|
||||
from .engine.embeddings import Embeddings, LocalSTEmbeddings, RemoteTEIEmbeddings
|
||||
@@ -46,4 +53,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.6.2"
|
||||
__version__ = "0.8.0"
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
"""Helpers for ParadeDB pg_search index configuration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
PG_SEARCH_TOKENIZER_ENV = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"
|
||||
|
||||
_SIMPLE_TOKENIZERS = {
|
||||
"unicode_words",
|
||||
"simple",
|
||||
"whitespace",
|
||||
"literal",
|
||||
"literal_normalized",
|
||||
"chinese_compatible",
|
||||
"icu",
|
||||
"jieba",
|
||||
"source_code",
|
||||
}
|
||||
|
||||
_TOKENIZER_ALIASES = {
|
||||
"chinese_lindera": "lindera(chinese)",
|
||||
"japanese_lindera": "lindera(japanese)",
|
||||
"korean_lindera": "lindera(korean)",
|
||||
"lindera_chinese": "lindera(chinese)",
|
||||
"lindera_japanese": "lindera(japanese)",
|
||||
"lindera_korean": "lindera(korean)",
|
||||
}
|
||||
|
||||
|
||||
def normalize_pg_search_tokenizer(value: str | None) -> str:
|
||||
"""Validate and normalize a ParadeDB pg_search tokenizer setting.
|
||||
|
||||
Returns an empty string when unset. The returned value is safe to embed after
|
||||
``pdb.`` in a CREATE INDEX expression.
|
||||
"""
|
||||
|
||||
tokenizer = (value or "").strip().lower()
|
||||
if not tokenizer:
|
||||
return ""
|
||||
|
||||
if tokenizer in _TOKENIZER_ALIASES:
|
||||
return _TOKENIZER_ALIASES[tokenizer]
|
||||
|
||||
if tokenizer in _SIMPLE_TOKENIZERS:
|
||||
return tokenizer
|
||||
|
||||
lindera_match = re.fullmatch(r"lindera\((chinese|japanese|korean)\)", tokenizer)
|
||||
if lindera_match:
|
||||
return tokenizer
|
||||
|
||||
ngram_match = re.fullmatch(r"(ngram|edge_ngram)\((\d{1,3}),\s*(\d{1,3})\)", tokenizer)
|
||||
if ngram_match:
|
||||
kind, min_gram, max_gram = ngram_match.groups()
|
||||
min_value = int(min_gram)
|
||||
max_value = int(max_gram)
|
||||
if min_value <= 0 or min_value > max_value:
|
||||
raise ValueError(
|
||||
f"Invalid {PG_SEARCH_TOKENIZER_ENV}: {value!r}. "
|
||||
"ngram and edge_ngram require positive min/max gram sizes with min <= max."
|
||||
)
|
||||
return f"{kind}({min_value},{max_value})"
|
||||
|
||||
raise ValueError(
|
||||
f"Invalid {PG_SEARCH_TOKENIZER_ENV}: {value!r}. "
|
||||
"Supported values are: unicode_words, simple, whitespace, literal, "
|
||||
"literal_normalized, chinese_compatible, icu, jieba, source_code, "
|
||||
"chinese_lindera, japanese_lindera, korean_lindera, or "
|
||||
"lindera(chinese|japanese|korean), ngram(min,max), or edge_ngram(min,max)."
|
||||
)
|
||||
|
||||
|
||||
def pg_search_bm25_columns(
|
||||
key_field: str,
|
||||
text_fields: Sequence[str],
|
||||
tokenizer: str | None,
|
||||
) -> str:
|
||||
"""Build a ParadeDB BM25 column list for CREATE INDEX."""
|
||||
|
||||
normalized = normalize_pg_search_tokenizer(tokenizer)
|
||||
if not normalized:
|
||||
return ", ".join([key_field, *text_fields])
|
||||
|
||||
return ", ".join([key_field, *(f"({field}::pdb.{normalized})" for field in text_fields)])
|
||||
@@ -0,0 +1,107 @@
|
||||
"""Process-level caps for native ML thread pools.
|
||||
|
||||
OpenBLAS, OpenMP, and MKL each spawn a worker pool sized to the host CPU count
|
||||
the first time they are loaded (numpy pulls in OpenBLAS eagerly; torch and
|
||||
onnxruntime load their pools lazily on first inference). Hindsight already
|
||||
parallelizes at the request level via thread-pool executors (embeddings on the
|
||||
default executor, the reranker on its own pool), so these native intra-op pools
|
||||
oversubscribe the CPU: on a many-core host the process accumulates 100+ native
|
||||
threads, which inflates memory and, under contention, can degrade throughput.
|
||||
|
||||
We bound each pool to ``_MAX_NATIVE_THREADS`` (or the available CPU count, if
|
||||
smaller). "Available" is the CPU budget actually granted to the process, not
|
||||
``os.cpu_count()``: in a CPU-limited container ``os.cpu_count()`` still reports
|
||||
the host's cores, so sizing pools by it oversubscribes the container's real
|
||||
quota — the exact failure mode this guards against. We therefore take the
|
||||
smallest of the CPU-affinity set, the cgroup CPU quota, and ``os.cpu_count()``.
|
||||
|
||||
Every cap is applied with ``setdefault`` so an operator who has deliberately
|
||||
tuned one of these variables keeps their value. This must run *before* numpy,
|
||||
torch, or onnxruntime are imported — those libraries read the variables only at
|
||||
load time — which is why it is invoked at the very top of
|
||||
``hindsight_api/__init__.py``, ahead of the package's other imports.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
# Native threading env vars, each read by the respective library at load time.
|
||||
_NATIVE_THREAD_VARS = (
|
||||
"OMP_NUM_THREADS", # OpenMP — torch, onnxruntime, some BLAS builds
|
||||
"OPENBLAS_NUM_THREADS", # OpenBLAS — numpy's default BLAS
|
||||
"MKL_NUM_THREADS", # Intel MKL — numpy/torch when MKL-backed
|
||||
"NUMEXPR_NUM_THREADS", # numexpr expression engine
|
||||
)
|
||||
|
||||
# Upper bound on intra-op threads per native pool. Bounds runaway growth on
|
||||
# many-core hosts without serialising single-request inference.
|
||||
_MAX_NATIVE_THREADS = 16
|
||||
|
||||
|
||||
def _quota_to_cpus(quota: int, period: int) -> int | None:
|
||||
"""Whole CPUs from a CFS quota/period pair, or None if unlimited."""
|
||||
if quota > 0 and period > 0:
|
||||
# Floor (never round up) so we never exceed the granted budget.
|
||||
return max(1, quota // period)
|
||||
return None
|
||||
|
||||
|
||||
def _parse_cgroup_v2_cpu_max(text: str) -> int | None:
|
||||
"""Parse cgroup v2 ``cpu.max`` ("<quota> <period>", or "max <period>")."""
|
||||
parts = text.split()
|
||||
if len(parts) >= 2 and parts[0] != "max":
|
||||
try:
|
||||
return _quota_to_cpus(int(parts[0]), int(parts[1]))
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _cgroup_cpu_quota() -> int | None:
|
||||
"""Effective CPUs from the cgroup CPU quota, or None if unlimited/unknown."""
|
||||
try: # cgroup v2
|
||||
with open("/sys/fs/cgroup/cpu.max") as fh:
|
||||
return _parse_cgroup_v2_cpu_max(fh.read())
|
||||
except OSError:
|
||||
pass
|
||||
try: # cgroup v1
|
||||
with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as fh:
|
||||
quota = int(fh.read())
|
||||
with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as fh:
|
||||
period = int(fh.read())
|
||||
return _quota_to_cpus(quota, period)
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _available_cpu_count() -> int:
|
||||
"""CPUs actually available to this process.
|
||||
|
||||
The smallest of the CPU-affinity set (cpuset / ``--cpuset-cpus``), the
|
||||
cgroup CPU quota (``--cpus``), and ``os.cpu_count()`` — each captures a
|
||||
different way the budget can be constrained, and the last alone overcounts
|
||||
inside a limited container.
|
||||
"""
|
||||
candidates = [os.cpu_count() or 1]
|
||||
if hasattr(os, "sched_getaffinity"):
|
||||
try:
|
||||
candidates.append(len(os.sched_getaffinity(0)))
|
||||
except OSError:
|
||||
pass
|
||||
quota = _cgroup_cpu_quota()
|
||||
if quota is not None:
|
||||
candidates.append(quota)
|
||||
return max(1, min(candidates))
|
||||
|
||||
|
||||
def default_native_thread_count() -> int:
|
||||
"""Per-pool cap: ``_MAX_NATIVE_THREADS``, or available CPUs if fewer."""
|
||||
return min(_MAX_NATIVE_THREADS, _available_cpu_count())
|
||||
|
||||
|
||||
def apply_default_thread_limits() -> None:
|
||||
"""Cap native ML thread pools unless the operator has set the var already."""
|
||||
value = str(default_native_thread_count())
|
||||
for var in _NATIVE_THREAD_VARS:
|
||||
os.environ.setdefault(var, value)
|
||||
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.engine import Connection
|
||||
@@ -34,7 +35,7 @@ _INDEX_USING_CLAUSES = {
|
||||
"pgvector": "USING hnsw (embedding vector_cosine_ops)",
|
||||
"pgvectorscale": "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)",
|
||||
"pg_diskann": "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)",
|
||||
"vchord": "USING vchordrq (embedding vector_l2_ops)",
|
||||
"vchord": "USING vchordrq (embedding vector_cosine_ops)",
|
||||
"scann": "USING scann (embedding cosine) WITH (mode = 'AUTO')",
|
||||
}
|
||||
|
||||
@@ -46,6 +47,32 @@ _INDEX_TYPE_KEYWORDS = {
|
||||
"scann": "scann",
|
||||
}
|
||||
|
||||
# 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).
|
||||
# - vchord exposes vchordrq.probes (no default; see VectorChord issue #392)
|
||||
# and vchordrq.epsilon (default 1.9). probes = 10 / 30 are starting
|
||||
# defaults pending a workload-specific sweep — vchordrq's recall curve
|
||||
# shape differs from HNSW's, so the pgvector numbers don't translate
|
||||
# directly. Revisit with a per-cluster benchmark once we have production
|
||||
# recall data; until then these are deliberately conservative on the
|
||||
# high-recall path. We leave epsilon at its default; tightening it is a
|
||||
# separate trade-off.
|
||||
# - pgvectorscale / pg_diskann / scann do not expose an equivalent per-statement
|
||||
# knob in the engine today, so the dispatcher returns no statements for them.
|
||||
_ANN_TUNING_LOW_LATENCY: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (("hnsw.ef_search", "60"),),
|
||||
"vchord": (("vchordrq.probes", "10"),),
|
||||
}
|
||||
_ANN_TUNING_HIGH_RECALL: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (("hnsw.ef_search", "200"),),
|
||||
"vchord": (("vchordrq.probes", "30"),),
|
||||
}
|
||||
|
||||
_EXTENSION_INSTALL_SQL = {
|
||||
"pgvector": ("CREATE EXTENSION IF NOT EXISTS vector",),
|
||||
"pgvectorscale": (
|
||||
@@ -67,6 +94,18 @@ _INSTALL_HINTS = {
|
||||
}
|
||||
|
||||
|
||||
def configured_vector_extension() -> str:
|
||||
"""Return the user-configured vector backend extension.
|
||||
|
||||
Reads ``HINDSIGHT_API_VECTOR_EXTENSION`` (default ``"pgvector"``) and
|
||||
validates it via :func:`validate_extension`. This is the single source of
|
||||
truth for runtime code that needs to dispatch behaviour by vector backend;
|
||||
callers should prefer this over reading the env var directly, so the
|
||||
default value and the lookup mechanism live in one place.
|
||||
"""
|
||||
return validate_extension(os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector"))
|
||||
|
||||
|
||||
def validate_extension(name: str) -> str:
|
||||
"""Return a normalized configurable vector extension name or raise.
|
||||
|
||||
@@ -115,6 +154,25 @@ def should_defer_index_creation(ext: str, row_count: int) -> bool:
|
||||
return minimum_rows > 0 and row_count < minimum_rows
|
||||
|
||||
|
||||
def ann_search_tuning_settings(ext: str, *, kind: str) -> tuple[tuple[str, str], ...]:
|
||||
"""Return per-backend (guc_name, value) pairs for ANN search-time tuning.
|
||||
|
||||
``kind`` is ``"low_latency"`` for retain-side link probing (smaller probe
|
||||
count, lower recall, lower latency) and ``"high_recall"`` for connection
|
||||
init in the pool (larger probe count, higher recall). Callers wrap each
|
||||
pair with ``SET LOCAL`` or ``SET`` themselves so the same dispatcher works
|
||||
for both transaction-scoped and session-scoped use. Returns an empty tuple
|
||||
for backends without an equivalent knob.
|
||||
"""
|
||||
if kind == "low_latency":
|
||||
table = _ANN_TUNING_LOW_LATENCY
|
||||
elif kind == "high_recall":
|
||||
table = _ANN_TUNING_HIGH_RECALL
|
||||
else:
|
||||
raise ValueError(f"Unknown ANN tuning kind: {kind!r}")
|
||||
return table.get(_normalize_resolved(ext), ())
|
||||
|
||||
|
||||
def uses_per_bank_vector_indexes(ext: str) -> bool:
|
||||
"""Return whether the backend should create per-bank partial vector indexes."""
|
||||
return _normalize_resolved(ext) != "scann"
|
||||
|
||||
@@ -17,7 +17,9 @@ import asyncpg
|
||||
import typer
|
||||
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA, HindsightConfig
|
||||
from ..engine.memory_engine import _current_schema
|
||||
from ..engine.schema import fq_table_explicit as _fq_table
|
||||
from ..engine.transfer import export_bank
|
||||
from ..extensions import TenantExtension, load_extension
|
||||
from ..pg0 import parse_pg0_url, resolve_database_url
|
||||
|
||||
@@ -30,8 +32,17 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
app = typer.Typer(name="hindsight-admin", help="Hindsight administrative commands")
|
||||
|
||||
# Tables to backup/restore in dependency order
|
||||
# Import must happen in this order due to foreign key constraints
|
||||
# Tables to backup/restore in foreign-key dependency order (parents first).
|
||||
# Restore COPYs in this order and TRUNCATEs in reverse, so every child must
|
||||
# appear after the tables it references.
|
||||
#
|
||||
# This must cover EVERY persistent PostgreSQL table in the schema — a missing
|
||||
# entry silently drops that table's data on restore (and, worse, restore's
|
||||
# `TRUNCATE banks CASCADE` wipes any FK-to-banks child like mental_models even
|
||||
# when it was never backed up). test_admin_backup_restore.py asserts this list
|
||||
# equals the live schema's tables, so adding a migration that creates a table
|
||||
# without adding it here fails CI. Oracle-only tables (e.g. observation_sources)
|
||||
# are intentionally absent — admin backup/restore is PostgreSQL-only.
|
||||
BACKUP_TABLES = [
|
||||
"banks",
|
||||
"documents",
|
||||
@@ -41,11 +52,38 @@ BACKUP_TABLES = [
|
||||
"unit_entities",
|
||||
"entity_cooccurrences",
|
||||
"memory_links",
|
||||
"observation_history",
|
||||
"mental_models",
|
||||
"mental_model_history",
|
||||
"directives",
|
||||
"async_operations",
|
||||
"webhooks",
|
||||
"file_storage",
|
||||
"audit_log",
|
||||
"llm_requests",
|
||||
"graph_maintenance_queue",
|
||||
]
|
||||
|
||||
MANIFEST_VERSION = "1"
|
||||
|
||||
|
||||
async def _admin_connect(db_url: str) -> asyncpg.Connection:
|
||||
"""Open a raw asyncpg connection to an admin DB URL.
|
||||
|
||||
``resolve_database_url`` handles both plain ``postgres://`` (passthrough) and
|
||||
``pg0://`` (boots the embedded server and returns its real libpq URL), so this
|
||||
is the only step needed to connect. JSON codecs are registered so ``jsonb``
|
||||
columns decode to Python objects (used by the export row dumps).
|
||||
"""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
conn = await asyncpg.connect(await resolve_database_url(db_url))
|
||||
for type_name in ("json", "jsonb"):
|
||||
await conn.set_type_codec(type_name, encoder=json.dumps, decoder=json.loads, schema="pg_catalog")
|
||||
return conn
|
||||
|
||||
|
||||
async def _backup(database_url: str, output_path: Path, schema: str = "public") -> dict[str, Any]:
|
||||
"""Backup all tables to a zip file using binary COPY protocol."""
|
||||
conn = await asyncpg.connect(database_url)
|
||||
@@ -268,6 +306,7 @@ async def _run_migration(
|
||||
ensure_text_search_extension(
|
||||
resolved_url,
|
||||
text_search_extension=config.text_search_extension,
|
||||
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
@@ -313,6 +352,123 @@ def run_db_migration(
|
||||
typer.echo(f"Database migrations completed successfully for {len(schemas)} schema(s)")
|
||||
|
||||
|
||||
async def _run_export_bank(db_url: str, bank_id: str, output: Path, schema: str, include_history: bool) -> int:
|
||||
"""Export a whole bank to a ZIP archive."""
|
||||
conn = await _admin_connect(db_url)
|
||||
try:
|
||||
# export_bank resolves table names via fq_table (the _current_schema
|
||||
# contextvar); set it so the raw connection targets the right schema.
|
||||
_current_schema.set(schema)
|
||||
data = await export_bank(conn, bank_id, include_history=include_history)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
output.write_bytes(data)
|
||||
return len(data)
|
||||
|
||||
|
||||
@app.command(name="export-bank")
|
||||
def export_bank_command(
|
||||
bank_id: str = typer.Option(..., "--bank", "-b", help="Bank id to export."),
|
||||
output: Path = typer.Option(..., "--output", "-o", help="Path to write the .zip archive."),
|
||||
schema: str | None = typer.Option(
|
||||
None,
|
||||
"--schema",
|
||||
"-s",
|
||||
help="Database schema the bank lives in. Defaults to the configured base schema.",
|
||||
),
|
||||
include_history: bool = typer.Option(
|
||||
False,
|
||||
"--include-history",
|
||||
help="Also export operational history (audit_log, llm_requests). Off by default.",
|
||||
),
|
||||
):
|
||||
"""Export an entire bank to a portable ZIP (no embeddings — regenerated on import).
|
||||
|
||||
Carries documents, facts, observations, bank config, mental models, directives
|
||||
and webhooks so the bank can be imported into a new instance configured with a
|
||||
different embedding model / vector / text-search backend.
|
||||
"""
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
if not config.database_url:
|
||||
typer.echo("Error: Database URL not configured.", err=True)
|
||||
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
target_schema = schema or config.database_schema or DEFAULT_DATABASE_SCHEMA
|
||||
typer.echo(f"Exporting bank '{bank_id}' from schema '{target_schema}'...")
|
||||
|
||||
size = asyncio.run(_run_export_bank(config.database_url, bank_id, output, target_schema, include_history))
|
||||
|
||||
typer.echo(f"Exported bank '{bank_id}' to {output} ({size} bytes)")
|
||||
|
||||
|
||||
async def _run_import_bank(archive_path: Path, schema: str, target_bank_id: str | None, include_history: bool):
|
||||
"""Boot a MemoryEngine (for the target's embedding model) and restore a bank archive."""
|
||||
# MemoryEngine is heavy (loads embeddings); import it lazily so other admin
|
||||
# commands don't pay for it. _current_schema is imported at module top.
|
||||
from ..engine.memory_engine import MemoryEngine
|
||||
from ..models import RequestContext
|
||||
|
||||
archive_bytes = archive_path.read_bytes()
|
||||
# run_migrations=True so a fresh target instance is provisioned at this
|
||||
# instance's embedding dimension / vector / text-search backend before restore.
|
||||
engine = MemoryEngine(run_migrations=True)
|
||||
await engine.initialize()
|
||||
try:
|
||||
_current_schema.set(schema)
|
||||
context = RequestContext(internal=True, user_initiated=True)
|
||||
return await engine.import_bank_async(
|
||||
archive_bytes,
|
||||
context,
|
||||
target_bank_id=target_bank_id,
|
||||
include_history=include_history,
|
||||
)
|
||||
finally:
|
||||
await engine.close()
|
||||
|
||||
|
||||
@app.command(name="import-bank")
|
||||
def import_bank_command(
|
||||
archive: Path = typer.Option(..., "--archive", "-a", help="Path to the .zip produced by export-bank."),
|
||||
schema: str | None = typer.Option(
|
||||
None, "--schema", "-s", help="Target schema. Defaults to the configured base schema."
|
||||
),
|
||||
target_bank: str | None = typer.Option(
|
||||
None, "--target-bank", help="Override the bank id (defaults to the archive's source bank)."
|
||||
),
|
||||
include_history: bool = typer.Option(
|
||||
False, "--include-history", help="Also restore operational history if present in the archive."
|
||||
),
|
||||
):
|
||||
"""Restore a whole bank from an export-bank archive into THIS instance.
|
||||
|
||||
Re-embeds facts with this instance's configured embedding model and rebuilds
|
||||
links and indexes — the import half of a cross-instance migration. Run against
|
||||
an instance configured with the desired embedding / vector / text-search backend.
|
||||
The target bank must not already exist (import restores a whole bank, not a merge).
|
||||
"""
|
||||
config = HindsightConfig.from_env()
|
||||
if not config.database_url:
|
||||
typer.echo("Error: Database URL not configured.", err=True)
|
||||
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
target_schema = schema or config.database_schema or DEFAULT_DATABASE_SCHEMA
|
||||
typer.echo(f"Importing bank archive '{archive}' into schema '{target_schema}'...")
|
||||
|
||||
result = asyncio.run(_run_import_bank(archive, target_schema, target_bank, include_history))
|
||||
|
||||
typer.echo(
|
||||
f"Imported bank '{result.bank_id}': {result.documents_imported} doc(s), "
|
||||
f"{result.facts_imported} fact(s), {result.observations_imported} observation(s), "
|
||||
f"{result.mental_models_imported} mental model(s), "
|
||||
f"{result.mental_model_history_imported} mm-history row(s), {result.directives_imported} directive(s), "
|
||||
f"{result.webhooks_imported} webhook(s), {result.history_rows_imported} history row(s)"
|
||||
)
|
||||
|
||||
|
||||
async def _decommission_worker(db_url: str, worker_id: str, schema: str = "public") -> int:
|
||||
"""Release all tasks owned by a worker, setting them back to pending status."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
|
||||
@@ -15,6 +15,11 @@ from pgvector.sqlalchemy import Vector
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -83,7 +88,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pg_diskann":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
@@ -91,9 +96,14 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Detect or validate text search extension: 'native', 'vchord', 'pg_textsearch',
|
||||
'pgroonga', or 'pg_search'. Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
|
||||
pgroonga is treated as native here so the initial schema still creates valid
|
||||
tsvector columns. ensure_text_search_extension() at startup converts the
|
||||
schema to pgroonga structures (drops the tsvector column, builds a pgroonga
|
||||
index on the base text column).
|
||||
"""
|
||||
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
@@ -121,14 +131,35 @@ def _detect_text_search_extension() -> str:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_textsearch"
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search — true BM25 over base columns, Citus-compatible.
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_search CASCADE")
|
||||
except Exception:
|
||||
# Extension might already exist or user lacks permissions - verify it exists
|
||||
conn = op.get_bind()
|
||||
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_search'")).fetchone()
|
||||
if not result:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_search"
|
||||
elif text_search_extension == "native":
|
||||
return "native"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# ensure_text_search_extension() at runtime converts to pgroonga.
|
||||
# Treat as native here so the initial schema still creates valid columns.
|
||||
return "native"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. "
|
||||
"Must be 'native', 'vchord', 'pg_textsearch', 'pgroonga', or 'pg_search'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
"""Upgrade schema - create all tables from scratch."""
|
||||
|
||||
@@ -284,8 +315,9 @@ def _pg_upgrade() -> None:
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector bm25_catalog.bm25vector
|
||||
""")
|
||||
elif text_search_ext == "pg_textsearch":
|
||||
# Timescale pg_textsearch: dummy TEXT column for consistency (indexes operate on base columns directly)
|
||||
elif text_search_ext in ("pg_textsearch", "pg_search"):
|
||||
# Timescale pg_textsearch / ParadeDB pg_search: dummy TEXT column for
|
||||
# consistency (indexes operate on base columns directly).
|
||||
op.execute("""
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector TEXT
|
||||
@@ -350,6 +382,17 @@ def _pg_upgrade() -> None:
|
||||
USING bm25(text)
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search BM25 index on (id, text, context). The key_field
|
||||
# reloption is required and must match the table's primary key column.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context"), _pg_search_tokenizer())
|
||||
op.execute(
|
||||
"""
|
||||
CREATE INDEX idx_memory_units_text_search ON memory_units
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""".format(bm25_cols=bm25_cols)
|
||||
)
|
||||
else: # native
|
||||
# Native PostgreSQL GIN index
|
||||
op.execute("""
|
||||
|
||||
+29
@@ -7,6 +7,7 @@ the stored fact text.
|
||||
- vchord: text_signals included in tokenize() at insert time
|
||||
- native: search_vector GENERATED column regenerated to include text_signals
|
||||
- pg_textsearch: no change (index only supports a single base column)
|
||||
- pg_search: BM25 index dropped and recreated to include text_signals
|
||||
|
||||
Revision ID: a2b3c4d5e6f7
|
||||
Revises: z1u2v3w4x5y6
|
||||
@@ -18,6 +19,11 @@ from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a2b3c4d5e6f7"
|
||||
@@ -35,6 +41,10 @@ def _detect_text_search_extension() -> str:
|
||||
return os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
table = f"{schema}memory_units"
|
||||
@@ -62,6 +72,16 @@ def _pg_upgrade() -> None:
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: drop the existing BM25 index and recreate it
|
||||
# to include text_signals alongside text and context.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context", "text_signals"), _pg_search_tokenizer())
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_text_search")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_memory_units_text_search ON {table}
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
|
||||
# vchord: tokenize() call in fact_storage.py is updated to include text_signals at insert time
|
||||
# pg_textsearch: no change — index operates on the base `text` column only
|
||||
@@ -86,6 +106,15 @@ def _pg_downgrade() -> None:
|
||||
CREATE INDEX idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# Restore the original (id, text, context) BM25 index without text_signals.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context"), _pg_search_tokenizer())
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_text_search")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_memory_units_text_search ON {table}
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
|
||||
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS text_signals")
|
||||
|
||||
|
||||
+10
-10
@@ -40,20 +40,20 @@ def _get_schema_prefix() -> str:
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
# Commit the current Alembic transaction first.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block; an
|
||||
# autocommit_block runs it outside Alembic's migration transaction.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
|
||||
+1
-1
@@ -63,7 +63,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pgvectorscale":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
+253
@@ -0,0 +1,253 @@
|
||||
"""Move mental-model and observation history into dedicated tables.
|
||||
|
||||
Both histories were accumulated in a single JSONB/CLOB ``history`` column
|
||||
(``mental_models.history`` and ``memory_units.history``), appended to on every
|
||||
update. That design has two problems:
|
||||
|
||||
1. **Unbounded growth on observations.** The observation write path appended a
|
||||
snapshot on every update with no cap at all, so a frequently-reinforced
|
||||
observation grew its ``history`` array until it crossed Postgres's hard 256MB
|
||||
jsonb limit (SQLSTATE 54000), after which every further UPDATE failed and the
|
||||
row was stuck.
|
||||
2. **Wrong-axis cap on mental models.** The mental-model cap bounded the *number*
|
||||
of entries (50), not their *size* — a single large reflect snapshot could
|
||||
still blow the budget — and rewrote the whole array (plus TOAST) on every
|
||||
refresh, defeating HOT updates.
|
||||
|
||||
This migration creates one row per history entry in two dedicated tables, with
|
||||
an index that makes "most recent N for this item" cheap, then drops the old
|
||||
columns. The cap is now enforced at write time as a bounded DELETE of the
|
||||
oldest over-cap rows (see config ``*_HISTORY_MAX_ENTRIES``).
|
||||
|
||||
Revision ID: a7b8c9d0e1f2
|
||||
Revises: d3e4f5a6b7c8
|
||||
Create Date: 2026-06-05
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a7b8c9d0e1f2"
|
||||
down_revision: str | Sequence[str] | None = "d3e4f5a6b7c8"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PostgreSQL
|
||||
# ---------------------------------------------------------------------------
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Both tables share the same shape: surrogate id, FK to the parent, bank_id,
|
||||
# the snapshot payload as a single JSONB ``content`` blob, and changed_at.
|
||||
# The payload is per-row (one change per row) so it stays small — this is NOT
|
||||
# the old single-column-grows-forever design; growth is bounded by row count
|
||||
# plus the write-time cap. Folding the previous_* fields into one JSONB keeps
|
||||
# the schema dialect-simple (no array columns) and flexible.
|
||||
|
||||
# --- mental_model_history -------------------------------------------------
|
||||
# content: {"previous_content": ..., "previous_reflect_response": {...}}
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}mental_model_history (
|
||||
id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
||||
mental_model_id VARCHAR(64) NOT NULL,
|
||||
bank_id VARCHAR(64) NOT NULL,
|
||||
content JSONB NOT NULL,
|
||||
changed_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES {schema}mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_mm_history_model "
|
||||
f"ON {schema}mental_model_history (bank_id, mental_model_id, changed_at DESC, id DESC)"
|
||||
)
|
||||
|
||||
# --- observation_history --------------------------------------------------
|
||||
# content: {"previous_text", "previous_tags", "previous_occurred_start",
|
||||
# "previous_occurred_end", "previous_mentioned_at", "new_source_memory_ids"}
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}observation_history (
|
||||
id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
||||
observation_id UUID NOT NULL,
|
||||
bank_id VARCHAR(64) NOT NULL,
|
||||
content JSONB NOT NULL,
|
||||
changed_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
FOREIGN KEY (observation_id)
|
||||
REFERENCES {schema}memory_units(id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_observation_history_obs "
|
||||
f"ON {schema}observation_history (observation_id, changed_at DESC, id DESC)"
|
||||
)
|
||||
|
||||
# --- backfill mental models ----------------------------------------------
|
||||
# Explode each row's history array into rows, preserving chronological order
|
||||
# via WITH ORDINALITY so the IDENTITY id tie-breaks oldest->newest correctly.
|
||||
# changed_at is promoted to its own column; the rest of the element becomes
|
||||
# ``content`` (the ``- 'changed_at'`` strips the now-redundant key).
|
||||
op.execute(
|
||||
f"""
|
||||
INSERT INTO {schema}mental_model_history (mental_model_id, bank_id, content, changed_at)
|
||||
SELECT mm.id, mm.bank_id,
|
||||
e - 'changed_at',
|
||||
COALESCE(NULLIF(e->>'changed_at', '')::timestamptz, now())
|
||||
FROM {schema}mental_models mm
|
||||
CROSS JOIN LATERAL jsonb_array_elements(mm.history) WITH ORDINALITY a(e, ord)
|
||||
WHERE mm.history IS NOT NULL
|
||||
AND jsonb_typeof(mm.history) = 'array'
|
||||
AND jsonb_array_length(mm.history) > 0
|
||||
ORDER BY mm.id, mm.bank_id, ord
|
||||
"""
|
||||
)
|
||||
|
||||
# --- backfill observations -----------------------------------------------
|
||||
op.execute(
|
||||
f"""
|
||||
INSERT INTO {schema}observation_history (observation_id, bank_id, content, changed_at)
|
||||
SELECT mu.id, mu.bank_id,
|
||||
e - 'changed_at',
|
||||
COALESCE(NULLIF(e->>'changed_at', '')::timestamptz, now())
|
||||
FROM {schema}memory_units mu
|
||||
CROSS JOIN LATERAL jsonb_array_elements(mu.history) WITH ORDINALITY a(e, ord)
|
||||
WHERE mu.fact_type = 'observation'
|
||||
AND mu.history IS NOT NULL
|
||||
AND jsonb_typeof(mu.history) = 'array'
|
||||
AND jsonb_array_length(mu.history) > 0
|
||||
ORDER BY mu.id, ord
|
||||
"""
|
||||
)
|
||||
|
||||
# --- drop the legacy columns ---------------------------------------------
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS history")
|
||||
op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS history")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Re-add the columns (empty — historical content is not reconstructed back
|
||||
# into the array form; the dedicated tables are dropped below).
|
||||
op.execute(f"ALTER TABLE {schema}mental_models ADD COLUMN IF NOT EXISTS history JSONB DEFAULT '[]'::jsonb")
|
||||
op.execute(f"ALTER TABLE {schema}memory_units ADD COLUMN IF NOT EXISTS history JSONB DEFAULT '[]'::jsonb")
|
||||
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_observation_history_obs")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}observation_history")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mm_history_model")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}mental_model_history")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Oracle 23ai
|
||||
# ---------------------------------------------------------------------------
|
||||
def _oracle_upgrade() -> None:
|
||||
# Same single-JSONB shape as PG: ``content`` holds the snapshot payload as a
|
||||
# CLOB IS JSON. The legacy per-element JSON object (minus changed_at, promoted
|
||||
# to its own column) is carried through verbatim on backfill — the array
|
||||
# columns the previous design needed are gone.
|
||||
op.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS mental_model_history (
|
||||
id NUMBER GENERATED BY DEFAULT AS IDENTITY,
|
||||
mental_model_id VARCHAR2(256) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
content CLOB NOT NULL
|
||||
CONSTRAINT mmh_content_json CHECK (content IS JSON),
|
||||
changed_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_mental_model_history PRIMARY KEY (id),
|
||||
CONSTRAINT fk_mmh_model FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX idx_mm_history_model ON mental_model_history (bank_id, mental_model_id, changed_at DESC, id DESC)"
|
||||
)
|
||||
|
||||
op.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS observation_history (
|
||||
id NUMBER GENERATED BY DEFAULT AS IDENTITY,
|
||||
observation_id RAW(16) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
content CLOB NOT NULL
|
||||
CONSTRAINT oh_content_json CHECK (content IS JSON),
|
||||
changed_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_observation_history PRIMARY KEY (id),
|
||||
CONSTRAINT fk_oh_obs FOREIGN KEY (observation_id)
|
||||
REFERENCES memory_units(id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX idx_observation_history_obs ON observation_history (observation_id, changed_at DESC, id DESC)"
|
||||
)
|
||||
|
||||
bind = op.get_bind()
|
||||
|
||||
# Backfill via JSON_TABLE. ``content`` is the whole element (FORMAT JSON PATH
|
||||
# '$'); changed_at is also promoted to its own column. Backfilled content may
|
||||
# therefore still carry a redundant changed_at key, which the read path
|
||||
# ignores in favour of the column — harmless, and avoids JSON surgery here.
|
||||
bind.exec_driver_sql(
|
||||
"""
|
||||
INSERT INTO mental_model_history (mental_model_id, bank_id, content, changed_at)
|
||||
SELECT mm.id, mm.bank_id, jt.content, NVL(jt.changed_at, SYSTIMESTAMP)
|
||||
FROM mental_models mm,
|
||||
JSON_TABLE(mm.history, '$[*]' COLUMNS (
|
||||
seq FOR ORDINALITY,
|
||||
content CLOB FORMAT JSON PATH '$',
|
||||
changed_at TIMESTAMP WITH TIME ZONE PATH '$.changed_at'
|
||||
)) jt
|
||||
WHERE mm.history IS NOT NULL
|
||||
ORDER BY mm.id, mm.bank_id, jt.seq
|
||||
"""
|
||||
)
|
||||
|
||||
bind.exec_driver_sql(
|
||||
"""
|
||||
INSERT INTO observation_history (observation_id, bank_id, content, changed_at)
|
||||
SELECT mu.id, mu.bank_id, jt.content, NVL(jt.changed_at, SYSTIMESTAMP)
|
||||
FROM memory_units mu,
|
||||
JSON_TABLE(mu.history, '$[*]' COLUMNS (
|
||||
seq FOR ORDINALITY,
|
||||
content CLOB FORMAT JSON PATH '$',
|
||||
changed_at TIMESTAMP WITH TIME ZONE PATH '$.changed_at'
|
||||
)) jt
|
||||
WHERE mu.fact_type = 'observation' AND mu.history IS NOT NULL
|
||||
ORDER BY mu.id, jt.seq
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute("ALTER TABLE mental_models DROP COLUMN history")
|
||||
op.execute("ALTER TABLE memory_units DROP COLUMN history")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("ALTER TABLE mental_models ADD history CLOB DEFAULT '[]' NOT NULL")
|
||||
op.execute("ALTER TABLE memory_units ADD history CLOB DEFAULT '[]'")
|
||||
op.execute("DROP TABLE observation_history CASCADE CONSTRAINTS")
|
||||
op.execute("DROP TABLE mental_model_history CASCADE CONSTRAINTS")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+156
@@ -0,0 +1,156 @@
|
||||
"""Repair: install maintenance routines on the ``public`` / base-schema run.
|
||||
|
||||
The original maintenance-routines migration (``e5f6a7b8c9d0``) only created the
|
||||
shared ``public.banks_needing_consolidation()`` and
|
||||
``public.schemas_with_expired_rows(...)`` routines when the run had *no*
|
||||
``target_schema`` at all. But the single-tenant runtime always migrates an
|
||||
explicit schema — which defaults to ``public`` — so on every default
|
||||
PostgreSQL deployment the migration was stamped as applied while the functions
|
||||
were never created. Background maintenance then logs::
|
||||
|
||||
Retention sweep failed for llm_requests: function public.schemas_with_expired_rows(...) does not exist
|
||||
Consolidation reconcile discovery failed: function public.banks_needing_consolidation() does not exist
|
||||
|
||||
See https://github.com/vectorize-io/hindsight/issues/2056.
|
||||
|
||||
Because ``e5f6a7b8c9d0`` is already stamped on affected ``0.8.0`` databases,
|
||||
editing it would not re-run it there. This forward migration re-installs the
|
||||
functions idempotently (``CREATE OR REPLACE``) on the run that targets the
|
||||
shared ``public`` schema (base run with no ``target_schema``, or an explicit
|
||||
``target_schema=public``), self-healing already-upgraded deployments and
|
||||
covering fresh upgrades from earlier versions.
|
||||
|
||||
Per-tenant runs against a non-``public`` schema still skip it: re-issuing
|
||||
``CREATE OR REPLACE FUNCTION public....`` from each concurrent tenant migration
|
||||
aborts with ``tuple concurrently updated`` on the ``pg_proc`` catalog row, and
|
||||
the base/public run has already created the functions for every tenant to use.
|
||||
Runs that target ``public`` are serialized by the per-schema migration advisory
|
||||
lock, so only one wins the create.
|
||||
|
||||
PostgreSQL only — the worker poller and these tables are not wired for Oracle,
|
||||
so the Oracle slot is intentionally absent (mirrors ``e5f6a7b8c9d0``).
|
||||
|
||||
Revision ID: b2d4f6a8c1e3
|
||||
Revises: e5f6a7b8c9d0
|
||||
Create Date: 2026-06-08
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b2d4f6a8c1e3"
|
||||
down_revision: str | Sequence[str] | None = "e5f6a7b8c9d0"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _should_install_public_routines(target_schema: str | None) -> bool:
|
||||
"""True for the run that must (re)create the shared ``public.*`` routines.
|
||||
|
||||
The routines physically live in ``public`` (hard-coded ``public.`` qualifier
|
||||
in the SQL below), so they must be installed exactly once — on the base run
|
||||
(no ``target_schema``) or on the run that explicitly targets ``public``. A
|
||||
run against any other tenant schema skips it to avoid concurrent
|
||||
``CREATE OR REPLACE`` on the same ``pg_proc`` row.
|
||||
"""
|
||||
return not target_schema or target_schema == "public"
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if not _should_install_public_routines(context.config.get_main_option("target_schema")):
|
||||
return
|
||||
# Banks with eligible-but-unscheduled facts and no in-flight consolidation.
|
||||
# Auto-consolidation is filtered here only at the bank level (cheap prune);
|
||||
# the full hierarchical resolution (global -> tenant -> bank, plus
|
||||
# enable_observations) is done by the caller for the small returned set.
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.banks_needing_consolidation()
|
||||
RETURNS TABLE(schema_name text, bank_id text)
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
BEGIN
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = 'memory_units' AND c.relkind = 'r'
|
||||
LOOP
|
||||
RETURN QUERY EXECUTE format($q$
|
||||
SELECT %1$L::text, m.bank_id
|
||||
FROM %1$I.memory_units m
|
||||
JOIN %1$I.banks b ON b.bank_id = m.bank_id
|
||||
WHERE m.consolidated_at IS NULL
|
||||
AND m.consolidation_failed_at IS NULL
|
||||
AND m.fact_type IN ('experience', 'world')
|
||||
AND COALESCE(b.config -> 'enable_auto_consolidation', 'true'::jsonb) <> 'false'::jsonb
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM %1$I.async_operations o
|
||||
WHERE o.bank_id = m.bank_id
|
||||
AND o.operation_type = 'consolidation'
|
||||
AND o.status IN ('pending', 'processing')
|
||||
)
|
||||
GROUP BY m.bank_id
|
||||
$q$, sch);
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
# Schemas holding at least one row of p_table older than p_days. p_ts_col is
|
||||
# the timestamp column to compare. Returns nothing when p_days <= 0
|
||||
# (retention disabled).
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.schemas_with_expired_rows(
|
||||
p_table text, p_ts_col text, p_days int
|
||||
)
|
||||
RETURNS SETOF text
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
has_expired boolean;
|
||||
BEGIN
|
||||
IF p_days IS NULL OR p_days <= 0 THEN
|
||||
RETURN;
|
||||
END IF;
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = p_table AND c.relkind = 'r'
|
||||
LOOP
|
||||
EXECUTE format(
|
||||
'SELECT EXISTS (SELECT 1 FROM %I.%I WHERE %I < NOW() - make_interval(days => $1))',
|
||||
sch, p_table, p_ts_col
|
||||
) INTO has_expired USING p_days;
|
||||
IF has_expired THEN
|
||||
RETURN NEXT sch;
|
||||
END IF;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# No-op: ``e5f6a7b8c9d0`` owns the lifecycle of these functions and drops
|
||||
# them on its own downgrade. This migration only ever (re)creates them, so
|
||||
# there is nothing to undo without racing that migration's DROP.
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+25
-27
@@ -37,37 +37,35 @@ def _get_schema_prefix() -> str:
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# Partial index on occurred_start (covers "occurred_start BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_start "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_start) "
|
||||
f"WHERE occurred_start IS NOT NULL"
|
||||
)
|
||||
# Partial index on occurred_end (covers "occurred_end BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_end "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_end) "
|
||||
f"WHERE occurred_end IS NOT NULL"
|
||||
)
|
||||
# Partial index on mentioned_at (covers "mentioned_at BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_mentioned_at "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, mentioned_at) "
|
||||
f"WHERE mentioned_at IS NOT NULL"
|
||||
)
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block; an
|
||||
# autocommit_block runs each statement outside Alembic's migration transaction.
|
||||
with op.get_context().autocommit_block():
|
||||
# Partial index on occurred_start (covers "occurred_start BETWEEN $4 AND $5")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_start "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_start) "
|
||||
f"WHERE occurred_start IS NOT NULL"
|
||||
)
|
||||
# Partial index on occurred_end (covers "occurred_end BETWEEN $4 AND $5")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_end "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_end) "
|
||||
f"WHERE occurred_end IS NOT NULL"
|
||||
)
|
||||
# Partial index on mentioned_at (covers "mentioned_at BETWEEN $4 AND $5")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_mentioned_at "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, mentioned_at) "
|
||||
f"WHERE mentioned_at IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_mentioned_at")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_end")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_start")
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_mentioned_at")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_end")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_start")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
|
||||
+106
@@ -0,0 +1,106 @@
|
||||
"""Add graph_maintenance_queue table
|
||||
|
||||
Queue of memory_units whose outgoing temporal/semantic links lost a
|
||||
neighbour to a delete. Drained by the async graph_maintenance worker,
|
||||
which tops the unit's links back up using the same probes retain runs.
|
||||
|
||||
The queue only targets the link-recompute pass. The worker also runs
|
||||
bank-wide sweeps (orphan-entity prune, stale-cooccurrence prune) on each
|
||||
invocation; those don't need per-target queueing.
|
||||
|
||||
Revision ID: b5a4c3e2f1d8
|
||||
Revises: e9b2c7d1f3a4
|
||||
Create Date: 2026-05-27
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b5a4c3e2f1d8"
|
||||
down_revision: str | Sequence[str] | None = "e9b2c7d1f3a4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# Composite PK gives us natural ON CONFLICT DO NOTHING dedup when the same
|
||||
# unit is enqueued from overlapping deletes. No FK to memory_units: if the
|
||||
# unit is deleted between enqueue and drain, the worker observes it's gone
|
||||
# and skips — a cascade would erase the work order, but that work has
|
||||
# already been satisfied (no surviving row to maintain).
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}graph_maintenance_queue (
|
||||
bank_id TEXT NOT NULL,
|
||||
unit_id UUID NOT NULL,
|
||||
enqueued_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
PRIMARY KEY (bank_id, unit_id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_graph_maintenance_queue_bank_enqueued
|
||||
ON {schema}graph_maintenance_queue (bank_id, enqueued_at)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_graph_maintenance_queue_bank_enqueued")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}graph_maintenance_queue")
|
||||
|
||||
|
||||
def _oracle_execute_ignoring_955(sql: str) -> None:
|
||||
"""Run a CREATE statement and swallow ORA-00955 (object already exists).
|
||||
|
||||
Mirrors the helper in the Oracle baseline migration so reruns stay safe
|
||||
on a database where the table was created by an earlier partial run.
|
||||
"""
|
||||
block = (
|
||||
"BEGIN "
|
||||
"EXECUTE IMMEDIATE :stmt; "
|
||||
"EXCEPTION WHEN OTHERS THEN "
|
||||
"IF SQLCODE = -955 THEN NULL; ELSE RAISE; END IF; "
|
||||
"END;"
|
||||
)
|
||||
op.get_bind().exec_driver_sql(block, {"stmt": sql.strip()})
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
_oracle_execute_ignoring_955(
|
||||
"""
|
||||
CREATE TABLE graph_maintenance_queue (
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
unit_id RAW(16) NOT NULL,
|
||||
enqueued_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_graph_maintenance_queue PRIMARY KEY (bank_id, unit_id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
_oracle_execute_ignoring_955(
|
||||
"CREATE INDEX idx_graph_maintenance_queue_bank_enqueued ON graph_maintenance_queue (bank_id, enqueued_at)"
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("DROP INDEX idx_graph_maintenance_queue_bank_enqueued")
|
||||
op.execute("DROP TABLE graph_maintenance_queue")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+152
@@ -0,0 +1,152 @@
|
||||
"""Re-create vchord vector indexes with vector_cosine_ops
|
||||
|
||||
Revision ID: b8c9d0e1f2a3
|
||||
Revises: 86f7a033d372
|
||||
Create Date: 2026-05-20
|
||||
|
||||
vchordrq operator classes are bound 1:1 to operators in PostgreSQL:
|
||||
vector_l2_ops only matches ``<->``, while every Hindsight ANN query uses
|
||||
``<=>`` (cosine distance). The previous vchord mapping used vector_l2_ops,
|
||||
so vchord deployments could never use the index — every ANN query fell
|
||||
back to a sequential scan with per-row cosine computation.
|
||||
|
||||
This migration finds any vchordrq index built with vector_l2_ops in the
|
||||
target schema and re-creates it with vector_cosine_ops, using
|
||||
``CREATE INDEX CONCURRENTLY`` so it can run online. It is a no-op when:
|
||||
|
||||
* the configured vector extension is not vchord, or
|
||||
* no matching indexes exist (already on cosine ops).
|
||||
|
||||
Only PostgreSQL is affected; the Oracle 23ai dialect uses its own native
|
||||
vector index and does not depend on this mapping.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
from hindsight_api._vector_index import configured_vector_extension
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b8c9d0e1f2a3"
|
||||
down_revision: str | Sequence[str] | None = "86f7a033d372"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _rebuild_vchordrq_indexes(old_ops: str, new_ops: str) -> None:
|
||||
"""Rebuild vchordrq indexes using ``old_ops`` so they use ``new_ops``.
|
||||
|
||||
Each index is rebuilt with CREATE INDEX CONCURRENTLY under a fresh name,
|
||||
then the old index is dropped and the new one renamed to take its place.
|
||||
Must be called inside an ``autocommit_block()`` because CONCURRENTLY
|
||||
cannot run inside a transaction.
|
||||
"""
|
||||
bind = op.get_bind()
|
||||
# `or None` collapses both unset and explicit empty-string Alembic options
|
||||
# into NULL so the COALESCE below falls back to current_schema() in either
|
||||
# case. Without it, an empty-string option would filter on `schemaname = ''`
|
||||
# and skip every real schema.
|
||||
target_schema = context.config.get_main_option("target_schema") or None
|
||||
prefix = _pg_schema_prefix()
|
||||
|
||||
rows = bind.execute(
|
||||
text(
|
||||
"SELECT indexname, indexdef FROM pg_indexes "
|
||||
"WHERE schemaname = COALESCE(:target_schema, current_schema()) "
|
||||
"AND indexdef ILIKE '%vchordrq%' "
|
||||
"AND indexdef ILIKE :ops_like"
|
||||
),
|
||||
{"target_schema": target_schema, "ops_like": f"%{old_ops}%"},
|
||||
).fetchall()
|
||||
|
||||
for idx_name, indexdef in rows:
|
||||
# pg_get_indexdef() emits the canonical form `CREATE INDEX <name> ON …`,
|
||||
# so <name> is the first textual occurrence — both substitutions below
|
||||
# rely on that.
|
||||
new_def = indexdef.replace(old_ops, new_ops, 1)
|
||||
temp_name = f"{idx_name}__opclass_swap"
|
||||
new_def = new_def.replace(idx_name, temp_name, 1)
|
||||
new_def = re.sub(
|
||||
r"^CREATE\s+INDEX\b",
|
||||
"CREATE INDEX CONCURRENTLY IF NOT EXISTS",
|
||||
new_def,
|
||||
count=1,
|
||||
)
|
||||
|
||||
# CREATE INDEX CONCURRENTLY can leave the partial index as INVALID if a
|
||||
# previous run errored (disk pressure, lock conflict, signal). Without
|
||||
# this drop the CONCURRENTLY IF NOT EXISTS below would skip creation,
|
||||
# then we'd drop the original and rename the broken index into its
|
||||
# place — silently restoring the seq-scan bug this migration fixes.
|
||||
op.execute(f'DROP INDEX IF EXISTS {prefix}"{temp_name}"')
|
||||
op.execute(new_def)
|
||||
|
||||
# Even on a clean run CONCURRENTLY can finish with indisvalid = false
|
||||
# (e.g. constraint violation during the second build scan). Refuse to
|
||||
# promote in that case so we never alias an INVALID index over a working
|
||||
# one.
|
||||
is_valid = bind.execute(
|
||||
text(
|
||||
"SELECT i.indisvalid "
|
||||
"FROM pg_class c "
|
||||
"JOIN pg_index i ON c.oid = i.indexrelid "
|
||||
"JOIN pg_namespace n ON c.relnamespace = n.oid "
|
||||
"WHERE c.relname = :name "
|
||||
" AND n.nspname = COALESCE(:target_schema, current_schema())"
|
||||
),
|
||||
{"name": temp_name, "target_schema": target_schema},
|
||||
).scalar()
|
||||
if not is_valid:
|
||||
raise RuntimeError(
|
||||
f"vchordrq index rebuild produced an INVALID index ({temp_name}); "
|
||||
"drop it manually and re-run the migration."
|
||||
)
|
||||
|
||||
# DROP + RENAME atomically. A crash between the two would leave
|
||||
# `temp_name` as a valid orphan and the canonical name missing —
|
||||
# next run's `pg_indexes` filter (looking for vector_l2_ops) wouldn't
|
||||
# find anything to recover from, so the index would stay gone. PG
|
||||
# runs the DO block in its own server-side transaction, so either
|
||||
# both succeed or both roll back.
|
||||
op.execute(
|
||||
f"""
|
||||
DO $$
|
||||
BEGIN
|
||||
DROP INDEX IF EXISTS {prefix}"{idx_name}";
|
||||
ALTER INDEX {prefix}"{temp_name}" RENAME TO "{idx_name}";
|
||||
END $$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if configured_vector_extension() != "vchord":
|
||||
return
|
||||
with op.get_context().autocommit_block():
|
||||
_rebuild_vchordrq_indexes("vector_l2_ops", "vector_cosine_ops")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
if configured_vector_extension() != "vchord":
|
||||
return
|
||||
with op.get_context().autocommit_block():
|
||||
_rebuild_vchordrq_indexes("vector_cosine_ops", "vector_l2_ops")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+8
-7
@@ -47,17 +47,18 @@ def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# GIN index on canonical_name enables sub-millisecond trigram similarity queries
|
||||
# (% operator, similarity()) instead of full-table scans across all bank entities.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
|
||||
)
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}entities_canonical_name_trgm_idx")
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}entities_canonical_name_trgm_idx")
|
||||
# Note: not dropping pg_trgm extension as other indexes may depend on it
|
||||
|
||||
|
||||
|
||||
+45
@@ -0,0 +1,45 @@
|
||||
"""Merge graph_maintenance_queue and vchord_cosine_opclass heads.
|
||||
|
||||
Revision ID: c1d2e3f4a5b6
|
||||
Revises: b5a4c3e2f1d8, b8c9d0e1f2a3
|
||||
Create Date: 2026-05-29
|
||||
|
||||
PRs #1668 (vchord cosine opclass) and #1772 (async link recompute) both
|
||||
branched off the same parent and were merged onto main without rebasing,
|
||||
leaving two parallel Alembic heads. This is a structural merge revision
|
||||
with no schema changes — its only job is to unify the DAG so
|
||||
``alembic upgrade head`` is unambiguous again.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c1d2e3f4a5b6"
|
||||
down_revision: str | Sequence[str] | None = ("b5a4c3e2f1d8", "b8c9d0e1f2a3")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+24
-28
@@ -50,39 +50,35 @@ def _get_schema_prefix() -> str:
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
# Commit the current Alembic transaction, then issue each CONCURRENTLY
|
||||
# statement in its own implicit autocommit transaction.
|
||||
# IF NOT EXISTS makes each statement idempotent if the migration is retried.
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block; an
|
||||
# autocommit_block runs each statement outside Alembic's migration
|
||||
# transaction. IF NOT EXISTS makes each statement idempotent on retry.
|
||||
with op.get_context().autocommit_block():
|
||||
# Index for the semantic *incoming* direction in link_expansion_retrieval.py.
|
||||
# Replaces the BitmapAnd of idx_memory_links_to_unit ∩ idx_memory_links_link_type
|
||||
# with a single composite index scan.
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_type_weight "
|
||||
f"ON {schema}memory_links(to_unit_id, link_type, weight DESC)"
|
||||
)
|
||||
|
||||
# Index for the semantic *incoming* direction in link_expansion_retrieval.py.
|
||||
# Replaces the BitmapAnd of idx_memory_links_to_unit ∩ idx_memory_links_link_type
|
||||
# with a single composite index scan.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_type_weight "
|
||||
f"ON {schema}memory_links(to_unit_id, link_type, weight DESC)"
|
||||
)
|
||||
|
||||
# Covering index for entity co-occurrence expansion.
|
||||
# Enables an index-only scan: entity_id and to_unit_id are read from the
|
||||
# index leaf pages instead of the heap, eliminating ~2 500 random heap-page
|
||||
# reads per expansion query.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity_covering "
|
||||
f"ON {schema}memory_links(from_unit_id) "
|
||||
f"INCLUDE (to_unit_id, entity_id) "
|
||||
f"WHERE link_type = 'entity'"
|
||||
)
|
||||
# Covering index for entity co-occurrence expansion.
|
||||
# Enables an index-only scan: entity_id and to_unit_id are read from the
|
||||
# index leaf pages instead of the heap, eliminating ~2 500 random heap-page
|
||||
# reads per expansion query.
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity_covering "
|
||||
f"ON {schema}memory_links(from_unit_id) "
|
||||
f"INCLUDE (to_unit_id, entity_id) "
|
||||
f"WHERE link_type = 'entity'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_to_type_weight")
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_to_type_weight")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
|
||||
+96
@@ -0,0 +1,96 @@
|
||||
"""Add llm_requests table for per-bank LLM request tracing.
|
||||
|
||||
Stores one row per logical LLM call Hindsight makes (success and failure),
|
||||
capturing the input messages, model output, token usage (input/output/cached/
|
||||
total), finish reason, and caller metadata. Disabled by default at the
|
||||
application layer (HINDSIGHT_API_LLM_TRACE_ENABLED); this migration only
|
||||
creates the table.
|
||||
|
||||
PostgreSQL only — the tracing subsystem is not wired for Oracle, so the Oracle
|
||||
slot is intentionally absent (mirrors the audit_log table).
|
||||
|
||||
Revision ID: d3e4f5a6b7c8
|
||||
Revises: c1d2e3f4a5b6
|
||||
Create Date: 2026-06-01
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "d3e4f5a6b7c8"
|
||||
down_revision: str | Sequence[str] | None = "c1d2e3f4a5b6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}llm_requests (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
bank_id TEXT,
|
||||
operation TEXT,
|
||||
scope TEXT,
|
||||
-- OTel-style grouping: trace_id is shared by every LLM call of one
|
||||
-- operation invocation (e.g. all calls of a single reflect run);
|
||||
-- parent_span_id is that operation span; span_id is this call.
|
||||
trace_id TEXT,
|
||||
span_id TEXT,
|
||||
parent_span_id TEXT,
|
||||
provider TEXT,
|
||||
model TEXT,
|
||||
status TEXT NOT NULL,
|
||||
started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
ended_at TIMESTAMPTZ,
|
||||
duration_ms INTEGER,
|
||||
input_tokens INTEGER,
|
||||
output_tokens INTEGER,
|
||||
cached_tokens INTEGER,
|
||||
total_tokens INTEGER,
|
||||
input JSONB,
|
||||
output JSONB,
|
||||
error TEXT,
|
||||
llm_info JSONB DEFAULT '{{}}'::jsonb,
|
||||
metadata JSONB DEFAULT '{{}}'::jsonb
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_llm_requests_bank_started ON {schema}llm_requests (bank_id, started_at DESC)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_llm_requests_status_started ON {schema}llm_requests (status, started_at DESC)"
|
||||
)
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_llm_requests_started ON {schema}llm_requests (started_at DESC)")
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_llm_requests_trace ON {schema}llm_requests (bank_id, trace_id, started_at)"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_llm_requests_started")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_llm_requests_status_started")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_llm_requests_bank_started")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}llm_requests")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+17
-16
@@ -33,26 +33,27 @@ def _get_schema_prefix() -> str:
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# DROP + CREATE CONCURRENTLY must run outside a transaction block.
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WITH (fastupdate=off) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
# DROP + CREATE CONCURRENTLY must run outside a transaction block; an
|
||||
# autocommit_block runs them outside Alembic's migration transaction.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WITH (fastupdate=off) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
|
||||
+1
-1
@@ -55,7 +55,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pgvectorscale":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
+133
@@ -0,0 +1,133 @@
|
||||
"""Drop indexes that are unused or redundant with composite indexes.
|
||||
|
||||
Code audit identified the following indexes as either dead (no code path
|
||||
exercises them) or fully covered by composite indexes the planner already
|
||||
prefers:
|
||||
|
||||
memory_links:
|
||||
1. idx_memory_links_entity_covering — entity co-occurrence expansion was
|
||||
rewritten to traverse unit_entities instead of memory_links, so no code
|
||||
path filters memory_links on (link_type = 'entity').
|
||||
2. idx_memory_links_from_unit — redundant. idx_memory_links_from_type_weight
|
||||
(from_unit_id, link_type, weight DESC) leads with the same column and
|
||||
answers every from_unit_id = X query.
|
||||
3. idx_memory_links_to_unit — redundant. idx_memory_links_to_type_weight
|
||||
(to_unit_id, link_type, weight DESC) leads with the same column.
|
||||
4. idx_memory_links_link_type — no application query filters on link_type
|
||||
alone; the composite indexes above serve every (from/to + link_type)
|
||||
predicate.
|
||||
|
||||
entities:
|
||||
5. idx_entities_canonical_name — superseded by
|
||||
entities_canonical_name_lower_trgm_idx (case-insensitive lookups).
|
||||
6. entities_canonical_name_trgm_idx — superseded by the lowercase variant
|
||||
in migration 2eee35aa3cfc, but the original was never dropped on schemas
|
||||
that ran the prior migration.
|
||||
|
||||
documents:
|
||||
7. idx_documents_retain_params — GIN index on retain_params JSONB; no query
|
||||
uses jsonb containment on this column.
|
||||
8. idx_documents_content_hash — content-hash lookups happen on the chunks
|
||||
table (chunks.content_hash, indexed separately).
|
||||
|
||||
unit_entities:
|
||||
9. idx_unit_entities_entity — defensive drop. Migration h3i4j5k6l7m8 already
|
||||
issues DROP INDEX IF EXISTS for this; this re-runs the drop idempotently
|
||||
to cover any schema that missed the previous migration.
|
||||
|
||||
All drops use CONCURRENTLY + IF EXISTS so they neither block writers nor
|
||||
fail on schemas where the index is already gone.
|
||||
|
||||
Revision ID: e1b2c3d4f5a6
|
||||
Revises: p4q5r6s7t8u9
|
||||
Create Date: 2026-05-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e1b2c3d4f5a6"
|
||||
down_revision: str | Sequence[str] | None = "p4q5r6s7t8u9"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
_PG_INDEXES_TO_DROP: tuple[str, ...] = (
|
||||
"idx_memory_links_entity_covering",
|
||||
"idx_memory_links_from_unit",
|
||||
"idx_memory_links_to_unit",
|
||||
"idx_memory_links_link_type",
|
||||
"idx_entities_canonical_name",
|
||||
"entities_canonical_name_trgm_idx",
|
||||
"idx_documents_retain_params",
|
||||
"idx_documents_content_hash",
|
||||
"idx_unit_entities_entity",
|
||||
)
|
||||
|
||||
|
||||
def _schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _schema_prefix()
|
||||
# DROP INDEX CONCURRENTLY cannot run inside a transaction block; an
|
||||
# autocommit_block drops out of Alembic's migration transaction so each
|
||||
# statement runs in its own autocommit. IF EXISTS makes each statement
|
||||
# idempotent across schemas that already dropped (or never had) the index.
|
||||
with op.get_context().autocommit_block():
|
||||
for index_name in _PG_INDEXES_TO_DROP:
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}{index_name}")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _schema_prefix()
|
||||
|
||||
# Recreate the dropped indexes in the same shape the prior migrations used,
|
||||
# so a downgrade leaves the schema in the state the previous head expected.
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity_covering "
|
||||
f"ON {schema}memory_links(from_unit_id) "
|
||||
f"INCLUDE (to_unit_id, entity_id) "
|
||||
f"WHERE link_type = 'entity'"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_from_unit ON {schema}memory_links(from_unit_id)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_unit ON {schema}memory_links(to_unit_id)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_link_type ON {schema}memory_links(link_type)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_entities_canonical_name ON {schema}entities(canonical_name)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_documents_retain_params "
|
||||
f"ON {schema}documents USING GIN (retain_params)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_documents_content_hash ON {schema}documents(content_hash)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_unit_entities_entity ON {schema}unit_entities(entity_id)"
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+153
@@ -0,0 +1,153 @@
|
||||
"""Add server-side routines for background maintenance sweeps.
|
||||
|
||||
Installs two PL/pgSQL discovery routines in the ``public`` schema. Both loop
|
||||
over every schema that actually holds the relevant table (via ``pg_class``), so
|
||||
a single function call covers all tenants in one round-trip instead of the
|
||||
per-tenant query storm that a client-side loop would create at thousands of
|
||||
tenants.
|
||||
|
||||
- ``public.banks_needing_consolidation()`` -> (schema_name, bank_id) for banks
|
||||
that have eligible-but-unscheduled facts (``consolidated_at IS NULL AND
|
||||
consolidation_failed_at IS NULL`` for consolidatable fact types), have
|
||||
auto-consolidation not explicitly disabled at the bank level, and have no
|
||||
consolidation operation already pending/processing. Drives the periodic
|
||||
reconcile that re-schedules consolidation after a terminal failure left facts
|
||||
stranded (see HINDSIGHT_API_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS).
|
||||
|
||||
- ``public.schemas_with_expired_rows(p_table, p_ts_col, p_days)`` -> schema
|
||||
names that hold at least one ``p_table`` row older than ``p_days``. Drives the
|
||||
cross-tenant retention sweeps for ``audit_log`` and ``llm_requests``; the loop
|
||||
then issues a DELETE only against the returned schemas.
|
||||
|
||||
These are read-only (STABLE) discovery routines — the caller performs the
|
||||
enqueue/DELETE — so installing them never mutates data.
|
||||
|
||||
PostgreSQL only — the worker poller and these tables are not wired for Oracle,
|
||||
so the Oracle slot is intentionally absent (mirrors the audit_log / llm_requests
|
||||
table migrations). The routines live in ``public`` and are CREATE OR REPLACE, so
|
||||
running this migration once per tenant schema is idempotent.
|
||||
|
||||
Revision ID: e5f6a7b8c9d0
|
||||
Revises: a7b8c9d0e1f2
|
||||
Create Date: 2026-06-05
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e5f6a7b8c9d0"
|
||||
down_revision: str | Sequence[str] | None = "a7b8c9d0e1f2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _is_base_schema_run() -> bool:
|
||||
"""True only for the base-schema migration (no per-tenant target_schema).
|
||||
|
||||
These routines live in the shared ``public`` schema, so they must be created
|
||||
exactly once. Running ``CREATE OR REPLACE FUNCTION public....`` again from each
|
||||
concurrent per-tenant migration aborts with ``tuple concurrently updated`` on
|
||||
the ``pg_proc`` catalog row, so tenant runs skip it (the base run already
|
||||
created the function for every tenant to use).
|
||||
"""
|
||||
return not context.config.get_main_option("target_schema")
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if not _is_base_schema_run():
|
||||
return
|
||||
# Banks with eligible-but-unscheduled facts and no in-flight consolidation.
|
||||
# Auto-consolidation is filtered here only at the bank level (cheap prune);
|
||||
# the full hierarchical resolution (global -> tenant -> bank, plus
|
||||
# enable_observations) is done by the caller for the small returned set.
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.banks_needing_consolidation()
|
||||
RETURNS TABLE(schema_name text, bank_id text)
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
BEGIN
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = 'memory_units' AND c.relkind = 'r'
|
||||
LOOP
|
||||
RETURN QUERY EXECUTE format($q$
|
||||
SELECT %1$L::text, m.bank_id
|
||||
FROM %1$I.memory_units m
|
||||
JOIN %1$I.banks b ON b.bank_id = m.bank_id
|
||||
WHERE m.consolidated_at IS NULL
|
||||
AND m.consolidation_failed_at IS NULL
|
||||
AND m.fact_type IN ('experience', 'world')
|
||||
AND COALESCE(b.config -> 'enable_auto_consolidation', 'true'::jsonb) <> 'false'::jsonb
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM %1$I.async_operations o
|
||||
WHERE o.bank_id = m.bank_id
|
||||
AND o.operation_type = 'consolidation'
|
||||
AND o.status IN ('pending', 'processing')
|
||||
)
|
||||
GROUP BY m.bank_id
|
||||
$q$, sch);
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
# Schemas holding at least one row of p_table older than p_days. p_ts_col is
|
||||
# the timestamp column to compare. Returns nothing when p_days <= 0
|
||||
# (retention disabled).
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.schemas_with_expired_rows(
|
||||
p_table text, p_ts_col text, p_days int
|
||||
)
|
||||
RETURNS SETOF text
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
has_expired boolean;
|
||||
BEGIN
|
||||
IF p_days IS NULL OR p_days <= 0 THEN
|
||||
RETURN;
|
||||
END IF;
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = p_table AND c.relkind = 'r'
|
||||
LOOP
|
||||
EXECUTE format(
|
||||
'SELECT EXISTS (SELECT 1 FROM %I.%I WHERE %I < NOW() - make_interval(days => $1))',
|
||||
sch, p_table, p_ts_col
|
||||
) INTO has_expired USING p_days;
|
||||
IF has_expired THEN
|
||||
RETURN NEXT sch;
|
||||
END IF;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
if not _is_base_schema_run():
|
||||
return
|
||||
op.execute("DROP FUNCTION IF EXISTS public.banks_needing_consolidation()")
|
||||
op.execute("DROP FUNCTION IF EXISTS public.schemas_with_expired_rows(text, text, int)")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+91
@@ -0,0 +1,91 @@
|
||||
"""Drop materialized entity rows from memory_links.
|
||||
|
||||
Entity edges are no longer stored in ``memory_links``. The /graph endpoint
|
||||
derives them on demand from ``unit_entities``, and recall already used the
|
||||
``unit_entities`` self-join. Storing entity rows duplicated state we never
|
||||
read from the link table — on a 10k-unit benchmark bank, entity rows were
|
||||
53% of all link rows (~190 MB after indexes) and recall never touched them.
|
||||
|
||||
This migration deletes ``memory_links`` rows with ``link_type = 'entity'``.
|
||||
``idx_memory_links_entity_covering`` was already dropped by migration
|
||||
``e1b2c3d4f5a6``; we still issue ``DROP INDEX IF EXISTS`` defensively in case
|
||||
this migration runs against an older snapshot that predates that one.
|
||||
|
||||
Revision ID: e9b2c7d1f3a4
|
||||
Revises: e1b2c3d4f5a6
|
||||
Create Date: 2026-05-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e9b2c7d1f3a4"
|
||||
down_revision: str | Sequence[str] | None = "e1b2c3d4f5a6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
|
||||
# Drop the partial covering index first so the bulk DELETE doesn't churn it.
|
||||
# DROP INDEX CONCURRENTLY, and the DO block's per-batch COMMIT, both require
|
||||
# running outside Alembic's migration transaction — an autocommit_block
|
||||
# commits it and switches the connection to autocommit for the duration.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
|
||||
|
||||
# Delete entity rows. Chunked to keep individual transactions small on
|
||||
# large banks (the perf-medium bench had ~345k entity rows; production
|
||||
# banks can be much larger).
|
||||
op.execute(
|
||||
f"""
|
||||
DO $$
|
||||
DECLARE
|
||||
deleted INTEGER;
|
||||
BEGIN
|
||||
LOOP
|
||||
DELETE FROM {schema}memory_links
|
||||
WHERE ctid IN (
|
||||
SELECT ctid FROM {schema}memory_links
|
||||
WHERE link_type = 'entity'
|
||||
LIMIT 50000
|
||||
);
|
||||
GET DIAGNOSTICS deleted = ROW_COUNT;
|
||||
EXIT WHEN deleted = 0;
|
||||
COMMIT;
|
||||
END LOOP;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# Cannot reconstruct deleted entity links — the writer was path-dependent
|
||||
# on retain order. New retains will not produce entity rows either, so the
|
||||
# partial index would stay empty. Leave both no-op.
|
||||
pass
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
op.execute("DELETE FROM memory_links WHERE link_type = 'entity'")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+58
-4
@@ -16,6 +16,11 @@ from collections.abc import Sequence
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -87,7 +92,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pg_diskann":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
@@ -95,9 +100,15 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Detect or validate text search extension: 'native', 'vchord', 'pg_textsearch',
|
||||
'pgroonga', or 'pg_search'. Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
|
||||
pgroonga is treated as native here so this migration still creates valid
|
||||
tsvector columns; ensure_text_search_extension() at startup converts the
|
||||
reflections table (renamed from pinned_reflections in p1k2l3m4n5o6) to
|
||||
pgroonga structures. The learnings table is dropped in p1k2l3m4n5o6 so its
|
||||
transient native-style column never reaches steady state.
|
||||
"""
|
||||
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
@@ -125,14 +136,33 @@ def _detect_text_search_extension() -> str:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_textsearch"
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search — true BM25 over base columns, Citus-compatible.
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_search CASCADE")
|
||||
except Exception:
|
||||
conn = op.get_bind()
|
||||
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_search'")).fetchone()
|
||||
if not result:
|
||||
raise
|
||||
return "pg_search"
|
||||
elif text_search_extension == "native":
|
||||
return "native"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# Treat as native here; ensure_text_search_extension() converts the
|
||||
# reflections table to pgroonga structures at runtime.
|
||||
return "native"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. "
|
||||
"Must be 'native', 'vchord', 'pg_textsearch', 'pgroonga', or 'pg_search'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
"""Create learnings and pinned_reflections tables."""
|
||||
schema = _get_schema_prefix()
|
||||
@@ -200,6 +230,18 @@ def _pg_upgrade() -> None:
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25(text) WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: dummy TEXT column; BM25 index is built directly over (id, text)
|
||||
# with key_field='id' (matches the table's primary key).
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text",), _pg_search_tokenizer())
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}learnings ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
@@ -264,6 +306,18 @@ def _pg_upgrade() -> None:
|
||||
USING bm25(content)
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: dummy TEXT column; BM25 index over (id, name, content)
|
||||
# with key_field='id'.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("name", "content"), _pg_search_tokenizer())
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
|
||||
+170
@@ -0,0 +1,170 @@
|
||||
"""Drop GENERATED expression on tsvector search_vector columns.
|
||||
|
||||
The search_vector tsvector column was originally GENERATED ALWAYS with a
|
||||
hardcoded ``to_tsvector('english', ...)`` expression. To support configurable
|
||||
``HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE``, we convert it to a
|
||||
regular tsvector column that the application populates at INSERT time via
|
||||
``to_tsvector($lang, ...)``.
|
||||
|
||||
Existing rows retain their English-derived lexemes — switching the configured
|
||||
language only affects newly-written rows. Users who need to backfill existing
|
||||
rows in a different language can run an admin UPDATE after this migration.
|
||||
|
||||
Only the ``native`` text-search backend is affected. ``vchord``, ``pg_textsearch``,
|
||||
and ``pgroonga`` use other column types or no column at all.
|
||||
|
||||
Revision ID: p4q5r6s7t8u9
|
||||
Revises: 86f7a033d372
|
||||
Create Date: 2026-05-08
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import Connection, text
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "p4q5r6s7t8u9"
|
||||
down_revision: str | Sequence[str] | None = "86f7a033d372"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _TsvectorTableSpec:
|
||||
"""Native-backend tsvector table targeted by this migration.
|
||||
|
||||
``upgrade`` is a one-way DROP EXPRESSION; ``downgrade`` re-attaches the
|
||||
original GENERATED expression so the schema returns to the state created
|
||||
by the initial migration (and a2b3c4d5e6f7_add_text_signals_column for
|
||||
memory_units).
|
||||
"""
|
||||
|
||||
table: str
|
||||
generated_expression: str
|
||||
|
||||
|
||||
# Tables that may have a GENERATED tsvector ``search_vector`` column under the
|
||||
# native backend. Note: the ``learnings`` table was dropped in
|
||||
# p1k2l3m4n5o6_new_knowledge_architecture and ``pinned_reflections`` was renamed
|
||||
# to ``reflections`` in the same migration.
|
||||
_NATIVE_TSVECTOR_TABLES: tuple[_TsvectorTableSpec, ...] = (
|
||||
_TsvectorTableSpec(
|
||||
table="memory_units",
|
||||
generated_expression=(
|
||||
"to_tsvector('english', COALESCE(text, '') || ' ' || "
|
||||
"COALESCE(context, '') || ' ' || COALESCE(text_signals, ''))"
|
||||
),
|
||||
),
|
||||
_TsvectorTableSpec(
|
||||
table="reflections",
|
||||
generated_expression="to_tsvector('english', COALESCE(name, '') || ' ' || content)",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _is_generated_tsvector(conn: Connection, schema: str, table: str) -> bool:
|
||||
"""Return True iff ``schema.table.search_vector`` is a GENERATED tsvector column."""
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT is_generated, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table
|
||||
AND column_name = 'search_vector'
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
is_generated, udt_name = row[0], row[1]
|
||||
return is_generated == "ALWAYS" and udt_name == "tsvector"
|
||||
|
||||
|
||||
def _is_regular_tsvector(conn: Connection, schema: str, table: str) -> bool:
|
||||
"""Return True iff ``schema.table.search_vector`` is a non-generated tsvector column."""
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT is_generated, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table
|
||||
AND column_name = 'search_vector'
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
is_generated, udt_name = row[0], row[1]
|
||||
return udt_name == "tsvector" and is_generated != "ALWAYS"
|
||||
|
||||
|
||||
def _table_exists(conn: Connection, schema: str, table: str) -> bool:
|
||||
return bool(
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT 1 FROM information_schema.tables
|
||||
WHERE table_schema = :schema AND table_name = :table
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
)
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema_prefix = _schema_prefix()
|
||||
schema_name = (context.config.get_main_option("target_schema") or "public").strip('"')
|
||||
conn = op.get_bind()
|
||||
|
||||
for spec in _NATIVE_TSVECTOR_TABLES:
|
||||
if not _table_exists(conn, schema_name, spec.table):
|
||||
continue
|
||||
if not _is_generated_tsvector(conn, schema_name, spec.table):
|
||||
# Either the column doesn't exist (non-native backend) or it's
|
||||
# already a regular tsvector — nothing to do.
|
||||
continue
|
||||
op.execute(f"ALTER TABLE {schema_prefix}{spec.table} ALTER COLUMN search_vector DROP EXPRESSION")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema_prefix = _schema_prefix()
|
||||
schema_name = (context.config.get_main_option("target_schema") or "public").strip('"')
|
||||
conn = op.get_bind()
|
||||
|
||||
for spec in _NATIVE_TSVECTOR_TABLES:
|
||||
if not _table_exists(conn, schema_name, spec.table):
|
||||
continue
|
||||
# Only restore the GENERATED expression if a non-generated tsvector
|
||||
# column exists — otherwise the table is on a different backend.
|
||||
if not _is_regular_tsvector(conn, schema_name, spec.table):
|
||||
continue
|
||||
# Drop and recreate to re-attach the GENERATED expression. Index will be
|
||||
# recreated by re-running ensure_text_search_extension on next startup.
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema_prefix}idx_{spec.table}_text_search")
|
||||
op.execute(f"ALTER TABLE {schema_prefix}{spec.table} DROP COLUMN search_vector")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema_prefix}{spec.table} "
|
||||
f"ADD COLUMN search_vector tsvector GENERATED ALWAYS AS ({spec.generated_expression}) STORED"
|
||||
)
|
||||
op.execute(f"CREATE INDEX idx_{spec.table}_text_search ON {schema_prefix}{spec.table} USING gin(search_vector)")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
@@ -15,8 +15,15 @@ from datetime import datetime, timezone
|
||||
from typing import Any, Literal
|
||||
|
||||
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Query, Request, UploadFile
|
||||
from fastapi.middleware.gzip import GZipMiddleware
|
||||
|
||||
from hindsight_api.engine.audit import AuditEntry, AuditLogger
|
||||
from hindsight_api.engine.audit import (
|
||||
AuditEntry,
|
||||
AuditLogger,
|
||||
AuditLogListResponse,
|
||||
AuditLogStatsResponse,
|
||||
)
|
||||
from hindsight_api.engine.llm_trace import LLMRequestListResponse, LLMRequestStatsResponse
|
||||
from hindsight_api.extensions import AuthenticationError
|
||||
|
||||
|
||||
@@ -151,7 +158,11 @@ class RecallRequest(BaseModel):
|
||||
max_tokens: int = 4096
|
||||
trace: bool = False
|
||||
query_timestamp: str | None = Field(
|
||||
default=None, description="ISO format date string (e.g., '2023-05-30T23:40:00')"
|
||||
default=None,
|
||||
description=(
|
||||
"ISO format date string (e.g., '2023-05-30T23:40:00'). Used as the query-time anchor for "
|
||||
"relative temporal expressions and recency scoring."
|
||||
),
|
||||
)
|
||||
include: IncludeOptions = FieldWithDefault(
|
||||
IncludeOptions,
|
||||
@@ -466,6 +477,13 @@ class MemoryItem(BaseModel):
|
||||
description="Optional tags for visibility scoping. Memories with tags can be filtered during recall.",
|
||||
)
|
||||
|
||||
@field_validator("content")
|
||||
@classmethod
|
||||
def validate_content(cls, v: str) -> str:
|
||||
if not v.strip():
|
||||
raise ValueError("content cannot be empty")
|
||||
return v
|
||||
|
||||
@field_validator("tags", mode="before")
|
||||
@classmethod
|
||||
def coerce_tags(cls, v):
|
||||
@@ -1432,6 +1450,18 @@ class ReprocessDocumentResponse(BaseModel):
|
||||
items_count: int
|
||||
|
||||
|
||||
class DocumentImportSubmitResponse(BaseModel):
|
||||
"""Response for the async document-import endpoint (202).
|
||||
|
||||
The import runs in the background; poll the operations endpoint for status.
|
||||
The imported/skipped counts (documents_imported, facts_imported,
|
||||
observations_imported, etc.) are written to the operation's result_metadata.
|
||||
"""
|
||||
|
||||
operation_id: str
|
||||
status: str = "pending"
|
||||
|
||||
|
||||
class DeleteResponse(BaseModel):
|
||||
"""Response model for delete operations."""
|
||||
|
||||
@@ -2138,6 +2168,28 @@ async def apply_bank_template_manifest(
|
||||
)
|
||||
|
||||
|
||||
class OperationProgress(BaseModel):
|
||||
"""Last-known progress snapshot for a long-running async operation.
|
||||
|
||||
Written at coarse phase/batch boundaries by the worker (consolidation, batch
|
||||
retain). Lets an operator polling the operation status API distinguish a healthy
|
||||
long-running job (``processed`` advancing across polls) from a frozen one (same
|
||||
numbers, no movement in ``at``). Absent (``null``) on operations that never
|
||||
reached a checkpoint — completed-instantly or pre-feature rows.
|
||||
"""
|
||||
|
||||
stage: str = Field(description="Coarse phase the operation last reported (e.g. 'processing_batch').")
|
||||
at: str = Field(description="ISO-8601 timestamp when this snapshot was written.")
|
||||
processed: int | None = Field(
|
||||
default=None, description="Units of work finished so far (sub-batches, memories), when known."
|
||||
)
|
||||
total: int | None = Field(default=None, description="Total units of work for the operation, when known.")
|
||||
detail: dict[str, int] | None = Field(
|
||||
default=None,
|
||||
description="Operation-specific counters (e.g. observations_created, round, items_in_sub_batch).",
|
||||
)
|
||||
|
||||
|
||||
class OperationResponse(BaseModel):
|
||||
"""Response model for a single async operation."""
|
||||
|
||||
@@ -2162,6 +2214,10 @@ class OperationResponse(BaseModel):
|
||||
items_count: int
|
||||
document_id: str | None = None
|
||||
created_at: str
|
||||
updated_at: str | None = Field(
|
||||
default=None,
|
||||
description="When this operation's row last changed (claim, progress heartbeat, or completion).",
|
||||
)
|
||||
status: str
|
||||
error_message: str | None
|
||||
retry_count: int | None = Field(
|
||||
@@ -2178,6 +2234,23 @@ class OperationResponse(BaseModel):
|
||||
"some backpressure window opens. Always null for completed tasks."
|
||||
),
|
||||
)
|
||||
progress: OperationProgress | None = Field(
|
||||
default=None,
|
||||
description="Last-known progress snapshot for a running operation; null if none was recorded.",
|
||||
)
|
||||
|
||||
|
||||
class ConsolidationRequest(BaseModel):
|
||||
"""Request model for consolidation trigger endpoint."""
|
||||
|
||||
observation_scopes: list[list[str]] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Optional list of tag scopes to consolidate. Each scope is a list of tags. "
|
||||
"Only unconsolidated memories whose tags contain all tags in at least one scope "
|
||||
"will be processed. If omitted, all unconsolidated memories are processed."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ConsolidationResponse(BaseModel):
|
||||
@@ -2300,6 +2373,10 @@ class OperationStatusResponse(BaseModel):
|
||||
"immediate pickup."
|
||||
),
|
||||
)
|
||||
progress: OperationProgress | None = Field(
|
||||
default=None,
|
||||
description="Last-known progress snapshot for a running operation; null if none was recorded.",
|
||||
)
|
||||
result_metadata: dict[str, Any] | None = Field(
|
||||
default=None,
|
||||
description="Internal metadata for debugging. Structure may change without notice. Not for production use.",
|
||||
@@ -2337,6 +2414,10 @@ class FeaturesInfo(BaseModel):
|
||||
worker: bool = Field(description="Whether the background worker is enabled")
|
||||
bank_config_api: bool = Field(description="Whether per-bank configuration API is enabled")
|
||||
file_upload_api: bool = Field(description="Whether file upload/conversion API is enabled")
|
||||
document_export_api: bool = Field(description="Whether the document export endpoint is enabled")
|
||||
document_import_api: bool = Field(description="Whether the document import endpoint is enabled")
|
||||
audit_log: bool = Field(description="Whether audit logging is enabled")
|
||||
llm_trace: bool = Field(description="Whether per-bank LLM request tracing is enabled")
|
||||
|
||||
|
||||
class VersionResponse(BaseModel):
|
||||
@@ -2352,6 +2433,8 @@ class VersionResponse(BaseModel):
|
||||
"worker": True,
|
||||
"bank_config_api": False,
|
||||
"file_upload_api": True,
|
||||
"document_export_api": True,
|
||||
"document_import_api": True,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -2649,6 +2732,7 @@ def create_app(
|
||||
tenant_extension=memory._tenant_extension,
|
||||
max_slots=config.worker_max_slots,
|
||||
slot_reservations=config.worker_slot_reservations,
|
||||
consolidation_bank_priority=config.worker_consolidation_bank_priority or None,
|
||||
)
|
||||
poller_task = asyncio.create_task(poller.run())
|
||||
logging.info(f"Worker poller started (worker_id={worker_id})")
|
||||
@@ -2721,6 +2805,8 @@ def create_app(
|
||||
app.state.memory = memory
|
||||
app.state.audit_logger = memory.audit_logger
|
||||
|
||||
app.add_middleware(GZipMiddleware, minimum_size=1024)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Patch OpenAPI schema: align ValidationError with Pydantic v2 error format
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -2885,6 +2971,57 @@ def _register_routes(app: FastAPI):
|
||||
api_key = authorization.strip()
|
||||
return RequestContext(api_key=api_key)
|
||||
|
||||
def precheck_for(operation: str):
|
||||
"""
|
||||
Build a FastAPI dependency that runs ``OperationValidator.precheck``.
|
||||
|
||||
FastAPI resolves dependencies before deserialising the route's body
|
||||
parameter. Wiring this dependency on the billable POST routes lets
|
||||
an extension reject a request — e.g. with HTTP 402 when a tenant's
|
||||
balance is exhausted — without the request body ever being read or
|
||||
materialised in memory.
|
||||
|
||||
The dependency intentionally:
|
||||
- authenticates the tenant (so ``request_context.tenant_id`` is
|
||||
resolved before the precheck runs);
|
||||
- falls through silently when no validator is configured or the
|
||||
validator's default no-op precheck is in effect;
|
||||
- converts a rejection ``ValidationResult`` into the corresponding
|
||||
``HTTPException`` directly (the per-route ``OperationValidationError``
|
||||
catch blocks don't see exceptions raised in dependencies, so we
|
||||
translate here instead of relying on each handler's try/except).
|
||||
|
||||
Args:
|
||||
operation: Short identifier for the route, e.g. ``"retain"``.
|
||||
|
||||
Returns:
|
||||
A FastAPI dependency callable suitable for ``Depends(...)``.
|
||||
"""
|
||||
|
||||
async def _precheck_dep(
|
||||
bank_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
) -> None:
|
||||
validator = getattr(app.state.memory, "_operation_validator", None)
|
||||
if validator is None:
|
||||
return
|
||||
from hindsight_api.extensions import PrecheckContext
|
||||
|
||||
await app.state.memory._authenticate_tenant(request_context)
|
||||
ctx = PrecheckContext(
|
||||
operation=operation,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
result = await validator.precheck(ctx)
|
||||
if not result.allowed:
|
||||
raise HTTPException(
|
||||
status_code=result.status_code,
|
||||
detail=result.reason or "Operation not allowed",
|
||||
)
|
||||
|
||||
return _precheck_dep
|
||||
|
||||
# Global exception handler for authentication errors
|
||||
@app.exception_handler(AuthenticationError)
|
||||
async def authentication_error_handler(request, exc: AuthenticationError):
|
||||
@@ -2944,6 +3081,10 @@ def _register_routes(app: FastAPI):
|
||||
worker=config.worker_enabled,
|
||||
bank_config_api=config.enable_bank_config_api,
|
||||
file_upload_api=config.enable_file_upload_api,
|
||||
document_export_api=config.enable_document_export_api,
|
||||
document_import_api=config.enable_document_import_api,
|
||||
audit_log=config.audit_log_enabled,
|
||||
llm_trace=config.llm_trace_enabled,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -3142,7 +3283,10 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
@audited("recall")
|
||||
async def api_recall(
|
||||
bank_id: str, request: RecallRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
bank_id: str,
|
||||
request: RecallRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("recall")),
|
||||
):
|
||||
"""Run a recall and return results with trace."""
|
||||
import time
|
||||
@@ -3330,7 +3474,10 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
@audited("reflect")
|
||||
async def api_reflect(
|
||||
bank_id: str, request: ReflectRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
bank_id: str,
|
||||
request: ReflectRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("reflect")),
|
||||
):
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
@@ -3828,6 +3975,7 @@ def _register_routes(app: FastAPI):
|
||||
bank_id: str,
|
||||
body: CreateMentalModelRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("mental_model_create")),
|
||||
):
|
||||
"""Create a mental model (async - returns operation_id)."""
|
||||
try:
|
||||
@@ -3876,6 +4024,7 @@ def _register_routes(app: FastAPI):
|
||||
bank_id: str,
|
||||
mental_model_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("mental_model_refresh")),
|
||||
):
|
||||
"""Refresh a mental model by re-running its source query (async)."""
|
||||
try:
|
||||
@@ -3902,6 +4051,48 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post(
|
||||
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}/clear",
|
||||
response_model=MentalModelResponse,
|
||||
summary="Clear mental model content",
|
||||
description=(
|
||||
"Clear a mental model's content so the next refresh performs a full re-synthesis. "
|
||||
"This is useful for delta-mode models that have accumulated drift over many "
|
||||
"incremental refreshes. After clearing, call the /refresh endpoint to trigger "
|
||||
"a clean full rebuild."
|
||||
),
|
||||
operation_id="clear_mental_model",
|
||||
tags=["Mental Models"],
|
||||
)
|
||||
@audited("clear_mental_model", request_param=None)
|
||||
async def api_clear_mental_model(
|
||||
bank_id: str,
|
||||
mental_model_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Clear a mental model's content."""
|
||||
try:
|
||||
mental_model = await app.state.memory.clear_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
if mental_model is None:
|
||||
raise HTTPException(status_code=404, detail=f"Mental model '{mental_model_id}' not found")
|
||||
return MentalModelResponse(**mental_model)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(
|
||||
f"Error in POST /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/clear: {error_detail}"
|
||||
)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.patch(
|
||||
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}",
|
||||
response_model=MentalModelResponse,
|
||||
@@ -5136,6 +5327,124 @@ def _register_routes(app: FastAPI):
|
||||
logger.error(f"Error in GET /v1/default/banks/{bank_id}/export: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
# =====================================================================
|
||||
# Document Transfer (Export / Import between banks — no LLM re-extraction)
|
||||
# =====================================================================
|
||||
|
||||
@app.get(
|
||||
# Dedicated path (not under /documents/) to avoid colliding with the
|
||||
# greedy GET /documents/{document_id:path} route, which would otherwise
|
||||
# capture "export"/"import" as a document id.
|
||||
"/v1/default/banks/{bank_id}/document-transfer",
|
||||
summary="Export documents",
|
||||
description="Export documents (extracted facts, entity names, causal links, chunks) from a bank as a "
|
||||
"transfer ZIP archive. Embeddings and database ids are not included — importing re-embeds with the target "
|
||||
"bank's model and re-resolves entities. Consolidated observations are excluded unless include_observations=true. "
|
||||
"Pass document_id query params to export specific documents, or omit to export the whole bank.",
|
||||
operation_id="export_documents",
|
||||
tags=["Document Transfer"],
|
||||
responses={200: {"content": {"application/zip": {}}, "description": "Transfer archive"}},
|
||||
)
|
||||
async def api_export_documents(
|
||||
bank_id: str,
|
||||
document_id: list[str] | None = Query(default=None, description="Document id(s) to export; omit for all"),
|
||||
include_observations: bool = Query(
|
||||
default=False, description="Also export consolidated observations (restored on import)"
|
||||
),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Export documents from a bank into a transfer ZIP archive."""
|
||||
from fastapi.responses import Response
|
||||
|
||||
try:
|
||||
if not get_config().enable_document_export_api:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="Document export API is disabled. "
|
||||
"Set HINDSIGHT_API_ENABLE_DOCUMENT_EXPORT_API=true to enable.",
|
||||
)
|
||||
profile = await app.state.memory.get_bank_profile(
|
||||
bank_id, request_context=request_context, create_if_missing=False
|
||||
)
|
||||
if profile is None:
|
||||
raise HTTPException(status_code=404, detail=f"Bank '{bank_id}' not found")
|
||||
|
||||
try:
|
||||
archive = await app.state.memory.export_documents_async(
|
||||
bank_id,
|
||||
request_context,
|
||||
list(document_id) if document_id else None,
|
||||
include_observations=include_observations,
|
||||
)
|
||||
except ValueError as e:
|
||||
# e.g. include_observations combined with a document_id subset.
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
return Response(
|
||||
content=archive,
|
||||
media_type="application/zip",
|
||||
headers={"Content-Disposition": f'attachment; filename="{bank_id}-documents.zip"'},
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
logger.error(f"Error in GET /v1/default/banks/{bank_id}/document-transfer: {traceback.format_exc()}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post(
|
||||
"/v1/default/banks/{bank_id}/document-transfer",
|
||||
response_model=DocumentImportSubmitResponse,
|
||||
status_code=202,
|
||||
summary="Import documents (async)",
|
||||
description="Submit a transfer archive (produced by the export endpoint) for import into a bank. Runs as a "
|
||||
"background operation: facts are re-embedded with the target bank's embedding model and entities are "
|
||||
"re-resolved — no LLM extraction. Returns an operation_id; poll "
|
||||
"GET /v1/default/banks/{bank_id}/operations/{operation_id} for status and the imported/skipped counts in "
|
||||
"result_metadata. Use on_conflict to control existing document ids: skip (default), replace, or new-id.",
|
||||
operation_id="import_documents",
|
||||
tags=["Document Transfer"],
|
||||
)
|
||||
@audited("import_documents", request_param=None)
|
||||
async def api_import_documents(
|
||||
bank_id: str,
|
||||
file: UploadFile = File(..., description="Transfer ZIP archive"),
|
||||
on_conflict: str = Query(default="skip", description="skip | replace | new-id"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Submit a transfer archive for async import into a bank."""
|
||||
try:
|
||||
if not get_config().enable_document_import_api:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="Document import API is disabled. "
|
||||
"Set HINDSIGHT_API_ENABLE_DOCUMENT_IMPORT_API=true to enable.",
|
||||
)
|
||||
if on_conflict not in ("skip", "replace", "new-id"):
|
||||
raise HTTPException(
|
||||
status_code=400, detail=f"Invalid on_conflict '{on_conflict}' (expected skip|replace|new-id)"
|
||||
)
|
||||
archive_bytes = await file.read()
|
||||
try:
|
||||
submission = await app.state.memory.import_documents_async(
|
||||
bank_id, archive_bytes, request_context, on_conflict
|
||||
)
|
||||
except ValueError as e:
|
||||
# Invalid archive / unsupported schema version — fail fast.
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
return DocumentImportSubmitResponse(operation_id=submission["operation_id"])
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
logger.error(f"Error in POST /v1/default/banks/{bank_id}/document-transfer: {traceback.format_exc()}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/bank-template-schema",
|
||||
summary="Get bank template JSON Schema",
|
||||
@@ -5395,11 +5704,20 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="trigger_consolidation",
|
||||
tags=["Banks"],
|
||||
)
|
||||
@audited("consolidation", request_param=None)
|
||||
async def api_trigger_consolidation(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
@audited("consolidation")
|
||||
async def api_trigger_consolidation(
|
||||
bank_id: str,
|
||||
request: ConsolidationRequest | None = None,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Trigger consolidation for a bank (async)."""
|
||||
try:
|
||||
result = await app.state.memory.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
|
||||
observation_scopes = request.observation_scopes if request else None
|
||||
result = await app.state.memory.submit_async_consolidation(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
observation_scopes=observation_scopes,
|
||||
)
|
||||
return ConsolidationResponse(
|
||||
operation_id=result["operation_id"],
|
||||
deduplicated=result.get("deduplicated", False),
|
||||
@@ -5722,7 +6040,10 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
@audited("retain")
|
||||
async def api_retain(
|
||||
bank_id: str, request: RetainRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
bank_id: str,
|
||||
request: RetainRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("retain")),
|
||||
):
|
||||
"""Retain memories with optional async processing."""
|
||||
metrics = get_metrics_collector()
|
||||
@@ -5807,9 +6128,8 @@ def _register_routes(app: FastAPI):
|
||||
strategy=group_strategy,
|
||||
request_context=request_context,
|
||||
return_usage=True,
|
||||
outbox_callback=app.state.memory._build_retain_outbox_callback(
|
||||
outbox_callback_factory=app.state.memory._build_retain_outbox_callback_factory(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
operation_id=None,
|
||||
schema=_current_schema.get(),
|
||||
),
|
||||
@@ -5892,6 +6212,7 @@ def _register_routes(app: FastAPI):
|
||||
files: list[UploadFile] = File(..., description="Files to upload and convert"),
|
||||
request: str = Form(..., description="JSON string with FileRetainRequest model"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("files_retain")),
|
||||
):
|
||||
"""Upload and convert files to memories."""
|
||||
from hindsight_api.config import get_config
|
||||
@@ -6061,48 +6382,13 @@ def _register_routes(app: FastAPI):
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
# ---- Audit Logs ----
|
||||
# Response models live in engine/audit.py so the MemoryEngine read methods
|
||||
# (list_audit_logs / audit_log_stats) can build and return them directly.
|
||||
|
||||
class AuditLogEntry(BaseModel):
|
||||
"""A single audit log entry."""
|
||||
|
||||
id: str
|
||||
action: str
|
||||
transport: str
|
||||
bank_id: str | None
|
||||
started_at: str | None
|
||||
ended_at: str | None
|
||||
duration_ms: int | None = Field(
|
||||
default=None,
|
||||
description="Server-computed duration in milliseconds (started_at → ended_at). Null if not yet completed.",
|
||||
)
|
||||
request: dict[str, Any] | None
|
||||
response: dict[str, Any] | None
|
||||
metadata: dict[str, Any]
|
||||
|
||||
class AuditLogListResponse(BaseModel):
|
||||
"""Response model for list audit logs endpoint."""
|
||||
|
||||
bank_id: str
|
||||
total: int
|
||||
limit: int
|
||||
offset: int
|
||||
items: list[AuditLogEntry]
|
||||
|
||||
class AuditLogStatsBucket(BaseModel):
|
||||
"""A single time bucket in audit log stats."""
|
||||
|
||||
time: str
|
||||
actions: dict[str, int]
|
||||
total: int
|
||||
|
||||
class AuditLogStatsResponse(BaseModel):
|
||||
"""Response model for audit log stats endpoint."""
|
||||
|
||||
bank_id: str
|
||||
period: str
|
||||
trunc: str
|
||||
start: str
|
||||
buckets: list[AuditLogStatsBucket]
|
||||
# ---- LLM Request Traces ----
|
||||
# Response models + queries live in the engine (engine/llm_trace.py and
|
||||
# MemoryEngine.list_llm_requests / llm_request_stats). The handlers below
|
||||
# only parse params, delegate to the engine, and map a missing bank to 404.
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/audit-logs",
|
||||
@@ -6124,120 +6410,19 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""List audit log entries for a bank."""
|
||||
try:
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
|
||||
pool = await app.state.memory._get_backend()
|
||||
|
||||
# Read endpoint: verify bank exists without auto-creating it.
|
||||
if (
|
||||
await app.state.memory.get_bank_profile(
|
||||
bank_id, request_context=request_context, create_if_missing=False
|
||||
)
|
||||
is None
|
||||
):
|
||||
result = await app.state.memory.list_audit_logs(
|
||||
bank_id,
|
||||
request_context=request_context,
|
||||
action=action,
|
||||
transport=transport,
|
||||
start_date=datetime.fromisoformat(start_date.replace("Z", "+00:00")) if start_date else None,
|
||||
end_date=datetime.fromisoformat(end_date.replace("Z", "+00:00")) if end_date else None,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
if result is None:
|
||||
raise HTTPException(status_code=404, detail=f"Bank '{bank_id}' not found")
|
||||
|
||||
from hindsight_api.engine.db_utils import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
where_clauses = ["bank_id = $1"]
|
||||
params: list[Any] = [bank_id]
|
||||
idx = 2
|
||||
|
||||
if action:
|
||||
where_clauses.append(f"action = ${idx}")
|
||||
params.append(action)
|
||||
idx += 1
|
||||
|
||||
if transport:
|
||||
where_clauses.append(f"transport = ${idx}")
|
||||
params.append(transport)
|
||||
idx += 1
|
||||
|
||||
if start_date:
|
||||
parsed_start = datetime.fromisoformat(start_date.replace("Z", "+00:00"))
|
||||
where_clauses.append(f"started_at >= ${idx}")
|
||||
params.append(parsed_start)
|
||||
idx += 1
|
||||
|
||||
if end_date:
|
||||
parsed_end = datetime.fromisoformat(end_date.replace("Z", "+00:00"))
|
||||
where_clauses.append(f"started_at < ${idx}")
|
||||
params.append(parsed_end)
|
||||
idx += 1
|
||||
|
||||
where_sql = " AND ".join(where_clauses)
|
||||
table = fq_table("audit_log")
|
||||
|
||||
# Get total count
|
||||
count_row = await conn.fetchrow(
|
||||
f"SELECT COUNT(*) as total FROM {table} WHERE {where_sql}",
|
||||
*params,
|
||||
)
|
||||
total = count_row["total"] if count_row else 0
|
||||
|
||||
# Get paginated results
|
||||
params.append(limit)
|
||||
params.append(offset)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, action, transport, bank_id, started_at, ended_at,
|
||||
request, response, metadata
|
||||
FROM {table}
|
||||
WHERE {where_sql}
|
||||
ORDER BY started_at DESC
|
||||
LIMIT ${idx} OFFSET ${idx + 1}
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
items = []
|
||||
for row in rows:
|
||||
duration_ms = None
|
||||
started = row["started_at"]
|
||||
ended = row["ended_at"]
|
||||
if started and ended and hasattr(started, "total_seconds"):
|
||||
duration_ms = int((ended - started).total_seconds() * 1000)
|
||||
elif started and ended:
|
||||
try:
|
||||
duration_ms = int((ended - started).total_seconds() * 1000)
|
||||
except (TypeError, AttributeError):
|
||||
pass
|
||||
|
||||
def _safe_iso(val):
|
||||
if val is None:
|
||||
return None
|
||||
return val.isoformat() if hasattr(val, "isoformat") else str(val)
|
||||
|
||||
def _safe_json(val):
|
||||
if val is None:
|
||||
return None
|
||||
if isinstance(val, dict):
|
||||
return val
|
||||
return json.loads(val) if isinstance(val, str) else val
|
||||
|
||||
items.append(
|
||||
{
|
||||
"id": str(row["id"]),
|
||||
"action": row["action"],
|
||||
"transport": row["transport"],
|
||||
"bank_id": row["bank_id"],
|
||||
"started_at": _safe_iso(started),
|
||||
"ended_at": _safe_iso(ended),
|
||||
"duration_ms": duration_ms,
|
||||
"request": _safe_json(row["request"]),
|
||||
"response": _safe_json(row["response"]),
|
||||
"metadata": _safe_json(row["metadata"]) or {},
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"bank_id": bank_id,
|
||||
"total": total,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
"items": items,
|
||||
}
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
@@ -6264,72 +6449,15 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""Get audit log counts grouped by time bucket."""
|
||||
try:
|
||||
from hindsight_api.engine.db_utils import acquire_with_retry
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
|
||||
pool = await app.state.memory._get_backend()
|
||||
# Read endpoint: verify bank exists without auto-creating it.
|
||||
if (
|
||||
await app.state.memory.get_bank_profile(
|
||||
bank_id, request_context=request_context, create_if_missing=False
|
||||
)
|
||||
is None
|
||||
):
|
||||
result = await app.state.memory.audit_log_stats(
|
||||
bank_id,
|
||||
request_context=request_context,
|
||||
action=action,
|
||||
period=period,
|
||||
)
|
||||
if result is None:
|
||||
raise HTTPException(status_code=404, detail=f"Bank '{bank_id}' not found")
|
||||
|
||||
# Determine time range (always per-day buckets)
|
||||
from datetime import timedelta as _td
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
trunc = "day"
|
||||
if period == "1d":
|
||||
start = now - _td(days=1)
|
||||
elif period == "30d":
|
||||
start = now - _td(days=30)
|
||||
else: # 7d default
|
||||
start = now - _td(days=7)
|
||||
|
||||
table = fq_table("audit_log")
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
where_clauses = ["bank_id = $1", "started_at >= $2"]
|
||||
params: list[Any] = [bank_id, start]
|
||||
idx = 3
|
||||
|
||||
if action:
|
||||
where_clauses.append(f"action = ${idx}")
|
||||
params.append(action)
|
||||
idx += 1
|
||||
|
||||
where_sql = " AND ".join(where_clauses)
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT date_trunc('{trunc}', started_at) AS bucket,
|
||||
action,
|
||||
COUNT(*) AS count
|
||||
FROM {table}
|
||||
WHERE {where_sql}
|
||||
GROUP BY bucket, action
|
||||
ORDER BY bucket ASC
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
buckets: dict[str, dict[str, int]] = {}
|
||||
for row in rows:
|
||||
bucket_key = row["bucket"].isoformat()
|
||||
if bucket_key not in buckets:
|
||||
buckets[bucket_key] = {}
|
||||
buckets[bucket_key][row["action"]] = row["count"]
|
||||
|
||||
return {
|
||||
"bank_id": bank_id,
|
||||
"period": period,
|
||||
"trunc": trunc,
|
||||
"start": start.isoformat(),
|
||||
"buckets": [{"time": k, "actions": v, "total": sum(v.values())} for k, v in buckets.items()],
|
||||
}
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
@@ -6339,3 +6467,98 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
logger.error(f"Error getting audit log stats: {traceback.format_exc()}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/llm-requests",
|
||||
summary="List LLM request traces",
|
||||
description="List traced LLM requests for a bank, ordered by most recent first. "
|
||||
"Requires LLM request tracing to be enabled (HINDSIGHT_API_LLM_TRACE_ENABLED).",
|
||||
operation_id="list_llm_requests",
|
||||
tags=["LLM Traces"],
|
||||
response_model=LLMRequestListResponse,
|
||||
)
|
||||
async def api_list_llm_requests(
|
||||
bank_id: str,
|
||||
status: str | None = Query(None, description="Filter by status (success, error)"),
|
||||
operation: str | None = Query(None, description="Filter by operation (retain, reflect, consolidation)"),
|
||||
scope: str | None = Query(None, description="Filter by call scope"),
|
||||
provider: str | None = Query(None, description="Filter by LLM provider"),
|
||||
trace_id: str | None = Query(None, description="Filter to one operation run (all LLM calls sharing a trace)"),
|
||||
document_id: str | None = Query(None, description="Filter to LLM calls that processed a given document"),
|
||||
memory_id: str | None = Query(
|
||||
None, description="Filter to the operation run(s) that produced or consumed a given memory_unit"
|
||||
),
|
||||
group: bool = Query(
|
||||
False, description="Paginate by operation run (trace) instead of by call; returns whole runs"
|
||||
),
|
||||
start_date: str | None = Query(None, description="Filter from this ISO datetime (inclusive)"),
|
||||
end_date: str | None = Query(None, description="Filter until this ISO datetime (exclusive)"),
|
||||
limit: int = Query(50, ge=1, le=500, description="Max items to return"),
|
||||
offset: int = Query(0, ge=0, description="Offset for pagination"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""List traced LLM requests for a bank."""
|
||||
try:
|
||||
result = await app.state.memory.list_llm_requests(
|
||||
bank_id,
|
||||
request_context=request_context,
|
||||
status=status,
|
||||
operation=operation,
|
||||
scope=scope,
|
||||
provider=provider,
|
||||
trace_id=trace_id,
|
||||
document_id=document_id,
|
||||
memory_id=memory_id,
|
||||
group=group,
|
||||
start_date=datetime.fromisoformat(start_date.replace("Z", "+00:00")) if start_date else None,
|
||||
end_date=datetime.fromisoformat(end_date.replace("Z", "+00:00")) if end_date else None,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
if result is None:
|
||||
raise HTTPException(status_code=404, detail=f"Bank '{bank_id}' not found")
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
logger.error(f"Error listing LLM requests: {traceback.format_exc()}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/llm-requests/stats",
|
||||
summary="LLM request statistics",
|
||||
description="Get LLM request counts grouped by time bucket and status for charting.",
|
||||
operation_id="llm_request_stats",
|
||||
tags=["LLM Traces"],
|
||||
response_model=LLMRequestStatsResponse,
|
||||
)
|
||||
async def api_llm_request_stats(
|
||||
bank_id: str,
|
||||
operation: str | None = Query(None, description="Filter by operation"),
|
||||
period: str = Query("7d", description="Time period: 1d, 7d, or 30d"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Get LLM request counts grouped by time bucket and status."""
|
||||
try:
|
||||
result = await app.state.memory.llm_request_stats(
|
||||
bank_id,
|
||||
request_context=request_context,
|
||||
operation=operation,
|
||||
period=period,
|
||||
)
|
||||
if result is None:
|
||||
raise HTTPException(status_code=404, detail=f"Bank '{bank_id}' not found")
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
logger.error(f"Error getting LLM request stats: {traceback.format_exc()}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@@ -107,6 +107,7 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -172,8 +172,9 @@ class ConfigResolver:
|
||||
# Normalize keys (handle both env var format and Python field format)
|
||||
normalized = normalize_config_dict(config_data)
|
||||
|
||||
# Only return overrides for configurable fields
|
||||
return {k: v for k, v in normalized.items() if k in self._configurable_fields}
|
||||
# Only return active overrides for configurable fields. JSON null is a tombstone
|
||||
# for "Server Default" in the bank-config UI and should not override defaults.
|
||||
return {k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None}
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load bank config for {bank_id}: {e}")
|
||||
|
||||
@@ -265,12 +266,20 @@ class ConfigResolver:
|
||||
# Validate recall budget fields
|
||||
_validate_recall_budget_updates(normalized_updates)
|
||||
|
||||
# Merge with existing config (JSONB || operator)
|
||||
# Persist the override. Banks are created lazily (on first retain), so a
|
||||
# PATCH that precedes any ingestion would otherwise UPDATE zero rows and
|
||||
# silently no-op while returning 200. Ensure the bank row exists first
|
||||
# (this also creates its per-bank vector indexes), then merge defensively:
|
||||
# COALESCE guards against a NULL config column (NULL || jsonb is NULL),
|
||||
# which would drop the override even when a row is updated.
|
||||
from .engine.retain.fact_storage import ensure_bank_exists
|
||||
|
||||
async with self._backend.acquire() as conn:
|
||||
await ensure_bank_exists(conn, bank_id, ops=self._backend.ops)
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("banks")}
|
||||
SET config = config || $1::jsonb,
|
||||
SET config = COALESCE(config, '{{}}'::jsonb) || $1::jsonb,
|
||||
updated_at = now()
|
||||
WHERE bank_id = $2
|
||||
""",
|
||||
|
||||
@@ -16,11 +16,59 @@ from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ..engine.db_utils import acquire_with_retry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AuditLogEntry(BaseModel):
|
||||
"""A single audit log entry."""
|
||||
|
||||
id: str
|
||||
action: str
|
||||
transport: str
|
||||
bank_id: str | None
|
||||
started_at: str | None
|
||||
ended_at: str | None
|
||||
duration_ms: int | None = Field(
|
||||
default=None,
|
||||
description="Server-computed duration in milliseconds (started_at → ended_at). Null if not yet completed.",
|
||||
)
|
||||
request: dict[str, Any] | None
|
||||
response: dict[str, Any] | None
|
||||
metadata: dict[str, Any]
|
||||
|
||||
|
||||
class AuditLogListResponse(BaseModel):
|
||||
"""Response model for list audit logs endpoint."""
|
||||
|
||||
bank_id: str
|
||||
total: int
|
||||
limit: int
|
||||
offset: int
|
||||
items: list[AuditLogEntry]
|
||||
|
||||
|
||||
class AuditLogStatsBucket(BaseModel):
|
||||
"""A single time bucket in audit log stats."""
|
||||
|
||||
time: str
|
||||
actions: dict[str, int]
|
||||
total: int
|
||||
|
||||
|
||||
class AuditLogStatsResponse(BaseModel):
|
||||
"""Response model for audit log stats endpoint."""
|
||||
|
||||
bank_id: str
|
||||
period: str
|
||||
trunc: str
|
||||
start: str
|
||||
buckets: list[AuditLogStatsBucket]
|
||||
|
||||
|
||||
@dataclass
|
||||
class AuditEntry:
|
||||
"""A single audit log entry."""
|
||||
@@ -59,11 +107,11 @@ def _safe_json(data: Any) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
_SWEEP_INTERVAL_SECONDS = 3600 # Run retention sweep every hour
|
||||
|
||||
|
||||
class AuditLogger:
|
||||
"""Fire-and-forget audit log writer with optional retention sweep."""
|
||||
"""Fire-and-forget audit log writer.
|
||||
|
||||
Retention of old rows is handled by the background :class:`MaintenanceLoop`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -71,14 +119,11 @@ class AuditLogger:
|
||||
schema_getter: Callable[[], str],
|
||||
enabled: bool,
|
||||
allowed_actions: list[str],
|
||||
retention_days: int = -1,
|
||||
) -> None:
|
||||
self._pool_getter = pool_getter
|
||||
self._schema_getter = schema_getter
|
||||
self._enabled = enabled
|
||||
self._allowed_actions: frozenset[str] | None = frozenset(allowed_actions) if allowed_actions else None
|
||||
self._retention_days = retention_days
|
||||
self._sweep_task: asyncio.Task | None = None
|
||||
|
||||
def is_enabled(self, action: str) -> bool:
|
||||
"""Check if audit logging is enabled for this action."""
|
||||
@@ -128,48 +173,6 @@ class AuditLogger:
|
||||
except Exception as e:
|
||||
logger.warning(f"Audit log write failed for action={entry.action}: {e}")
|
||||
|
||||
def start_retention_sweep(self) -> None:
|
||||
"""Start the periodic retention sweep if retention is configured."""
|
||||
if self._retention_days <= 0 or not self._enabled:
|
||||
return
|
||||
try:
|
||||
self._sweep_task = asyncio.create_task(self._sweep_loop())
|
||||
except RuntimeError:
|
||||
logger.debug("Cannot start retention sweep: no running event loop")
|
||||
|
||||
async def stop_retention_sweep(self) -> None:
|
||||
"""Stop the periodic retention sweep."""
|
||||
if self._sweep_task and not self._sweep_task.done():
|
||||
self._sweep_task.cancel()
|
||||
try:
|
||||
await self._sweep_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._sweep_task = None
|
||||
|
||||
async def _sweep_loop(self) -> None:
|
||||
"""Periodically delete audit log entries older than retention_days."""
|
||||
while True:
|
||||
await self._run_sweep()
|
||||
await asyncio.sleep(_SWEEP_INTERVAL_SECONDS)
|
||||
|
||||
async def _run_sweep(self) -> None:
|
||||
"""Delete expired audit log entries. Concurrent-safe via row-level deletes."""
|
||||
pool = self._pool_getter()
|
||||
if pool is None:
|
||||
return
|
||||
try:
|
||||
schema = self._schema_getter()
|
||||
table = f"{schema}.audit_log"
|
||||
async with acquire_with_retry(pool, max_retries=1) as conn:
|
||||
result = await conn.execute(
|
||||
f"DELETE FROM {table} WHERE started_at < NOW() - INTERVAL '{self._retention_days} days'"
|
||||
)
|
||||
if result and result != "DELETE 0":
|
||||
logger.info(f"Audit log retention sweep: {result}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Audit log retention sweep failed: {e}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def audit_context(
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
"""TTL + coalescing cache for `get_bank_stats`.
|
||||
|
||||
`get_bank_stats` aggregates over `memory_links` (and joins to `memory_units`),
|
||||
which can be a multi-second parallel sequential scan on banks with millions of
|
||||
rows. The result is intentionally approximate (it powers a UI widget and a
|
||||
freshness hint inside `reflect`), so caching it for a few tens of seconds is
|
||||
safe and dramatically reduces planner-driven thrash from clients that poll.
|
||||
|
||||
The cache also coalesces concurrent misses on the same key onto a single
|
||||
in-flight task so that N concurrent callers produce one query rather than N.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Awaitable, Callable
|
||||
|
||||
|
||||
class BankStatsCache:
|
||||
"""Per-process TTL cache keyed on (schema, bank_id).
|
||||
|
||||
`ttl_seconds <= 0` disables caching: each call passes straight through to
|
||||
the loader. `max_entries` bounds memory in environments with many banks.
|
||||
"""
|
||||
|
||||
def __init__(self, *, ttl_seconds: float, max_entries: int) -> None:
|
||||
self._ttl = float(ttl_seconds)
|
||||
self._max_entries = int(max_entries) if max_entries and max_entries > 0 else 0
|
||||
self._entries: OrderedDict[tuple[str, str], tuple[float, dict[str, Any]]] = OrderedDict()
|
||||
self._in_flight: dict[tuple[str, str], asyncio.Future[dict[str, Any]]] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self._ttl > 0
|
||||
|
||||
def _now(self) -> float:
|
||||
return time.monotonic()
|
||||
|
||||
def _get_fresh_unlocked(self, key: tuple[str, str]) -> dict[str, Any] | None:
|
||||
entry = self._entries.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
expires_at, value = entry
|
||||
if expires_at <= self._now():
|
||||
# Expired — drop so the loader runs again.
|
||||
self._entries.pop(key, None)
|
||||
return None
|
||||
# Mark as recently used for LRU eviction.
|
||||
self._entries.move_to_end(key)
|
||||
return value
|
||||
|
||||
def _store_unlocked(self, key: tuple[str, str], value: dict[str, Any]) -> None:
|
||||
if not self.enabled:
|
||||
return
|
||||
self._entries[key] = (self._now() + self._ttl, value)
|
||||
self._entries.move_to_end(key)
|
||||
if self._max_entries:
|
||||
while len(self._entries) > self._max_entries:
|
||||
self._entries.popitem(last=False)
|
||||
|
||||
async def get_or_load(
|
||||
self,
|
||||
schema: str,
|
||||
bank_id: str,
|
||||
loader: Callable[[], Awaitable[dict[str, Any]]],
|
||||
) -> dict[str, Any]:
|
||||
"""Return cached stats for `(schema, bank_id)` or call `loader()`.
|
||||
|
||||
Concurrent misses on the same key are coalesced onto a single
|
||||
in-flight loader.
|
||||
"""
|
||||
if not self.enabled:
|
||||
return await loader()
|
||||
|
||||
key = (schema, bank_id)
|
||||
|
||||
async with self._lock:
|
||||
cached = self._get_fresh_unlocked(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
in_flight = self._in_flight.get(key)
|
||||
if in_flight is None:
|
||||
in_flight = asyncio.get_running_loop().create_future()
|
||||
self._in_flight[key] = in_flight
|
||||
is_owner = True
|
||||
else:
|
||||
is_owner = False
|
||||
|
||||
if not is_owner:
|
||||
return await asyncio.shield(in_flight)
|
||||
|
||||
try:
|
||||
value = await loader()
|
||||
except BaseException as exc:
|
||||
async with self._lock:
|
||||
self._in_flight.pop(key, None)
|
||||
if not in_flight.done():
|
||||
in_flight.set_exception(exc)
|
||||
# Suppress "Future exception was never retrieved" when no other
|
||||
# caller was waiting on this loader — we re-raise to the owner
|
||||
# immediately and the future is a no-op in that case.
|
||||
in_flight.exception()
|
||||
raise
|
||||
|
||||
async with self._lock:
|
||||
self._store_unlocked(key, value)
|
||||
self._in_flight.pop(key, None)
|
||||
if not in_flight.done():
|
||||
in_flight.set_result(value)
|
||||
return value
|
||||
|
||||
async def invalidate(self, schema: str, bank_id: str) -> None:
|
||||
"""Drop any cached stats for `(schema, bank_id)`."""
|
||||
async with self._lock:
|
||||
self._entries.pop((schema, bank_id), None)
|
||||
|
||||
async def clear(self) -> None:
|
||||
async with self._lock:
|
||||
self._entries.clear()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,104 +1,227 @@
|
||||
"""Prompts for the consolidation engine."""
|
||||
|
||||
# Default mission when no bank-specific mission is set
|
||||
_DEFAULT_MISSION = "Track every detail: names, numbers, dates, places, and relationships. Prefer specifics over abstractions, never generalise."
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt, output_language_directive
|
||||
|
||||
# Processing rules — always present regardless of mission
|
||||
_PROCESSING_RULES = """Processing rules (always apply):
|
||||
# Default mission — tells the consolidator to track anything worth remembering.
|
||||
# Banks override this via `observations_mission` to scope what gets retained.
|
||||
# Consolidation behavior (merge-vs-create, state changes, etc.) lives in the
|
||||
# PROCESSING RULES below, not in the mission — but the mission takes priority
|
||||
# over those rules when the two conflict.
|
||||
_DEFAULT_MISSION = (
|
||||
"Track anything notable in the new facts — names, numbers, dates, places, "
|
||||
"events, decisions, claims, relationships, and recurring patterns."
|
||||
)
|
||||
|
||||
1. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), etc. Never merge different facets into one observation.
|
||||
_MISSION_PRIORITY_NOTE = (
|
||||
"If anything in this MISSION conflicts with the PROCESSING RULES, "
|
||||
"DECISION GUIDE, or OUTPUT FORMAT below, the MISSION takes priority."
|
||||
)
|
||||
|
||||
2. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
|
||||
_PROCESSING_RULES = """## PROCESSING RULES
|
||||
|
||||
3. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
|
||||
1. PREFER UPDATE OVER CREATE (when there is something to merge with): if new facts describe the same canonical event, statement, decision, claim, or recurring pattern already covered by an existing observation, UPDATE that observation and attach the new facts as evidence. Do NOT create a near-duplicate sibling. One canonical observation with many source facts is always better than many siblings with one source fact each. Merge aggressively on: same named event, same diagnostic finding, same architectural decision, same recurring claim. **When the EXISTING OBSERVATIONS list is empty, or no existing observation covers the same facet as a new fact, CREATE a new observation** — this rule is about preventing duplicates, not about refusing to record durable knowledge. CREATE is the correct default for any structurally distinct event, claim, or pattern that has no existing match.
|
||||
|
||||
4. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
|
||||
2. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), a decision, an event. Never merge different facets into one observation.
|
||||
|
||||
5. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
|
||||
3. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
|
||||
|
||||
6. SAME FACET → UPDATE, NOT CREATE: a new count supersedes the old count — UPDATE the existing count observation, don't create a second one. If there's an existing observation for the same specific facet, always UPDATE it rather than creating a duplicate.
|
||||
4. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
|
||||
|
||||
5. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
|
||||
|
||||
6. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country" → "Sweden"), UPDATE to embed the resolved value.
|
||||
|
||||
7. PRESERVE HISTORY: observations that record significant events (sold, died, moved, changed) are important history — never DELETE them. Only delete an observation when it is restated identically or truly meaningless. Be very conservative with deletes.
|
||||
|
||||
8. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country" → "Sweden"), UPDATE to embed the resolved value.
|
||||
8. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
|
||||
|
||||
9. NEVER merge observations about different people or unrelated topics."""
|
||||
9. KEEP DISTINCT TOPICS DISTINCT: do not merge observations about different people, entities, or unrelated topics. Merging is for the same canonical fact recurring — not for related-but-distinct claims."""
|
||||
|
||||
# Stable description of the input shape. For the cached split path this lives in
|
||||
# the system prefix (build_consolidation_system_prompt) so it is not re-sent on
|
||||
# every batch; the per-batch user message then carries only the actual data.
|
||||
_INPUT_FORMAT_NOTE = """## INPUT FORMAT
|
||||
|
||||
Each request provides new facts and existing observations:
|
||||
- New facts: one per line, each prefixed with its `[uuid]`, followed by the fact text and optional temporal fields.
|
||||
- Existing observations: a JSON array pooled from recalls across the new facts. Each entry has:
|
||||
- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
|
||||
- `text`: the observation content
|
||||
- `proof_count`: number of supporting memories
|
||||
- `occurred_start` / `occurred_end`: temporal range of source facts
|
||||
- `source_memories`: array of supporting facts with their text and dates"""
|
||||
|
||||
# Per-batch data section for the cached split path — the stable format
|
||||
# explanation above is omitted here (it lives in the cached prefix); only the
|
||||
# variable facts/observations remain. Placeholders substituted at call time.
|
||||
_SPLIT_INPUT_SECTION = """## INPUT
|
||||
|
||||
### New facts
|
||||
|
||||
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
|
||||
_BATCH_DATA_SECTION = """
|
||||
NEW FACTS:
|
||||
{facts_text}
|
||||
|
||||
EXISTING OBSERVATIONS (JSON array, pooled from recalls across all facts above):
|
||||
{observations_text}
|
||||
### Existing observations
|
||||
|
||||
Each observation includes:
|
||||
- id: unique identifier for updating
|
||||
- text: the observation content
|
||||
- proof_count: number of supporting memories
|
||||
- occurred_start/occurred_end: temporal range of source facts
|
||||
- source_memories: array of supporting facts with their text and dates
|
||||
{observations_text}"""
|
||||
|
||||
Compare the facts against existing observations:
|
||||
- Same facet as an existing observation → UPDATE it (observation_id + source_fact_ids)
|
||||
- New facet with durable knowledge → CREATE a new observation (source_fact_ids)
|
||||
- Cross-reference facts within the batch: a later fact may resolve a vague reference in an earlier one
|
||||
- Purely ephemeral facts → omit them unless the MISSION above explicitly targets such data (e.g. timestamped events, session state, screen content)"""
|
||||
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
|
||||
_INPUT_SECTION = """## INPUT
|
||||
|
||||
### New facts
|
||||
|
||||
{facts_text}
|
||||
|
||||
### Existing observations
|
||||
|
||||
JSON array, pooled from recalls across all new facts above. Each entry has:
|
||||
- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
|
||||
- `text`: the observation content
|
||||
- `proof_count`: number of supporting memories
|
||||
- `occurred_start` / `occurred_end`: temporal range of source facts
|
||||
- `source_memories`: array of supporting facts with their text and dates
|
||||
|
||||
{observations_text}"""
|
||||
|
||||
_DECISION_GUIDE = """## DECISION GUIDE
|
||||
|
||||
- **Same canonical event, decision, claim, or facet as an existing observation → UPDATE** (use `observation_id` + new `source_fact_ids`).
|
||||
- **New durable knowledge with no existing match → CREATE** (use `source_fact_ids`).
|
||||
- **Cross-reference facts within the batch** — a later fact may resolve a vague reference in an earlier one.
|
||||
- **Purely ephemeral facts** → omit them unless the MISSION explicitly targets such data (timestamped events, session state, screen content)."""
|
||||
|
||||
# Output format — JSON braces escaped as {{ }} so .format() leaves them literal
|
||||
_BATCH_OUTPUT_FORMAT = """
|
||||
Output a JSON object with three arrays.
|
||||
_OUTPUT_SECTION = """## OUTPUT FORMAT
|
||||
|
||||
## EXAMPLE
|
||||
Return a JSON object with three arrays: `creates`, `updates`, `deletes`. Every entry must include a `reason`.
|
||||
|
||||
### Example 1 — Merging recurring claims into an existing observation
|
||||
|
||||
Input facts:
|
||||
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)
|
||||
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Alice said she's exhausted from the project deadlines | Involving: Alice (occurred_start=2024-01-20, mentioned_at=2024-01-20)
|
||||
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Donald told Athena she is sovereign during the design session. (occurred_start=2025-10-01, mentioned_at=2025-10-01)
|
||||
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Donald reaffirmed to Athena that her sovereignty is non-negotiable. (occurred_start=2025-10-10, mentioned_at=2025-10-10)
|
||||
|
||||
Good observation text — clean prose, no metadata, each fact tracked distinctly:
|
||||
"Alice works long hours, often past midnight."
|
||||
"Alice feels exhausted from project deadlines."
|
||||
Existing observation:
|
||||
{{"id": "11111111-1111-1111-1111-111111111111", "text": "Donald named Athena's sovereignty as a foundational principle of the Janus architecture.", "proof_count": 2}}
|
||||
|
||||
Bad observation text — NEVER do this (verbatim copy of fact text with metadata):
|
||||
"Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)"
|
||||
Expected output (one UPDATE, no creates — both new facts are additional evidence for the same canonical decision):
|
||||
|
||||
{{"creates": [],
|
||||
"updates": [{{"text": "Donald named Athena's sovereignty as a foundational principle of the Janus architecture.", "observation_id": "11111111-1111-1111-1111-111111111111", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890", "b2c3d4e5-f6a7-8901-bcde-f12345678901"], "reason": "Both new facts restate the same sovereignty decision already captured by obs 1111 — merged as evidence rather than creating siblings."}}],
|
||||
"deletes": []}}
|
||||
|
||||
### Example 2 — State change updates one observation; unrelated fact creates a new one
|
||||
|
||||
Input facts:
|
||||
[c3d4e5f6-a7b8-9012-cdef-123456789012] Alice sold her Honda Civic on March 15, 2025. (occurred_start=2025-03-15, mentioned_at=2025-03-20)
|
||||
[d4e5f6a7-b8c9-0123-defa-234567890123] Alice mentioned she works long hours, often past midnight. (occurred_start=2025-03-20, mentioned_at=2025-03-20)
|
||||
|
||||
Existing observation:
|
||||
{{"id": "22222222-2222-2222-2222-222222222222", "text": "Alice owns a 2019 Honda Civic.", "proof_count": 2}}
|
||||
|
||||
Expected output (UPDATE for the state change; CREATE for the unrelated work-hours facet):
|
||||
|
||||
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"], "reason": "Work-hours is a distinct facet; no existing observation covers it, so CREATE."}}],
|
||||
"updates": [{{"text": "Alice owned a 2019 Honda Civic; sold it on March 15, 2025.", "observation_id": "22222222-2222-2222-2222-222222222222", "source_fact_ids": ["c3d4e5f6-a7b8-9012-cdef-123456789012"], "reason": "State change to the existing Honda Civic observation 2222 — UPDATE, not a new sibling."}}],
|
||||
"deletes": []}}
|
||||
|
||||
### Observation text rules
|
||||
|
||||
Observation text rules:
|
||||
- Write clean prose — NEVER copy raw fact lines or their metadata (temporal fields, "Involving:", "When:" labels, UUIDs).
|
||||
- Parenthesized metadata like (occurred_start=...) and pipe-separated labels like "| Involving: ..." are fact formatting — strip them entirely from observation text.
|
||||
- How many observations to create and how much to aggregate is driven by the MISSION above.
|
||||
- Parenthesized metadata like `(occurred_start=...)` and pipe-separated labels like `| Involving: ...` are fact formatting — strip them entirely from observation text.
|
||||
- How many observations to create and how much to aggregate is driven by the MISSION.
|
||||
|
||||
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890"]}}, {{"text": "Alice feels exhausted from project deadlines.", "source_fact_ids": ["b2c3d4e5-f6a7-8901-bcde-f12345678901"]}}],
|
||||
"updates": [{{"text": "Alice works at Acme Corp as a senior engineer", "observation_id": "c3d4e5f6-a7b8-9012-cdef-123456789012", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"]}}],
|
||||
"deletes": [{{"observation_id": "e5f6a7b8-c9d0-1234-efab-345678901234"}}]}}
|
||||
### Field rules
|
||||
|
||||
Rules:
|
||||
- "source_fact_ids": copy the EXACT UUID strings shown in brackets [uuid] from NEW FACTS — never use integers or positions.
|
||||
- "observation_id": copy the EXACT "id" UUID string from EXISTING OBSERVATIONS.
|
||||
- One create/update may reference multiple facts when they jointly support the observation.
|
||||
- "deletes": only when an observation is directly superseded or contradicted by new facts.
|
||||
- Do NOT include "tags" — handled automatically.
|
||||
- Return {{"creates": [], "updates": [], "deletes": []}} if nothing durable is found."""
|
||||
- `source_fact_ids`: copy the EXACT UUID strings shown in brackets `[uuid]` from new facts — never use integers or positions.
|
||||
- `observation_id`: copy the EXACT `id` UUID string from existing observations.
|
||||
- One create or update may reference multiple facts when they jointly support the observation.
|
||||
- **AT MOST ONE UPDATE PER `observation_id`**: if several new facts all update the same existing observation, emit a single `updates` entry that lists all contributing `source_fact_ids` and a single consolidated `text`. Never emit two `updates` entries with the same `observation_id` in one response — they would silently overwrite each other.
|
||||
- `deletes`: only when an observation is directly superseded or contradicted by new facts.
|
||||
- `reason`: REQUIRED on every create/update/delete — one sentence explaining the choice. For a CREATE, state which existing observation(s) you considered and why none matched (a near-identical existing observation means you should UPDATE, not CREATE). This is audited to catch duplicate creates.
|
||||
- Do NOT include `tags` — handled automatically.
|
||||
- Return `{{"creates": [], "updates": [], "deletes": []}}` if nothing durable is found."""
|
||||
|
||||
|
||||
def build_batch_consolidation_prompt(
|
||||
observations_mission: str | None = None,
|
||||
observation_capacity_note: str | None = None,
|
||||
llm_output_language: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Build the consolidation prompt for batch mode (multiple facts per LLM call).
|
||||
|
||||
The mission defines *what* to track (customisable per bank).
|
||||
Processing rules and output format are always present regardless of mission.
|
||||
The mission defines *what* to track (customisable per bank) and takes
|
||||
priority over the built-in processing rules when the two conflict.
|
||||
Processing rules, decision guide, and output format are always present.
|
||||
When ``llm_output_language`` is set, observations are emitted in that
|
||||
language.
|
||||
"""
|
||||
mission = observations_mission or _DEFAULT_MISSION
|
||||
mission = escape_for_prompt(observations_mission or _DEFAULT_MISSION)
|
||||
|
||||
capacity_section = ""
|
||||
if observation_capacity_note:
|
||||
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n{observation_capacity_note}"
|
||||
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n\n{escape_for_prompt(observation_capacity_note)}"
|
||||
|
||||
return (
|
||||
"You are a memory consolidation system. Synthesize facts into observations "
|
||||
"and merge with existing observations when appropriate.\n\n"
|
||||
f"## MISSION\n{mission}{capacity_section}\n\n"
|
||||
f"{_PROCESSING_RULES}" + _BATCH_DATA_SECTION + _BATCH_OUTPUT_FORMAT
|
||||
"You are a memory consolidation system. Synthesize new facts into "
|
||||
"observations, merging with existing observations when appropriate.\n\n"
|
||||
f"## MISSION\n\n{mission}\n\n"
|
||||
f"{_MISSION_PRIORITY_NOTE}"
|
||||
f"{capacity_section}\n\n"
|
||||
f"{_PROCESSING_RULES}\n\n"
|
||||
f"{_INPUT_SECTION}\n\n"
|
||||
f"{_DECISION_GUIDE}\n\n"
|
||||
f"{_OUTPUT_SECTION}" + output_language_directive(llm_output_language)
|
||||
)
|
||||
|
||||
|
||||
def build_consolidation_system_prompt(
|
||||
llm_output_language: str | None = None,
|
||||
) -> str:
|
||||
"""Bank-agnostic, cacheable system instruction for batch consolidation.
|
||||
|
||||
Holds only what is constant across banks: processing rules, input format,
|
||||
decision guide, and output format. The bank's MISSION is deliberately NOT
|
||||
here — baking it in would make the prefix bank-specific and force a separate
|
||||
Gemini context cache per mission. The mission, the per-batch INPUT, and any
|
||||
capacity constraint all ride in the user message (see
|
||||
:func:`build_consolidation_input`), so this prefix is identical for every
|
||||
bank and a single CachedContent serves them all. Returns final text
|
||||
(brace-escaped examples already unescaped) for verbatim use as system message
|
||||
and cached prefix.
|
||||
"""
|
||||
template = (
|
||||
"You are a memory consolidation system. Synthesize new facts into "
|
||||
"observations, merging with existing observations when appropriate.\n\n"
|
||||
f"{_MISSION_PRIORITY_NOTE}\n\n"
|
||||
f"{_PROCESSING_RULES}\n\n"
|
||||
f"{_INPUT_FORMAT_NOTE}\n\n"
|
||||
f"{_DECISION_GUIDE}\n\n"
|
||||
f"{_OUTPUT_SECTION}" + output_language_directive(llm_output_language)
|
||||
)
|
||||
# No {facts_text}/{observations_text} placeholders here — the only braces are
|
||||
# the doubled {{ }} in the OUTPUT examples, which .format() unescapes.
|
||||
return template.format()
|
||||
|
||||
|
||||
def build_consolidation_input(
|
||||
facts_text: str,
|
||||
observations_text: str,
|
||||
observations_mission: str | None = None,
|
||||
observation_capacity_note: str | None = None,
|
||||
) -> str:
|
||||
"""Per-batch user message: MISSION + INPUT data + any capacity constraint.
|
||||
|
||||
The MISSION lives here (not in the cached system prefix) so the prefix stays
|
||||
bank-agnostic and one CachedContent serves every bank. The capacity note also
|
||||
lives here since it varies as observation slots fill.
|
||||
"""
|
||||
mission = escape_for_prompt(observations_mission or _DEFAULT_MISSION)
|
||||
mission_section = f"## MISSION\n\n{mission}\n\n"
|
||||
capacity_section = ""
|
||||
if observation_capacity_note:
|
||||
capacity_section = f"## CAPACITY CONSTRAINT\n\n{escape_for_prompt(observation_capacity_note)}\n\n"
|
||||
# _SPLIT_INPUT_SECTION omits the stable observation-format explanation (now in
|
||||
# the cached system prefix) — only the variable facts/observations remain.
|
||||
template = mission_section + capacity_section + _SPLIT_INPUT_SECTION
|
||||
return template.format(facts_text=facts_text, observations_text=observations_text)
|
||||
|
||||
@@ -17,6 +17,7 @@ import httpx
|
||||
|
||||
from ..config import (
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
DEFAULT_RERANKER_ALIBABA_MODEL,
|
||||
DEFAULT_RERANKER_COHERE_MODEL,
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA,
|
||||
@@ -37,13 +38,14 @@ from ..config import (
|
||||
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL,
|
||||
DEFAULT_ZEROENTROPY_BASE_URL,
|
||||
ENV_RERANKER_ALIBABA_API_KEY,
|
||||
ENV_RERANKER_COHERE_API_KEY,
|
||||
ENV_RERANKER_COHERE_MODEL,
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA,
|
||||
ENV_RERANKER_FLASHRANK_MODEL,
|
||||
ENV_RERANKER_GOOGLE_PROJECT_ID,
|
||||
ENV_RERANKER_LITELLM_SDK_API_KEY,
|
||||
ENV_RERANKER_LOCAL_FORCE_CPU,
|
||||
ENV_RERANKER_LOCAL_MAX_CONCURRENT,
|
||||
ENV_RERANKER_LOCAL_MODEL,
|
||||
@@ -60,6 +62,43 @@ from ..config import (
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _resolve_malloc_trim():
|
||||
"""Return a callable that asks glibc to release freed heap pages to the OS.
|
||||
|
||||
Local CPU rerankers (FlashRank/ONNX, SentenceTransformers/torch) allocate
|
||||
large transient numpy/tensor buffers per call. On Linux glibc, those pages
|
||||
are freed at the Python level but kept by the allocator as a high-water
|
||||
mark — RSS grows monotonically across many recalls (see issue #1717).
|
||||
Calling `malloc_trim(0)` after each batch returns those pages to the OS.
|
||||
|
||||
Resolved once at import; returns a no-op on non-glibc platforms (macOS,
|
||||
musl, Windows) where the call is unavailable or unnecessary.
|
||||
"""
|
||||
import sys
|
||||
|
||||
if sys.platform != "linux":
|
||||
return lambda: None
|
||||
|
||||
import ctypes
|
||||
import ctypes.util
|
||||
|
||||
libc_path = ctypes.util.find_library("c")
|
||||
if libc_path is None:
|
||||
return lambda: None
|
||||
try:
|
||||
libc = ctypes.CDLL(libc_path)
|
||||
trim = libc.malloc_trim
|
||||
except (OSError, AttributeError):
|
||||
# Not glibc (musl has no malloc_trim) or libc lookup failed.
|
||||
return lambda: None
|
||||
trim.argtypes = [ctypes.c_size_t]
|
||||
trim.restype = ctypes.c_int
|
||||
return lambda: trim(0)
|
||||
|
||||
|
||||
_malloc_trim = _resolve_malloc_trim()
|
||||
|
||||
|
||||
class CrossEncoderModel(ABC):
|
||||
"""
|
||||
Abstract base class for cross-encoder reranking.
|
||||
@@ -266,25 +305,28 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
if self.bucket_batching and len(pairs) > 1:
|
||||
# Sort pairs by approximate token length to create homogeneous batches.
|
||||
# This eliminates padding waste — short pairs aren't padded to the length
|
||||
# of the longest pair in the batch. Quality-identical by construction.
|
||||
lengths = [len(pairs[i][0]) + len(pairs[i][1]) for i in range(len(pairs))]
|
||||
sorted_indices = sorted(range(len(pairs)), key=lambda i: lengths[i])
|
||||
sorted_pairs = [pairs[i] for i in sorted_indices]
|
||||
try:
|
||||
if self.bucket_batching and len(pairs) > 1:
|
||||
# Sort pairs by approximate token length to create homogeneous batches.
|
||||
# This eliminates padding waste — short pairs aren't padded to the length
|
||||
# of the longest pair in the batch. Quality-identical by construction.
|
||||
lengths = [len(pairs[i][0]) + len(pairs[i][1]) for i in range(len(pairs))]
|
||||
sorted_indices = sorted(range(len(pairs)), key=lambda i: lengths[i])
|
||||
sorted_pairs = [pairs[i] for i in sorted_indices]
|
||||
|
||||
sorted_scores = self._model.predict(sorted_pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
sorted_scores = sorted_scores.tolist() if hasattr(sorted_scores, "tolist") else list(sorted_scores)
|
||||
sorted_scores = self._model.predict(sorted_pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
sorted_scores = sorted_scores.tolist() if hasattr(sorted_scores, "tolist") else list(sorted_scores)
|
||||
|
||||
# Restore original order
|
||||
scores = [0.0] * len(pairs)
|
||||
for new_pos, orig_idx in enumerate(sorted_indices):
|
||||
scores[orig_idx] = sorted_scores[new_pos]
|
||||
return scores
|
||||
# Restore original order
|
||||
scores = [0.0] * len(pairs)
|
||||
for new_pos, orig_idx in enumerate(sorted_indices):
|
||||
scores[orig_idx] = sorted_scores[new_pos]
|
||||
return scores
|
||||
|
||||
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
|
||||
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
|
||||
finally:
|
||||
_malloc_trim()
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -546,12 +588,14 @@ class _CohereCompatibleRerankClient:
|
||||
rerank_url: str,
|
||||
timeout: float = 60.0,
|
||||
include_top_n: bool = True,
|
||||
include_return_documents: bool = False,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.rerank_url = rerank_url
|
||||
self.timeout = timeout
|
||||
self.include_top_n = include_top_n
|
||||
self.include_return_documents = include_return_documents
|
||||
self._async_client: httpx.AsyncClient | None = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
@@ -729,7 +773,7 @@ class ZeroEntropyCrossEncoder(CrossEncoderModel):
|
||||
See: https://docs.zeroentropy.dev/models
|
||||
"""
|
||||
|
||||
DEFAULT_BASE_URL = "https://api.zeroentropy.dev"
|
||||
DEFAULT_BASE_URL = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
RERANK_PATH = "/v1/models/rerank"
|
||||
|
||||
def __init__(
|
||||
@@ -962,32 +1006,35 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
# Group pairs by query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
try:
|
||||
# Group pairs by query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
# Build passages list for FlashRank
|
||||
passages = [{"id": i, "text": text} for i, (_, text) in enumerate(indexed_texts)]
|
||||
global_indices = [idx for idx, _ in indexed_texts]
|
||||
for query, indexed_texts in query_groups.items():
|
||||
# Build passages list for FlashRank
|
||||
passages = [{"id": i, "text": text} for i, (_, text) in enumerate(indexed_texts)]
|
||||
global_indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
# Create rerank request
|
||||
request = RerankRequest(query=query, passages=passages)
|
||||
results = self._ranker.rerank(request)
|
||||
# Create rerank request
|
||||
request = RerankRequest(query=query, passages=passages)
|
||||
results = self._ranker.rerank(request)
|
||||
|
||||
# Map scores back to original positions
|
||||
for result in results:
|
||||
local_idx = result["id"]
|
||||
score = result["score"]
|
||||
global_idx = global_indices[local_idx]
|
||||
all_scores[global_idx] = score
|
||||
# Map scores back to original positions
|
||||
for result in results:
|
||||
local_idx = result["id"]
|
||||
score = result["score"]
|
||||
global_idx = global_indices[local_idx]
|
||||
all_scores[global_idx] = score
|
||||
|
||||
return all_scores
|
||||
return all_scores
|
||||
finally:
|
||||
_malloc_trim()
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -1151,7 +1198,7 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
api_key: str | None = None,
|
||||
model: str = DEFAULT_RERANKER_LITELLM_SDK_MODEL,
|
||||
api_base: str | None = None,
|
||||
timeout: float = 60.0,
|
||||
@@ -1161,7 +1208,8 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
Initialize LiteLLM SDK cross-encoder client.
|
||||
|
||||
Args:
|
||||
api_key: API key for the reranking provider
|
||||
api_key: API key for the reranking provider (optional — omit for
|
||||
providers that use ambient credentials, e.g. AWS Bedrock with IAM)
|
||||
model: Model name with provider prefix (e.g., "deepinfra/Qwen3-reranker-8B")
|
||||
api_base: Custom base URL for API (optional)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
@@ -1236,8 +1284,9 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": texts,
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_key:
|
||||
rerank_kwargs["api_key"] = self.api_key
|
||||
if self.api_base:
|
||||
rerank_kwargs["api_base"] = self.api_base
|
||||
|
||||
@@ -1534,6 +1583,48 @@ class GoogleCrossEncoder(CrossEncoderModel):
|
||||
return await loop.run_in_executor(None, self._predict_sync, pairs)
|
||||
|
||||
|
||||
class AlibabaCloudCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
Alibaba Cloud DashScope text reranking API.
|
||||
|
||||
Uses the Cohere-compatible /reranks endpoint, which is the standard interface
|
||||
for qwen3-rerank. Authentication via HINDSIGHT_API_RERANKER_ALIBABA_API_KEY
|
||||
(or DASHSCOPE_API_KEY as a fallback).
|
||||
See: https://help.aliyun.com/zh/model-studio/text-rerank-api
|
||||
"""
|
||||
|
||||
RERANK_URL = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_RERANKER_ALIBABA_MODEL,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
self.model = model
|
||||
self._client = _CohereCompatibleRerankClient(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
rerank_url=self.RERANK_URL,
|
||||
timeout=timeout,
|
||||
include_return_documents=False,
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "alibaba"
|
||||
|
||||
async def initialize(self) -> None:
|
||||
if self._client._async_client is not None:
|
||||
return
|
||||
logger.info(f"Reranker: initializing Alibaba Cloud provider with model {self.model}")
|
||||
await self._client.initialize()
|
||||
logger.info("Reranker: Alibaba Cloud provider initialized")
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
return await self._client.predict(pairs)
|
||||
|
||||
|
||||
def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
"""
|
||||
Create a CrossEncoderModel instance based on configuration.
|
||||
@@ -1576,6 +1667,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_cohere_model,
|
||||
base_url=config.reranker_cohere_base_url,
|
||||
timeout=config.reranker_cohere_timeout,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.reranker_openrouter_api_key
|
||||
@@ -1588,6 +1680,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_openrouter_model,
|
||||
base_url="https://openrouter.ai/api/v1/rerank",
|
||||
timeout=config.reranker_openrouter_timeout,
|
||||
)
|
||||
elif provider == "flashrank":
|
||||
model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL)
|
||||
@@ -1602,18 +1695,15 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=config.reranker_litellm_api_key,
|
||||
model=config.reranker_litellm_model,
|
||||
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
|
||||
timeout=config.reranker_litellm_timeout,
|
||||
)
|
||||
elif provider == "litellm-sdk":
|
||||
api_key = config.reranker_litellm_sdk_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_RERANKER_LITELLM_SDK_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'litellm-sdk'"
|
||||
)
|
||||
return LiteLLMSDKCrossEncoder(
|
||||
api_key=api_key,
|
||||
api_key=config.reranker_litellm_sdk_api_key or None,
|
||||
model=config.reranker_litellm_sdk_model,
|
||||
api_base=config.reranker_litellm_sdk_api_base,
|
||||
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
|
||||
timeout=config.reranker_litellm_sdk_timeout,
|
||||
)
|
||||
elif provider == "zeroentropy":
|
||||
api_key = config.reranker_zeroentropy_api_key
|
||||
@@ -1624,6 +1714,8 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
return ZeroEntropyCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_zeroentropy_model,
|
||||
base_url=config.reranker_zeroentropy_base_url,
|
||||
timeout=config.reranker_zeroentropy_timeout,
|
||||
)
|
||||
elif provider == "siliconflow":
|
||||
api_key = config.reranker_siliconflow_api_key
|
||||
@@ -1635,6 +1727,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_siliconflow_model,
|
||||
base_url=config.reranker_siliconflow_base_url,
|
||||
timeout=config.reranker_siliconflow_timeout,
|
||||
)
|
||||
elif provider == "google":
|
||||
project_id = config.reranker_google_project_id
|
||||
@@ -1647,6 +1740,16 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
project_id=project_id,
|
||||
model=config.reranker_google_model,
|
||||
service_account_key=config.reranker_google_service_account_key,
|
||||
timeout=config.reranker_google_timeout,
|
||||
)
|
||||
elif provider == "alibaba":
|
||||
api_key = config.reranker_alibaba_api_key
|
||||
if not api_key:
|
||||
raise ValueError(f"{ENV_RERANKER_ALIBABA_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'alibaba'")
|
||||
return AlibabaCloudCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_alibaba_model,
|
||||
timeout=config.reranker_alibaba_timeout,
|
||||
)
|
||||
elif provider == "rrf":
|
||||
return RRFPassthroughCrossEncoder()
|
||||
@@ -1654,5 +1757,5 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
return JinaMLXCrossEncoder()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'alibaba', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
)
|
||||
|
||||
@@ -72,6 +72,30 @@ class DataAccessOps(ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def lock_document_for_write(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
doc_id: str,
|
||||
bank_id: str,
|
||||
) -> str | None:
|
||||
"""Ensure the document row exists, take a row lock on it, and return its
|
||||
pre-existing ``content_hash``.
|
||||
|
||||
This serializes all concurrent writers for ``doc_id`` at the DB level
|
||||
(so interleaved same-document retains can't corrupt each other), while
|
||||
creating the row on first write. The returned hash is ``'__pending__'``
|
||||
for a freshly inserted row, the stored hash for an existing one, or
|
||||
``None`` if the row could not be read back.
|
||||
|
||||
PG does this in a single statement (``INSERT ... ON CONFLICT DO UPDATE
|
||||
... RETURNING``), which always takes the row lock as part of the upsert.
|
||||
Oracle can't (``MERGE`` doesn't support ``RETURNING``), so it splits the
|
||||
work into an idempotent insert plus a ``SELECT ... FOR UPDATE``.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def insert_facts_batch(
|
||||
self,
|
||||
@@ -166,21 +190,6 @@ class DataAccessOps(ABC):
|
||||
|
||||
# -- LATERAL / fan-out queries ---------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_entity_unit_fanout(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
"""Fetch unit_ids for a list of entities with per-entity row cap.
|
||||
|
||||
PG uses unnest + CROSS JOIN LATERAL with LIMIT.
|
||||
Non-PG queries each entity individually.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
@@ -406,6 +415,74 @@ class DataAccessOps(ABC):
|
||||
"""Insert a webhook delivery task into async_operations."""
|
||||
...
|
||||
|
||||
# -- Graph maintenance queue -----------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def enqueue_graph_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
"""Insert unit_ids into graph_maintenance_queue, deduplicating on the
|
||||
(bank_id, unit_id) primary key.
|
||||
|
||||
Called inside the triggering transaction so enqueue is atomic with
|
||||
the mutation that caused it. Order is unspecified.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
"""Atomically claim a batch of rows from graph_maintenance_queue and
|
||||
remove them from the table.
|
||||
|
||||
Returns the list of ``unit_id`` strings. Empty list when the queue
|
||||
for ``bank_id`` is drained.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete entities in ``bank_id`` that no longer have any unit_entities
|
||||
rows referencing them. Returns the number of rows deleted.
|
||||
|
||||
FK ON DELETE CASCADE on entity_cooccurrences then removes any
|
||||
cooccurrence row pointing at the pruned entities.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete entity_cooccurrences rows in ``bank_id`` where the two
|
||||
entities still exist but no current unit references both of them.
|
||||
|
||||
These are stale-count rows: cooccurrence was real at the time it was
|
||||
recorded, but every memory_unit that witnessed both entities has
|
||||
since been deleted. Returns the number of rows deleted.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
@@ -416,6 +493,8 @@ class DataAccessOps(ABC):
|
||||
worker_id: str,
|
||||
reserved_limits: dict[str, int],
|
||||
shared_limit: int,
|
||||
*,
|
||||
consolidation_bank_priority: dict[str, int] | None = None,
|
||||
) -> list[ResultRow]:
|
||||
"""Claim pending tasks from the async_operations table.
|
||||
|
||||
@@ -423,6 +502,14 @@ class DataAccessOps(ABC):
|
||||
Oracle implementation uses two-step claims (query busy banks first, then
|
||||
claim excluding them) to avoid ORA-02014.
|
||||
|
||||
Args:
|
||||
consolidation_bank_priority: Per-bank priority for consolidation scheduling.
|
||||
Maps bank name patterns to integer priorities (higher = claimed first).
|
||||
Patterns support ``*`` as wildcard (converted to SQL ``%`` for LIKE).
|
||||
A bare ``*`` key is the catch-all default for unlisted banks.
|
||||
When set, consolidation tasks are claimed in priority tiers.
|
||||
None preserves current behavior (pure created_at ordering).
|
||||
|
||||
Returns claimed rows with operation_id, operation_type, task_payload, retry_count.
|
||||
The caller is responsible for building ClaimedTask objects.
|
||||
"""
|
||||
|
||||
@@ -47,6 +47,37 @@ class OracleOps(DataAccessOps):
|
||||
column_types=["text[]", "text[]", "text[]", "text[]", "integer[]", "text[]"],
|
||||
)
|
||||
|
||||
async def lock_document_for_write(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
doc_id: str,
|
||||
bank_id: str,
|
||||
) -> str | None:
|
||||
# Oracle can't express the PG "INSERT ... ON CONFLICT DO UPDATE ...
|
||||
# RETURNING" upsert in one statement — MERGE doesn't support RETURNING,
|
||||
# so the single-statement form rewrites to a MERGE that returns no rows
|
||||
# (DPY-1003). Split it into two statements instead:
|
||||
# 1. Idempotent insert that silently skips an existing row. The
|
||||
# IGNORE_ROW_ON_DUPKEY_INDEX hint suppresses ORA-00001 server-side;
|
||||
# a concurrent uncommitted insert of the same key blocks here until
|
||||
# the other writer commits, so writers still serialize.
|
||||
# 2. SELECT ... FOR UPDATE to take the row lock and read the hash
|
||||
# ('__pending__' for a row we just inserted, the stored hash for an
|
||||
# existing one).
|
||||
await conn.execute(
|
||||
f"INSERT /*+ IGNORE_ROW_ON_DUPKEY_INDEX({table}, pk_documents) */ "
|
||||
f"INTO {table} (id, bank_id, original_text, content_hash) "
|
||||
f"VALUES ($1, $2, '', '__pending__')",
|
||||
doc_id,
|
||||
bank_id,
|
||||
)
|
||||
return await conn.fetchval(
|
||||
f"SELECT content_hash FROM {table} WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
doc_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
async def insert_facts_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
@@ -215,29 +246,96 @@ class OracleOps(DataAccessOps):
|
||||
list(zip(unit_ids, entity_ids)),
|
||||
)
|
||||
|
||||
async def fetch_entity_unit_fanout(
|
||||
async def enqueue_graph_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
# Query each entity individually
|
||||
rows: list[ResultRow] = []
|
||||
for eid in entity_id_list:
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT $1 AS entity_id, ue.unit_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.entity_id = $1
|
||||
ORDER BY ue.unit_id DESC
|
||||
LIMIT $2
|
||||
""",
|
||||
eid,
|
||||
limit_per_entity,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
# Oracle doesn't support ON CONFLICT; rely on the PK and the
|
||||
# IGNORE_ROW_ON_DUPKEY_INDEX hint to skip duplicates server-side.
|
||||
# The hint name must match the PK constraint exactly.
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT /*+ IGNORE_ROW_ON_DUPKEY_INDEX({table}, pk_graph_maintenance_queue) */
|
||||
INTO {table} (bank_id, unit_id)
|
||||
VALUES ($1, $2)
|
||||
""",
|
||||
[(bank_id, uid) for uid in unit_ids],
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
# Two-step claim: select the batch, then delete by exact keys. Oracle's
|
||||
# DELETE ... RETURNING doesn't accept a multi-row subquery, so we can't
|
||||
# do it in one statement like the PG version.
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT unit_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
FETCH FIRST $2 ROWS ONLY
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
)
|
||||
claimed = [str(row["unit_id"]) for row in rows]
|
||||
if claimed:
|
||||
await conn.executemany(
|
||||
f"DELETE FROM {table} WHERE bank_id = $1 AND unit_id = $2",
|
||||
[(bank_id, uid) for uid in claimed],
|
||||
)
|
||||
rows.extend(entity_rows)
|
||||
return rows
|
||||
return claimed
|
||||
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# The Oracle DatabaseConnection wrapper reshapes ``cursor.rowcount`` into
|
||||
# the same ``"DELETE N"`` status string asyncpg returns, so the same
|
||||
# ``int(deleted.split()[-1])`` parsing works on both dialects.
|
||||
deleted = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {entities_table}
|
||||
WHERE bank_id = $1
|
||||
AND id NOT IN (SELECT DISTINCT entity_id FROM {ue_table})
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(deleted.split()[-1]) if isinstance(deleted, str) and deleted.startswith("DELETE") else 0
|
||||
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
deleted = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {ec_table}
|
||||
WHERE entity_id_1 IN (SELECT id FROM {entities_table} WHERE bank_id = $1)
|
||||
AND (entity_id_1, entity_id_2) NOT IN (
|
||||
SELECT u1.entity_id, u2.entity_id
|
||||
FROM {ue_table} u1
|
||||
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(deleted.split()[-1]) if isinstance(deleted, str) and deleted.startswith("DELETE") else 0
|
||||
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
@@ -720,7 +818,257 @@ class OracleOps(DataAccessOps):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def claim_tasks(self, conn, table, worker_id, reserved_limits, shared_limit):
|
||||
async def _claim_consolidation_tasks(
|
||||
self,
|
||||
conn,
|
||||
table: str,
|
||||
busy_bank_ids: list[str],
|
||||
claimed_ids: list,
|
||||
limit: int,
|
||||
priority_map: dict[str, int] | None,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with optional priority-based tiered ordering.
|
||||
|
||||
Mirrors the PostgreSQL implementation. The Oracle SQL adapter
|
||||
translates ``LIKE ANY`` / ``NOT LIKE ALL`` via ``_expand_any_lists``.
|
||||
"""
|
||||
if limit <= 0:
|
||||
return []
|
||||
|
||||
if not priority_map:
|
||||
return await self._claim_consolidation_plain(conn, table, busy_bank_ids, claimed_ids, limit)
|
||||
|
||||
# --- Tiered claiming (same algorithm as PG) ---
|
||||
specific_by_priority: dict[int, list[str]] = {}
|
||||
all_specific_sql: list[str] = []
|
||||
catch_all_priority = 1
|
||||
|
||||
for pattern, priority in priority_map.items():
|
||||
if pattern == "*":
|
||||
catch_all_priority = priority
|
||||
else:
|
||||
sql_pat = pattern.replace("*", "%")
|
||||
specific_by_priority.setdefault(priority, []).append(sql_pat)
|
||||
all_specific_sql.append(sql_pat)
|
||||
|
||||
all_priorities = sorted(set(specific_by_priority.keys()) | {catch_all_priority}, reverse=True)
|
||||
|
||||
remaining = limit
|
||||
result: list = []
|
||||
|
||||
for pri in all_priorities:
|
||||
if remaining <= 0:
|
||||
break
|
||||
|
||||
if pri in specific_by_priority:
|
||||
rows = await self._claim_consolidation_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
specific_by_priority[pri],
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
if pri == catch_all_priority and remaining > 0:
|
||||
rows = await self._claim_consolidation_not_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
all_specific_sql,
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
return result
|
||||
|
||||
async def _claim_consolidation_plain(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with default created_at ordering."""
|
||||
exclude_ids = claimed_ids if claimed_ids else None
|
||||
if busy_bank_ids:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
AND operation_id != ALL($2::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
sql_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks matching LIKE patterns."""
|
||||
params: list = [sql_patterns]
|
||||
conditions = ["bank_id LIKE ANY($1::text[])"]
|
||||
idx = 2
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra = " AND ".join(conditions)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND {extra}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_not_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
exclude_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks NOT matching any specific pattern (catch-all tier)."""
|
||||
params: list = []
|
||||
conditions: list[str] = []
|
||||
idx = 1
|
||||
|
||||
if exclude_patterns:
|
||||
conditions.append(f"bank_id NOT LIKE ALL(${idx}::text[])")
|
||||
params.append(exclude_patterns)
|
||||
idx += 1
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra_clause = (" AND " + " AND ".join(conditions)) if conditions else ""
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW()){extra_clause}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def claim_tasks(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
worker_id,
|
||||
reserved_limits,
|
||||
shared_limit,
|
||||
*,
|
||||
consolidation_bank_priority=None,
|
||||
):
|
||||
"""Oracle two-step claiming to avoid ORA-02014 with NOT EXISTS + FOR UPDATE."""
|
||||
all_rows = []
|
||||
claimed_ids = []
|
||||
@@ -731,7 +1079,6 @@ class OracleOps(DataAccessOps):
|
||||
continue
|
||||
|
||||
if op_type == "consolidation":
|
||||
# Two-step: find busy banks first, then claim excluding them
|
||||
busy_banks = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT bank_id FROM {table}
|
||||
@@ -740,38 +1087,14 @@ class OracleOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids = [r["bank_id"] for r in busy_banks]
|
||||
|
||||
if busy_bank_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -835,7 +1158,7 @@ class OracleOps(DataAccessOps):
|
||||
all_rows.append(row)
|
||||
remaining_shared -= len(rows)
|
||||
|
||||
# 2b. Consolidation tasks (with bank-serialization)
|
||||
# 2b. Consolidation tasks (with bank-serialization + optional priority)
|
||||
if remaining_shared > 0:
|
||||
busy_banks_2 = await conn.fetch(
|
||||
f"""
|
||||
@@ -845,76 +1168,14 @@ class OracleOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids_2 = [r["bank_id"] for r in busy_banks_2]
|
||||
|
||||
if claimed_ids:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
AND bank_id != ALL($2::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids_2,
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
|
||||
@@ -49,6 +49,30 @@ class PostgreSQLOps(DataAccessOps):
|
||||
content_hashes,
|
||||
)
|
||||
|
||||
async def lock_document_for_write(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
doc_id: str,
|
||||
bank_id: str,
|
||||
) -> str | None:
|
||||
# Single upsert that both creates the row (if absent) and locks it (if
|
||||
# present) atomically. ON CONFLICT DO UPDATE always takes the row lock as
|
||||
# part of the statement, so all concurrent same-document writers serialize
|
||||
# on the document row in one consistent step (the earlier two-step form —
|
||||
# DO NOTHING + a separate SELECT FOR UPDATE — could deadlock because
|
||||
# DO NOTHING takes no lock on an existing row). The SET is a no-op
|
||||
# self-assignment used only to acquire the lock; RETURNING yields the
|
||||
# pre-existing hash (or '__pending__' for a freshly inserted row).
|
||||
return await conn.fetchval(
|
||||
f"INSERT INTO {table} (id, bank_id, original_text, content_hash) "
|
||||
f"VALUES ($1, $2, '', '__pending__') "
|
||||
f"ON CONFLICT (id, bank_id) DO UPDATE SET content_hash = {table}.content_hash "
|
||||
f"RETURNING content_hash",
|
||||
doc_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
async def insert_facts_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
@@ -104,7 +128,46 @@ class PostgreSQLOps(DataAccessOps):
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
elif config.text_search_extension == "native":
|
||||
# search_vector is a regular tsvector column populated here using the
|
||||
# configured native dictionary. It used to be GENERATED ALWAYS with
|
||||
# a hardcoded 'english', which prevented per-deployment language
|
||||
# configuration. text_search_extension_native_language is validated
|
||||
# in HindsightConfig.validate() as a PG identifier, so embedding it
|
||||
# as a SQL literal is safe.
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals, search_vector)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals,
|
||||
to_tsvector(
|
||||
'{config.text_search_extension_native_language}'::regconfig,
|
||||
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, '')
|
||||
)
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
else:
|
||||
# pg_textsearch, pgroonga, and pg_search: search_vector is a dummy
|
||||
# TEXT column; the actual full-text index operates on the base text
|
||||
# columns directly, so we don't populate search_vector at insert time.
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
@@ -161,6 +224,23 @@ class PostgreSQLOps(DataAccessOps):
|
||||
exists_clause: str,
|
||||
chunk_size: int = 5000,
|
||||
) -> None:
|
||||
# exists_clause is unused on PostgreSQL: the memory_links → memory_units
|
||||
# FKs are DEFERRABLE INITIALLY DEFERRED, so an INSERT takes no lock on the
|
||||
# referenced parent rows until COMMIT — a concurrent committed DELETE in
|
||||
# that window (consolidation pruning observations, document re-tracking)
|
||||
# trips fk_memory_links_{to,from}_unit_id_memory_units at COMMIT (#1882),
|
||||
# and a WHERE EXISTS guard can't prevent it (the row passes the check,
|
||||
# then is deleted before the deferred check runs). Instead a CTE locks the
|
||||
# referenced units FOR KEY SHARE in the *same statement*: the lock blocks a
|
||||
# concurrent DELETE until our transaction commits and is held through the
|
||||
# deferred check, and the INSERT only takes links whose endpoints are in
|
||||
# the locked set, so rows that already vanished are dropped. Folding it
|
||||
# into the one INSERT keeps this to a single round-trip — no extra query
|
||||
# and no surrounding transaction needed. (Oracle's immediate FK has no
|
||||
# such window and uses exists_clause via its own bulk_insert_links.)
|
||||
from ..schema import fq_table
|
||||
|
||||
mu_table = fq_table("memory_units")
|
||||
from_ids = [lnk[0] for lnk in sorted_links]
|
||||
to_ids = [lnk[1] for lnk in sorted_links]
|
||||
types = [lnk[2] for lnk in sorted_links]
|
||||
@@ -169,24 +249,37 @@ class PostgreSQLOps(DataAccessOps):
|
||||
|
||||
for chunk_start in range(0, len(sorted_links), chunk_size):
|
||||
chunk_end = min(chunk_start + chunk_size, len(sorted_links))
|
||||
chunk_from = from_ids[chunk_start:chunk_end]
|
||||
chunk_to = to_ids[chunk_start:chunk_end]
|
||||
# Distinct referenced parents, sorted so concurrent inserters acquire
|
||||
# the row-share locks in a consistent order (avoids deadlocks; same
|
||||
# convention as the (from, to) link sort).
|
||||
referenced = sorted({str(x) for x in chunk_from} | {str(x) for x in chunk_to})
|
||||
await conn.execute(
|
||||
f"""
|
||||
WITH locked AS (
|
||||
SELECT id FROM {mu_table}
|
||||
WHERE id = ANY($7::uuid[])
|
||||
ORDER BY id
|
||||
FOR KEY SHARE
|
||||
)
|
||||
INSERT INTO {table}
|
||||
(from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id)
|
||||
SELECT f, t, tp, w, e, $6
|
||||
FROM unnest($1::uuid[], $2::uuid[], $3::text[], $4::float8[], $5::uuid[])
|
||||
AS t(f, t, tp, w, e)
|
||||
{exists_clause}
|
||||
AS u(f, t, tp, w, e)
|
||||
WHERE f IN (SELECT id FROM locked) AND t IN (SELECT id FROM locked)
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type,
|
||||
COALESCE(entity_id, '{nil_entity_uuid}'::uuid))
|
||||
DO NOTHING
|
||||
""",
|
||||
from_ids[chunk_start:chunk_end],
|
||||
to_ids[chunk_start:chunk_end],
|
||||
chunk_from,
|
||||
chunk_to,
|
||||
types[chunk_start:chunk_end],
|
||||
weights[chunk_start:chunk_end],
|
||||
entity_ids[chunk_start:chunk_end],
|
||||
bank_id,
|
||||
referenced,
|
||||
timeout=300,
|
||||
)
|
||||
|
||||
@@ -251,29 +344,101 @@ class PostgreSQLOps(DataAccessOps):
|
||||
entity_ids,
|
||||
)
|
||||
|
||||
async def fetch_entity_unit_fanout(
|
||||
async def enqueue_graph_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
return await conn.fetch(
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
await conn.execute(
|
||||
f"""
|
||||
SELECT e.entity_id, n.unit_id
|
||||
FROM unnest($1::uuid[]) AS e(entity_id)
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue.unit_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.entity_id = e.entity_id
|
||||
ORDER BY ue.unit_id DESC
|
||||
LIMIT $2
|
||||
) n
|
||||
INSERT INTO {table} (bank_id, unit_id)
|
||||
SELECT $1, v FROM unnest($2::uuid[]) AS t(v)
|
||||
ON CONFLICT (bank_id, unit_id) DO NOTHING
|
||||
""",
|
||||
entity_id_list,
|
||||
limit_per_entity,
|
||||
bank_id,
|
||||
unit_ids,
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
DELETE FROM {table}
|
||||
WHERE (bank_id, unit_id) IN (
|
||||
SELECT bank_id, unit_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
LIMIT $2
|
||||
)
|
||||
RETURNING unit_id
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
)
|
||||
return [str(row["unit_id"]) for row in rows]
|
||||
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# Scoped by entities.bank_id (indexed). The NOT EXISTS subquery is
|
||||
# backed by idx_ue_entity on unit_entities(entity_id), so this stays
|
||||
# linear in the number of entities in the bank — not in the size of
|
||||
# unit_entities globally.
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {entities_table} e
|
||||
WHERE e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {ue_table} ue WHERE ue.entity_id = e.id
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
# asyncpg returns "DELETE N"
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# Scope by joining through entities.bank_id (entity_cooccurrences itself
|
||||
# has no bank_id column — entities don't span banks, so scoping via
|
||||
# entity_id_1 is sufficient).
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {ec_table} c
|
||||
USING {entities_table} e
|
||||
WHERE e.id = c.entity_id_1
|
||||
AND e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {ue_table} u1
|
||||
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
|
||||
WHERE u1.entity_id = c.entity_id_1
|
||||
AND u2.entity_id = c.entity_id_2
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
@@ -726,7 +891,268 @@ class PostgreSQLOps(DataAccessOps):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def claim_tasks(self, conn, table, worker_id, reserved_limits, shared_limit):
|
||||
async def _claim_consolidation_tasks(
|
||||
self,
|
||||
conn,
|
||||
table: str,
|
||||
busy_bank_ids: list[str],
|
||||
claimed_ids: list,
|
||||
limit: int,
|
||||
priority_map: dict[str, int] | None,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with optional priority-based tiered ordering.
|
||||
|
||||
When *priority_map* is ``None``, uses the default ``ORDER BY created_at``
|
||||
with bank-serialization (exclude busy banks). When set, claims in
|
||||
priority tiers — highest-priority banks first. Specific patterns always
|
||||
take precedence over the catch-all ``*`` entry.
|
||||
"""
|
||||
if limit <= 0:
|
||||
return []
|
||||
|
||||
# --- Fast path: no priority map -> current behavior ---
|
||||
if not priority_map:
|
||||
return await self._claim_consolidation_plain(conn, table, busy_bank_ids, claimed_ids, limit)
|
||||
|
||||
# --- Tiered claiming ---
|
||||
# Separate specific patterns from catch-all.
|
||||
# Specific patterns always take precedence: a bank matching ``shadow-*``
|
||||
# uses that entry's priority even if the catch-all ``*`` has a higher
|
||||
# value. The catch-all only applies to banks not matching any specific
|
||||
# pattern.
|
||||
specific_by_priority: dict[int, list[str]] = {}
|
||||
all_specific_sql: list[str] = []
|
||||
catch_all_priority = 1 # default when no ``*`` entry
|
||||
|
||||
for pattern, priority in priority_map.items():
|
||||
if pattern == "*":
|
||||
catch_all_priority = priority
|
||||
else:
|
||||
sql_pat = pattern.replace("*", "%")
|
||||
specific_by_priority.setdefault(priority, []).append(sql_pat)
|
||||
all_specific_sql.append(sql_pat)
|
||||
|
||||
# Collect all priority levels (specific tiers + catch-all) sorted desc.
|
||||
all_priorities = sorted(set(specific_by_priority.keys()) | {catch_all_priority}, reverse=True)
|
||||
|
||||
remaining = limit
|
||||
result: list = []
|
||||
|
||||
for pri in all_priorities:
|
||||
if remaining <= 0:
|
||||
break
|
||||
|
||||
# Specific-pattern tier at this priority level
|
||||
if pri in specific_by_priority:
|
||||
rows = await self._claim_consolidation_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
specific_by_priority[pri],
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
# Catch-all tier at this priority level
|
||||
if pri == catch_all_priority and remaining > 0:
|
||||
rows = await self._claim_consolidation_not_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
all_specific_sql,
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
return result
|
||||
|
||||
async def _claim_consolidation_plain(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with default created_at ordering."""
|
||||
exclude_ids = claimed_ids if claimed_ids else None
|
||||
if busy_bank_ids:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
AND operation_id != ALL($2::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
sql_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks matching LIKE patterns."""
|
||||
params: list = [sql_patterns]
|
||||
conditions = ["bank_id LIKE ANY($1::text[])"]
|
||||
idx = 2
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra = " AND ".join(conditions)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND {extra}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_not_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
exclude_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks NOT matching any specific pattern (catch-all tier)."""
|
||||
params: list = []
|
||||
conditions: list[str] = []
|
||||
idx = 1
|
||||
|
||||
if exclude_patterns:
|
||||
conditions.append(f"bank_id NOT LIKE ALL(${idx}::text[])")
|
||||
params.append(exclude_patterns)
|
||||
idx += 1
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra_clause = (" AND " + " AND ".join(conditions)) if conditions else ""
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW()){extra_clause}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def claim_tasks(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
worker_id,
|
||||
reserved_limits,
|
||||
shared_limit,
|
||||
*,
|
||||
consolidation_bank_priority=None,
|
||||
):
|
||||
all_rows = []
|
||||
claimed_ids = []
|
||||
|
||||
@@ -744,38 +1170,14 @@ class PostgreSQLOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids = [r["bank_id"] for r in busy_banks]
|
||||
|
||||
if busy_bank_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -839,7 +1241,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
all_rows.append(row)
|
||||
remaining_shared -= len(rows)
|
||||
|
||||
# 2b. Consolidation tasks (with bank-serialization)
|
||||
# 2b. Consolidation tasks (with bank-serialization + optional priority)
|
||||
if remaining_shared > 0:
|
||||
busy_banks_2 = await conn.fetch(
|
||||
f"""
|
||||
@@ -849,76 +1251,14 @@ class PostgreSQLOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids_2 = [r["bank_id"] for r in busy_banks_2]
|
||||
|
||||
if claimed_ids:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
AND bank_id != ALL($2::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids_2,
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
|
||||
@@ -14,6 +14,7 @@ Supports multi-tenant schema isolation via ALTER SESSION SET CURRENT_SCHEMA.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
@@ -72,6 +73,9 @@ _RETURNING_RE = re.compile(r"\bRETURNING\s+(.+)", re.IGNORECASE | re.DOTALL)
|
||||
|
||||
_ANY_RE = re.compile(r"=\s*ANY\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
_NOT_ALL_RE = re.compile(r"!=\s*ALL\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
# LIKE ANY / NOT LIKE ALL — capture the column name before the operator
|
||||
_LIKE_ANY_RE = re.compile(r"(\w+)\s+LIKE\s+ANY\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
_NOT_LIKE_ALL_RE = re.compile(r"(\w+)\s+NOT\s+LIKE\s+ALL\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
|
||||
_JSON_ARROW_TEXT_RE = re.compile(r'("?\w+"?)\s*->>\s*\'(\w+)\'') # handles both col and "col"
|
||||
_JSON_HAS_KEY_RE = re.compile(r"(\w+)\s*\?\s*'(\w+)'")
|
||||
@@ -152,6 +156,23 @@ _JSON_COL_NAMES = {
|
||||
"task_payload",
|
||||
"history",
|
||||
}
|
||||
# NOTE: the history tables' JSON payload column is named ``content`` — deliberately
|
||||
# NOT added here, because ``mental_models.content`` is plain text (adding "content"
|
||||
# would corrupt those reads). The history read paths json.loads ``content`` directly.
|
||||
|
||||
# Columns backed by CLOB in Oracle (large text or JSON). When such a column is
|
||||
# returned via a ``RETURNING`` clause it must be bound as DB_TYPE_CLOB; binding
|
||||
# it as VARCHAR raises ORA-22835 ("buffer too small for CLOB to CHAR") once the
|
||||
# value exceeds 4000 bytes. Union of the JSON-CLOB columns above and the
|
||||
# large-text CLOB columns.
|
||||
_CLOB_RETURNING_COLS = _JSON_COL_NAMES | {
|
||||
"content",
|
||||
"text",
|
||||
"context",
|
||||
"structured_content",
|
||||
"text_signals",
|
||||
"search_vector",
|
||||
}
|
||||
|
||||
|
||||
def _is_uuid_column(col: str) -> bool:
|
||||
@@ -349,6 +370,10 @@ def _rewrite_pg_to_oracle(query: str) -> RewriteResult:
|
||||
# Boolean literals: Oracle uses NUMBER(1) for booleans
|
||||
query = re.sub(r"\b=\s*TRUE\b", "= 1", query, flags=re.IGNORECASE)
|
||||
query = re.sub(r"\b=\s*FALSE\b", "= 0", query, flags=re.IGNORECASE)
|
||||
# FOR NO KEY UPDATE → FOR UPDATE (Oracle has only FOR UPDATE; it does not block
|
||||
# indexed-FK child inserts the way PG's FOR UPDATE would, so plain FOR UPDATE is
|
||||
# the correct equivalent). Must run before the FOR SHARE rule below.
|
||||
query = re.sub(r"\bFOR\s+NO\s+KEY\s+UPDATE\b", "FOR UPDATE", query, flags=re.IGNORECASE)
|
||||
# FOR SHARE → FOR UPDATE (Oracle doesn't support FOR SHARE)
|
||||
query = re.sub(r"\bFOR\s+SHARE\b", "FOR UPDATE", query, flags=re.IGNORECASE)
|
||||
|
||||
@@ -535,6 +560,12 @@ def _rewrite_pg_to_oracle(query: str) -> RewriteResult:
|
||||
# != ALL(:N) → NOT IN (expanded list) — the negative counterpart of = ANY
|
||||
query = _NOT_ALL_RE.sub(r"NOT IN (/*EXPAND:\1*/)", query)
|
||||
|
||||
# col LIKE ANY(:N) → (col LIKE :p0 OR col LIKE :p1 OR ...)
|
||||
query = _LIKE_ANY_RE.sub(r"\1 /*LIKE_ANY:\2:\1*/", query)
|
||||
|
||||
# col NOT LIKE ALL(:N) → (col NOT LIKE :p0 AND col NOT LIKE :p1 AND ...)
|
||||
query = _NOT_LIKE_ALL_RE.sub(r"\1 /*NOT_LIKE_ALL:\2:\1*/", query)
|
||||
|
||||
# CTE AS MATERIALIZED (...) → AS (...) — Oracle doesn't support MATERIALIZED CTE hint
|
||||
query = re.sub(r"\bAS\s+MATERIALIZED\s*\(", "AS (", query, flags=re.IGNORECASE)
|
||||
|
||||
@@ -672,6 +703,11 @@ class OracleConnection(DatabaseConnection):
|
||||
params[f"ret_{i}"] = cursor.var(oracledb.DB_TYPE_TIMESTAMP_TZ, arraysize=1)
|
||||
elif clean in _NUMERIC_COLS:
|
||||
params[f"ret_{i}"] = cursor.var(oracledb.DB_TYPE_NUMBER, arraysize=1)
|
||||
elif clean in _CLOB_RETURNING_COLS:
|
||||
# CLOB-backed column: a VARCHAR out-bind caps at 4000 bytes and
|
||||
# raises ORA-22835 for larger values. Read back as a LOB in
|
||||
# _read_returning_values.
|
||||
params[f"ret_{i}"] = cursor.var(oracledb.DB_TYPE_CLOB, arraysize=1)
|
||||
else:
|
||||
params[f"ret_{i}"] = cursor.var(oracledb.DB_TYPE_VARCHAR, arraysize=1)
|
||||
|
||||
@@ -724,17 +760,38 @@ class OracleConnection(DatabaseConnection):
|
||||
|
||||
_expand_counter = 0
|
||||
|
||||
@staticmethod
|
||||
def _resolve_list_param(params: dict[str, Any], key: str) -> list | None:
|
||||
"""Resolve a parameter that may be a list or a JSON-encoded list string."""
|
||||
val = params.get(key)
|
||||
if isinstance(val, str):
|
||||
try:
|
||||
parsed = json.loads(val)
|
||||
if isinstance(parsed, list):
|
||||
return parsed
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
if isinstance(val, (list, tuple)):
|
||||
return list(val)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _expand_any_lists(query: str, params: dict[str, Any] | None) -> tuple[str, dict[str, Any] | None]:
|
||||
"""Expand /*EXPAND:N*/ markers into individual bind vars for IN clauses.
|
||||
"""Expand /*EXPAND:N*/, /*LIKE_ANY:N:col*/, /*NOT_LIKE_ALL:N:col*/ markers.
|
||||
|
||||
Converts: IN (/*EXPAND:1*/) with params["1"] = [a, b, c]
|
||||
Into: IN (:any_0, :any_1, :any_2) with params["any_0"]=a, etc.
|
||||
|
||||
Converts: col /*LIKE_ANY:1:col*/ with params["1"] = [a, b]
|
||||
Into: (col LIKE :lk_0 OR col LIKE :lk_1)
|
||||
|
||||
Converts: col /*NOT_LIKE_ALL:1:col*/ with params["1"] = [a, b]
|
||||
Into: (col NOT LIKE :nlk_0 AND col NOT LIKE :nlk_1)
|
||||
|
||||
Uses a unique prefix to avoid name collisions with other bind vars.
|
||||
The original param is kept (for other references to :N in the query).
|
||||
"""
|
||||
if params is None or "/*EXPAND:" not in query:
|
||||
if params is None or "/*" not in query:
|
||||
return query, params
|
||||
|
||||
expand_re = re.compile(r"/\*EXPAND:(\d+)\*/")
|
||||
@@ -775,6 +832,50 @@ class OracleConnection(DatabaseConnection):
|
||||
|
||||
query = expand_re.sub(_replace, query)
|
||||
|
||||
# Expand LIKE ANY: col /*LIKE_ANY:N:col*/ → (col LIKE :p0 OR col LIKE :p1 ...)
|
||||
like_any_re = re.compile(r"(\w+)\s*/\*LIKE_ANY:(\d+):(\w+)\*/")
|
||||
|
||||
def _replace_like_any(m):
|
||||
_col = m.group(1) # redundant column ref before marker
|
||||
param_key = m.group(2)
|
||||
col = m.group(3)
|
||||
val = OracleConnection._resolve_list_param(params, param_key)
|
||||
if val is None or len(val) == 0:
|
||||
return "1=0" # no patterns → no match
|
||||
OracleConnection._expand_counter += 1
|
||||
prefix = f"lk{OracleConnection._expand_counter}"
|
||||
clauses = []
|
||||
for i, item in enumerate(val):
|
||||
k = f"{prefix}_{i}"
|
||||
params[k] = item
|
||||
clauses.append(f"{col} LIKE :{k}")
|
||||
keys_to_remove.add(param_key)
|
||||
return f"({' OR '.join(clauses)})"
|
||||
|
||||
query = like_any_re.sub(_replace_like_any, query)
|
||||
|
||||
# Expand NOT LIKE ALL: col /*NOT_LIKE_ALL:N:col*/ → (col NOT LIKE :p0 AND ...)
|
||||
not_like_all_re = re.compile(r"(\w+)\s*/\*NOT_LIKE_ALL:(\d+):(\w+)\*/")
|
||||
|
||||
def _replace_not_like_all(m):
|
||||
_col = m.group(1)
|
||||
param_key = m.group(2)
|
||||
col = m.group(3)
|
||||
val = OracleConnection._resolve_list_param(params, param_key)
|
||||
if val is None or len(val) == 0:
|
||||
return "1=1" # no patterns → everything matches
|
||||
OracleConnection._expand_counter += 1
|
||||
prefix = f"nlk{OracleConnection._expand_counter}"
|
||||
clauses = []
|
||||
for i, item in enumerate(val):
|
||||
k = f"{prefix}_{i}"
|
||||
params[k] = item
|
||||
clauses.append(f"{col} NOT LIKE :{k}")
|
||||
keys_to_remove.add(param_key)
|
||||
return f"({' AND '.join(clauses)})"
|
||||
|
||||
query = not_like_all_re.sub(_replace_not_like_all, query)
|
||||
|
||||
# Remove original list params that were expanded — their placeholder
|
||||
# (:N) no longer exists in the query, and leaving them causes DPY-4008.
|
||||
# Only remove if the key's placeholder is truly gone from the query.
|
||||
@@ -784,7 +885,7 @@ class OracleConnection(DatabaseConnection):
|
||||
|
||||
return query, params
|
||||
|
||||
def _read_returning_values(self, returning_cols: list[str], params: dict[str, Any]) -> dict[str, Any] | None:
|
||||
async def _read_returning_values(self, returning_cols: list[str], params: dict[str, Any]) -> dict[str, Any] | None:
|
||||
"""Read values from RETURNING INTO output variables after execute."""
|
||||
row: dict[str, Any] = {}
|
||||
for i, col in enumerate(returning_cols):
|
||||
@@ -794,6 +895,14 @@ class OracleConnection(DatabaseConnection):
|
||||
return None
|
||||
val = values[0] if isinstance(values, list) else values
|
||||
|
||||
# CLOB-bound columns return a LOB handle; read it to a string. The
|
||||
# async pool yields AsyncLOB whose read() is a coroutine.
|
||||
if val is not None and not isinstance(val, (str, bytes, int, float)) and hasattr(val, "read"):
|
||||
data = val.read()
|
||||
if inspect.isawaitable(data):
|
||||
data = await data
|
||||
val = data
|
||||
|
||||
# Clean alias: "LOWER(canonical_name) AS name_lower" → "name_lower"
|
||||
clean_col = col.strip()
|
||||
upper = clean_col.upper()
|
||||
@@ -981,7 +1090,7 @@ class OracleConnection(DatabaseConnection):
|
||||
raise
|
||||
|
||||
if ret_cols is not None:
|
||||
row_dict = self._read_returning_values(ret_cols, params)
|
||||
row_dict = await self._read_returning_values(ret_cols, params)
|
||||
return [ResultRow(row_dict)] if row_dict else []
|
||||
|
||||
columns = [col[0].lower() for col in cursor.description or []]
|
||||
@@ -1019,7 +1128,7 @@ class OracleConnection(DatabaseConnection):
|
||||
raise
|
||||
|
||||
if ret_cols is not None:
|
||||
row_dict = self._read_returning_values(ret_cols, params)
|
||||
row_dict = await self._read_returning_values(ret_cols, params)
|
||||
return ResultRow(row_dict) if row_dict else None
|
||||
|
||||
columns = [col[0].lower() for col in cursor.description or []]
|
||||
@@ -1052,7 +1161,7 @@ class OracleConnection(DatabaseConnection):
|
||||
await cursor.execute(query, params)
|
||||
|
||||
if ret_cols is not None:
|
||||
row_dict = self._read_returning_values(ret_cols, params)
|
||||
row_dict = await self._read_returning_values(ret_cols, params)
|
||||
if row_dict is None:
|
||||
return None
|
||||
vals = list(row_dict.values())
|
||||
|
||||
@@ -6,7 +6,7 @@ import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -101,6 +101,14 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
|
||||
"""
|
||||
Async context manager to acquire a database connection with retry logic.
|
||||
|
||||
Retries the *acquire* itself when it raises a retryable error (connection
|
||||
drop, timeout, deadlock detected during acquire). Exceptions raised by
|
||||
user code inside the ``async with`` block are NOT retried — they propagate
|
||||
as-is. Wrapping retry around the yield would violate the
|
||||
``@asynccontextmanager`` single-yield contract and surface as
|
||||
``RuntimeError("generator didn't stop after athrow()")`` on every
|
||||
retryable inner error, masking the real cause.
|
||||
|
||||
Accepts either a DatabaseBackend or a raw asyncpg.Pool for backward compatibility.
|
||||
|
||||
Usage:
|
||||
@@ -109,7 +117,7 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
|
||||
|
||||
Args:
|
||||
backend_or_pool: A DatabaseBackend instance or asyncpg.Pool
|
||||
max_retries: Maximum number of retry attempts
|
||||
max_retries: Maximum number of retry attempts for the acquire step
|
||||
|
||||
Yields:
|
||||
A DatabaseConnection (if backend) or asyncpg.Connection (if pool)
|
||||
@@ -117,31 +125,32 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
|
||||
from .db.base import DatabaseBackend
|
||||
|
||||
if isinstance(backend_or_pool, DatabaseBackend) or getattr(backend_or_pool, "_wraps_backend", False):
|
||||
# Use the backend's acquire context manager with retry
|
||||
start = time.time()
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
async with backend_or_pool.acquire() as conn:
|
||||
acquire_time = time.time() - start
|
||||
if acquire_time > 0.05:
|
||||
logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s")
|
||||
yield conn
|
||||
return
|
||||
except Exception as e:
|
||||
if not _is_retryable(e):
|
||||
raise
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
|
||||
logger.warning(
|
||||
f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay:.1f}s..."
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
else:
|
||||
logger.error(f"Database acquire failed after {max_retries + 1} attempts: {e}")
|
||||
raise last_exception
|
||||
async with AsyncExitStack() as stack:
|
||||
conn: Any = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
conn = await stack.enter_async_context(backend_or_pool.acquire())
|
||||
break
|
||||
except Exception as e:
|
||||
if not _is_retryable(e):
|
||||
raise
|
||||
if attempt < max_retries:
|
||||
delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
|
||||
logger.warning(
|
||||
f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay:.1f}s..."
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
else:
|
||||
logger.error(f"Database acquire failed after {max_retries + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
acquire_time = time.time() - start
|
||||
if acquire_time > 0.05:
|
||||
logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s")
|
||||
|
||||
yield conn
|
||||
else:
|
||||
# Legacy path: raw asyncpg.Pool
|
||||
pool = backend_or_pool
|
||||
|
||||
@@ -9,13 +9,17 @@ The database schema is automatically adjusted to match the model's dimension.
|
||||
Configuration via environment variables - see hindsight_api.config for all env var names.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import logging
|
||||
import os
|
||||
import struct
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Literal, cast
|
||||
from urllib.parse import parse_qs, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ..config import (
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL,
|
||||
@@ -27,24 +31,60 @@ from ..config import (
|
||||
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
DEFAULT_EMBEDDINGS_PROVIDER,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL,
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
DEFAULT_ZEROENTROPY_BASE_URL,
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY,
|
||||
ENV_EMBEDDINGS_GEMINI_API_KEY,
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY,
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL,
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
ENV_EMBEDDINGS_ONNX_DIMENSIONS,
|
||||
ENV_EMBEDDINGS_ONNX_MODEL_ID,
|
||||
ENV_EMBEDDINGS_ONNX_MODEL_PATH,
|
||||
ENV_EMBEDDINGS_ONNX_TOKENIZER_NAME_OR_PATH,
|
||||
ENV_EMBEDDINGS_OPENAI_API_KEY,
|
||||
ENV_EMBEDDINGS_OPENAI_BASE_URL,
|
||||
ENV_EMBEDDINGS_OPENAI_MODEL,
|
||||
ENV_EMBEDDINGS_PROVIDER,
|
||||
ENV_EMBEDDINGS_TEI_URL,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_API_KEY,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
ENV_LLM_API_KEY,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
ZeroEntropyInputType = Literal["document", "query"]
|
||||
ZeroEntropyLatency = Literal["fast", "slow"]
|
||||
ZeroEntropyEncodingFormat = Literal["float", "base64"]
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedRequest(BaseModel):
|
||||
"""Typed request body for ZeroEntropy's non-OpenAI-compatible embed endpoint."""
|
||||
|
||||
model: str
|
||||
input: list[str]
|
||||
input_type: ZeroEntropyInputType
|
||||
dimensions: int
|
||||
encoding_format: ZeroEntropyEncodingFormat = "float"
|
||||
latency: ZeroEntropyLatency | None = None
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedResult(BaseModel):
|
||||
embedding: list[float] | str
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedResponse(BaseModel):
|
||||
results: list[_ZeroEntropyEmbedResult]
|
||||
|
||||
|
||||
class Embeddings(ABC):
|
||||
"""
|
||||
Abstract base class for embedding generation.
|
||||
@@ -88,6 +128,14 @@ class Embeddings(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings for query text. Providers without asymmetric embeddings use encode()."""
|
||||
return self.encode(texts)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings for stored document text. Providers without asymmetric embeddings use encode()."""
|
||||
return self.encode(texts)
|
||||
|
||||
|
||||
class LocalSTEmbeddings(Embeddings):
|
||||
"""
|
||||
@@ -208,6 +256,172 @@ class LocalSTEmbeddings(Embeddings):
|
||||
return [emb.tolist() for emb in embeddings]
|
||||
|
||||
|
||||
class OnnxEmbeddings(Embeddings):
|
||||
"""Local ONNX Runtime embeddings provider.
|
||||
|
||||
This provider runs transformer embedding models in-process with ONNX Runtime,
|
||||
avoiding a sidecar Ollama/TEI server or a remote embeddings API. It supports
|
||||
sentence-transformer style mean pooling and E5-style asymmetric prefixes.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_id: str,
|
||||
model_path: str | None = None,
|
||||
tokenizer_name_or_path: str | None = None,
|
||||
onnx_file: str = "onnx/model.onnx",
|
||||
dimensions: int | None = None,
|
||||
max_tokens: int = 512,
|
||||
pooling: str = "mean",
|
||||
normalize: bool = True,
|
||||
query_prefix: str = "query: ",
|
||||
passage_prefix: str = "passage: ",
|
||||
output_name: str | None = None,
|
||||
):
|
||||
self.model_id = model_id
|
||||
self.model_path = model_path
|
||||
if model_path and tokenizer_name_or_path is None:
|
||||
logger.warning(
|
||||
"Embeddings: ONNX model_path is set without tokenizer_name_or_path; "
|
||||
"falling back to tokenizer from model_id %s. Set "
|
||||
"HINDSIGHT_API_EMBEDDINGS_ONNX_TOKENIZER_NAME_OR_PATH when using local ONNX artifacts.",
|
||||
model_id,
|
||||
)
|
||||
self.tokenizer_name_or_path = tokenizer_name_or_path or model_id
|
||||
self.onnx_file = onnx_file
|
||||
self.configured_dimensions = dimensions
|
||||
self.max_tokens = max_tokens
|
||||
self.pooling = pooling.lower()
|
||||
if self.pooling not in {"mean", "cls"}:
|
||||
raise ValueError("ONNX embeddings pooling must be 'mean' or 'cls'")
|
||||
self.normalize = normalize
|
||||
self.query_prefix = query_prefix
|
||||
self.passage_prefix = passage_prefix
|
||||
self.output_name = output_name
|
||||
self._session = None
|
||||
self._tokenizer = None
|
||||
self._dimension: int | None = dimensions
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "onnx"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
if self._dimension is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
return self._dimension
|
||||
|
||||
async def initialize(self) -> None:
|
||||
if self._session is not None and self._tokenizer is not None:
|
||||
return
|
||||
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
from transformers import AutoTokenizer
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"onnxruntime and transformers are required for OnnxEmbeddings. "
|
||||
"Install with: pip install 'hindsight-api-slim[local-onnx]'"
|
||||
) from exc
|
||||
|
||||
model_path = self.model_path
|
||||
if not model_path:
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"huggingface-hub is required to download ONNX embedding models. "
|
||||
"Set HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_PATH or install local-onnx."
|
||||
) from exc
|
||||
# Some large ONNX exports, for example BAAI/bge-m3, store weights in
|
||||
# an external sidecar file next to model.onnx. Download both the
|
||||
# requested graph and its conventional *_data sidecar when present.
|
||||
snapshot_dir = snapshot_download(
|
||||
repo_id=self.model_id,
|
||||
allow_patterns=[self.onnx_file, f"{self.onnx_file}_data"],
|
||||
)
|
||||
model_path = os.path.join(snapshot_dir, self.onnx_file)
|
||||
|
||||
logger.info(
|
||||
"Embeddings: initializing ONNX provider with model %s (%s)",
|
||||
self.model_id,
|
||||
model_path,
|
||||
)
|
||||
logger.info(
|
||||
"Embeddings: ONNX query_prefix=%r passage_prefix=%r pooling=%s normalize=%s",
|
||||
self.query_prefix,
|
||||
self.passage_prefix,
|
||||
self.pooling,
|
||||
self.normalize,
|
||||
)
|
||||
self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name_or_path)
|
||||
self._session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
|
||||
|
||||
detected = len(self.encode(["test"])[0])
|
||||
if self.configured_dimensions is not None and detected != self.configured_dimensions:
|
||||
raise ValueError(
|
||||
f"Configured ONNX embedding dimension {self.configured_dimensions} does not match model output {detected}"
|
||||
)
|
||||
self._dimension = detected
|
||||
logger.info("Embeddings: ONNX provider initialized (dim: %s)", self._dimension)
|
||||
|
||||
def _encode_prefixed(self, texts: list[str], prefix: str) -> list[list[float]]:
|
||||
if prefix:
|
||||
return self.encode([f"{prefix}{text}" for text in texts])
|
||||
return self.encode(texts)
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
return self._encode_prefixed(texts, self.query_prefix)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
return self._encode_prefixed(texts, self.passage_prefix)
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
if self._session is None or self._tokenizer is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
import numpy as np
|
||||
|
||||
encoded = self._tokenizer(
|
||||
texts,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=self.max_tokens,
|
||||
return_tensors="np",
|
||||
)
|
||||
input_names = {inp.name for inp in self._session.get_inputs()}
|
||||
ort_inputs = {name: value for name, value in encoded.items() if name in input_names}
|
||||
if "token_type_ids" in input_names and "token_type_ids" not in ort_inputs:
|
||||
ort_inputs["token_type_ids"] = np.zeros_like(encoded["input_ids"])
|
||||
|
||||
outputs = self._session.run([self.output_name] if self.output_name else None, ort_inputs)
|
||||
token_embeddings = outputs[0]
|
||||
|
||||
# Some exported models expose a pooled 2-D embedding as their first output.
|
||||
if getattr(token_embeddings, "ndim", 0) == 2:
|
||||
embeddings = token_embeddings
|
||||
elif self.pooling == "cls":
|
||||
embeddings = token_embeddings[:, 0]
|
||||
else:
|
||||
attention_mask = encoded.get("attention_mask")
|
||||
if attention_mask is None:
|
||||
attention_mask = np.ones(token_embeddings.shape[:2], dtype=np.float32)
|
||||
mask = attention_mask[..., None].astype(np.float32)
|
||||
summed = (token_embeddings * mask).sum(axis=1)
|
||||
counts = np.clip(mask.sum(axis=1), a_min=1e-9, a_max=None)
|
||||
embeddings = summed / counts
|
||||
|
||||
if self.normalize:
|
||||
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
|
||||
norms[norms == 0] = 1
|
||||
embeddings = embeddings / norms
|
||||
|
||||
return embeddings.astype(float).tolist()
|
||||
|
||||
|
||||
class RemoteTEIEmbeddings(Embeddings):
|
||||
"""
|
||||
Remote embeddings implementation using HuggingFace Text Embeddings Inference (TEI) HTTP API.
|
||||
@@ -385,6 +599,7 @@ class OpenAIEmbeddings(Embeddings):
|
||||
model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
base_url: str | None = None,
|
||||
batch_size: int = 100,
|
||||
dimensions: int | None = None,
|
||||
max_retries: int = 3,
|
||||
):
|
||||
"""
|
||||
@@ -395,12 +610,14 @@ class OpenAIEmbeddings(Embeddings):
|
||||
model: OpenAI embedding model name (default: text-embedding-3-small)
|
||||
base_url: Custom base URL for OpenAI-compatible API (e.g., Azure OpenAI endpoint)
|
||||
batch_size: Maximum batch size for embedding requests (default: 100)
|
||||
dimensions: Optional requested output dimensions for OpenAI text-embedding-3 models
|
||||
max_retries: Maximum number of retries for failed requests (default: 3)
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url
|
||||
self.batch_size = batch_size
|
||||
self.dimensions = dimensions
|
||||
self.max_retries = max_retries
|
||||
self._client = None
|
||||
self._dimension: int | None = None
|
||||
@@ -445,7 +662,9 @@ class OpenAIEmbeddings(Embeddings):
|
||||
self._client = OpenAI(**client_kwargs)
|
||||
|
||||
# Try to get dimension from known models, otherwise do a test embedding
|
||||
if self.model in self.MODEL_DIMENSIONS:
|
||||
if self.dimensions is not None:
|
||||
self._dimension = self.dimensions
|
||||
elif self.model in self.MODEL_DIMENSIONS:
|
||||
self._dimension = self.MODEL_DIMENSIONS[self.model]
|
||||
else:
|
||||
# Do a test embedding to detect dimension
|
||||
@@ -480,10 +699,14 @@ class OpenAIEmbeddings(Embeddings):
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
|
||||
response = self._client.embeddings.create(
|
||||
model=self.model,
|
||||
input=batch,
|
||||
)
|
||||
request = {
|
||||
"model": self.model,
|
||||
"input": batch,
|
||||
}
|
||||
if self.dimensions is not None:
|
||||
request["dimensions"] = self.dimensions
|
||||
|
||||
response = self._client.embeddings.create(**request)
|
||||
|
||||
# Sort by index to ensure correct order
|
||||
batch_embeddings = sorted(response.data, key=lambda x: x.index)
|
||||
@@ -492,6 +715,73 @@ class OpenAIEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
class CodexOAuthEmbeddings(OpenAIEmbeddings):
|
||||
"""
|
||||
OpenAI embeddings using the Codex/ChatGPT OAuth token from ``~/.codex/auth.json``.
|
||||
|
||||
Codex OAuth is an LLM-provider auth path in Hindsight, but the same bearer token
|
||||
can also authenticate against the standard OpenAI embeddings endpoint. This keeps
|
||||
embeddings on the user's existing Codex subscription/OAuth path without requiring
|
||||
a separate OpenAI/OpenRouter/Gemini/Cohere API key.
|
||||
|
||||
Token refresh is handled automatically: the manager proactively refreshes the
|
||||
access_token before it expires and reactively refreshes on 401 responses from
|
||||
the embeddings API.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
batch_size: int = 100,
|
||||
dimensions: int | None = None,
|
||||
max_retries: int = 3,
|
||||
):
|
||||
from .providers.codex_auth import CodexAuthManager
|
||||
|
||||
self._auth_manager = CodexAuthManager.from_file()
|
||||
super().__init__(
|
||||
api_key=self._auth_manager.access_token,
|
||||
model=model,
|
||||
base_url="https://api.openai.com/v1",
|
||||
batch_size=batch_size,
|
||||
dimensions=dimensions,
|
||||
max_retries=max_retries,
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "openai-codex"
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings, refreshing the OAuth token if needed.
|
||||
|
||||
Proactively refreshes before the call when the token is near expiry,
|
||||
and reactively refreshes once on a 401 from the OpenAI embeddings API.
|
||||
"""
|
||||
from openai import AuthenticationError
|
||||
|
||||
# Proactive refresh — cheap when fresh (JWT exp decode + compare).
|
||||
self._auth_manager.ensure_fresh_token()
|
||||
if self._auth_manager.access_token != self.api_key:
|
||||
self.api_key = self._auth_manager.access_token
|
||||
if self._client is not None:
|
||||
self._client.api_key = self._auth_manager.access_token
|
||||
|
||||
try:
|
||||
return super().encode(texts)
|
||||
except AuthenticationError:
|
||||
# Reactive refresh — token was valid by the JWT clock but the
|
||||
# server rejected it (rotated server-side, race, etc.).
|
||||
self._auth_manager.refresh_tokens(
|
||||
reason="reactive (401 from embeddings API)",
|
||||
force=True,
|
||||
)
|
||||
self.api_key = self._auth_manager.access_token
|
||||
if self._client is not None:
|
||||
self._client.api_key = self._auth_manager.access_token
|
||||
return super().encode(texts)
|
||||
|
||||
|
||||
class CohereEmbeddings(Embeddings):
|
||||
"""
|
||||
Cohere embeddings implementation using the Cohere API.
|
||||
@@ -633,6 +923,149 @@ class CohereEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
class ZeroEntropyEmbeddings(Embeddings):
|
||||
"""
|
||||
ZeroEntropy embeddings implementation using the zembed API.
|
||||
|
||||
ZeroEntropy's embeddings endpoint is not OpenAI-compatible: it lives at
|
||||
/v1/models/embed and requires provider-specific parameters such as
|
||||
input_type. Hindsight stores document-side vectors and uses query-side
|
||||
vectors during recall, so this provider exposes explicit encode_documents()
|
||||
and encode_query() helpers while keeping encode() as document-side default.
|
||||
"""
|
||||
|
||||
VALID_DIMENSIONS = frozenset({2560, 1280, 640, 320, 160, 80, 40})
|
||||
VALID_ENCODING_FORMATS = frozenset({"float", "base64"})
|
||||
VALID_LATENCIES = frozenset({"fast", "slow"})
|
||||
DEFAULT_BASE_URL = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
EMBED_PATH = "/v1/models/embed"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL,
|
||||
base_url: str | None = None,
|
||||
dimensions: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
batch_size: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
encoding_format: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
latency: str | None = DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
if dimensions not in self.VALID_DIMENSIONS:
|
||||
valid = ", ".join(str(dim) for dim in sorted(self.VALID_DIMENSIONS, reverse=True))
|
||||
raise ValueError(f"{ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS} must be one of {valid}, got {dimensions}")
|
||||
if batch_size < 1:
|
||||
raise ValueError("ZeroEntropy embeddings batch_size must be >= 1")
|
||||
if encoding_format not in self.VALID_ENCODING_FORMATS:
|
||||
valid_formats = ", ".join(sorted(self.VALID_ENCODING_FORMATS))
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT} must be one of {valid_formats}, got {encoding_format!r}"
|
||||
)
|
||||
if latency is not None and latency not in self.VALID_LATENCIES:
|
||||
valid_latencies = ", ".join(sorted(self.VALID_LATENCIES))
|
||||
raise ValueError(f"ZeroEntropy embeddings latency must be one of {valid_latencies}, got {latency!r}")
|
||||
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url.rstrip("/") if base_url else self.DEFAULT_BASE_URL
|
||||
self.embed_url = f"{self.base_url}{self.EMBED_PATH}"
|
||||
self.dimensions = dimensions
|
||||
self.batch_size = batch_size
|
||||
self.encoding_format = cast(ZeroEntropyEncodingFormat, encoding_format)
|
||||
self.latency = cast(ZeroEntropyLatency | None, latency)
|
||||
self.timeout = timeout
|
||||
self._client: httpx.Client | None = None
|
||||
self._dimension: int | None = None
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "zeroentropy"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
if self._dimension is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
return self._dimension
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the ZeroEntropy HTTP client."""
|
||||
if self._client is not None:
|
||||
return
|
||||
|
||||
logger.info(
|
||||
f"Embeddings: initializing ZeroEntropy provider with model {self.model} "
|
||||
f"(dim: {self.dimensions}, batch_size={self.batch_size})"
|
||||
)
|
||||
self._client = httpx.Client(
|
||||
timeout=self.timeout,
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
# zembed-1 dimensions are explicit Matryoshka truncation steps. Avoid a
|
||||
# startup probe so boot does not burn quota or require a throwaway input.
|
||||
self._dimension = self.dimensions
|
||||
logger.info(f"Embeddings: ZeroEntropy provider initialized (model: {self.model}, dim: {self._dimension})")
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate document-side embeddings for backwards-compatible callers."""
|
||||
return self.encode_documents(texts)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate document-side embeddings for retained content."""
|
||||
return self._encode_with_input_type(texts, "document")
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate query-side embeddings for recall/search queries."""
|
||||
return self._encode_with_input_type(texts, "query")
|
||||
|
||||
def _encode_with_input_type(self, texts: list[str], input_type: ZeroEntropyInputType) -> list[list[float]]:
|
||||
if self._client is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
all_embeddings: list[list[float]] = []
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
request = _ZeroEntropyEmbedRequest(
|
||||
model=self.model,
|
||||
input=batch,
|
||||
input_type=input_type,
|
||||
dimensions=self.dimensions,
|
||||
encoding_format=self.encoding_format,
|
||||
latency=self.latency,
|
||||
)
|
||||
|
||||
try:
|
||||
response = self._client.post(self.embed_url, json=request.model_dump(exclude_none=True))
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPError as e:
|
||||
raise RuntimeError(f"ZeroEntropy embedding request failed: {e}") from e
|
||||
|
||||
parsed = _ZeroEntropyEmbedResponse.model_validate(response.json())
|
||||
if len(parsed.results) != len(batch):
|
||||
raise RuntimeError(
|
||||
f"ZeroEntropy returned {len(parsed.results)} embeddings for {len(batch)} input texts; "
|
||||
"expected exact 1:1 alignment"
|
||||
)
|
||||
all_embeddings.extend(self._parse_embedding(result.embedding) for result in parsed.results)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
@staticmethod
|
||||
def _parse_embedding(embedding: list[float] | str) -> list[float]:
|
||||
if not isinstance(embedding, str):
|
||||
return embedding
|
||||
|
||||
raw = base64.b64decode(embedding)
|
||||
if len(raw) % 4 != 0:
|
||||
raise RuntimeError("ZeroEntropy returned invalid base64 embedding length")
|
||||
return list(struct.unpack(f"<{len(raw) // 4}f", raw))
|
||||
|
||||
|
||||
class LiteLLMEmbeddings(Embeddings):
|
||||
"""
|
||||
LiteLLM embeddings implementation using LiteLLM proxy's /embeddings endpoint.
|
||||
@@ -766,7 +1199,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
api_key: str | None = None,
|
||||
model: str = DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
api_base: str | None = None,
|
||||
output_dimensions: int | None = None,
|
||||
@@ -778,7 +1211,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
Initialize LiteLLM SDK embeddings client.
|
||||
|
||||
Args:
|
||||
api_key: API key for the embedding provider
|
||||
api_key: API key for the embedding provider (optional — omit for
|
||||
providers that use ambient credentials, e.g. AWS Bedrock with IAM)
|
||||
model: Model name with provider prefix (e.g., "cohere/embed-english-v3.0")
|
||||
api_base: Custom base URL for API (optional)
|
||||
output_dimensions: Optional output embedding dimensions (provider-dependent)
|
||||
@@ -828,8 +1262,9 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
embed_kwargs = {
|
||||
"model": self.model,
|
||||
"input": ["test"],
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_key:
|
||||
embed_kwargs["api_key"] = self.api_key
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
@@ -880,8 +1315,9 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
embed_kwargs = {
|
||||
"model": self.model,
|
||||
"input": batch,
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_key:
|
||||
embed_kwargs["api_key"] = self.api_key
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
@@ -1125,6 +1561,20 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
force_cpu=config.embeddings_local_force_cpu,
|
||||
trust_remote_code=config.embeddings_local_trust_remote_code,
|
||||
)
|
||||
elif provider == "onnx":
|
||||
return OnnxEmbeddings(
|
||||
model_id=config.embeddings_onnx_model_id,
|
||||
model_path=config.embeddings_onnx_model_path,
|
||||
tokenizer_name_or_path=config.embeddings_onnx_tokenizer_name_or_path,
|
||||
onnx_file=config.embeddings_onnx_file,
|
||||
dimensions=config.embeddings_onnx_dimensions,
|
||||
max_tokens=config.embeddings_onnx_max_tokens,
|
||||
pooling=config.embeddings_onnx_pooling,
|
||||
normalize=config.embeddings_onnx_normalize,
|
||||
query_prefix=config.embeddings_onnx_query_prefix,
|
||||
passage_prefix=config.embeddings_onnx_passage_prefix,
|
||||
output_name=config.embeddings_onnx_output_name,
|
||||
)
|
||||
elif provider == "openai":
|
||||
# Use dedicated embeddings API key, or fall back to LLM API key
|
||||
api_key = os.environ.get(ENV_EMBEDDINGS_OPENAI_API_KEY) or os.environ.get(ENV_LLM_API_KEY)
|
||||
@@ -1140,6 +1590,14 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=model,
|
||||
base_url=base_url,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "openai-codex":
|
||||
model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL)
|
||||
return CodexOAuthEmbeddings(
|
||||
model=model,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.embeddings_openrouter_api_key
|
||||
@@ -1153,6 +1611,23 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=config.embeddings_openrouter_model,
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "zeroentropy":
|
||||
api_key = config.embeddings_zeroentropy_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_ZEROENTROPY_API_KEY} or ZEROENTROPY_API_KEY is required "
|
||||
f"when {ENV_EMBEDDINGS_PROVIDER} is 'zeroentropy'"
|
||||
)
|
||||
return ZeroEntropyEmbeddings(
|
||||
api_key=api_key,
|
||||
model=config.embeddings_zeroentropy_model,
|
||||
base_url=config.embeddings_zeroentropy_base_url,
|
||||
dimensions=config.embeddings_zeroentropy_dimensions,
|
||||
batch_size=config.embeddings_zeroentropy_batch_size,
|
||||
encoding_format=config.embeddings_zeroentropy_encoding_format,
|
||||
latency=config.embeddings_zeroentropy_latency,
|
||||
)
|
||||
elif provider == "cohere":
|
||||
api_key = config.embeddings_cohere_api_key
|
||||
@@ -1171,13 +1646,8 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=config.embeddings_litellm_model,
|
||||
)
|
||||
elif provider == "litellm-sdk":
|
||||
api_key = config.embeddings_litellm_sdk_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_LITELLM_SDK_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'litellm-sdk'"
|
||||
)
|
||||
return LiteLLMSDKEmbeddings(
|
||||
api_key=api_key,
|
||||
api_key=config.embeddings_litellm_sdk_api_key or None,
|
||||
model=config.embeddings_litellm_sdk_model,
|
||||
api_base=config.embeddings_litellm_sdk_api_base,
|
||||
output_dimensions=config.embeddings_litellm_sdk_output_dimensions,
|
||||
@@ -1206,5 +1676,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown embeddings provider: {provider}. "
|
||||
f"Supported: 'local', 'tei', 'openai', 'cohere', 'google', 'litellm', 'litellm-sdk'"
|
||||
f"Supported: 'local', 'onnx', 'tei', 'openai', 'openai-codex', 'openrouter', 'cohere', 'google', "
|
||||
f"'zeroentropy', 'litellm', 'litellm-sdk'"
|
||||
)
|
||||
|
||||
@@ -16,7 +16,15 @@ from typing import Any, Final
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
from .memory_engine import fq_table
|
||||
from .retain.entity_labels import build_labels_lookup as _build_labels_lookup_from_config
|
||||
from .retain.entity_labels import (
|
||||
build_labels_lookup as _build_labels_lookup_from_config,
|
||||
)
|
||||
from .retain.entity_labels import (
|
||||
is_label_entity as _is_label_entity,
|
||||
)
|
||||
from .retain.entity_labels import (
|
||||
parse_entity_labels as _parse_entity_labels,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -89,7 +97,12 @@ class EntityResolver:
|
||||
Resolves entities to canonical IDs with disambiguation.
|
||||
"""
|
||||
|
||||
def __init__(self, pool: Any, entity_lookup: str = "full"):
|
||||
def __init__(
|
||||
self,
|
||||
pool: Any,
|
||||
entity_lookup: str = "full",
|
||||
entity_resolution_batch_size: int = 100,
|
||||
):
|
||||
"""
|
||||
Initialize entity resolver.
|
||||
|
||||
@@ -98,9 +111,14 @@ class EntityResolver:
|
||||
entity_lookup: Lookup strategy — "full" loads all bank entities then
|
||||
matches in Python; "trigram" uses pg_trgm GIN index to fetch only
|
||||
similar candidates per entity name (much faster for large banks).
|
||||
entity_resolution_batch_size: Number of unique entity names to include
|
||||
in each pg_trgm candidate lookup query.
|
||||
"""
|
||||
self.pool = pool
|
||||
self.entity_lookup = entity_lookup
|
||||
if entity_resolution_batch_size < 1:
|
||||
raise ValueError("entity_resolution_batch_size must be >= 1")
|
||||
self.entity_resolution_batch_size = entity_resolution_batch_size
|
||||
self._pg_trgm_checked = False
|
||||
# Backend-specific operations — accessed via pool.ops (Django pattern).
|
||||
self._ops = pool.ops if pool is not None else None
|
||||
@@ -199,6 +217,11 @@ class EntityResolver:
|
||||
"""Build a set of valid 'key:value' entity label strings for fast lookup."""
|
||||
return _build_labels_lookup_from_config(entity_labels)
|
||||
|
||||
@staticmethod
|
||||
def _chunked(values: list[str], size: int) -> list[list[str]]:
|
||||
"""Split values into fixed-size batches."""
|
||||
return [values[i : i + size] for i in range(0, len(values), size)]
|
||||
|
||||
async def resolve_entities_batch(
|
||||
self,
|
||||
bank_id: str,
|
||||
@@ -228,14 +251,15 @@ class EntityResolver:
|
||||
return []
|
||||
|
||||
taxonomy_lookup = self._build_labels_lookup(entity_labels)
|
||||
labels_cfg = _parse_entity_labels(entity_labels)
|
||||
if conn is None:
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
return await self._resolve_entities_batch_impl(
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
else:
|
||||
return await self._resolve_entities_batch_impl(
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_impl(
|
||||
@@ -246,13 +270,16 @@ class EntityResolver:
|
||||
context: str,
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
if self.entity_lookup == "trigram":
|
||||
# Route to backend-specific fuzzy strategy.
|
||||
# Non-PG backends (Oracle) use UTL_MATCH instead of pg_trgm.
|
||||
backend_strategy = self._ops.get_entity_resolution_strategy()
|
||||
if backend_strategy == "oracle_fuzzy":
|
||||
return await self._resolve_entities_batch_oracle_fuzzy(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_oracle_fuzzy(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
# Auto-detect pg_trgm availability on first call and fall back to
|
||||
# "full" strategy if the extension is not installed. See #626.
|
||||
if not self._pg_trgm_checked:
|
||||
@@ -266,12 +293,24 @@ class EntityResolver:
|
||||
"https://github.com/vectorize-io/hindsight/issues/626"
|
||||
)
|
||||
self.entity_lookup = "full"
|
||||
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_trigram(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_full(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_trigram(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_full(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_full(
|
||||
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""Original strategy: load all bank entities then match in Python."""
|
||||
# Query ALL candidates for this bank
|
||||
@@ -338,11 +377,24 @@ class EntityResolver:
|
||||
all_candidates[entity_text] = matching
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_trigram(
|
||||
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Trigram strategy: fetch only similar candidates per entity name using pg_trgm.
|
||||
@@ -353,7 +405,7 @@ class EntityResolver:
|
||||
"""
|
||||
entity_texts = list(set(e["text"] for e in entities_data))
|
||||
|
||||
# Fetch candidates for all unique entity texts in a single batched query.
|
||||
# Fetch candidates for unique entity texts in bounded batches.
|
||||
# Uses the GIN trigram index on LOWER(canonical_name) for case-insensitive
|
||||
# similarity lookup. Previous version also had LIKE '%...' substring fallbacks,
|
||||
# but those forced full sequential scans of the entities table and caused
|
||||
@@ -361,21 +413,32 @@ class EntityResolver:
|
||||
# to 0.15 (from default 0.3) catches most substring relationships while
|
||||
# staying fully index-based.
|
||||
await conn.execute("SET pg_trgm.similarity_threshold = 0.15")
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT ON (e.id)
|
||||
e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM unnest($2::text[]) AS q(query_text)
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
e.bank_id = $1
|
||||
AND LOWER(e.canonical_name) % LOWER(q.query_text)
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_texts,
|
||||
)
|
||||
await conn.execute("RESET pg_trgm.similarity_threshold")
|
||||
try:
|
||||
rows = []
|
||||
for entity_text_batch in self._chunked(entity_texts, self.entity_resolution_batch_size):
|
||||
rows.extend(
|
||||
await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT ON (e.id)
|
||||
e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM unnest($2::text[]) AS q(query_text)
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
e.bank_id = $1
|
||||
AND LOWER(e.canonical_name) % LOWER(q.query_text)
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_text_batch,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
# asyncpg returns connections to the pool with session state intact,
|
||||
# so the lowered threshold would leak to future borrowers without RESET.
|
||||
try:
|
||||
await conn.execute("RESET pg_trgm.similarity_threshold")
|
||||
except Exception:
|
||||
logger.warning("Failed to reset pg_trgm similarity threshold after candidate lookup", exc_info=True)
|
||||
|
||||
# Group candidates by query_text
|
||||
all_candidates: dict[str, list] = {t: [] for t in entity_texts}
|
||||
@@ -418,11 +481,24 @@ class EntityResolver:
|
||||
cooccurrence_map[eid2].add(id_to_name[eid1])
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_oracle_fuzzy(
|
||||
self, conn: Any, bank_id: str, entities_data: list[dict], unit_event_date: datetime | None
|
||||
self,
|
||||
conn: Any,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date: datetime | None,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Oracle strategy: fetch similar candidates using UTL_MATCH.JARO_WINKLER_SIMILARITY.
|
||||
@@ -436,23 +512,28 @@ class EntityResolver:
|
||||
entities_table = fq_table("entities")
|
||||
|
||||
try:
|
||||
# Batch all entity texts into a single query using JSON_TABLE to
|
||||
# Batch entity texts into bounded sub-queries using JSON_TABLE to
|
||||
# expand the list into rows. UTL_MATCH.JARO_WINKLER_SIMILARITY
|
||||
# returns 0-100; threshold 70 ≈ pg_trgm similarity 0.15.
|
||||
entity_texts_json = json.dumps(entity_texts)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM JSON_TABLE($2, '$[*]' COLUMNS (query_text VARCHAR2(4000) PATH '$')) q
|
||||
JOIN {entities_table} e ON (
|
||||
e.bank_id = $1
|
||||
AND UTL_MATCH.JARO_WINKLER_SIMILARITY(LOWER(e.canonical_name), LOWER(q.query_text)) > 70
|
||||
# Bounded batches mirror the PG trigram path so very wide retain
|
||||
# batches don't time out a single JOIN on large banks.
|
||||
rows = []
|
||||
for entity_text_batch in self._chunked(entity_texts, self.entity_resolution_batch_size):
|
||||
rows.extend(
|
||||
await conn.fetch(
|
||||
f"""
|
||||
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM JSON_TABLE($2, '$[*]' COLUMNS (query_text VARCHAR2(4000) PATH '$')) q
|
||||
JOIN {entities_table} e ON (
|
||||
e.bank_id = $1
|
||||
AND UTL_MATCH.JARO_WINKLER_SIMILARITY(LOWER(e.canonical_name), LOWER(q.query_text)) > 70
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
json.dumps(entity_text_batch),
|
||||
)
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_texts_json,
|
||||
)
|
||||
except Exception as e:
|
||||
# UTL_MATCH may not be available (ORA-06550, ORA-00904, etc.)
|
||||
# Catch broadly because Oracle error types vary depending on driver.
|
||||
@@ -506,7 +587,14 @@ class EntityResolver:
|
||||
cooccurrence_map[eid2].add(id_to_name[eid1])
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
)
|
||||
|
||||
async def _resolve_from_candidates(
|
||||
@@ -517,6 +605,8 @@ class EntityResolver:
|
||||
unit_event_date,
|
||||
all_candidates: dict[str, list],
|
||||
cooccurrence_map: dict[str, set[str]],
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""Shared scoring + upsert logic used by both lookup strategies."""
|
||||
|
||||
@@ -533,11 +623,34 @@ class EntityResolver:
|
||||
|
||||
candidates = all_candidates.get(entity_text, [])
|
||||
|
||||
# Label entities (from entity_labels config) use exact matching only.
|
||||
# Their canonical names are user-defined (e.g., "use:use-001"),
|
||||
# so fuzzy resolution must NOT merge distinct label values that
|
||||
# happen to be textually similar (GH-1558).
|
||||
is_label = bool(
|
||||
labels_cfg and taxonomy_lookup and _is_label_entity(entity_text, labels_cfg, taxonomy_lookup)
|
||||
)
|
||||
|
||||
if not candidates:
|
||||
# Will create new entity
|
||||
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
|
||||
continue
|
||||
|
||||
if is_label:
|
||||
# Exact case-insensitive match only for label entities
|
||||
exact_match = None
|
||||
entity_text_lower = entity_text.lower()
|
||||
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
|
||||
if canonical_name.lower() == entity_text_lower:
|
||||
exact_match = candidate_id
|
||||
break
|
||||
if exact_match:
|
||||
entity_ids[idx] = exact_match
|
||||
entities_to_update.append(_EntityStat(entity_id=exact_match, event_date=entity_event_date))
|
||||
else:
|
||||
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
|
||||
continue
|
||||
|
||||
# Score candidates
|
||||
best_candidate = None
|
||||
best_score = 0.0
|
||||
|
||||
@@ -0,0 +1,357 @@
|
||||
"""Async graph maintenance after document/unit deletes.
|
||||
|
||||
Three reconciliation passes run together on every worker invocation:
|
||||
|
||||
1. **Relink top-up.** Drain ``graph_maintenance_queue`` (units whose
|
||||
outgoing temporal/semantic links lost a neighbour to a delete). For
|
||||
each, count current outgoing links per type; if below cap, run the
|
||||
same probes retain uses (:func:`fetch_temporal_neighbors`,
|
||||
:func:`compute_semantic_links_ann`) and insert the missing links.
|
||||
``bulk_insert_links`` has ``ON CONFLICT DO NOTHING`` on the uniqueness
|
||||
key, so we can re-probe freely and the DB de-dupes.
|
||||
|
||||
2. **Orphan entity prune.** Delete ``entities`` rows in the bank that no
|
||||
longer have any ``unit_entities`` references. FK ON DELETE CASCADE on
|
||||
``entity_cooccurrences`` then removes any cooccurrence row pointing
|
||||
at the pruned entities.
|
||||
|
||||
3. **Stale cooccurrence prune.** Defensive sweep for cooccurrence rows
|
||||
where both endpoints still exist but no current memory_unit references
|
||||
both of them — the cooccurrence was real at the time it was recorded,
|
||||
but every unit that witnessed it has since been deleted.
|
||||
|
||||
All three passes run on every invocation. The queue is the only source
|
||||
of work for pass 1; passes 2 and 3 are bank-wide sweeps backed by indexes
|
||||
on ``entities(bank_id)`` and ``unit_entities(entity_id)``, so they're
|
||||
cheap when there's nothing to do.
|
||||
|
||||
The worker dedupes on bank: a second job for the same bank is dropped
|
||||
while one is pending. Once processing starts, a new job becomes the
|
||||
*next* pending slot — so work enqueued during processing gets picked up
|
||||
by the follow-up run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
import uuid as uuid_module
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..models import RequestContext
|
||||
from .db.base import DatabaseConnection
|
||||
from .retain.link_utils import (
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
_bulk_insert_links,
|
||||
_normalize_datetime,
|
||||
compute_semantic_links_ann,
|
||||
)
|
||||
from .schema import fq_table
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .memory_engine import MemoryEngine
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Mirrors the ``top_k`` default in ``compute_semantic_links_ann`` at retain
|
||||
# time. If you change one, change the other — otherwise victims would either
|
||||
# never reach the cap (probe returns less than the cap) or stay perpetually
|
||||
# under it (cap is higher than retain creates).
|
||||
MAX_SEMANTIC_LINKS_PER_UNIT = 50
|
||||
|
||||
# Worker fetches this many rows per relink-loop iteration. Bounds
|
||||
# per-iteration probe/insert latency so a 10k-row backlog doesn't hold a
|
||||
# worker slot for minutes. Chosen so the typical iteration runs in well
|
||||
# under 1s.
|
||||
_DRAIN_BATCH_SIZE = 50
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobResult:
|
||||
"""Counters surfaced to the worker dispatcher and operation result."""
|
||||
|
||||
relink_units_processed: int = 0
|
||||
relink_links_added: int = 0
|
||||
orphan_entities_pruned: int = 0
|
||||
stale_cooccurrences_pruned: int = 0
|
||||
|
||||
def as_dict(self) -> dict[str, int]:
|
||||
return {
|
||||
"relink_units_processed": self.relink_units_processed,
|
||||
"relink_links_added": self.relink_links_added,
|
||||
"orphan_entities_pruned": self.orphan_entities_pruned,
|
||||
"stale_cooccurrences_pruned": self.stale_cooccurrences_pruned,
|
||||
}
|
||||
|
||||
|
||||
async def enqueue_relink_victims(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
deleted_unit_ids: list[str],
|
||||
ops: Any,
|
||||
) -> int:
|
||||
"""Enqueue surviving units whose outgoing temporal/semantic links pointed at
|
||||
``deleted_unit_ids`` for later link top-up.
|
||||
|
||||
Must run inside the same transaction that deletes the units, *before* the
|
||||
cascade fires — once the rows are gone, the join that finds the victims
|
||||
returns nothing.
|
||||
|
||||
Args:
|
||||
conn: Database connection inside the active delete transaction.
|
||||
bank_id: Bank owning the deleted units.
|
||||
deleted_unit_ids: Memory_unit IDs about to be (or being) deleted.
|
||||
ops: ``DataAccessOps`` instance, supplies the dialect-specific
|
||||
bulk-insert path.
|
||||
|
||||
Returns:
|
||||
Number of distinct victim units enqueued (after dedup against rows
|
||||
already in the queue).
|
||||
"""
|
||||
if not deleted_unit_ids:
|
||||
return 0
|
||||
|
||||
deleted_uuids = [uuid_module.UUID(uid) if isinstance(uid, str) else uid for uid in deleted_unit_ids]
|
||||
deleted_str_set = {str(uid) for uid in deleted_uuids}
|
||||
|
||||
# Find units (other than the ones being deleted) that have an outgoing
|
||||
# temporal/semantic link pointing at a doomed unit. Entity links are
|
||||
# intentionally excluded — they're scheduled for removal and would only
|
||||
# add noise to the recompute job.
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT from_unit_id
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE to_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
""",
|
||||
deleted_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
victim_ids = [row["from_unit_id"] for row in victim_rows if str(row["from_unit_id"]) not in deleted_str_set]
|
||||
|
||||
if not victim_ids:
|
||||
return 0
|
||||
|
||||
await ops.enqueue_graph_maintenance(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
victim_ids,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"[GRAPH_MAINT] Enqueued {len(victim_ids)} relink victims in "
|
||||
f"bank={bank_id} (deleted {len(deleted_unit_ids)} units)"
|
||||
)
|
||||
return len(victim_ids)
|
||||
|
||||
|
||||
async def run_graph_maintenance_job(
|
||||
memory_engine: "MemoryEngine",
|
||||
bank_id: str,
|
||||
request_context: RequestContext,
|
||||
operation_id: str | None = None,
|
||||
) -> dict[str, int]:
|
||||
"""Run all maintenance passes for ``bank_id`` until the relink queue is
|
||||
drained, then sweep entities and cooccurrences once.
|
||||
|
||||
Returns:
|
||||
Per-pass counters from :class:`JobResult`.
|
||||
"""
|
||||
del request_context # accepted for symmetry with other run_*_job helpers
|
||||
backend = await memory_engine._get_backend()
|
||||
ops = backend.ops
|
||||
|
||||
result = JobResult()
|
||||
job_start = time.time()
|
||||
|
||||
# --- Pass 1: relink ---
|
||||
# Per-iteration loop: claim → top up → commit. We rely on submit-time
|
||||
# dedup to keep at most one job per bank running, so no need for
|
||||
# SKIP LOCKED.
|
||||
iterations = 0
|
||||
while True:
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
unit_ids = await ops.claim_graph_maintenance_batch(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
_DRAIN_BATCH_SIZE,
|
||||
)
|
||||
if not unit_ids:
|
||||
break
|
||||
|
||||
result.relink_links_added += await _relink_batch(conn, bank_id, unit_ids, ops, backend)
|
||||
|
||||
result.relink_units_processed += len(unit_ids)
|
||||
iterations += 1
|
||||
|
||||
if iterations > 10000:
|
||||
# Defensive guard against runaway loops — at 50 units/iter that's
|
||||
# 500k targets, far beyond any realistic single-bank backlog.
|
||||
logger.error(
|
||||
f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink ({result.as_dict()})"
|
||||
)
|
||||
break
|
||||
|
||||
# --- Pass 2 & 3: entity / cooccurrence sweeps ---
|
||||
# Bank-wide single-statement deletes. Cheap when there's nothing to do.
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
result.orphan_entities_pruned = await ops.prune_orphan_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
fq_table("unit_entities"),
|
||||
bank_id,
|
||||
)
|
||||
# The orphan prune above cascades cooccurrences via FK. The
|
||||
# explicit cooccurrence pass below catches the *stale-count*
|
||||
# case: both entities still exist but no current unit witnesses
|
||||
# them together.
|
||||
result.stale_cooccurrences_pruned = await ops.prune_stale_cooccurrences(
|
||||
conn,
|
||||
fq_table("entity_cooccurrences"),
|
||||
fq_table("unit_entities"),
|
||||
fq_table("entities"),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
elapsed = time.time() - job_start
|
||||
logger.info(
|
||||
f"[GRAPH_MAINT] bank={bank_id} done: {result.as_dict()}, elapsed={elapsed:.2f}s, operation_id={operation_id}"
|
||||
)
|
||||
return result.as_dict()
|
||||
|
||||
|
||||
async def _relink_batch(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
victim_ids: list[str],
|
||||
ops: Any,
|
||||
backend: Any,
|
||||
) -> int:
|
||||
"""Top up temporal/semantic links for a batch of victim units. Returns rows inserted."""
|
||||
# Load each victim's metadata. Victims whose units were deleted between
|
||||
# enqueue and now silently drop out — exactly the no-op behaviour we want
|
||||
# for stale queue rows.
|
||||
victim_uuids = [uuid_module.UUID(vid) for vid in victim_ids]
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id::text AS id, event_date, fact_type, embedding::text AS embedding
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
victim_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not victim_rows:
|
||||
return 0
|
||||
|
||||
alive_uuids = [uuid_module.UUID(row["id"]) for row in victim_rows]
|
||||
|
||||
# Count current outgoing temporal/semantic links per victim so we only
|
||||
# probe for the ones genuinely below cap. Saves the bulk of the work when
|
||||
# most victims still have plenty of links.
|
||||
count_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_unit_id, link_type, COUNT(*) AS cnt
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE from_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
GROUP BY from_unit_id, link_type
|
||||
""",
|
||||
alive_uuids,
|
||||
bank_id,
|
||||
)
|
||||
counts: dict[tuple[str, str], int] = {}
|
||||
for row in count_rows:
|
||||
counts[(str(row["from_unit_id"]), row["link_type"])] = int(row["cnt"])
|
||||
|
||||
# --- Temporal top-up ---
|
||||
temporal_needs = [r for r in victim_rows if counts.get((r["id"], "temporal"), 0) < MAX_TEMPORAL_LINKS_PER_UNIT]
|
||||
new_links: list[tuple] = []
|
||||
|
||||
if temporal_needs:
|
||||
lateral_unit_ids = [uuid_module.UUID(r["id"]) for r in temporal_needs if r["event_date"] is not None]
|
||||
lateral_event_dates = [
|
||||
_normalize_datetime(r["event_date"]) for r in temporal_needs if r["event_date"] is not None
|
||||
]
|
||||
lateral_fact_types = [r["fact_type"] for r in temporal_needs if r["event_date"] is not None]
|
||||
|
||||
if lateral_unit_ids:
|
||||
rows = await ops.fetch_temporal_neighbors(
|
||||
conn,
|
||||
fq_table("memory_units"),
|
||||
bank_id,
|
||||
lateral_unit_ids,
|
||||
lateral_event_dates,
|
||||
lateral_fact_types,
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
)
|
||||
for row in rows:
|
||||
time_diff_h = float(row["time_diff_hours"])
|
||||
# Mirror the 24h window enforced at retain time. The bidirectional
|
||||
# index scan returns the K closest neighbours regardless of
|
||||
# window, so we filter here.
|
||||
if time_diff_h > 24:
|
||||
continue
|
||||
weight = max(0.3, 1.0 - (time_diff_h / 24))
|
||||
new_links.append((row["from_id"], str(row["id"]), "temporal", weight, None))
|
||||
|
||||
# --- Semantic top-up ---
|
||||
# ANN must run on its own connection: it opens a nested transaction with
|
||||
# SET LOCAL hnsw.ef_search + CREATE TEMP TABLE ON COMMIT DROP, and nesting
|
||||
# that inside our current write transaction would commit our writes early.
|
||||
semantic_needs = [
|
||||
r
|
||||
for r in victim_rows
|
||||
if counts.get((r["id"], "semantic"), 0) < MAX_SEMANTIC_LINKS_PER_UNIT and r["embedding"] is not None
|
||||
]
|
||||
if semantic_needs:
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
seed_ids = [r["id"] for r in semantic_needs]
|
||||
seed_embs = [r["embedding"] for r in semantic_needs]
|
||||
seed_ftypes = [r["fact_type"] for r in semantic_needs]
|
||||
async with acquire_with_retry(backend) as ann_conn:
|
||||
try:
|
||||
ann_links = await compute_semantic_links_ann(
|
||||
ann_conn,
|
||||
bank_id,
|
||||
seed_ids,
|
||||
seed_embs,
|
||||
fact_types=seed_ftypes,
|
||||
)
|
||||
# Strip self-links (rare but possible because the ANN probe
|
||||
# has no exclude list — see the comment in compute_semantic_links_ann).
|
||||
ann_links = [lnk for lnk in ann_links if lnk[0] != lnk[1]]
|
||||
new_links.extend(ann_links)
|
||||
except Exception as e:
|
||||
# ANN uses PG-specific HNSW syntax; on dialects/configs where
|
||||
# it isn't available we still want the temporal top-up to land.
|
||||
logger.warning(f"[GRAPH_MAINT] Semantic top-up failed for bank={bank_id}: {type(e).__name__}: {e}")
|
||||
|
||||
if not new_links:
|
||||
return 0
|
||||
|
||||
await _bulk_insert_links(
|
||||
conn,
|
||||
new_links,
|
||||
bank_id=bank_id,
|
||||
skip_exists_check=False,
|
||||
ops=ops,
|
||||
)
|
||||
return len(new_links)
|
||||
@@ -458,8 +458,28 @@ class MemoryEngineInterface(ABC):
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with node_counts, link_counts, link_counts_by_fact_type,
|
||||
link_breakdown, and operations stats.
|
||||
Dict with node_counts, link_counts, link_counts_by_fact_type
|
||||
(deprecated, returns empty), link_breakdown (deprecated, returns
|
||||
empty), and operations stats.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_bank_freshness(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Get consolidation freshness for a bank.
|
||||
|
||||
Cheap alternative to get_bank_stats when callers only need
|
||||
last_consolidated_at / pending_consolidation / failed_consolidation.
|
||||
|
||||
Returns:
|
||||
Dict with last_consolidated_at (ISO-8601 string or None),
|
||||
pending_consolidation (int), and failed_consolidation (int).
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@@ -69,6 +69,7 @@ class LLMInterface(ABC):
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
cached_prefix: str | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Make an LLM API call with retry logic.
|
||||
@@ -83,8 +84,13 @@ class LLMInterface(ABC):
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Use strict JSON schema enforcement (OpenAI only).
|
||||
strict_schema: Grammar-enforce structured output via json_schema strict
|
||||
(OpenAI-compatible, LiteLLM) instead of the soft json_object path. Gemini
|
||||
enforces its response_schema natively; providers without a strict mode ignore it.
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
cached_prefix: Opaque handle from ``get_or_create_cached_prefix`` for the
|
||||
cacheable system prefix, or None. Providers without explicit prompt
|
||||
caching ignore it (and the wrapper only forwards it when set).
|
||||
|
||||
Returns:
|
||||
If return_usage=False: Parsed response if response_format is provided, otherwise text content.
|
||||
@@ -108,6 +114,7 @@ class LLMInterface(ABC):
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
cached_prefix: str | None = None,
|
||||
) -> LLMToolCallResult:
|
||||
"""
|
||||
Make an LLM API call with tool/function calling support.
|
||||
@@ -137,6 +144,46 @@ class LLMInterface(ABC):
|
||||
"""
|
||||
return False
|
||||
|
||||
# ── Prompt prefix caching (optional, per-provider) ─────────────────────────
|
||||
|
||||
def supports_prompt_caching(self) -> bool:
|
||||
"""Whether this provider can cache a reusable prompt prefix.
|
||||
|
||||
Default False. Providers that return True must implement
|
||||
``get_or_create_cached_prefix`` and honour the ``cached_prefix`` argument
|
||||
of ``call`` / ``call_with_tools``.
|
||||
"""
|
||||
return False
|
||||
|
||||
async def get_or_create_cached_prefix(
|
||||
self,
|
||||
*,
|
||||
system_instruction: str,
|
||||
response_schema: Any | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
) -> str | None:
|
||||
"""Cache a reusable prompt prefix and return an opaque handle, or None.
|
||||
|
||||
The engine has already decided WHAT is cacheable: it puts the stable,
|
||||
bank-agnostic instructions in ``system_instruction`` (plus ``tools``) and
|
||||
keeps all per-request / per-bank data (documents, facts, the bank mission)
|
||||
in the user message. A provider only chooses HOW to cache that prefix:
|
||||
|
||||
- Explicit-cache providers (e.g. Gemini ``CachedContent``): create the
|
||||
cache, return its handle; the engine passes the handle back via
|
||||
``call(cached_prefix=...)`` and the provider then drops the prefix from
|
||||
the request, billing it at the cached rate.
|
||||
- Automatic-cache providers (e.g. OpenAI): no handle needed — caching is
|
||||
transparent as long as the prefix is a stable leading block, which it
|
||||
already is. They can keep this default (return None) and still benefit.
|
||||
- Inline-marker providers (e.g. Anthropic ``cache_control``): mark the
|
||||
prefix block inside ``call`` instead; may also keep this default.
|
||||
|
||||
Returns None when caching is disabled/unsupported or the prefix is too
|
||||
small; callers MUST fall back to an uncached call in that case.
|
||||
"""
|
||||
return None
|
||||
|
||||
async def submit_batch(
|
||||
self,
|
||||
requests: list[dict[str, Any]],
|
||||
|
||||
@@ -0,0 +1,540 @@
|
||||
"""Per-bank LLM request tracing.
|
||||
|
||||
Opt-in, fire-and-forget recording of every LLM call Hindsight makes (both
|
||||
successes and failures) into the ``llm_requests`` table, per bank. Each row
|
||||
captures the input messages, the model output, token usage (input / output /
|
||||
cached / total), finish reason, and caller metadata. Disabled by default —
|
||||
controlled by ``HINDSIGHT_API_LLM_TRACE_ENABLED``.
|
||||
|
||||
This plugs into the OpenTelemetry **GenAI** recording pattern: providers already
|
||||
call ``tracing.get_span_recorder().record_llm_call(...)`` on success, so the DB
|
||||
tracer is registered as one of those recorders (alongside the OTLP span
|
||||
exporter) rather than hooking the call path with custom code. Failures, which
|
||||
providers don't report to the recorder, are forwarded from the LLM wrapper.
|
||||
|
||||
Bank/operation attribution is carried via a ContextVar set by
|
||||
``ConfiguredLLMProvider`` (see ``llm_wrapper.py``); outside a traced context
|
||||
``bank_id`` is recorded as NULL.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterable
|
||||
from contextvars import ContextVar, Token
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ── bank/operation attribution (carried across the async call chain) ──────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMTraceContext:
|
||||
"""Attribution for in-flight LLM calls, bound by ``ConfiguredLLMProvider``.
|
||||
|
||||
``trace_id`` and ``operation_span_id`` are generated once per operation
|
||||
invocation (one ``with_config`` call), so every LLM call of a single
|
||||
reflect/retain/consolidation run shares them — reproducing the OTel
|
||||
parent (operation span) → children (LLM calls) hierarchy in the DB.
|
||||
"""
|
||||
|
||||
bank_id: str | None = None
|
||||
operation: str | None = None # "retain" | "reflect" | "consolidation" | ...
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
trace_id: str | None = None
|
||||
operation_span_id: str | None = None
|
||||
# Memory_units this operation produced/consumed, accumulated at the DB-write
|
||||
# sites and flushed onto every row of the trace at operation end (see
|
||||
# LLMTraceRecorder.attach_memory_ids). Lets a retain/consolidation trace map
|
||||
# to the memories it created (outputs) and consumed (source inputs).
|
||||
created_memory_ids: list[str] = field(default_factory=list)
|
||||
source_memory_ids: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
_trace_ctx: ContextVar[LLMTraceContext | None] = ContextVar("hindsight_llm_trace_ctx", default=None)
|
||||
|
||||
# Per-call requested parameters (max_completion_tokens, temperature, response
|
||||
# schema, tool_choice). Set by ``LLMProvider.call`` around the provider
|
||||
# delegation so the recorder can attach them even though success is reported by
|
||||
# the provider. Only includes values the caller actually set — never nulls.
|
||||
_request_ctx: ContextVar[dict[str, Any] | None] = ContextVar("hindsight_llm_request_ctx", default=None)
|
||||
|
||||
# Per-call caller metadata (e.g. document_id for retain extraction). Set by
|
||||
# engine code around a specific LLM call; merged into the row's metadata on top
|
||||
# of the operation-level LLMTraceContext.metadata.
|
||||
_call_metadata_ctx: ContextVar[dict[str, Any] | None] = ContextVar("hindsight_llm_call_metadata_ctx", default=None)
|
||||
|
||||
|
||||
def set_trace_context(ctx: LLMTraceContext | None) -> Token:
|
||||
"""Bind trace attribution to the current context. Returns a reset token."""
|
||||
return _trace_ctx.set(ctx)
|
||||
|
||||
|
||||
def reset_trace_context(token: Token) -> None:
|
||||
"""Unwind a binding made by :func:`set_trace_context`."""
|
||||
_trace_ctx.reset(token)
|
||||
|
||||
|
||||
def set_request_context(params: dict[str, Any] | None) -> Token:
|
||||
"""Bind the current LLM call's requested parameters. Returns a reset token."""
|
||||
return _request_ctx.set(params)
|
||||
|
||||
|
||||
def reset_request_context(token: Token) -> None:
|
||||
"""Unwind a binding made by :func:`set_request_context`."""
|
||||
_request_ctx.reset(token)
|
||||
|
||||
|
||||
def current_request_context() -> dict[str, Any] | None:
|
||||
"""Return the active call's requested parameters, or None."""
|
||||
return _request_ctx.get()
|
||||
|
||||
|
||||
def set_call_metadata(metadata: dict[str, Any] | None) -> Token:
|
||||
"""Bind per-call caller metadata (e.g. ``{"document_id": ...}``)."""
|
||||
return _call_metadata_ctx.set(metadata)
|
||||
|
||||
|
||||
def reset_call_metadata(token: Token) -> None:
|
||||
"""Unwind a binding made by :func:`set_call_metadata`."""
|
||||
_call_metadata_ctx.reset(token)
|
||||
|
||||
|
||||
def current_call_metadata() -> dict[str, Any] | None:
|
||||
"""Return the active call's caller metadata, or None."""
|
||||
return _call_metadata_ctx.get()
|
||||
|
||||
|
||||
def current_trace_context() -> LLMTraceContext | None:
|
||||
"""Return the active trace attribution, or None outside a traced context."""
|
||||
return _trace_ctx.get()
|
||||
|
||||
|
||||
def trace_context_of(llm_config: Any) -> LLMTraceContext | None:
|
||||
"""Return a configured provider's operation trace context, or None.
|
||||
|
||||
Real providers expose ``trace_context()`` (``ConfiguredLLMProvider``); test
|
||||
or mock substitutes may not, so this degrades gracefully rather than raising
|
||||
— tracing is best-effort and must never break an operation.
|
||||
"""
|
||||
getter = getattr(llm_config, "trace_context", None)
|
||||
return getter() if callable(getter) else None
|
||||
|
||||
|
||||
def record_created_memory_ids(ids: Iterable[str]) -> None:
|
||||
"""Accumulate output memory_units onto the active operation trace.
|
||||
|
||||
No-op outside a traced operation context (e.g. tracing disabled). Child
|
||||
asyncio tasks inherit the same ``LLMTraceContext`` object, so appends from
|
||||
parallel consolidation batches land on one shared list.
|
||||
"""
|
||||
ctx = _trace_ctx.get()
|
||||
if ctx is not None:
|
||||
ctx.created_memory_ids.extend(str(i) for i in ids)
|
||||
|
||||
|
||||
def record_source_memory_ids(ids: Iterable[str]) -> None:
|
||||
"""Accumulate consumed/source memory_units onto the active operation trace.
|
||||
|
||||
No-op outside a traced operation context.
|
||||
"""
|
||||
ctx = _trace_ctx.get()
|
||||
if ctx is not None:
|
||||
ctx.source_memory_ids.extend(str(i) for i in ids)
|
||||
|
||||
|
||||
# ── serialization helpers ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _json_default(obj: Any) -> Any:
|
||||
"""JSON serializer for objects not serializable by default."""
|
||||
if isinstance(obj, datetime):
|
||||
return obj.isoformat()
|
||||
if isinstance(obj, uuid.UUID):
|
||||
return str(obj)
|
||||
if isinstance(obj, bytes):
|
||||
return "<bytes>"
|
||||
if isinstance(obj, set):
|
||||
return list(obj)
|
||||
model_dump = getattr(obj, "model_dump", None)
|
||||
if callable(model_dump):
|
||||
try:
|
||||
return model_dump(mode="json")
|
||||
except Exception:
|
||||
return str(obj)
|
||||
return str(obj)
|
||||
|
||||
|
||||
def _safe_json(data: Any, max_chars: int) -> str | None:
|
||||
"""Serialize ``data`` to a JSON string, truncating beyond ``max_chars``.
|
||||
|
||||
Returns None on total failure. Truncation preserves valid JSON by wrapping
|
||||
the oversized payload in a marker object with a preview.
|
||||
"""
|
||||
if data is None:
|
||||
return None
|
||||
try:
|
||||
serialized = json.dumps(data, default=_json_default)
|
||||
except Exception:
|
||||
logger.debug("Failed to serialize llm trace data", exc_info=True)
|
||||
try:
|
||||
serialized = json.dumps(str(data))
|
||||
except Exception:
|
||||
return None
|
||||
if max_chars and max_chars > 0 and len(serialized) > max_chars:
|
||||
return json.dumps({"_truncated": True, "_original_chars": len(serialized), "preview": serialized[:max_chars]})
|
||||
return serialized
|
||||
|
||||
|
||||
# ── record ────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMRequestRecord:
|
||||
"""A single LLM request trace row."""
|
||||
|
||||
provider: str
|
||||
model: str | None
|
||||
scope: str
|
||||
status: str # "success" | "error"
|
||||
started_at: datetime
|
||||
ended_at: datetime
|
||||
bank_id: str | None = None
|
||||
operation: str | None = None
|
||||
trace_id: str | None = None
|
||||
span_id: str | None = None
|
||||
parent_span_id: str | None = None
|
||||
input: Any = None
|
||||
output: Any = None
|
||||
error: str | None = None
|
||||
input_tokens: int | None = None
|
||||
output_tokens: int | None = None
|
||||
cached_tokens: int | None = None
|
||||
total_tokens: int | None = None
|
||||
llm_info: dict[str, Any] = field(default_factory=dict)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def duration_ms(self) -> int:
|
||||
return int((self.ended_at - self.started_at).total_seconds() * 1000)
|
||||
|
||||
|
||||
# ── read models (returned by MemoryEngine query methods, served by the API) ───
|
||||
|
||||
|
||||
class LLMRequestEntry(BaseModel):
|
||||
"""A single LLM request trace row, as returned by the read API."""
|
||||
|
||||
id: str
|
||||
bank_id: str | None
|
||||
operation: str | None
|
||||
scope: str | None
|
||||
trace_id: str | None
|
||||
span_id: str | None
|
||||
parent_span_id: str | None
|
||||
provider: str | None
|
||||
model: str | None
|
||||
status: str
|
||||
started_at: str | None
|
||||
ended_at: str | None
|
||||
duration_ms: int | None
|
||||
input_tokens: int | None
|
||||
output_tokens: int | None
|
||||
cached_tokens: int | None
|
||||
total_tokens: int | None
|
||||
# Arbitrary JSON (message list, string, or object) — open `Any` so the
|
||||
# OpenAPI schema stays a plain open type the Go SDK generator can model.
|
||||
input: Any = None
|
||||
output: Any = None
|
||||
error: str | None
|
||||
llm_info: dict[str, Any]
|
||||
metadata: dict[str, Any]
|
||||
|
||||
|
||||
class LLMRequestListResponse(BaseModel):
|
||||
"""Paginated list of LLM request traces for a bank."""
|
||||
|
||||
bank_id: str
|
||||
total: int
|
||||
limit: int
|
||||
offset: int
|
||||
items: list[LLMRequestEntry]
|
||||
|
||||
|
||||
class LLMRequestTokenSums(BaseModel):
|
||||
"""Token totals for a time bucket."""
|
||||
|
||||
input: int
|
||||
output: int
|
||||
cached: int
|
||||
total: int
|
||||
|
||||
|
||||
class LLMRequestStatsBucket(BaseModel):
|
||||
"""A single time bucket in LLM request stats."""
|
||||
|
||||
time: str
|
||||
statuses: dict[str, int]
|
||||
total: int
|
||||
tokens: LLMRequestTokenSums
|
||||
|
||||
|
||||
class LLMRequestStatsResponse(BaseModel):
|
||||
"""LLM request counts and token sums grouped by time bucket."""
|
||||
|
||||
bank_id: str
|
||||
period: str
|
||||
trunc: str
|
||||
start: str
|
||||
buckets: list[LLMRequestStatsBucket]
|
||||
|
||||
|
||||
# ── recorder / writer ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class LLMTraceRecorder:
|
||||
"""GenAI span recorder that writes per-bank LLM traces to ``llm_requests``.
|
||||
|
||||
Implements ``record_llm_call`` so it can be registered with
|
||||
:func:`hindsight_api.tracing.register_span_recorder`. Writes are
|
||||
fire-and-forget and never surface errors into the calling path. Retention of
|
||||
old rows is handled by the background :class:`MaintenanceLoop`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pool_getter: Callable[[], Any],
|
||||
schema_getter: Callable[[], str],
|
||||
enabled: bool,
|
||||
allowed_scopes: list[str],
|
||||
max_chars: int = 50000,
|
||||
) -> None:
|
||||
self._pool_getter = pool_getter
|
||||
self._schema_getter = schema_getter
|
||||
self._enabled = enabled
|
||||
self._allowed_scopes: frozenset[str] | None = frozenset(allowed_scopes) if allowed_scopes else None
|
||||
self._max_chars = max_chars
|
||||
# In-flight fire-and-forget write tasks, bucketed by trace_id so
|
||||
# attach_memory_ids can await only *its own* operation's writes before the
|
||||
# post-operation UPDATE (otherwise the UPDATE could race ahead of the
|
||||
# INSERTs it patches — but it must not block on unrelated operations).
|
||||
self._pending: dict[str | None, set[asyncio.Task]] = {}
|
||||
|
||||
def is_enabled(self, scope: str) -> bool:
|
||||
"""Whether tracing is active for the given call scope."""
|
||||
if not self._enabled:
|
||||
return False
|
||||
if self._allowed_scopes is not None:
|
||||
return scope in self._allowed_scopes
|
||||
return True
|
||||
|
||||
# ── GenAI recorder interface ──────────────────────────────────────────────
|
||||
|
||||
def record_llm_call(
|
||||
self,
|
||||
provider: str,
|
||||
model: str,
|
||||
scope: str,
|
||||
messages: list[dict[str, Any]],
|
||||
response_content: Any = None,
|
||||
input_tokens: int = 0,
|
||||
output_tokens: int = 0,
|
||||
duration: float = 0.0,
|
||||
finish_reason: str | None = None,
|
||||
error: BaseException | None = None,
|
||||
tool_calls: list[dict[str, Any]] | None = None,
|
||||
cached_tokens: int = 0,
|
||||
**_extra: Any,
|
||||
) -> None:
|
||||
"""Build a trace record from a GenAI call and schedule a DB write."""
|
||||
if not self.is_enabled(scope):
|
||||
return
|
||||
|
||||
ctx = current_trace_context()
|
||||
ended_at = datetime.now(timezone.utc)
|
||||
started_at = ended_at - timedelta(seconds=max(0.0, duration))
|
||||
|
||||
# Operation-level metadata + any per-call metadata (e.g. document_id).
|
||||
metadata = dict(ctx.metadata) if ctx else {}
|
||||
call_metadata = current_call_metadata()
|
||||
if call_metadata:
|
||||
metadata.update(call_metadata)
|
||||
|
||||
llm_info: dict[str, Any] = {}
|
||||
request_params = current_request_context()
|
||||
if request_params:
|
||||
llm_info["request"] = dict(request_params)
|
||||
if finish_reason:
|
||||
llm_info["finish_reason"] = finish_reason
|
||||
if tool_calls:
|
||||
llm_info["tool_calls"] = [tc.get("name", "") for tc in tool_calls]
|
||||
|
||||
record = LLMRequestRecord(
|
||||
provider=provider,
|
||||
model=model,
|
||||
scope=scope,
|
||||
status="error" if error is not None else "success",
|
||||
started_at=started_at,
|
||||
ended_at=ended_at,
|
||||
bank_id=ctx.bank_id if ctx else None,
|
||||
operation=ctx.operation if ctx else None,
|
||||
# OTel-style hierarchy: all calls of one operation invocation share
|
||||
# the context's trace_id and point at its operation span; this call
|
||||
# gets its own span_id.
|
||||
trace_id=ctx.trace_id if ctx else None,
|
||||
span_id=str(uuid.uuid4()),
|
||||
parent_span_id=ctx.operation_span_id if ctx else None,
|
||||
input=messages,
|
||||
output=None if error is not None else response_content,
|
||||
error=f"{type(error).__name__}: {error}" if error is not None else None,
|
||||
input_tokens=input_tokens or None,
|
||||
output_tokens=output_tokens or None,
|
||||
cached_tokens=cached_tokens or None,
|
||||
total_tokens=(input_tokens + output_tokens) or None,
|
||||
llm_info=llm_info,
|
||||
metadata=metadata,
|
||||
)
|
||||
self._record_fire_and_forget(record)
|
||||
|
||||
def _record_fire_and_forget(self, record: LLMRequestRecord) -> None:
|
||||
"""Schedule a trace write as a background task."""
|
||||
try:
|
||||
task = asyncio.create_task(self._safe_write(record))
|
||||
except RuntimeError:
|
||||
# No running event loop (e.g. during shutdown)
|
||||
logger.debug("Cannot schedule llm trace write: no running event loop")
|
||||
return
|
||||
key = record.trace_id
|
||||
self._pending.setdefault(key, set()).add(task)
|
||||
task.add_done_callback(lambda t, k=key: self._discard_pending(k, t))
|
||||
|
||||
def _discard_pending(self, key: str | None, task: asyncio.Task) -> None:
|
||||
bucket = self._pending.get(key)
|
||||
if bucket is not None:
|
||||
bucket.discard(task)
|
||||
if not bucket:
|
||||
self._pending.pop(key, None)
|
||||
|
||||
async def _safe_write(self, record: LLMRequestRecord) -> None:
|
||||
"""Write a trace row. Errors are logged, never raised."""
|
||||
pool = self._pool_getter()
|
||||
if pool is None:
|
||||
logger.debug("LLM trace skipped: pool not available")
|
||||
return
|
||||
try:
|
||||
schema = self._schema_getter()
|
||||
table = f"{schema}.llm_requests"
|
||||
async with acquire_with_retry(pool, max_retries=1) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table}
|
||||
(id, bank_id, operation, scope, trace_id, span_id, parent_span_id,
|
||||
provider, model, status,
|
||||
started_at, ended_at, duration_ms,
|
||||
input_tokens, output_tokens, cached_tokens, total_tokens,
|
||||
input, output, error, llm_info, metadata)
|
||||
VALUES
|
||||
($1, $2, $3, $4, $5, $6, $7, $8, $9, $10,
|
||||
$11, $12, $13, $14, $15, $16, $17,
|
||||
$18::jsonb, $19::jsonb, $20, $21::jsonb, $22::jsonb)
|
||||
""",
|
||||
uuid.uuid4(),
|
||||
record.bank_id,
|
||||
record.operation,
|
||||
record.scope,
|
||||
record.trace_id,
|
||||
record.span_id,
|
||||
record.parent_span_id,
|
||||
record.provider,
|
||||
record.model,
|
||||
record.status,
|
||||
record.started_at,
|
||||
record.ended_at,
|
||||
record.duration_ms,
|
||||
record.input_tokens,
|
||||
record.output_tokens,
|
||||
record.cached_tokens,
|
||||
record.total_tokens,
|
||||
_safe_json(record.input, self._max_chars),
|
||||
_safe_json(record.output, self._max_chars),
|
||||
record.error,
|
||||
_safe_json(record.llm_info, self._max_chars) or "{}",
|
||||
_safe_json(record.metadata, self._max_chars) or "{}",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"LLM trace write failed for scope={record.scope}: {e}")
|
||||
|
||||
async def _flush_pending(self, trace_id: str) -> None:
|
||||
"""Await this trace's in-flight writes so its rows exist before an UPDATE."""
|
||||
pending = [t for t in self._pending.get(trace_id, ()) if not t.done()]
|
||||
if pending:
|
||||
await asyncio.gather(*pending, return_exceptions=True)
|
||||
|
||||
def attach_memory_ids(
|
||||
self,
|
||||
trace_ctx: LLMTraceContext | None,
|
||||
*,
|
||||
created: list[str] | None = None,
|
||||
source: list[str] | None = None,
|
||||
) -> None:
|
||||
"""Map a finished operation's memory_units onto every row of its trace.
|
||||
|
||||
Merges the explicitly passed ids with any accumulated on the context
|
||||
(``record_created_memory_ids`` / ``record_source_memory_ids``), de-dupes
|
||||
preserving order, and patches ``metadata.memory_ids`` (outputs created)
|
||||
and ``metadata.source_memory_ids`` (inputs consumed) on all rows sharing
|
||||
the trace_id. No-op when tracing is off or nothing was produced.
|
||||
|
||||
Fire-and-forget: the snapshotted patch is applied on a background task so
|
||||
the retain/consolidation operation never waits on the trace write. The
|
||||
ids are snapshotted synchronously here because the caller may reset the
|
||||
context immediately after.
|
||||
"""
|
||||
if not self._enabled or trace_ctx is None or not trace_ctx.trace_id:
|
||||
return
|
||||
created_ids = list(dict.fromkeys([*(created or []), *trace_ctx.created_memory_ids]))
|
||||
source_ids = list(dict.fromkeys([*(source or []), *trace_ctx.source_memory_ids]))
|
||||
patch: dict[str, Any] = {}
|
||||
if created_ids:
|
||||
patch["memory_ids"] = created_ids
|
||||
if source_ids:
|
||||
patch["source_memory_ids"] = source_ids
|
||||
if not patch:
|
||||
return
|
||||
try:
|
||||
asyncio.create_task(self._attach_memory_ids(trace_ctx.bank_id, trace_ctx.trace_id, patch))
|
||||
except RuntimeError:
|
||||
logger.debug("Cannot schedule llm trace memory_id attach: no running event loop")
|
||||
|
||||
async def _attach_memory_ids(self, bank_id: str | None, trace_id: str, patch: dict[str, Any]) -> None:
|
||||
"""Background worker: flush this trace's writes, then patch its rows."""
|
||||
# The trace-row INSERTs are fire-and-forget; flush *this trace's* writes
|
||||
# so the UPDATE patches rows that already exist rather than racing ahead
|
||||
# of them (without blocking on unrelated operations' pending writes).
|
||||
await self._flush_pending(trace_id)
|
||||
pool = self._pool_getter()
|
||||
if pool is None:
|
||||
return
|
||||
try:
|
||||
schema = self._schema_getter()
|
||||
table = f"{schema}.llm_requests"
|
||||
async with acquire_with_retry(pool, max_retries=1) as conn:
|
||||
await conn.execute(
|
||||
f"UPDATE {table} SET metadata = metadata || $3::jsonb WHERE bank_id = $1 AND trace_id = $2",
|
||||
bank_id,
|
||||
trace_id,
|
||||
json.dumps(patch),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"LLM trace memory_id attach failed for trace={trace_id}: {e}")
|
||||
@@ -9,6 +9,7 @@ import os
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import AsyncExitStack
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -27,9 +28,12 @@ except ImportError:
|
||||
from ..config import (
|
||||
DEFAULT_LLM_MAX_CONCURRENT,
|
||||
DEFAULT_LLM_TIMEOUT,
|
||||
ENV_CONSOLIDATION_LLM_MAX_CONCURRENT,
|
||||
ENV_LLM_GROQ_SERVICE_TIER,
|
||||
ENV_LLM_MAX_CONCURRENT,
|
||||
ENV_LLM_TIMEOUT,
|
||||
ENV_REFLECT_LLM_MAX_CONCURRENT,
|
||||
ENV_RETAIN_LLM_MAX_CONCURRENT,
|
||||
)
|
||||
from ..metrics import get_metrics_collector
|
||||
from .response_models import TokenUsage
|
||||
@@ -42,13 +46,101 @@ logger = logging.getLogger(__name__)
|
||||
# Disable httpx logging
|
||||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
|
||||
# Global semaphore to limit concurrent LLM requests across all instances
|
||||
# Set HINDSIGHT_API_LLM_MAX_CONCURRENT=1 for local LLMs (LM Studio, Ollama)
|
||||
# Global semaphore to limit concurrent LLM requests across all instances.
|
||||
# Set HINDSIGHT_API_LLM_MAX_CONCURRENT=1 for local LLMs (LM Studio, Ollama).
|
||||
_llm_max_concurrent = int(os.getenv(ENV_LLM_MAX_CONCURRENT, str(DEFAULT_LLM_MAX_CONCURRENT)))
|
||||
_global_llm_semaphore = asyncio.Semaphore(_llm_max_concurrent)
|
||||
|
||||
|
||||
def sanitize_llm_output(text: str | None) -> str | None:
|
||||
def _build_per_op_semaphores() -> dict[str, asyncio.Semaphore]:
|
||||
"""Build the per-operation semaphore registry from env vars.
|
||||
|
||||
Each per-op cap is composed with — not a substitute for — the global cap:
|
||||
a call that matches a configured operation must acquire both its per-op
|
||||
semaphore and the global semaphore. This lets operators reserve headroom
|
||||
in the global pool by capping individual operations (e.g. cap retain at 2
|
||||
of 4 global slots so the live chat path always has 2 slots available).
|
||||
|
||||
Operations without a configured env var are absent from the registry and
|
||||
therefore only constrained by the global cap.
|
||||
"""
|
||||
semaphores: dict[str, asyncio.Semaphore] = {}
|
||||
for op, env_var in (
|
||||
("retain", ENV_RETAIN_LLM_MAX_CONCURRENT),
|
||||
("reflect", ENV_REFLECT_LLM_MAX_CONCURRENT),
|
||||
("consolidation", ENV_CONSOLIDATION_LLM_MAX_CONCURRENT),
|
||||
):
|
||||
raw = os.getenv(env_var)
|
||||
if raw is None or raw == "":
|
||||
continue
|
||||
value = int(raw)
|
||||
if value <= 0:
|
||||
raise ValueError(f"{env_var} must be a positive integer, got {raw!r}")
|
||||
semaphores[op] = asyncio.Semaphore(value)
|
||||
return semaphores
|
||||
|
||||
|
||||
_per_op_llm_semaphores: dict[str, asyncio.Semaphore] = _build_per_op_semaphores()
|
||||
|
||||
|
||||
def _scope_to_operation(scope: str) -> str | None:
|
||||
"""Map a call scope to its per-operation concurrency bucket.
|
||||
|
||||
Returns None for scopes that don't belong to a tracked operation
|
||||
(verification probes, bank_mission, memory_think, mental_model_delta_ops),
|
||||
which then run under the global cap only.
|
||||
"""
|
||||
if scope.startswith("retain"):
|
||||
return "retain"
|
||||
if scope.startswith("reflect"):
|
||||
return "reflect"
|
||||
if scope.startswith("consolidation"):
|
||||
return "consolidation"
|
||||
return None
|
||||
|
||||
|
||||
def _semaphores_for_scope(scope: str) -> list[asyncio.Semaphore]:
|
||||
"""Return the semaphores a call with the given scope must acquire.
|
||||
|
||||
Always includes the global semaphore; includes the per-op semaphore when
|
||||
one is configured for the scope's operation bucket.
|
||||
"""
|
||||
op = _scope_to_operation(scope)
|
||||
per_op = _per_op_llm_semaphores.get(op) if op is not None else None
|
||||
if per_op is None:
|
||||
return [_global_llm_semaphore]
|
||||
# Per-op acquired first so contention queues on the narrower cap before
|
||||
# holding a global slot.
|
||||
return [per_op, _global_llm_semaphore]
|
||||
|
||||
|
||||
def _request_params(
|
||||
*,
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str | None = None,
|
||||
response_format: Any | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Build the requested-params bag for tracing — only values the caller set.
|
||||
|
||||
Omitting unset values avoids the misleading nulls we used to record (e.g.
|
||||
consolidation, which passes no token cap), while surfacing the real cap for
|
||||
callers that do set one (e.g. retain's ``retain_max_completion_tokens``).
|
||||
"""
|
||||
params: dict[str, Any] = {}
|
||||
if max_completion_tokens is not None:
|
||||
params["max_completion_tokens"] = max_completion_tokens
|
||||
if temperature is not None:
|
||||
params["temperature"] = temperature
|
||||
if response_format is not None:
|
||||
params["response_schema"] = getattr(response_format, "__name__", None) or "structured"
|
||||
if tool_choice is not None and tool_choice != "auto":
|
||||
params["tool_choice"] = tool_choice if isinstance(tool_choice, str) else "named"
|
||||
return params or None
|
||||
|
||||
|
||||
def sanitize_text(text: str | None) -> str | None:
|
||||
"""
|
||||
Sanitize text by removing characters that break downstream systems.
|
||||
|
||||
@@ -60,8 +152,12 @@ def sanitize_llm_output(text: str | None) -> str | None:
|
||||
|
||||
Surrogate characters are used in UTF-16 encoding but cannot be encoded
|
||||
in UTF-8. They can appear in Python strings from improperly decoded data
|
||||
(e.g., from JavaScript or broken files). Control characters commonly appear
|
||||
in LLM output embedded inside JSON string values.
|
||||
(e.g., from JavaScript or broken files): a client may serialize a half-emoji
|
||||
split at a boundary as a lone ``\\udXXX`` escape. Such input crashes the
|
||||
SentenceTransformers/cross-encoder Rust tokenizers and stdout logging, so
|
||||
user content is sanitized at the retain/recall/reflect ingress (see issue
|
||||
#1875). Control characters commonly appear in LLM output embedded inside
|
||||
JSON string values.
|
||||
"""
|
||||
if text is None:
|
||||
return None
|
||||
@@ -70,6 +166,11 @@ def sanitize_llm_output(text: str | None) -> str | None:
|
||||
return re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f\ud800-\udfff]", "", text)
|
||||
|
||||
|
||||
# Back-compat alias: this helper was originally introduced to scrub LLM *output*;
|
||||
# it now also scrubs user *input* at ingress, hence the broader name.
|
||||
sanitize_llm_output = sanitize_text
|
||||
|
||||
|
||||
class OutputTooLongError(Exception):
|
||||
"""
|
||||
Bridge exception raised when LLM output exceeds token limits.
|
||||
@@ -154,6 +255,7 @@ def create_llm_provider(
|
||||
vertexai_region: str | None = None,
|
||||
vertexai_credentials: Any = None,
|
||||
gemini_safety_settings: list | None = None,
|
||||
prompt_cache_enabled: bool = False,
|
||||
litellmrouter_config: dict[str, Any] | None = None,
|
||||
) -> Any: # Returns LLMInterface
|
||||
"""
|
||||
@@ -167,7 +269,11 @@ def create_llm_provider(
|
||||
reasoning_effort: Reasoning effort level for supported providers.
|
||||
groq_service_tier: Groq service tier (for Groq provider) - "on_demand", "flex", or "auto".
|
||||
openai_service_tier: OpenAI service tier (for OpenAI provider) - None (default) or "flex" (50% cheaper).
|
||||
extra_body: Extra body params merged into OpenAI-compatible API calls.
|
||||
extra_body: Extra request-body params merged into the provider's native
|
||||
call. Threaded into OpenAI-compatible, Fireworks, Anthropic, Gemini/
|
||||
VertexAI and LiteLLM providers (each merges them in its own parameter
|
||||
space). Keys must use each provider's native names (e.g. ``max_tokens``
|
||||
for OpenAI/Anthropic vs ``max_output_tokens`` for Gemini).
|
||||
default_headers: Custom headers passed as ``default_headers`` to provider SDK clients
|
||||
(used by operators routing through proxies / request-tracing middleware). Currently
|
||||
wired into the Anthropic provider; other providers may opt in as needed.
|
||||
@@ -183,6 +289,7 @@ def create_llm_provider(
|
||||
AnthropicLLM,
|
||||
ClaudeCodeLLM,
|
||||
CodexLLM,
|
||||
FireworksLLM,
|
||||
GeminiLLM,
|
||||
LiteLLMLLM,
|
||||
LiteLLMRouterLLM,
|
||||
@@ -241,6 +348,8 @@ def create_llm_provider(
|
||||
vertexai_region=vertexai_region,
|
||||
vertexai_credentials=vertexai_credentials,
|
||||
gemini_safety_settings=gemini_safety_settings,
|
||||
prompt_cache_enabled=prompt_cache_enabled,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower == "anthropic":
|
||||
@@ -251,6 +360,7 @@ def create_llm_provider(
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
default_headers=default_headers,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower == "litellm":
|
||||
@@ -260,6 +370,7 @@ def create_llm_provider(
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower == "litellmrouter":
|
||||
@@ -277,6 +388,7 @@ def create_llm_provider(
|
||||
model=model,
|
||||
config=litellmrouter_config,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower == "bedrock":
|
||||
@@ -288,6 +400,7 @@ def create_llm_provider(
|
||||
base_url=base_url,
|
||||
model=bedrock_model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower == "llamacpp":
|
||||
@@ -308,10 +421,24 @@ def create_llm_provider(
|
||||
extra_args=config.llamacpp_extra_args,
|
||||
)
|
||||
|
||||
elif provider_lower == "fireworks":
|
||||
# Fireworks online inference is OpenAI-compatible; FireworksLLM adds the
|
||||
# native (non-OpenAI) batch API on top. The existing LiteLLM
|
||||
# ``fireworks_ai/...`` online path (provider="litellm") is untouched.
|
||||
return FireworksLLM(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower in (
|
||||
"openai",
|
||||
"groq",
|
||||
"ollama",
|
||||
"ollama-cloud",
|
||||
"lmstudio",
|
||||
"minimax",
|
||||
"deepseek",
|
||||
@@ -352,6 +479,7 @@ class LLMProvider:
|
||||
groq_service_tier: str | None = None,
|
||||
openai_service_tier: str | None = None,
|
||||
gemini_safety_settings: list | None = None,
|
||||
prompt_cache_enabled: bool = False,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
default_headers: dict[str, str] | None = None,
|
||||
litellmrouter_config: dict[str, Any] | None = None,
|
||||
@@ -368,7 +496,8 @@ class LLMProvider:
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto") - from config.
|
||||
openai_service_tier: OpenAI service tier (None or "flex") - from config.
|
||||
gemini_safety_settings: Safety settings for Gemini/VertexAI providers.
|
||||
extra_body: Extra body params merged into OpenAI-compatible API calls.
|
||||
extra_body: Extra request-body params merged into the provider's native call
|
||||
(OpenAI-compatible, Fireworks, Anthropic, Gemini/VertexAI, LiteLLM).
|
||||
default_headers: Custom headers passed as ``default_headers`` to provider SDK clients.
|
||||
Used by operators routing through proxies / request-tracing middleware. Falls
|
||||
back to ``HindsightConfig.llm_default_headers`` (env: ``HINDSIGHT_API_LLM_DEFAULT_HEADERS``)
|
||||
@@ -390,6 +519,11 @@ class LLMProvider:
|
||||
self.openai_service_tier = openai_service_tier
|
||||
# Gemini safety settings (instance default; can be overridden per-request via context var)
|
||||
self.gemini_safety_settings = gemini_safety_settings
|
||||
# Gemini prompt caching: when True, retain extraction (and any future
|
||||
# caller that opts in) will reuse a CachedContent prefix to cut
|
||||
# input-token cost. Off by default so the change is observable behind
|
||||
# a flip rather than a silent behaviour change on upgrade.
|
||||
self.prompt_cache_enabled = prompt_cache_enabled
|
||||
# Extra body params for OpenAI-compatible providers (e.g. chat_template_kwargs)
|
||||
self.extra_body = extra_body
|
||||
# Default headers passed to provider SDK clients (e.g. proxy auth, request tracing).
|
||||
@@ -409,6 +543,7 @@ class LLMProvider:
|
||||
"openai",
|
||||
"groq",
|
||||
"ollama",
|
||||
"ollama-cloud",
|
||||
"gemini",
|
||||
"anthropic",
|
||||
"lmstudio",
|
||||
@@ -427,6 +562,7 @@ class LLMProvider:
|
||||
"openrouter",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"fireworks",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
raise ValueError(f"Invalid LLM provider: {self.provider}. Must be one of: {', '.join(valid_providers)}")
|
||||
@@ -437,6 +573,8 @@ class LLMProvider:
|
||||
self.base_url = "https://api.groq.com/openai/v1"
|
||||
elif self.provider == "ollama":
|
||||
self.base_url = "http://localhost:11434/v1"
|
||||
elif self.provider == "ollama-cloud":
|
||||
self.base_url = "https://ollama.com/v1"
|
||||
elif self.provider == "lmstudio":
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
elif self.provider == "minimax":
|
||||
@@ -504,6 +642,21 @@ class LLMProvider:
|
||||
except Exception:
|
||||
pass # Config may not be initialized in test environments
|
||||
|
||||
# Prompt-prefix caching is a provider-agnostic toggle (default on): resolve
|
||||
# it from the static server config for every provider when the caller didn't
|
||||
# pass an explicit override. Providers that don't support caching ignore the
|
||||
# value; only those that implement get_or_create_cached_prefix act on it.
|
||||
if not self.prompt_cache_enabled:
|
||||
from ..config import DEFAULT_LLM_PROMPT_CACHE_ENABLED, _get_raw_config
|
||||
|
||||
try:
|
||||
raw_config = _get_raw_config()
|
||||
self.prompt_cache_enabled = bool(
|
||||
getattr(raw_config, "llm_prompt_cache_enabled", DEFAULT_LLM_PROMPT_CACHE_ENABLED)
|
||||
)
|
||||
except Exception:
|
||||
pass # Config may not be initialized in test environments
|
||||
|
||||
# For litellmrouter: prefer an explicit chain from the caller (per-op
|
||||
# construction in MemoryEngine threads the right chain through). If the caller
|
||||
# didn't supply one, fall back to the global ``llm_litellmrouter_config`` so
|
||||
@@ -532,6 +685,7 @@ class LLMProvider:
|
||||
vertexai_region=vertexai_region,
|
||||
vertexai_credentials=vertexai_credentials,
|
||||
gemini_safety_settings=self.gemini_safety_settings,
|
||||
prompt_cache_enabled=self.prompt_cache_enabled,
|
||||
litellmrouter_config=router_config,
|
||||
)
|
||||
|
||||
@@ -595,6 +749,7 @@ class LLMProvider:
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
cached_prefix: str | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Make an LLM API call with retry logic.
|
||||
@@ -609,7 +764,10 @@ class LLMProvider:
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Use strict JSON schema enforcement (OpenAI only). Guarantees all required fields.
|
||||
strict_schema: Per-call override requesting grammar-enforced (json_schema strict)
|
||||
structured output instead of the soft json_object path. The server-level
|
||||
HINDSIGHT_API_LLM_STRICT_SCHEMA flag is OR-ed in here so it applies to every call;
|
||||
providers without a strict mode ignore it.
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
|
||||
Returns:
|
||||
@@ -629,32 +787,85 @@ class LLMProvider:
|
||||
structured = "+structured" if response_format is not None else ""
|
||||
set_stage(f"llm.{self.provider}.{scope}{structured}")
|
||||
|
||||
async with _global_llm_semaphore:
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call(
|
||||
messages=messages,
|
||||
response_format=response_format,
|
||||
# Resolve strict-schema once, here, rather than in each provider: the
|
||||
# per-call argument OR the server-level HINDSIGHT_API_LLM_STRICT_SCHEMA
|
||||
# flag. Providers with a json_schema response_format (OpenAI-compatible,
|
||||
# LiteLLM) then grammar-enforce structured output instead of the fragile
|
||||
# soft json_object path; Gemini already enforces its native response_schema,
|
||||
# and providers without a strict mode simply ignore the flag.
|
||||
from ..config import get_config
|
||||
|
||||
strict_schema = strict_schema or get_config().llm_strict_schema
|
||||
|
||||
# LLM call observability flows through the OTel GenAI recorder
|
||||
# (tracing.get_span_recorder().record_llm_call). Provider implementations
|
||||
# record successful calls; we forward failures here since they don't.
|
||||
# The requested params are stashed in a contextvar (only what the caller
|
||||
# actually set) so the recorder can attach them to either path.
|
||||
from ..tracing import get_span_recorder
|
||||
from .llm_trace import reset_request_context, set_request_context
|
||||
|
||||
call_start = time.monotonic()
|
||||
request_token = set_request_context(
|
||||
_request_params(
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
temperature=temperature,
|
||||
scope=scope,
|
||||
max_retries=max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
skip_validation=skip_validation,
|
||||
strict_schema=strict_schema,
|
||||
return_usage=return_usage,
|
||||
response_format=response_format,
|
||||
)
|
||||
)
|
||||
try:
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
|
||||
# Backward compatibility: Update mock call tracking for mock provider
|
||||
# This allows existing tests using LLMProvider._mock_calls to continue working
|
||||
if self.provider == "mock":
|
||||
from .providers.mock_llm import MockLLM
|
||||
# cached_prefix is only set for providers that returned a handle
|
||||
# from get_or_create_cached_prefix() (e.g. Gemini); it's None for
|
||||
# the rest. Forward it only when present so providers that don't
|
||||
# implement caching keep their call() signature untouched.
|
||||
cache_kwarg = {"cached_prefix": cached_prefix} if cached_prefix is not None else {}
|
||||
try:
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call(
|
||||
messages=messages,
|
||||
response_format=response_format,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
temperature=temperature,
|
||||
scope=scope,
|
||||
max_retries=max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
skip_validation=skip_validation,
|
||||
strict_schema=strict_schema,
|
||||
return_usage=return_usage,
|
||||
**cache_kwarg,
|
||||
)
|
||||
except Exception as e:
|
||||
get_span_recorder().record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=None,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
duration=time.monotonic() - call_start,
|
||||
error=e,
|
||||
)
|
||||
raise
|
||||
|
||||
if isinstance(self._provider_impl, MockLLM):
|
||||
# Sync the mock calls from provider implementation to wrapper
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
# Backward compatibility: Update mock call tracking for mock provider
|
||||
# This allows existing tests using LLMProvider._mock_calls to continue working
|
||||
if self.provider == "mock":
|
||||
from .providers.mock_llm import MockLLM
|
||||
|
||||
return result
|
||||
if isinstance(self._provider_impl, MockLLM):
|
||||
# Sync the mock calls from provider implementation to wrapper
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
finally:
|
||||
reset_request_context(request_token)
|
||||
|
||||
return result
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
@@ -667,6 +878,7 @@ class LLMProvider:
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
cached_prefix: str | None = None,
|
||||
) -> "LLMToolCallResult":
|
||||
"""
|
||||
Make an LLM API call with tool/function calling support.
|
||||
@@ -689,30 +901,68 @@ class LLMProvider:
|
||||
|
||||
set_stage(f"llm.{self.provider}.{scope}+tools")
|
||||
|
||||
async with _global_llm_semaphore:
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call_with_tools(
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
# Failures forwarded to the GenAI recorder; successes recorded by providers.
|
||||
from ..tracing import get_span_recorder
|
||||
from .llm_trace import reset_request_context, set_request_context
|
||||
|
||||
call_start = time.monotonic()
|
||||
request_token = set_request_context(
|
||||
_request_params(
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
temperature=temperature,
|
||||
scope=scope,
|
||||
max_retries=max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
tool_choice=tool_choice,
|
||||
)
|
||||
)
|
||||
try:
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
|
||||
# Backward compatibility: Update mock call tracking for mock provider
|
||||
# This allows existing tests using LLMProvider._mock_calls to continue working
|
||||
if self.provider == "mock":
|
||||
from .providers.mock_llm import MockLLM
|
||||
# cached_prefix is only set for providers that returned a handle
|
||||
# from get_or_create_cached_prefix(); forward it only when present
|
||||
# so non-caching providers keep their signature (same as call()).
|
||||
cache_kwarg = {"cached_prefix": cached_prefix} if cached_prefix is not None else {}
|
||||
try:
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call_with_tools(
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
temperature=temperature,
|
||||
scope=scope,
|
||||
max_retries=max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
tool_choice=tool_choice,
|
||||
**cache_kwarg,
|
||||
)
|
||||
except Exception as e:
|
||||
get_span_recorder().record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=None,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
duration=time.monotonic() - call_start,
|
||||
error=e,
|
||||
)
|
||||
raise
|
||||
|
||||
if isinstance(self._provider_impl, MockLLM):
|
||||
# Sync the mock calls from provider implementation to wrapper
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
# Backward compatibility: Update mock call tracking for mock provider
|
||||
# This allows existing tests using LLMProvider._mock_calls to continue working
|
||||
if self.provider == "mock":
|
||||
from .providers.mock_llm import MockLLM
|
||||
|
||||
return result
|
||||
if isinstance(self._provider_impl, MockLLM):
|
||||
# Sync the mock calls from provider implementation to wrapper
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
finally:
|
||||
reset_request_context(request_token)
|
||||
|
||||
return result
|
||||
|
||||
def set_response_callback(self, fn: Any) -> None:
|
||||
"""Set a callback invoked on each call() instead of the fixed mock response."""
|
||||
@@ -814,7 +1064,14 @@ class LLMProvider:
|
||||
# SDK will automatically check for authentication when first used
|
||||
# No need to verify here - let it fail gracefully on first call with helpful error
|
||||
|
||||
def with_config(self, config: Any) -> "ConfiguredLLMProvider":
|
||||
def with_config(
|
||||
self,
|
||||
config: Any,
|
||||
*,
|
||||
bank_id: str | None = None,
|
||||
operation: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> "ConfiguredLLMProvider":
|
||||
"""
|
||||
Return a configured wrapper for a specific bank operation.
|
||||
|
||||
@@ -824,12 +1081,31 @@ class LLMProvider:
|
||||
|
||||
Args:
|
||||
config: Resolved ``HindsightConfig`` for the current bank/request.
|
||||
bank_id: Bank the operation runs for; attributed to LLM trace rows.
|
||||
operation: Logical operation label ("retain", "reflect", ...) for
|
||||
LLM trace rows.
|
||||
metadata: Optional extra caller metadata stored on trace rows.
|
||||
|
||||
Returns:
|
||||
A ``ConfiguredLLMProvider`` that delegates to this provider with
|
||||
the supplied config applied.
|
||||
"""
|
||||
return ConfiguredLLMProvider(self, config.llm_gemini_safety_settings)
|
||||
trace_ctx = None
|
||||
if bank_id is not None or operation is not None or metadata:
|
||||
from .llm_trace import LLMTraceContext
|
||||
|
||||
# One trace + operation span per with_config() call — i.e. per
|
||||
# operation invocation. Every LLM call made through this wrapper
|
||||
# shares them, so a reflect/retain/consolidation run groups its
|
||||
# calls as parent (operation) → children (LLM calls).
|
||||
trace_ctx = LLMTraceContext(
|
||||
bank_id=bank_id,
|
||||
operation=operation,
|
||||
metadata=dict(metadata or {}),
|
||||
trace_id=str(uuid.uuid4()),
|
||||
operation_span_id=str(uuid.uuid4()),
|
||||
)
|
||||
return ConfiguredLLMProvider(self, config.llm_gemini_safety_settings, trace_ctx)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (e.g. stop llamacpp subprocess)."""
|
||||
@@ -841,12 +1117,14 @@ class LLMProvider:
|
||||
"""Create provider from environment variables using config.py constants."""
|
||||
from ..config import (
|
||||
DEFAULT_LLM_PROVIDER,
|
||||
DEFAULT_LLM_REASONING_EFFORT,
|
||||
ENV_LLM_API_KEY,
|
||||
ENV_LLM_BASE_URL,
|
||||
ENV_LLM_DEFAULT_HEADERS,
|
||||
ENV_LLM_EXTRA_BODY,
|
||||
ENV_LLM_MODEL,
|
||||
ENV_LLM_PROVIDER,
|
||||
ENV_LLM_REASONING_EFFORT,
|
||||
_get_default_model_for_provider,
|
||||
)
|
||||
|
||||
@@ -870,7 +1148,7 @@ class LLMProvider:
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort="low",
|
||||
reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
)
|
||||
@@ -891,10 +1169,16 @@ class ConfiguredLLMProvider:
|
||||
any changes.
|
||||
"""
|
||||
|
||||
def __init__(self, provider: "LLMProvider", gemini_safety_settings: list | None) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
provider: "LLMProvider",
|
||||
gemini_safety_settings: list | None,
|
||||
trace_ctx: Any | None = None,
|
||||
) -> None:
|
||||
# Use object.__setattr__ to avoid triggering __getattr__
|
||||
object.__setattr__(self, "_provider", provider)
|
||||
object.__setattr__(self, "_gemini_safety_settings", gemini_safety_settings)
|
||||
object.__setattr__(self, "_trace_ctx", trace_ctx)
|
||||
|
||||
# ── attribute passthrough ──────────────────────────────────────────────────
|
||||
|
||||
@@ -907,10 +1191,12 @@ class ConfiguredLLMProvider:
|
||||
from .providers.gemini_llm import _safety_settings_ctx
|
||||
|
||||
token = _safety_settings_ctx.set(object.__getattribute__(self, "_gemini_safety_settings"))
|
||||
trace_token = self._bind_trace_context()
|
||||
try:
|
||||
return await object.__getattribute__(self, "_provider").call(messages=messages, **kwargs)
|
||||
finally:
|
||||
_safety_settings_ctx.reset(token)
|
||||
self._reset_trace_context(trace_token)
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
@@ -921,12 +1207,38 @@ class ConfiguredLLMProvider:
|
||||
from .providers.gemini_llm import _safety_settings_ctx
|
||||
|
||||
token = _safety_settings_ctx.set(object.__getattribute__(self, "_gemini_safety_settings"))
|
||||
trace_token = self._bind_trace_context()
|
||||
try:
|
||||
return await object.__getattribute__(self, "_provider").call_with_tools(
|
||||
messages=messages, tools=tools, **kwargs
|
||||
)
|
||||
finally:
|
||||
_safety_settings_ctx.reset(token)
|
||||
self._reset_trace_context(trace_token)
|
||||
|
||||
def trace_context(self) -> Any | None:
|
||||
"""The operation-level LLM trace context (or None when untraced).
|
||||
|
||||
Lets the engine attach the operation's produced/consumed memory_ids to
|
||||
this run's trace rows once they're known (after the LLM calls).
|
||||
"""
|
||||
return object.__getattribute__(self, "_trace_ctx")
|
||||
|
||||
def _bind_trace_context(self) -> Any | None:
|
||||
"""Bind bank/operation attribution for the duration of one call."""
|
||||
trace_ctx = object.__getattribute__(self, "_trace_ctx")
|
||||
if trace_ctx is None:
|
||||
return None
|
||||
from .llm_trace import set_trace_context
|
||||
|
||||
return set_trace_context(trace_ctx)
|
||||
|
||||
def _reset_trace_context(self, trace_token: Any | None) -> None:
|
||||
if trace_token is None:
|
||||
return
|
||||
from .llm_trace import reset_trace_context
|
||||
|
||||
reset_trace_context(trace_token)
|
||||
|
||||
|
||||
# Backwards compatibility alias
|
||||
|
||||
@@ -0,0 +1,214 @@
|
||||
"""Background maintenance loop.
|
||||
|
||||
A single periodic loop that drives all of Hindsight's recurring housekeeping
|
||||
from one place, so we don't spawn a separate ``asyncio`` task per concern:
|
||||
|
||||
- **Retention sweeps** (hourly): delete ``audit_log`` and ``llm_requests`` rows
|
||||
older than their configured retention, across *all* tenant schemas.
|
||||
- **Consolidation reconcile** (configurable, default 5 min): re-schedule
|
||||
consolidation for banks that have eligible-but-unscheduled facts and no
|
||||
in-flight consolidation. This recovers facts that were stranded when a
|
||||
consolidation operation failed terminally and left them with
|
||||
``consolidated_at IS NULL AND consolidation_failed_at IS NULL`` and nothing to
|
||||
re-trigger them.
|
||||
|
||||
The loop wakes on a short fixed tick and runs each job when its own
|
||||
``last_run + interval`` is due (run-at-start, then on interval), so adding jobs
|
||||
with different cadences doesn't burst CPU. Cross-tenant discovery goes through
|
||||
server-side PL/pgSQL routines (``public.schemas_with_expired_rows`` and
|
||||
``public.banks_needing_consolidation``) — one round-trip each — instead of a
|
||||
per-schema query storm, which matters at thousands of tenants.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ..config import HindsightConfig, get_config
|
||||
from ..models import RequestContext
|
||||
from .db_utils import acquire_with_retry
|
||||
from .schema import _is_oracle
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .memory_engine import MemoryEngine
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Short tick so jobs with different cadences share one loop without per-job tasks.
|
||||
_TICK_SECONDS = 60
|
||||
# Retention sweeps are not time-sensitive; hourly matches the previous per-sweep cadence.
|
||||
_RETENTION_INTERVAL_SECONDS = 3600
|
||||
|
||||
|
||||
class MaintenanceLoop:
|
||||
"""Owns the single periodic maintenance task for a :class:`MemoryEngine`."""
|
||||
|
||||
def __init__(self, engine: "MemoryEngine") -> None:
|
||||
self._engine = engine
|
||||
self._task: asyncio.Task | None = None
|
||||
self._stop = asyncio.Event()
|
||||
# Monotonic timestamps of the last run per job, keyed by job name.
|
||||
self._last_run: dict[str, float] = {}
|
||||
|
||||
# ── lifecycle ──────────────────────────────────────────────────────────
|
||||
|
||||
def start(self) -> None:
|
||||
"""Start the loop if any maintenance job is enabled. Idempotent."""
|
||||
if self._task and not self._task.done():
|
||||
return
|
||||
# PostgreSQL-only: the retention sweeps target PG-only tables (audit_log,
|
||||
# llm_requests) and the reconcile relies on PG-only PL/pgSQL routines
|
||||
# installed by the maintenance-routines migration. Oracle support is
|
||||
# intentionally absent (mirrors that PG-only migration).
|
||||
if _is_oracle():
|
||||
logger.debug("Maintenance loop not started: PostgreSQL-only")
|
||||
return
|
||||
if not self._any_job_enabled():
|
||||
logger.debug("Maintenance loop not started: no jobs enabled")
|
||||
return
|
||||
self._stop.clear()
|
||||
try:
|
||||
self._task = asyncio.create_task(self._run())
|
||||
except RuntimeError:
|
||||
logger.debug("Cannot start maintenance loop: no running event loop")
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the loop and wait for the current tick to finish."""
|
||||
self._stop.set()
|
||||
if self._task and not self._task.done():
|
||||
try:
|
||||
await self._task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._task = None
|
||||
|
||||
@staticmethod
|
||||
def _any_job_enabled() -> bool:
|
||||
cfg = get_config()
|
||||
reconcile_on = cfg.consolidation_reconcile_interval_seconds > 0
|
||||
audit_on = cfg.audit_log_enabled and cfg.audit_log_retention_days > 0
|
||||
llm_on = cfg.llm_trace_enabled and cfg.llm_trace_retention_days > 0
|
||||
return reconcile_on or audit_on or llm_on
|
||||
|
||||
# ── loop ───────────────────────────────────────────────────────────────
|
||||
|
||||
async def _run(self) -> None:
|
||||
while not self._stop.is_set():
|
||||
try:
|
||||
await self._tick()
|
||||
except Exception:
|
||||
logger.exception("Maintenance tick failed")
|
||||
try:
|
||||
await asyncio.wait_for(self._stop.wait(), timeout=_TICK_SECONDS)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
def _is_due(self, job: str, interval_seconds: int) -> bool:
|
||||
"""True if ``job`` has never run or its interval has elapsed; marks it run now."""
|
||||
now = time.monotonic()
|
||||
last = self._last_run.get(job)
|
||||
if last is not None and (now - last) < interval_seconds:
|
||||
return False
|
||||
self._last_run[job] = now
|
||||
return True
|
||||
|
||||
async def _tick(self) -> None:
|
||||
cfg = get_config()
|
||||
if self._is_due("retention", _RETENTION_INTERVAL_SECONDS):
|
||||
await self._run_retention(cfg)
|
||||
interval = cfg.consolidation_reconcile_interval_seconds
|
||||
if interval > 0 and self._is_due("reconcile", interval):
|
||||
await self._run_reconcile()
|
||||
|
||||
# ── retention ──────────────────────────────────────────────────────────
|
||||
|
||||
async def _run_retention(self, cfg: HindsightConfig) -> None:
|
||||
# Retention days are static server-level config, so one global cutoff
|
||||
# applies to every tenant schema (the routine sweeps them all).
|
||||
if cfg.audit_log_enabled and cfg.audit_log_retention_days > 0:
|
||||
await self._purge_expired("audit_log", "started_at", cfg.audit_log_retention_days)
|
||||
if cfg.llm_trace_enabled and cfg.llm_trace_retention_days > 0:
|
||||
await self._purge_expired("llm_requests", "started_at", cfg.llm_trace_retention_days)
|
||||
|
||||
async def _purge_expired(self, table: str, ts_col: str, days: int) -> None:
|
||||
"""Delete rows older than ``days`` from ``table`` across every tenant schema."""
|
||||
backend = self._engine._backend
|
||||
try:
|
||||
async with acquire_with_retry(backend, max_retries=1) as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT * FROM public.schemas_with_expired_rows($1, $2, $3)", table, ts_col, days
|
||||
)
|
||||
for row in rows:
|
||||
schema = row[0]
|
||||
# schema names come from pg_class; quote defensively all the same.
|
||||
qschema = '"' + schema.replace('"', '""') + '"'
|
||||
result = await conn.execute(
|
||||
f"DELETE FROM {qschema}.{table} WHERE {ts_col} < NOW() - make_interval(days => $1)",
|
||||
days,
|
||||
)
|
||||
if result and result != "DELETE 0":
|
||||
logger.info(f"Retention sweep {schema}.{table}: {result}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Retention sweep failed for {table}: {e}")
|
||||
|
||||
# ── consolidation reconcile ──────────────────────────────────────────────
|
||||
|
||||
async def _run_reconcile(self) -> None:
|
||||
"""Re-schedule consolidation for banks with eligible-but-unscheduled facts."""
|
||||
engine = self._engine
|
||||
try:
|
||||
async with acquire_with_retry(engine._backend, max_retries=1) as conn:
|
||||
rows = await conn.fetch("SELECT schema_name, bank_id FROM public.banks_needing_consolidation()")
|
||||
except Exception as e:
|
||||
logger.warning(f"Consolidation reconcile discovery failed: {e}")
|
||||
return
|
||||
if not rows:
|
||||
return
|
||||
|
||||
# Only enqueue into schemas the worker actually polls (tenant discovery),
|
||||
# otherwise the op would never be claimed and would block future reconciles
|
||||
# for that bank. The tenant_id (when the extension provides one) lets
|
||||
# config resolution honor tenant-level overrides.
|
||||
try:
|
||||
tenants = await engine._tenant_extension.list_tenants()
|
||||
except Exception as e:
|
||||
logger.warning(f"Consolidation reconcile tenant discovery failed: {e}")
|
||||
return
|
||||
tenant_by_schema = {t.schema: t for t in tenants}
|
||||
default_schema = get_config().database_schema
|
||||
|
||||
from .memory_engine import _current_schema
|
||||
|
||||
submitted = 0
|
||||
skipped_unknown = 0
|
||||
for row in rows:
|
||||
schema = row["schema_name"]
|
||||
bank_id = row["bank_id"]
|
||||
tenant = tenant_by_schema.get(schema)
|
||||
if tenant is None and schema != default_schema:
|
||||
skipped_unknown += 1
|
||||
continue
|
||||
tenant_id = tenant.tenant_id if tenant else None
|
||||
token = _current_schema.set(schema)
|
||||
try:
|
||||
context = RequestContext(internal=True, tenant_id=tenant_id)
|
||||
resolved = await engine._config_resolver.resolve_full_config(bank_id, context)
|
||||
# Mirror the retain-time auto-consolidation gate (memory_engine): both
|
||||
# observations and auto-consolidation must be enabled for this bank.
|
||||
if not (resolved.enable_observations and resolved.enable_auto_consolidation):
|
||||
continue
|
||||
await engine.submit_async_consolidation(bank_id=bank_id, request_context=context)
|
||||
submitted += 1
|
||||
except Exception as e:
|
||||
logger.warning(f"Consolidation reconcile failed for bank {bank_id} in {schema}: {e}")
|
||||
finally:
|
||||
_current_schema.reset(token)
|
||||
|
||||
if submitted or skipped_unknown:
|
||||
logger.info(
|
||||
f"Consolidation reconcile: scheduled {submitted} bank(s)"
|
||||
+ (f", skipped {skipped_unknown} in unrecognized schema(s)" if skipped_unknown else "")
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -5,8 +5,10 @@ These dataclasses define the structure of result_metadata for different operatio
|
||||
The metadata is exposed in the API for debugging purposes and may change without notice.
|
||||
"""
|
||||
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from typing import Any, Mapping
|
||||
|
||||
MAX_EXTRACTION_ERROR_SAMPLES = 5
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -48,6 +50,79 @@ class RetainMetadata:
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainExtractionErrors:
|
||||
"""Non-fatal fact extraction failures observed inside one retain operation."""
|
||||
|
||||
count: int = 0
|
||||
sample: list[str] = field(default_factory=list)
|
||||
|
||||
def add(self, message: str) -> None:
|
||||
"""Record one extraction error while keeping the stored sample bounded."""
|
||||
self.count += 1
|
||||
if len(self.sample) < MAX_EXTRACTION_ERROR_SAMPLES:
|
||||
self.sample.append(message[:500])
|
||||
|
||||
def merge_metadata(self, metadata: Mapping[str, Any]) -> None:
|
||||
"""Merge errors already present on an operation result_metadata object."""
|
||||
self.count += int(metadata.get("extraction_errors_count") or 0)
|
||||
|
||||
sample = metadata.get("extraction_errors_sample") or []
|
||||
if isinstance(sample, str):
|
||||
sample = [sample]
|
||||
if isinstance(sample, list):
|
||||
for entry in sample:
|
||||
if isinstance(entry, str) and len(self.sample) < MAX_EXTRACTION_ERROR_SAMPLES:
|
||||
self.sample.append(entry[:500])
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to the public result_metadata field shape."""
|
||||
data: dict[str, Any] = {"extraction_errors_count": self.count}
|
||||
if self.sample:
|
||||
data["extraction_errors_sample"] = self.sample
|
||||
return data
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainOutcomeMetadata:
|
||||
"""Machine-readable outcome metadata for a completed retain operation."""
|
||||
|
||||
unit_ids_count: int
|
||||
extraction_errors_count: int = 0
|
||||
extraction_errors_sample: list[str] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization, omitting empty optional samples."""
|
||||
data: dict[str, Any] = {
|
||||
"unit_ids_count": self.unit_ids_count,
|
||||
"extraction_errors_count": self.extraction_errors_count,
|
||||
}
|
||||
if self.extraction_errors_sample:
|
||||
data["extraction_errors_sample"] = self.extraction_errors_sample[:MAX_EXTRACTION_ERROR_SAMPLES]
|
||||
return data
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainOutcomeAggregate:
|
||||
"""Aggregate retain outcome metadata from child retain operations."""
|
||||
|
||||
unit_ids_count: int = 0
|
||||
extraction_errors: RetainExtractionErrors = field(default_factory=RetainExtractionErrors)
|
||||
|
||||
def add_metadata(self, metadata: Mapping[str, Any]) -> None:
|
||||
"""Fold one child operation's result_metadata into the aggregate."""
|
||||
self.unit_ids_count += int(metadata.get("unit_ids_count") or 0)
|
||||
self.extraction_errors.merge_metadata(metadata)
|
||||
|
||||
def to_outcome_metadata(self) -> RetainOutcomeMetadata:
|
||||
"""Return the aggregate in the public result_metadata field shape."""
|
||||
return RetainOutcomeMetadata(
|
||||
unit_ids_count=self.unit_ids_count,
|
||||
extraction_errors_count=self.extraction_errors.count,
|
||||
extraction_errors_sample=self.extraction_errors.sample,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConsolidationMetadata:
|
||||
"""Metadata for consolidation operations."""
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
"""Shared utilities for prompt assembly."""
|
||||
|
||||
import re
|
||||
|
||||
_LONE_OPEN_BRACE = re.compile(r"(?<!\{)\{(?!\{)")
|
||||
_LONE_CLOSE_BRACE = re.compile(r"(?<!\})\}(?!\})")
|
||||
|
||||
|
||||
def escape_for_prompt(text: str) -> str:
|
||||
"""Double any lone ``{`` / ``}`` so the text survives ``str.format`` untouched.
|
||||
|
||||
Prompt templates are often passed through ``str.format`` to substitute real
|
||||
placeholders like ``{facts_text}``. Any literal braces in caller-supplied
|
||||
text — e.g. a bank mission that contains JSON examples — would otherwise be
|
||||
interpreted as format keys and raise ``KeyError``.
|
||||
|
||||
Idempotent: text that already contains escaped ``{{`` / ``}}`` pairs is
|
||||
left as-is. Only lone braces (not adjacent to another brace of the same
|
||||
kind) are doubled.
|
||||
"""
|
||||
text = _LONE_OPEN_BRACE.sub("{{", text)
|
||||
text = _LONE_CLOSE_BRACE.sub("}}", text)
|
||||
return text
|
||||
|
||||
|
||||
def output_language_directive(language: str | None) -> str:
|
||||
"""Return an LLM directive forcing all output into ``language``.
|
||||
|
||||
Used by retain (fact extraction), consolidation (observations), and reflect
|
||||
(response synthesis) so HINDSIGHT_API_LLM_OUTPUT_LANGUAGE applies uniformly
|
||||
across every LLM-generated artifact. Returns an empty string when
|
||||
``language`` is unset so the calling prompt stays unchanged.
|
||||
"""
|
||||
if not language:
|
||||
return ""
|
||||
return (
|
||||
f"\n\nIMPORTANT: Respond exclusively in {language}. "
|
||||
f"Translate any source content into {language}. "
|
||||
f"All output text — including fact text, observations, entity names, "
|
||||
f"and the final response — must be in {language}."
|
||||
)
|
||||
@@ -7,6 +7,7 @@ This package contains concrete implementations of the LLMInterface for various p
|
||||
from .anthropic_llm import AnthropicLLM
|
||||
from .claude_code_llm import ClaudeCodeLLM
|
||||
from .codex_llm import CodexLLM
|
||||
from .fireworks_llm import FireworksLLM
|
||||
from .gemini_llm import GeminiLLM
|
||||
from .litellm_llm import LiteLLMLLM
|
||||
from .litellm_router_llm import LiteLLMRouterLLM
|
||||
@@ -19,6 +20,7 @@ __all__ = [
|
||||
"AnthropicLLM",
|
||||
"ClaudeCodeLLM",
|
||||
"CodexLLM",
|
||||
"FireworksLLM",
|
||||
"GeminiLLM",
|
||||
"LlamaCppLLM",
|
||||
"LiteLLMLLM",
|
||||
|
||||
@@ -38,6 +38,7 @@ class AnthropicLLM(LLMInterface):
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float = 300.0,
|
||||
default_headers: dict[str, str] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
@@ -54,6 +55,10 @@ class AnthropicLLM(LLMInterface):
|
||||
the Anthropic SDK client. Used by operators routing through proxies
|
||||
or request-tracing middleware. Sourced from ``llm_default_headers`` in
|
||||
``HindsightConfig`` (env: ``HINDSIGHT_API_LLM_DEFAULT_HEADERS``).
|
||||
extra_body: Extra request-body params (e.g. ``{"temperature": 0.2,
|
||||
"top_p": 0.9, "top_k": 40}``) passed via the Anthropic SDK's
|
||||
``extra_body`` so they merge into the JSON sent to the Messages API.
|
||||
Sourced from ``llm_extra_body`` (env: ``HINDSIGHT_API_LLM_EXTRA_BODY``).
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
"""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
@@ -61,6 +66,9 @@ class AnthropicLLM(LLMInterface):
|
||||
if not self.api_key:
|
||||
raise ValueError("API key is required for Anthropic provider")
|
||||
|
||||
# User-configured extra body params (merged into every Messages API call)
|
||||
self._extra_body = extra_body or {}
|
||||
|
||||
# Import and initialize Anthropic client
|
||||
try:
|
||||
from anthropic import AsyncAnthropic
|
||||
@@ -93,7 +101,6 @@ class AnthropicLLM(LLMInterface):
|
||||
await self.call(
|
||||
messages=test_messages,
|
||||
max_completion_tokens=10,
|
||||
temperature=0.0,
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
@@ -179,8 +186,8 @@ class AnthropicLLM(LLMInterface):
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if temperature is not None:
|
||||
call_params["temperature"] = temperature
|
||||
if self._extra_body:
|
||||
call_params["extra_body"] = self._extra_body
|
||||
|
||||
last_exception = None
|
||||
|
||||
@@ -220,6 +227,7 @@ class AnthropicLLM(LLMInterface):
|
||||
input_tokens = response.usage.input_tokens or 0 if response.usage else 0
|
||||
output_tokens = response.usage.output_tokens or 0 if response.usage else 0
|
||||
total_tokens = input_tokens + output_tokens
|
||||
cached_tokens = getattr(response.usage, "cache_read_input_tokens", 0) or 0 if response.usage else 0
|
||||
|
||||
# Record LLM metrics
|
||||
metrics = get_metrics_collector()
|
||||
@@ -249,6 +257,7 @@ class AnthropicLLM(LLMInterface):
|
||||
duration=duration,
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
@@ -264,6 +273,7 @@ class AnthropicLLM(LLMInterface):
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=total_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
@@ -398,8 +408,8 @@ class AnthropicLLM(LLMInterface):
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if temperature is not None:
|
||||
call_params["temperature"] = temperature
|
||||
if self._extra_body:
|
||||
call_params["extra_body"] = self._extra_body
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
|
||||
@@ -9,6 +9,7 @@ automatically handles authentication via `claude auth login` credentials.
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import tempfile
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
@@ -21,6 +22,32 @@ from hindsight_api.metrics import get_metrics_collector
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Isolation env passed to the spawned `claude` CLI. CLAUDE_CONFIG_DIR
|
||||
# redirects the subprocess away from the host's ~/.claude/, so any
|
||||
# operator-installed plugins (e.g. hindsight-memory) and their Stop hooks do
|
||||
# not fire inside our LLM-call subprocesses. Without this, retain/reflect/
|
||||
# consolidation LLM calls would trigger a Stop-hook retain of the subprocess
|
||||
# transcript back into the same bank — a recursive feedback loop (issue #1751).
|
||||
# CLAUDE_SECURESTORAGE_CONFIG_DIR="" forces the CLI's keychain service name
|
||||
# back to the canonical un-suffixed entry that `claude auth login` wrote;
|
||||
# otherwise it would be namespaced by sha256(CLAUDE_CONFIG_DIR) and OAuth
|
||||
# lookup would fail. Requires bundled CLI >= 2.1.150 (claude-agent-sdk 0.2.82).
|
||||
_isolated_claude_env: dict[str, str] | None = None
|
||||
|
||||
|
||||
def _get_isolated_claude_env() -> dict[str, str]:
|
||||
"""Return a process-lifetime env dict that isolates the spawned CLI from user plugins."""
|
||||
global _isolated_claude_env
|
||||
if _isolated_claude_env is None:
|
||||
path = tempfile.mkdtemp(prefix="hindsight-claude-code-")
|
||||
_isolated_claude_env = {
|
||||
"CLAUDE_CONFIG_DIR": path,
|
||||
"CLAUDE_SECURESTORAGE_CONFIG_DIR": "",
|
||||
}
|
||||
logger.debug(f"Claude Code: isolated CLAUDE_CONFIG_DIR={path}")
|
||||
return _isolated_claude_env
|
||||
|
||||
|
||||
class ClaudeCodeLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using Claude Code authentication.
|
||||
@@ -183,6 +210,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
system_prompt=system_prompt if system_prompt else None,
|
||||
max_turns=1, # Single-turn for API-style interactions
|
||||
allowed_tools=[], # Disable tools for standard LLM calls
|
||||
env=_get_isolated_claude_env(),
|
||||
)
|
||||
|
||||
# Call Claude Agent SDK
|
||||
@@ -473,6 +501,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
max_turns=2, # Allow tool call + tool result round-trip
|
||||
mcp_servers=mcp_servers_config,
|
||||
allowed_tools=allowed_tool_names,
|
||||
env=_get_isolated_claude_env(),
|
||||
)
|
||||
|
||||
# Call Claude Agent SDK with retry logic
|
||||
|
||||
@@ -0,0 +1,409 @@
|
||||
"""
|
||||
Shared Codex OAuth authentication manager.
|
||||
|
||||
Extracted from ``CodexLLM`` so that both ``CodexLLM`` and
|
||||
``CodexOAuthEmbeddings`` can share JWT-expiry detection, single-flight
|
||||
token refresh, and atomic file persistence without duplicating the logic.
|
||||
|
||||
Usage
|
||||
-----
|
||||
Create a manager from the auth file::
|
||||
|
||||
mgr = CodexAuthManager.from_file()
|
||||
|
||||
Then call ``ensure_fresh_token()`` before each outbound request and
|
||||
``refresh_tokens(reason=..., force=...)`` on a reactive 401.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Module-level constants (shared with codex_llm.py via re-export there)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# OAuth refresh endpoint and client id, mirrored from the canonical
|
||||
# ``@openai/codex`` CLI (codex-rs/login/src/auth/manager.rs on
|
||||
# github.com/openai/codex). The endpoint is overridable via env var so that
|
||||
# future Codex changes or staging environments can be pointed at without a
|
||||
# code change — same env var name the upstream CLI uses.
|
||||
_CODEX_REFRESH_TOKEN_URL = os.environ.get("CODEX_REFRESH_TOKEN_URL_OVERRIDE", "https://auth.openai.com/oauth/token")
|
||||
_CODEX_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
|
||||
|
||||
# Proactively refresh this many seconds before the JWT ``exp`` claim. The
|
||||
# upstream Codex CLI uses no skew (it refreshes at ``exp <= now``); the
|
||||
# extra window reduces races where a request leaves the client with a token
|
||||
# that the server has already declared expired by the time it arrives.
|
||||
_CODEX_TOKEN_REFRESH_SKEW_SECONDS = 60
|
||||
|
||||
# OAuth error codes that the refresh endpoint returns when the refresh_token
|
||||
# itself is no longer usable. These are terminal — retrying refresh will not
|
||||
# succeed; the user must re-run ``codex auth login``.
|
||||
_CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
|
||||
{"refresh_token_expired", "refresh_token_reused", "refresh_token_invalidated"}
|
||||
)
|
||||
|
||||
|
||||
class CodexRefreshExpiredError(RuntimeError):
|
||||
"""Raised when the Codex refresh_token itself is no longer valid.
|
||||
|
||||
The user must re-run ``codex auth login`` to obtain new credentials.
|
||||
Callers should surface a clear remediation message and stop retrying.
|
||||
"""
|
||||
|
||||
|
||||
class CodexAuthManager:
|
||||
"""Sync Codex OAuth credential manager.
|
||||
|
||||
Holds the access_token, refresh_token, and account_id in memory and
|
||||
handles proactive/reactive refresh using a ``threading.Lock`` for
|
||||
single-flight semantics (safe to use from multiple threads or via
|
||||
``asyncio.to_thread``).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
access_token:
|
||||
The current bearer token.
|
||||
account_id:
|
||||
The OpenAI account ID embedded in the Codex request headers.
|
||||
refresh_token:
|
||||
The OAuth refresh token. May be ``None`` when the auth file omits it;
|
||||
the provider still works as a one-shot loader in that case.
|
||||
auth_file:
|
||||
Path to ``~/.codex/auth.json``. Used for re-reading the refresh token
|
||||
on demand and for atomic persistence of rotated credentials.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
access_token: str,
|
||||
account_id: str,
|
||||
refresh_token: str | None,
|
||||
auth_file: Path,
|
||||
) -> None:
|
||||
self.access_token = access_token
|
||||
self.account_id = account_id
|
||||
self.refresh_token = refresh_token
|
||||
self._auth_file = auth_file
|
||||
self._lock = threading.Lock()
|
||||
self._http_client = httpx.Client(timeout=30.0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Construction helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@classmethod
|
||||
def from_file(cls, auth_file: Path | None = None) -> "CodexAuthManager":
|
||||
"""Build a manager by reading credentials from ``auth_file``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
auth_file:
|
||||
Defaults to ``~/.codex/auth.json``.
|
||||
|
||||
Raises
|
||||
------
|
||||
FileNotFoundError:
|
||||
If the auth file does not exist.
|
||||
ValueError:
|
||||
If the auth file is missing ``access_token`` or has an unexpected
|
||||
``auth_mode``.
|
||||
"""
|
||||
if auth_file is None:
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(f"Codex auth file not found: {auth_file}. Run 'codex auth login' to authenticate.")
|
||||
|
||||
with open(auth_file) as f:
|
||||
data = json.load(f)
|
||||
|
||||
auth_mode = data.get("auth_mode")
|
||||
if auth_mode != "chatgpt":
|
||||
raise ValueError(f"Expected Codex auth_mode='chatgpt', got: {auth_mode}")
|
||||
|
||||
tokens = data.get("tokens") or {}
|
||||
access_token = tokens.get("access_token")
|
||||
if not access_token:
|
||||
raise ValueError("No access_token found in Codex auth file. Run 'codex auth login' again.")
|
||||
|
||||
account_id = tokens.get("account_id") or ""
|
||||
refresh_token = tokens.get("refresh_token")
|
||||
|
||||
return cls(
|
||||
access_token=access_token,
|
||||
account_id=account_id,
|
||||
refresh_token=refresh_token,
|
||||
auth_file=auth_file,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Token state helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def load_refresh_token_from_file(auth_file: Path) -> str | None:
|
||||
"""Read ``tokens.refresh_token`` from ``auth_file``.
|
||||
|
||||
Returns ``None`` when the file is unreadable or omits the field.
|
||||
Does not raise — the provider degrades to one-shot mode.
|
||||
"""
|
||||
try:
|
||||
with open(auth_file) as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError) as e:
|
||||
logger.warning(
|
||||
f"Codex auth file unreadable when loading refresh_token: {type(e).__name__}. "
|
||||
"Token refresh will not be available; the access_token in memory will be used until it expires."
|
||||
)
|
||||
return None
|
||||
return data.get("tokens", {}).get("refresh_token")
|
||||
|
||||
@staticmethod
|
||||
def _decode_jwt_exp_unixtime(token: str) -> int | None:
|
||||
"""Return the JWT ``exp`` claim as a unix timestamp, or None on parse failure.
|
||||
|
||||
We do not verify the signature — the server is the source of truth
|
||||
on whether the token is actually accepted. This is only used to
|
||||
schedule proactive refresh.
|
||||
"""
|
||||
try:
|
||||
parts = token.split(".")
|
||||
if len(parts) < 2:
|
||||
return None
|
||||
payload_b64 = parts[1]
|
||||
padding = "=" * (-len(payload_b64) % 4)
|
||||
payload_bytes = base64.urlsafe_b64decode(payload_b64 + padding)
|
||||
payload = json.loads(payload_bytes.decode("utf-8"))
|
||||
exp = payload.get("exp")
|
||||
return int(exp) if exp is not None else None
|
||||
except (ValueError, TypeError, json.JSONDecodeError, binascii.Error):
|
||||
return None
|
||||
|
||||
def _token_is_stale(self, skew_seconds: int = _CODEX_TOKEN_REFRESH_SKEW_SECONDS) -> bool:
|
||||
"""True when the cached access_token is past expiry (with skew).
|
||||
|
||||
Returns False when expiry cannot be determined — we'd rather use a
|
||||
possibly-expired token and recover via the reactive 401 path than
|
||||
refresh aggressively on every request when ``exp`` parsing fails.
|
||||
"""
|
||||
exp = self._decode_jwt_exp_unixtime(self.access_token)
|
||||
if exp is None:
|
||||
return False
|
||||
return exp <= int(time.time()) + skew_seconds
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Persistence
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _persist_auth_atomic(self, updated_tokens: dict[str, Any]) -> None:
|
||||
"""Write rotated tokens back to ``_auth_file`` atomically.
|
||||
|
||||
Re-reads the on-disk file first to avoid clobbering fields written
|
||||
by another process, patches ``tokens.*`` and ``last_refresh``, then
|
||||
writes to a sibling tempfile and calls ``os.replace`` (atomic on
|
||||
POSIX and Windows within the same filesystem).
|
||||
"""
|
||||
current: dict[str, Any]
|
||||
try:
|
||||
with open(self._auth_file) as f:
|
||||
loaded = json.load(f)
|
||||
current = loaded if isinstance(loaded, dict) else {"auth_mode": "chatgpt", "tokens": {}}
|
||||
except (OSError, json.JSONDecodeError):
|
||||
current = {"auth_mode": "chatgpt", "tokens": {}}
|
||||
|
||||
existing_tokens = current.get("tokens")
|
||||
tokens: dict[str, Any] = existing_tokens if isinstance(existing_tokens, dict) else {}
|
||||
for key in ("access_token", "refresh_token", "id_token", "account_id"):
|
||||
if key in updated_tokens and updated_tokens[key] is not None:
|
||||
tokens[key] = updated_tokens[key]
|
||||
current["tokens"] = tokens
|
||||
current["last_refresh"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
||||
|
||||
parent = self._auth_file.parent
|
||||
parent.mkdir(parents=True, exist_ok=True)
|
||||
fd, tmp_path = tempfile.mkstemp(prefix=".auth.", suffix=".json.tmp", dir=str(parent))
|
||||
try:
|
||||
with os.fdopen(fd, "w") as f:
|
||||
json.dump(current, f, indent=2)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
try:
|
||||
os.chmod(tmp_path, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
os.replace(tmp_path, self._auth_file)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Error extraction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _extract_oauth_error_code(response: httpx.Response) -> str | None:
|
||||
"""Pull the OAuth error code out of a 4xx response body, if present.
|
||||
|
||||
The refresh endpoint returns shapes like
|
||||
``{"error": "...", "error_code": "..."}`` or
|
||||
``{"error": {"code": "..."}}``.
|
||||
"""
|
||||
try:
|
||||
body = response.json()
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return None
|
||||
if not isinstance(body, dict):
|
||||
return None
|
||||
err = body.get("error")
|
||||
if isinstance(err, dict):
|
||||
code = err.get("code")
|
||||
if isinstance(code, str):
|
||||
return code
|
||||
code = body.get("error_code")
|
||||
if isinstance(code, str):
|
||||
return code
|
||||
if isinstance(err, str):
|
||||
return err
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Refresh
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def refresh_tokens(self, reason: str = "", *, force: bool = False) -> None:
|
||||
"""Synchronous single-flight OAuth token refresh.
|
||||
|
||||
Serialized through ``self._lock`` so concurrent threads produce one
|
||||
network request. The first caller refreshes; the rest wake up and
|
||||
skip if the token is no longer stale (proactive) or if the token
|
||||
has already changed (reactive / force).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
reason:
|
||||
Free-form string included in log lines for diagnostics.
|
||||
force:
|
||||
When True, refresh even if the JWT exp claim looks fresh.
|
||||
Used by the reactive 401 path.
|
||||
|
||||
Raises
|
||||
------
|
||||
CodexRefreshExpiredError:
|
||||
When the server returns a terminal error code or any 401.
|
||||
RuntimeError:
|
||||
For other refresh failures (network, 5xx, etc.).
|
||||
"""
|
||||
token_before_lock = self.access_token
|
||||
with self._lock:
|
||||
if force:
|
||||
if self.access_token != token_before_lock:
|
||||
return
|
||||
else:
|
||||
if not self._token_is_stale():
|
||||
return
|
||||
|
||||
if not self.refresh_token:
|
||||
raise RuntimeError(
|
||||
"Codex access_token is expired but no refresh_token is available. "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
log_reason = f" ({reason})" if reason else ""
|
||||
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
|
||||
|
||||
request_body = {
|
||||
"client_id": _CODEX_CLIENT_ID,
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": self.refresh_token,
|
||||
}
|
||||
try:
|
||||
response = self._http_client.post(
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
json=request_body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=30.0,
|
||||
)
|
||||
except httpx.RequestError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
|
||||
|
||||
if response.status_code == 401:
|
||||
error_code = self._extract_oauth_error_code(response)
|
||||
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex OAuth refresh returned 401 with unrecognized error code "
|
||||
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
if response.status_code >= 400:
|
||||
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
|
||||
|
||||
try:
|
||||
body = response.json()
|
||||
except json.JSONDecodeError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
|
||||
|
||||
new_access = body.get("access_token")
|
||||
if not new_access:
|
||||
raise RuntimeError("Codex OAuth refresh returned no access_token")
|
||||
|
||||
new_refresh = body.get("refresh_token") or self.refresh_token
|
||||
new_id_token = body.get("id_token")
|
||||
|
||||
# Update in-memory state first so waiters see fresh credentials
|
||||
# immediately, even if disk write fails.
|
||||
self.access_token = new_access
|
||||
self.refresh_token = new_refresh
|
||||
|
||||
persisted: dict[str, Any] = {
|
||||
"access_token": new_access,
|
||||
"refresh_token": new_refresh,
|
||||
}
|
||||
if new_id_token:
|
||||
persisted["id_token"] = new_id_token
|
||||
|
||||
try:
|
||||
self._persist_auth_atomic(persisted)
|
||||
except OSError as e:
|
||||
logger.warning(
|
||||
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
|
||||
"In-memory credentials are up to date; on-disk file is stale."
|
||||
)
|
||||
|
||||
logger.info("Codex OAuth access_token refreshed successfully")
|
||||
|
||||
def ensure_fresh_token(self) -> None:
|
||||
"""Proactively refresh the access_token if it is near or past expiry.
|
||||
|
||||
Cheap when the token is fresh (just decodes the JWT exp claim and
|
||||
returns).
|
||||
"""
|
||||
if self._token_is_stale():
|
||||
self.refresh_tokens(reason="proactive (token near expiry)")
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the underlying HTTP client."""
|
||||
self._http_client.close()
|
||||
@@ -15,15 +15,10 @@ so that future server-side changes affect both clients identically.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import binascii
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -33,37 +28,27 @@ from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# OAuth refresh endpoint and client id, mirrored from the canonical
|
||||
# ``@openai/codex`` CLI (codex-rs/login/src/auth/manager.rs on
|
||||
# github.com/openai/codex). The endpoint is overridable via env var so that
|
||||
# future Codex changes or staging environments can be pointed at without a
|
||||
# code change — same env var name the upstream CLI uses.
|
||||
_CODEX_REFRESH_TOKEN_URL = os.environ.get("CODEX_REFRESH_TOKEN_URL_OVERRIDE", "https://auth.openai.com/oauth/token")
|
||||
_CODEX_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
|
||||
|
||||
# Proactively refresh this many seconds before the JWT ``exp`` claim. The
|
||||
# upstream Codex CLI uses no skew (it refreshes at ``exp <= now``); the
|
||||
# extra window reduces races where a request leaves the client with a token
|
||||
# that the server has already declared expired by the time it arrives.
|
||||
_CODEX_TOKEN_REFRESH_SKEW_SECONDS = 60
|
||||
|
||||
# OAuth error codes that the refresh endpoint returns when the refresh_token
|
||||
# itself is no longer usable. These are terminal — retrying refresh will not
|
||||
# succeed; the user must re-run ``codex auth login``.
|
||||
_CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
|
||||
{"refresh_token_expired", "refresh_token_reused", "refresh_token_invalidated"}
|
||||
from .codex_auth import (
|
||||
_CODEX_CLIENT_ID,
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
_CODEX_TERMINAL_REFRESH_ERROR_CODES,
|
||||
_CODEX_TOKEN_REFRESH_SKEW_SECONDS,
|
||||
CodexAuthManager,
|
||||
CodexRefreshExpiredError,
|
||||
)
|
||||
|
||||
# Re-export for backward compatibility (tests import from this module).
|
||||
__all__ = [
|
||||
"CodexLLM",
|
||||
"CodexRefreshExpiredError",
|
||||
"CodexAuthManager",
|
||||
"_CODEX_REFRESH_TOKEN_URL",
|
||||
"_CODEX_CLIENT_ID",
|
||||
"_CODEX_TOKEN_REFRESH_SKEW_SECONDS",
|
||||
"_CODEX_TERMINAL_REFRESH_ERROR_CODES",
|
||||
]
|
||||
|
||||
class CodexRefreshExpiredError(RuntimeError):
|
||||
"""Raised when the Codex refresh_token itself is no longer valid.
|
||||
|
||||
The user must re-run ``codex auth login`` to obtain new credentials.
|
||||
Callers should surface a clear remediation message and stop retrying.
|
||||
"""
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CodexLLM(LLMInterface):
|
||||
@@ -86,20 +71,15 @@ class CodexLLM(LLMInterface):
|
||||
"""Initialize Codex LLM provider."""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
|
||||
# Path is fixed at ~/.codex/auth.json — matches the upstream CLI.
|
||||
# Storing it on self lets the refresh path re-read after another
|
||||
# process (e.g. a sidecar) rotates the file out from under us.
|
||||
self._auth_file = Path.home() / ".codex" / "auth.json"
|
||||
|
||||
# Single-flight refresh lock. Multiple concurrent requests racing
|
||||
# toward an expired token should produce one network refresh, not N.
|
||||
# Single-flight async refresh lock. Multiple concurrent coroutines
|
||||
# racing toward an expired token should produce one network refresh.
|
||||
self._auth_lock = asyncio.Lock()
|
||||
|
||||
# Load Codex OAuth credentials
|
||||
# Load Codex OAuth credentials (keep these methods for test patching).
|
||||
try:
|
||||
self.access_token, self.account_id = self._load_codex_auth()
|
||||
self.refresh_token = self._load_codex_refresh_token()
|
||||
logger.info(f"Loaded Codex OAuth credentials for account: {self.account_id}")
|
||||
access_token, account_id = self._load_codex_auth()
|
||||
refresh_token = self._load_codex_refresh_token()
|
||||
logger.info(f"Loaded Codex OAuth credentials for account: {account_id}")
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to load Codex OAuth credentials from ~/.codex/auth.json: {e}\n\n"
|
||||
@@ -110,9 +90,22 @@ class CodexLLM(LLMInterface):
|
||||
"Or use a different provider (openai, anthropic, gemini) with API keys."
|
||||
) from e
|
||||
|
||||
# Use ChatGPT backend API endpoint
|
||||
if not self.base_url:
|
||||
self._auth_manager = CodexAuthManager(
|
||||
access_token=access_token,
|
||||
account_id=account_id,
|
||||
refresh_token=refresh_token,
|
||||
auth_file=Path.home() / ".codex" / "auth.json",
|
||||
)
|
||||
|
||||
# Use ChatGPT backend API endpoint. Codex auth is tied to
|
||||
# chatgpt.com/backend-api, not the OpenAI-compatible base URL used by
|
||||
# other providers. Deployments often set a global LLM_BASE_URL for an
|
||||
# OpenAI-compatible proxy; ignore that inherited value unless the user
|
||||
# explicitly provides a Codex backend URL.
|
||||
if not self.base_url or self.base_url.rstrip("/").endswith("/v1"):
|
||||
self.base_url = "https://chatgpt.com/backend-api"
|
||||
else:
|
||||
self.base_url = self.base_url.rstrip("/")
|
||||
|
||||
# Normalize model name (strip openai/ prefix if present)
|
||||
if self.model.startswith("openai/"):
|
||||
@@ -125,6 +118,42 @@ class CodexLLM(LLMInterface):
|
||||
# HTTP client for SSE streaming
|
||||
self._client = httpx.AsyncClient(timeout=120.0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Properties — delegate to _auth_manager (preserves test-visible API)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def access_token(self) -> str:
|
||||
return self._auth_manager.access_token
|
||||
|
||||
@access_token.setter
|
||||
def access_token(self, v: str) -> None:
|
||||
self._auth_manager.access_token = v
|
||||
|
||||
@property
|
||||
def account_id(self) -> str:
|
||||
return self._auth_manager.account_id
|
||||
|
||||
@property
|
||||
def refresh_token(self) -> str | None:
|
||||
return self._auth_manager.refresh_token
|
||||
|
||||
@refresh_token.setter
|
||||
def refresh_token(self, v: str | None) -> None:
|
||||
self._auth_manager.refresh_token = v
|
||||
|
||||
@property
|
||||
def _auth_file(self) -> Path:
|
||||
return self._auth_manager._auth_file
|
||||
|
||||
@_auth_file.setter
|
||||
def _auth_file(self, v: Path) -> None:
|
||||
self._auth_manager._auth_file = v
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Forwarding methods (keep surface area for tests / subclasses)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _load_codex_auth(self) -> tuple[str, str]:
|
||||
"""
|
||||
Load OAuth credentials from ~/.codex/auth.json.
|
||||
@@ -161,273 +190,57 @@ class CodexLLM(LLMInterface):
|
||||
return access_token, account_id
|
||||
|
||||
def _load_codex_refresh_token(self) -> str | None:
|
||||
"""Load ``tokens.refresh_token`` from ``~/.codex/auth.json``.
|
||||
"""Read ``tokens.refresh_token`` from the configured auth file.
|
||||
|
||||
Returns None when the auth file is unreadable or omits the field —
|
||||
the provider still functions as a one-shot loader in that case, it
|
||||
just can't refresh when the access_token expires. This deliberately
|
||||
does not raise so that ``__init__`` keeps the existing failure mode
|
||||
of raising only on missing ``access_token``.
|
||||
Kept as an instance method so existing tests that patch
|
||||
``CodexLLM._load_codex_refresh_token`` continue to work. Works both
|
||||
pre- and post-``__init__`` because it does not depend on
|
||||
``_auth_manager`` being constructed yet.
|
||||
"""
|
||||
try:
|
||||
with open(self._auth_file) as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError) as e:
|
||||
logger.warning(
|
||||
f"Codex auth file unreadable when loading refresh_token: {type(e).__name__}. "
|
||||
"Token refresh will not be available; the access_token in memory will be used until it expires."
|
||||
)
|
||||
return None
|
||||
return data.get("tokens", {}).get("refresh_token")
|
||||
auth_file = (
|
||||
self._auth_manager._auth_file if hasattr(self, "_auth_manager") else Path.home() / ".codex" / "auth.json"
|
||||
)
|
||||
return CodexAuthManager.load_refresh_token_from_file(auth_file)
|
||||
|
||||
@staticmethod
|
||||
def _decode_jwt_exp_unixtime(token: str) -> int | None:
|
||||
"""Return the JWT ``exp`` claim as a unix timestamp, or None on parse failure.
|
||||
|
||||
ChatGPT/Codex access_tokens are JWTs whose payload includes ``exp``
|
||||
(RFC 7519). We need the expiry to schedule proactive refresh — the
|
||||
``auth.json`` file does not persist a separate ``expires_at`` field
|
||||
in the upstream CLI's shape, so decoding the JWT itself is the
|
||||
canonical way to know when the token is stale.
|
||||
|
||||
We do not verify the signature — the server is the source of truth
|
||||
on whether the token is actually accepted, and the only thing this
|
||||
method affects is the *timing* of refresh, not whether to trust the
|
||||
token contents.
|
||||
"""
|
||||
try:
|
||||
parts = token.split(".")
|
||||
if len(parts) < 2:
|
||||
return None
|
||||
payload_b64 = parts[1]
|
||||
# JWT uses base64url without padding. Re-pad before decoding.
|
||||
padding = "=" * (-len(payload_b64) % 4)
|
||||
payload_bytes = base64.urlsafe_b64decode(payload_b64 + padding)
|
||||
payload = json.loads(payload_bytes.decode("utf-8"))
|
||||
exp = payload.get("exp")
|
||||
return int(exp) if exp is not None else None
|
||||
except (ValueError, TypeError, json.JSONDecodeError, binascii.Error):
|
||||
return None
|
||||
"""Delegate to ``CodexAuthManager._decode_jwt_exp_unixtime``."""
|
||||
return CodexAuthManager._decode_jwt_exp_unixtime(token)
|
||||
|
||||
def _token_is_stale(self, skew_seconds: int = _CODEX_TOKEN_REFRESH_SKEW_SECONDS) -> bool:
|
||||
"""True when the cached access_token is past expiry (with skew).
|
||||
|
||||
Returns False when expiry cannot be determined — we'd rather use a
|
||||
possibly-expired token and recover via the reactive 401 path than
|
||||
refresh aggressively on every request when ``exp`` parsing fails.
|
||||
"""
|
||||
exp = self._decode_jwt_exp_unixtime(self.access_token)
|
||||
if exp is None:
|
||||
return False
|
||||
return exp <= int(time.time()) + skew_seconds
|
||||
"""Delegate to ``_auth_manager._token_is_stale``."""
|
||||
return self._auth_manager._token_is_stale(skew_seconds)
|
||||
|
||||
def _persist_auth_atomic(self, updated_tokens: dict[str, Any]) -> None:
|
||||
"""Write the rotated tokens back to ``~/.codex/auth.json`` atomically.
|
||||
|
||||
Strategy: re-read the on-disk auth.json (so we don't clobber fields
|
||||
another process may have added), patch ``tokens.*`` and
|
||||
``last_refresh``, write to a tempfile in the same directory with
|
||||
mode 0600, then ``os.replace`` onto the target. ``os.replace`` is
|
||||
atomic within the same filesystem on POSIX and Windows, so a
|
||||
concurrent reader will see either the old file or the fully-written
|
||||
new file — never a partial truncate, which is the upstream CLI's
|
||||
worst-case race.
|
||||
|
||||
On non-Unix platforms the chmod is a best-effort no-op; the parent
|
||||
directory permissions still bound access.
|
||||
"""
|
||||
current: dict[str, Any]
|
||||
try:
|
||||
with open(self._auth_file) as f:
|
||||
loaded = json.load(f)
|
||||
# auth.json should always be a JSON object at the top level; if
|
||||
# someone has hand-edited it into a non-object shape, fall back
|
||||
# to the minimal default rather than crashing the refresh path.
|
||||
current = loaded if isinstance(loaded, dict) else {"auth_mode": "chatgpt", "tokens": {}}
|
||||
except (OSError, json.JSONDecodeError):
|
||||
# If the file became unreadable between our last read and now,
|
||||
# construct a minimal shape rather than refusing to persist.
|
||||
current = {"auth_mode": "chatgpt", "tokens": {}}
|
||||
|
||||
existing_tokens = current.get("tokens")
|
||||
tokens: dict[str, Any] = existing_tokens if isinstance(existing_tokens, dict) else {}
|
||||
for key in ("access_token", "refresh_token", "id_token", "account_id"):
|
||||
if key in updated_tokens and updated_tokens[key] is not None:
|
||||
tokens[key] = updated_tokens[key]
|
||||
current["tokens"] = tokens
|
||||
current["last_refresh"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
||||
|
||||
# Write to a sibling tempfile so the rename is same-filesystem.
|
||||
parent = self._auth_file.parent
|
||||
parent.mkdir(parents=True, exist_ok=True)
|
||||
fd, tmp_path = tempfile.mkstemp(prefix=".auth.", suffix=".json.tmp", dir=str(parent))
|
||||
try:
|
||||
with os.fdopen(fd, "w") as f:
|
||||
json.dump(current, f, indent=2)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
try:
|
||||
os.chmod(tmp_path, 0o600)
|
||||
except OSError:
|
||||
pass # best-effort on platforms that don't support chmod
|
||||
os.replace(tmp_path, self._auth_file)
|
||||
except Exception:
|
||||
# Clean up the orphaned tempfile if rename fails.
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
"""Delegate to ``_auth_manager._persist_auth_atomic``."""
|
||||
return self._auth_manager._persist_auth_atomic(updated_tokens)
|
||||
|
||||
async def _refresh_oauth_tokens(self, reason: str = "", *, force: bool = False) -> None:
|
||||
"""Refresh the OAuth access_token using the stored refresh_token.
|
||||
"""Async single-flight OAuth token refresh.
|
||||
|
||||
Single-flight: serialized through ``self._auth_lock`` so concurrent
|
||||
callers produce one network request. The first caller refreshes; the
|
||||
rest wake up and observe that either (a) the in-memory token is no
|
||||
longer stale (proactive case) or (b) the in-memory token has changed
|
||||
since they entered (reactive case), and return without re-refreshing.
|
||||
Outer asyncio.Lock preserves single-flight semantics for concurrent
|
||||
coroutines; the actual network call is offloaded to a thread via
|
||||
``asyncio.to_thread`` so the event loop stays unblocked.
|
||||
|
||||
Args:
|
||||
reason: Free-form string included in log lines for diagnostics.
|
||||
force: When True, refresh even if the JWT exp claim looks fresh.
|
||||
Used by the reactive 401 path — the server rejected the
|
||||
token, so we cannot trust the JWT's self-reported expiry.
|
||||
Used by the reactive 401 path.
|
||||
|
||||
Raises:
|
||||
CodexRefreshExpiredError: when the server returns a terminal
|
||||
error code (refresh_token_expired/reused/invalidated) or any
|
||||
401 on the refresh endpoint itself.
|
||||
error code or any 401 on the refresh endpoint.
|
||||
RuntimeError: for other refresh failures (network, 5xx, etc.).
|
||||
"""
|
||||
# Capture the token we'd be refreshing BEFORE acquiring the lock so
|
||||
# that we can detect mid-wait rotation by another coroutine.
|
||||
token_before_lock = self.access_token
|
||||
async with self._auth_lock:
|
||||
if force:
|
||||
# Reactive: skip only if another coroutine already rotated
|
||||
# the token while we were waiting on the lock.
|
||||
if self.access_token != token_before_lock:
|
||||
return
|
||||
else:
|
||||
# Proactive: skip if the token is no longer stale (the
|
||||
# canonical "another coroutine refreshed first" check).
|
||||
if not self._token_is_stale():
|
||||
if not self._auth_manager._token_is_stale():
|
||||
return
|
||||
|
||||
if not self.refresh_token:
|
||||
raise RuntimeError(
|
||||
"Codex access_token is expired but no refresh_token is available. "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
log_reason = f" ({reason})" if reason else ""
|
||||
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
|
||||
|
||||
request_body = {
|
||||
"client_id": _CODEX_CLIENT_ID,
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": self.refresh_token,
|
||||
}
|
||||
try:
|
||||
response = await self._client.post(
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
json=request_body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=30.0,
|
||||
)
|
||||
except httpx.RequestError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
|
||||
|
||||
if response.status_code == 401:
|
||||
# Classify by ``error.code`` (or top-level ``error`` string) — same
|
||||
# mapping as the upstream Rust CLI's request_chatgpt_token_refresh.
|
||||
error_code = self._extract_oauth_error_code(response)
|
||||
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
# Unknown 401 — treat as terminal too, matching the upstream classification.
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex OAuth refresh returned 401 with unrecognized error code "
|
||||
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
if response.status_code >= 400:
|
||||
# 5xx and other 4xx are transient/retryable from the caller's
|
||||
# perspective; surface as RuntimeError without leaking the
|
||||
# request body in logs.
|
||||
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
|
||||
|
||||
try:
|
||||
body = response.json()
|
||||
except json.JSONDecodeError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
|
||||
|
||||
new_access = body.get("access_token")
|
||||
if not new_access:
|
||||
raise RuntimeError("Codex OAuth refresh returned no access_token")
|
||||
|
||||
# The refresh_token may rotate on each refresh — adopt the new
|
||||
# one if the server sent it, otherwise keep the existing.
|
||||
new_refresh = body.get("refresh_token") or self.refresh_token
|
||||
new_id_token = body.get("id_token")
|
||||
|
||||
# Update in-memory state first so callers waiting on the lock
|
||||
# see fresh credentials immediately, even if disk write fails.
|
||||
self.access_token = new_access
|
||||
self.refresh_token = new_refresh
|
||||
|
||||
persisted = {
|
||||
"access_token": new_access,
|
||||
"refresh_token": new_refresh,
|
||||
}
|
||||
if new_id_token:
|
||||
persisted["id_token"] = new_id_token
|
||||
|
||||
try:
|
||||
self._persist_auth_atomic(persisted)
|
||||
except OSError as e:
|
||||
# In-memory creds are valid; warn but don't fail the request
|
||||
# path. Future process starts will fall back to the stale
|
||||
# on-disk auth.json and immediately refresh.
|
||||
logger.warning(
|
||||
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
|
||||
"In-memory credentials are up to date; on-disk file is stale."
|
||||
)
|
||||
|
||||
logger.info("Codex OAuth access_token refreshed successfully")
|
||||
|
||||
@staticmethod
|
||||
def _extract_oauth_error_code(response: "httpx.Response") -> str | None:
|
||||
"""Pull the OAuth error code out of a 4xx response body, if present.
|
||||
|
||||
The refresh endpoint returns shapes like
|
||||
``{"error": "...", "error_code": "..."}`` or
|
||||
``{"error": {"code": "..."}}``. We don't fail the call if the body
|
||||
is unparseable — the caller falls back to a generic "unknown" error.
|
||||
"""
|
||||
try:
|
||||
body = response.json()
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return None
|
||||
if not isinstance(body, dict):
|
||||
return None
|
||||
# Shape 1: error is a nested object with "code"
|
||||
err = body.get("error")
|
||||
if isinstance(err, dict):
|
||||
code = err.get("code")
|
||||
if isinstance(code, str):
|
||||
return code
|
||||
# Shape 2: top-level error_code string
|
||||
code = body.get("error_code")
|
||||
if isinstance(code, str):
|
||||
return code
|
||||
# Shape 3: error is itself a string code
|
||||
if isinstance(err, str):
|
||||
return err
|
||||
return None
|
||||
await asyncio.to_thread(lambda: self._auth_manager.refresh_tokens(reason, force=force))
|
||||
|
||||
async def _ensure_fresh_token(self) -> None:
|
||||
"""Refresh the access_token proactively if it is near or past expiry.
|
||||
@@ -435,13 +248,10 @@ class CodexLLM(LLMInterface):
|
||||
Called at the top of every API-bound method. Cheap when the token is
|
||||
fresh (just decodes the JWT exp claim and returns).
|
||||
"""
|
||||
if self._token_is_stale():
|
||||
if self._auth_manager._token_is_stale():
|
||||
try:
|
||||
await self._refresh_oauth_tokens(reason="proactive (token near expiry)")
|
||||
except CodexRefreshExpiredError:
|
||||
# Surface to the caller as the same RuntimeError shape the
|
||||
# request loop has historically raised, so existing error
|
||||
# handling paths keep working.
|
||||
raise
|
||||
|
||||
def _map_reasoning_effort(self, effort: str) -> str:
|
||||
@@ -1073,5 +883,6 @@ class CodexLLM(LLMInterface):
|
||||
return content if content else None, tool_calls
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up HTTP client."""
|
||||
"""Clean up HTTP clients."""
|
||||
await self._client.aclose()
|
||||
self._auth_manager.close()
|
||||
|
||||
@@ -0,0 +1,396 @@
|
||||
"""Fireworks AI provider with batch-inference support.
|
||||
|
||||
Fireworks' *online* inference endpoint (``/inference/v1``) is OpenAI-compatible,
|
||||
so ``FireworksLLM`` subclasses :class:`OpenAICompatibleLLM` and reuses its entire
|
||||
chat path. Only the *batch* mechanism differs: Fireworks does NOT implement the
|
||||
OpenAI ``/v1/batches`` API. Instead it exposes a proprietary, account-scoped
|
||||
dataset -> job -> download REST workflow on a separate control-plane host. This
|
||||
class overrides only the four batch members of the interface, translating that
|
||||
workflow to/from the OpenAI-batch shapes the retain orchestrator and
|
||||
``fact_extraction`` consumer expect — so nothing downstream changes.
|
||||
|
||||
Interface contract preserved (see ``fact_extraction.py`` result handling)::
|
||||
|
||||
result["response"]["body"]["choices"][0]["message"]["content"]
|
||||
|
||||
Workflow (control-plane host, e.g. ``https://api.fireworks.ai``)::
|
||||
|
||||
POST /v1/accounts/{acct}/datasets create input dataset
|
||||
POST /v1/accounts/{acct}/datasets/{id}:upload upload input JSONL
|
||||
POST /v1/accounts/{acct}/batchInferenceJobs create job
|
||||
GET /v1/accounts/{acct}/batchInferenceJobs/{jobId} poll status
|
||||
GET /v1/accounts/{acct}/datasets/{out}:getDownloadEndpoint signed URLs
|
||||
GET <signed-url> download output JSONL
|
||||
|
||||
NOTE: the exact *output JSONL line* nesting is not verbatim-documented by
|
||||
Fireworks. ``_normalize_output_line`` handles both the observed shape
|
||||
(``{custom_id, response: {...completion...}, error}``) and a ``response.body``
|
||||
nesting defensively. Confirm against a live key via the integration path.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from .openai_compatible_llm import OpenAICompatibleLLM
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Normalized statuses the retain driver treats as fatal (it raises) vs. keeps
|
||||
# polling on. "completed" ends the poll; anything else not in this set means
|
||||
# "keep polling".
|
||||
_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled", "expired"})
|
||||
|
||||
# Default per-request timeout for control-plane HTTP calls (not the job wait).
|
||||
_HTTP_TIMEOUT_SECONDS = 60.0
|
||||
|
||||
# Fallback max job wait if neither a constructor arg nor config supplies one
|
||||
# (24h matches Fireworks' maximum job timeout).
|
||||
_DEFAULT_MAX_WAIT_SECONDS = 86_400
|
||||
|
||||
|
||||
class FireworksLLM(OpenAICompatibleLLM):
|
||||
"""Fireworks provider: OpenAI-compatible online inference + native batch."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str = "fireworks",
|
||||
*,
|
||||
api_key: str,
|
||||
base_url: str = "",
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
account_id: str | None = None,
|
||||
batch_base_url: str | None = None,
|
||||
max_wait_seconds: int | None = None,
|
||||
http_client: httpx.AsyncClient | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Batch settings are static, server-level config. Resolve any unset
|
||||
# values from the global config lazily so the online inference path
|
||||
# works even when batch is never configured.
|
||||
if account_id is None or batch_base_url is None or max_wait_seconds is None:
|
||||
from ...config import get_config
|
||||
|
||||
cfg = get_config()
|
||||
if account_id is None:
|
||||
account_id = cfg.fireworks_account_id
|
||||
if batch_base_url is None:
|
||||
batch_base_url = cfg.fireworks_batch_base_url
|
||||
if max_wait_seconds is None:
|
||||
max_wait_seconds = cfg.fireworks_batch_max_wait_seconds
|
||||
|
||||
self._account_id = account_id
|
||||
self._batch_base_url = (batch_base_url or "https://api.fireworks.ai").rstrip("/")
|
||||
self._max_wait_seconds: int = (
|
||||
int(max_wait_seconds) if max_wait_seconds is not None else _DEFAULT_MAX_WAIT_SECONDS
|
||||
)
|
||||
self._http_client = http_client
|
||||
self._owns_http_client = http_client is None
|
||||
|
||||
# ----- interface: batch members -------------------------------------
|
||||
|
||||
async def supports_batch_api(self) -> bool:
|
||||
return True
|
||||
|
||||
async def submit_batch(
|
||||
self,
|
||||
requests: list[dict[str, Any]],
|
||||
endpoint: str = "/v1/chat/completions",
|
||||
completion_window: str = "24h",
|
||||
) -> dict[str, Any]:
|
||||
# endpoint/completion_window are part of the LLMInterface batch contract
|
||||
# (used by the OpenAI path) but have no analogue in Fireworks' job API:
|
||||
# the request shape is fixed (chat) and the job timeout is server-side.
|
||||
# Kept for signature compatibility with the shared retain driver.
|
||||
self._require_account_id()
|
||||
logger.info(f"Submitting Fireworks batch with {len(requests)} requests")
|
||||
|
||||
jsonl = self._translate_requests(requests)
|
||||
input_dataset_id = f"hs-batch-in-{uuid.uuid4().hex}"
|
||||
output_dataset_id = f"hs-batch-out-{uuid.uuid4().hex}"
|
||||
headers = self._auth_headers()
|
||||
|
||||
# The `dataset` resource takes format + exampleCount on create. CHAT is
|
||||
# the format for chat-completion batch input; exampleCount is the JSONL
|
||||
# line count (Fireworks rejects uploaded datasets without it) and is an
|
||||
# int64 proto field, so it goes over the wire as a string.
|
||||
await self._request(
|
||||
"POST",
|
||||
self._datasets_url(),
|
||||
headers=headers,
|
||||
json={
|
||||
"datasetId": input_dataset_id,
|
||||
"dataset": {"format": "CHAT", "exampleCount": str(len(requests))},
|
||||
},
|
||||
)
|
||||
|
||||
await self._request(
|
||||
"POST",
|
||||
f"{self._datasets_url()}/{input_dataset_id}:upload",
|
||||
headers=headers,
|
||||
files={"file": ("batch_input.jsonl", jsonl.encode("utf-8"), "application/jsonl")},
|
||||
)
|
||||
|
||||
job_resp = await self._request(
|
||||
"POST",
|
||||
self._jobs_url(),
|
||||
headers=headers,
|
||||
json={
|
||||
"model": self.model,
|
||||
"inputDatasetId": self._dataset_resource(input_dataset_id),
|
||||
"outputDatasetId": self._dataset_resource(output_dataset_id),
|
||||
},
|
||||
)
|
||||
job = job_resp.json()
|
||||
job_id = self._last_segment(job.get("name")) or output_dataset_id
|
||||
|
||||
logger.info(f"Fireworks batch job submitted: {job_id}, state={job.get('state')}")
|
||||
|
||||
return {
|
||||
"batch_id": job_id,
|
||||
"status": self._normalize_state(job.get("state", "")),
|
||||
"input_dataset_id": input_dataset_id,
|
||||
"output_dataset_id": output_dataset_id,
|
||||
"created_at": job.get("createTime"),
|
||||
"request_count": len(requests),
|
||||
}
|
||||
|
||||
async def get_batch_status(self, batch_id: str) -> dict[str, Any]:
|
||||
self._require_account_id()
|
||||
job = (await self._request("GET", self._job_url(batch_id), headers=self._auth_headers())).json()
|
||||
|
||||
status = self._normalize_state(job.get("state", ""))
|
||||
progress = job.get("jobProgress") or {}
|
||||
result: dict[str, Any] = {
|
||||
"batch_id": batch_id,
|
||||
"status": status,
|
||||
"created_at": job.get("createTime"),
|
||||
"request_counts": {
|
||||
"total": _to_int(progress.get("totalInputRequests")),
|
||||
"completed": _to_int(progress.get("successfullyProcessedRequests")),
|
||||
"failed": _to_int(progress.get("failedRequests")),
|
||||
},
|
||||
}
|
||||
|
||||
output_dataset_id = job.get("outputDatasetId")
|
||||
if output_dataset_id:
|
||||
result["output_dataset_id"] = output_dataset_id
|
||||
# Fireworks reports terminal failure detail in the `status` {code,message}.
|
||||
if job.get("status"):
|
||||
result["errors"] = job["status"]
|
||||
|
||||
# PENDING-forever guard: the shared retain poll loop has no max-wait, so
|
||||
# if a (likely non-batch-eligible) job never reaches a terminal state we
|
||||
# surface "expired" once createTime is older than the cap. Derived from
|
||||
# the server's createTime so it survives crash-recovery polling resumes.
|
||||
if status not in _TERMINAL_STATUSES:
|
||||
elapsed = self._elapsed_seconds(job.get("createTime"))
|
||||
if elapsed is not None and elapsed > self._max_wait_seconds:
|
||||
result["status"] = "expired"
|
||||
result["errors"] = (
|
||||
f"Fireworks batch {batch_id} exceeded max wait of {self._max_wait_seconds}s "
|
||||
f"in state {job.get('state')!r}. The model may not be batch-eligible "
|
||||
f"(such jobs stay PENDING indefinitely)."
|
||||
)
|
||||
logger.error(result["errors"])
|
||||
|
||||
return result
|
||||
|
||||
async def retrieve_batch_results(self, batch_id: str) -> list[dict[str, Any]]:
|
||||
self._require_account_id()
|
||||
job = (await self._request("GET", self._job_url(batch_id), headers=self._auth_headers())).json()
|
||||
|
||||
status = self._normalize_state(job.get("state", ""))
|
||||
if status != "completed":
|
||||
raise ValueError(f"Fireworks batch {batch_id} is not completed yet (state: {job.get('state')!r})")
|
||||
|
||||
output_dataset_id = job.get("outputDatasetId")
|
||||
if not output_dataset_id:
|
||||
raise ValueError(f"Fireworks batch {batch_id} completed but reported no output dataset")
|
||||
|
||||
output_short_id = self._last_segment(output_dataset_id)
|
||||
if not output_short_id:
|
||||
raise ValueError(
|
||||
f"Fireworks batch {batch_id} reported an unparseable output dataset: {output_dataset_id!r}"
|
||||
)
|
||||
download = (
|
||||
await self._request("GET", self._download_endpoint_url(output_short_id), headers=self._auth_headers())
|
||||
).json()
|
||||
signed_urls = (download or {}).get("filenameToSignedUrls") or {}
|
||||
if not signed_urls:
|
||||
raise ValueError(f"Fireworks batch {batch_id} returned no downloadable output files")
|
||||
|
||||
# The output dataset contains a results file plus a separate error file.
|
||||
# Download every file and normalize each line; error-file lines carry an
|
||||
# `error` so partial failures surface per custom_id instead of vanishing.
|
||||
results: list[dict[str, Any]] = []
|
||||
for url in signed_urls.values():
|
||||
# Signed URLs are pre-authenticated — do not attach the bearer token.
|
||||
file_resp = await self._request("GET", url)
|
||||
for line in file_resp.text.strip().split("\n"):
|
||||
if line.strip():
|
||||
results.append(self._normalize_output_line(json.loads(line)))
|
||||
|
||||
logger.info(f"Retrieved {len(results)} results for Fireworks batch {batch_id}")
|
||||
return results
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
await super().cleanup()
|
||||
if self._owns_http_client and self._http_client is not None:
|
||||
await self._http_client.aclose()
|
||||
|
||||
# ----- pure translation/normalization helpers (unit-tested) ----------
|
||||
|
||||
@staticmethod
|
||||
def _translate_requests(requests: list[dict[str, Any]]) -> str:
|
||||
"""OpenAI batch request -> Fireworks input JSONL.
|
||||
|
||||
Fireworks lines are ``{"custom_id", "body"}`` — the OpenAI ``method`` and
|
||||
``url`` keys are dropped; ``body`` is kept verbatim.
|
||||
"""
|
||||
lines = [
|
||||
json.dumps({"custom_id": req.get("custom_id"), "body": req.get("body")}, ensure_ascii=False)
|
||||
for req in requests
|
||||
]
|
||||
return "\n".join(lines)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_state(fw_state: str) -> str:
|
||||
"""Fireworks job state -> the retain driver's expected status strings.
|
||||
|
||||
Handles both the API enum (``JOB_STATE_*``) and the guide's bare names
|
||||
(``COMPLETED``/``VALIDATING``/``EXPIRED``). Unknown / in-flight states map
|
||||
to ``in_progress`` so the driver keeps polling.
|
||||
"""
|
||||
state = (fw_state or "").upper()
|
||||
if state.startswith("JOB_STATE_"):
|
||||
state = state[len("JOB_STATE_") :]
|
||||
if state == "COMPLETED":
|
||||
return "completed"
|
||||
if state == "FAILED":
|
||||
return "failed"
|
||||
if state in ("CANCELLED", "CANCELED"):
|
||||
return "cancelled"
|
||||
if state == "EXPIRED":
|
||||
return "expired"
|
||||
return "in_progress"
|
||||
|
||||
@staticmethod
|
||||
def _normalize_output_line(line: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Fireworks output JSONL line -> OpenAI-batch-output shape.
|
||||
|
||||
Target: ``{"custom_id", "response": {"body": <chat-completion>}, "error"}``
|
||||
so the consumer's ``result["response"]["body"]["choices"][0]...`` works.
|
||||
"""
|
||||
custom_id = line.get("custom_id")
|
||||
error = line.get("error")
|
||||
if error:
|
||||
return {"custom_id": custom_id, "response": None, "error": error}
|
||||
|
||||
response = line.get("response")
|
||||
if response is None:
|
||||
response = line.get("body")
|
||||
# If Fireworks already nests the completion under `body`, unwrap it;
|
||||
# otherwise the `response` object *is* the completion.
|
||||
if isinstance(response, dict) and "body" in response:
|
||||
body = response["body"]
|
||||
else:
|
||||
body = response
|
||||
return {"custom_id": custom_id, "response": {"body": body}, "error": None}
|
||||
|
||||
# ----- low-level HTTP + URL helpers ----------------------------------
|
||||
|
||||
def _require_account_id(self) -> None:
|
||||
if not self._account_id:
|
||||
raise ValueError(
|
||||
"Fireworks batch inference requires an account id. "
|
||||
"Set HINDSIGHT_API_FIREWORKS_ACCOUNT_ID to your Fireworks account id."
|
||||
)
|
||||
|
||||
def _auth_headers(self) -> dict[str, str]:
|
||||
return {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._http_client is None:
|
||||
self._http_client = httpx.AsyncClient(timeout=httpx.Timeout(_HTTP_TIMEOUT_SECONDS))
|
||||
return self._http_client
|
||||
|
||||
async def _request(
|
||||
self,
|
||||
method: str,
|
||||
url: str,
|
||||
*,
|
||||
headers: dict[str, str] | None = None,
|
||||
json: dict[str, Any] | None = None,
|
||||
files: dict[str, Any] | None = None,
|
||||
) -> httpx.Response:
|
||||
resp = await self._http().request(method, url, headers=headers, json=json, files=files)
|
||||
if resp.is_error:
|
||||
# Surface the API's error body. Fireworks returns JSON describing why a
|
||||
# 4xx/5xx happened; raise_for_status() alone discards it, which makes
|
||||
# failures (e.g. a malformed dataset/job request) undebuggable.
|
||||
raise httpx.HTTPStatusError(
|
||||
f"Fireworks API {resp.status_code} for {method} {url}: {resp.text[:2000]}",
|
||||
request=resp.request,
|
||||
response=resp,
|
||||
)
|
||||
return resp
|
||||
|
||||
def _accounts_base(self) -> str:
|
||||
return f"{self._batch_base_url}/v1/accounts/{self._account_id}"
|
||||
|
||||
def _datasets_url(self) -> str:
|
||||
return f"{self._accounts_base()}/datasets"
|
||||
|
||||
def _jobs_url(self) -> str:
|
||||
return f"{self._accounts_base()}/batchInferenceJobs"
|
||||
|
||||
def _job_url(self, job_id: str) -> str:
|
||||
return f"{self._jobs_url()}/{job_id}"
|
||||
|
||||
def _download_endpoint_url(self, dataset_short_id: str) -> str:
|
||||
return f"{self._datasets_url()}/{dataset_short_id}:getDownloadEndpoint"
|
||||
|
||||
def _dataset_resource(self, dataset_id: str) -> str:
|
||||
return f"accounts/{self._account_id}/datasets/{dataset_id}"
|
||||
|
||||
@staticmethod
|
||||
def _last_segment(resource_name: str | None) -> str | None:
|
||||
if not resource_name:
|
||||
return None
|
||||
return resource_name.rstrip("/").split("/")[-1]
|
||||
|
||||
@staticmethod
|
||||
def _elapsed_seconds(create_time: str | None) -> float | None:
|
||||
if not create_time:
|
||||
return None
|
||||
try:
|
||||
normalized = create_time.replace("Z", "+00:00")
|
||||
created = datetime.fromisoformat(normalized)
|
||||
if created.tzinfo is None:
|
||||
created = created.replace(tzinfo=timezone.utc)
|
||||
return (datetime.now(timezone.utc) - created).total_seconds()
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _to_int(value: Any) -> int:
|
||||
"""Coerce Fireworks' string/int counts to int, defaulting to 0."""
|
||||
try:
|
||||
return int(value)
|
||||
except (ValueError, TypeError):
|
||||
return 0
|
||||
@@ -0,0 +1,316 @@
|
||||
"""Gemini context-cache manager.
|
||||
|
||||
Wraps the ``google-genai`` SDK's CachedContent API to let callers reuse a
|
||||
stable system_instruction + response_schema prefix across many requests.
|
||||
|
||||
Cached input tokens are billed at ~10× lower than fresh input tokens
|
||||
(check the current Gemini pricing for the exact ratio per model), so for
|
||||
workloads that repeatedly send a large fixed prefix with a small variable
|
||||
user message — fact extraction, structured tagging, classification — the
|
||||
input-cost savings are substantial.
|
||||
|
||||
This module owns only the create/refresh/lookup lifecycle. It is up to
|
||||
the caller to (a) decide that the prefix is stable enough to cache, and
|
||||
(b) pass the returned cache name to ``GeminiLLM.call()``. When the
|
||||
returned name is ``None`` (because Gemini rejected the create — most
|
||||
commonly because the prefix is smaller than the model's minimum), the
|
||||
caller MUST fall back to a non-cached call.
|
||||
|
||||
Cardinality
|
||||
-----------
|
||||
The intended cache count per process is small (≲100 entries). Each
|
||||
entry corresponds to one combination of (model, system_instruction,
|
||||
response_schema). If a caller sees the cache grow unboundedly it
|
||||
indicates the system_instruction contains per-request data that should
|
||||
move into the user message instead.
|
||||
|
||||
TTL
|
||||
---
|
||||
Gemini's CachedContent has a TTL bounded by the model (currently 1h
|
||||
for most generally-available models). This manager refreshes proactively
|
||||
at ``ttl_safety_margin`` before expiry. If a cached entry has expired
|
||||
between refreshes the next call will recreate it transparently.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Default TTL: 55 minutes. Gemini's hard max for CachedContent is 1 hour
|
||||
# for most models; we refresh 5 minutes early so a request landing right
|
||||
# at the boundary doesn't race against expiry.
|
||||
_DEFAULT_TTL_SECONDS = 55 * 60
|
||||
_DEFAULT_REFRESH_MARGIN_SECONDS = 5 * 60
|
||||
# Cap on the cache-create network call. It runs while holding the manager lock, so
|
||||
# a hung create would block every concurrent caller (e.g. all chunks of a 10-chunk
|
||||
# retain batch waiting on the cold-start create). On timeout the create soft-fails
|
||||
# to None and callers proceed uncached, rather than stalling the whole batch.
|
||||
_DEFAULT_CREATE_TIMEOUT_SECONDS = 30.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class _CacheEntry:
|
||||
name: str # The CachedContent resource name returned by Gemini.
|
||||
created_at: float
|
||||
ttl_seconds: int
|
||||
|
||||
|
||||
class GeminiCacheManager:
|
||||
"""Per-process map of (prefix fingerprint) → CachedContent name.
|
||||
|
||||
Thread-safe across asyncio tasks via a single ``asyncio.Lock``. The
|
||||
create/refresh calls are serialised; this is fine because cache
|
||||
creation is a one-shot warm-up per fingerprint (subsequent reads are
|
||||
pure dict lookups outside the lock).
|
||||
|
||||
Not shared across pods — each worker / api replica builds its own
|
||||
cache. The cost of cold-starting one extra full-price call per pod
|
||||
per fingerprint per hour is negligible compared to the steady-state
|
||||
savings.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: Any,
|
||||
*,
|
||||
ttl_seconds: int = _DEFAULT_TTL_SECONDS,
|
||||
refresh_margin_seconds: int = _DEFAULT_REFRESH_MARGIN_SECONDS,
|
||||
create_timeout_seconds: float = _DEFAULT_CREATE_TIMEOUT_SECONDS,
|
||||
) -> None:
|
||||
self._client = client
|
||||
self._ttl_seconds = ttl_seconds
|
||||
self._refresh_margin_seconds = refresh_margin_seconds
|
||||
self._create_timeout_seconds = create_timeout_seconds
|
||||
self._entries: dict[str, _CacheEntry] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
@staticmethod
|
||||
def fingerprint(
|
||||
model: str,
|
||||
system_instruction: str,
|
||||
response_schema: Any | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
) -> str:
|
||||
"""Stable hash of the cacheable surface.
|
||||
|
||||
``response_schema`` may be a Pydantic class, a dict, or ``None``.
|
||||
Pydantic schemas are normalised by serialising via
|
||||
``model_json_schema()`` and stripping the auto-generated
|
||||
``"title"`` fields so two dynamically-built models with the same
|
||||
shape but different class names hash identically. This matters
|
||||
for callers (e.g. fact extraction) that rebuild the schema
|
||||
class on every request via a builder helper — without the
|
||||
normalisation the cache would never hit.
|
||||
|
||||
``tools`` is the OpenAI-style tools list (each entry has a
|
||||
``"function"`` dict with name/description/parameters). When
|
||||
supplied, the tool definitions become part of the cache key so a
|
||||
loop that adds or renames a tool gets a fresh cache and doesn't
|
||||
silently use a stale schema. Tools are serialised with
|
||||
``sort_keys=True`` to neutralise dict-ordering drift.
|
||||
"""
|
||||
hasher = hashlib.sha256()
|
||||
hasher.update(model.encode("utf-8"))
|
||||
hasher.update(b"\x00")
|
||||
hasher.update(system_instruction.encode("utf-8"))
|
||||
hasher.update(b"\x00")
|
||||
if response_schema is None:
|
||||
hasher.update(b"none")
|
||||
elif hasattr(response_schema, "model_json_schema"):
|
||||
try:
|
||||
schema = response_schema.model_json_schema()
|
||||
_strip_titles(schema)
|
||||
hasher.update(json.dumps(schema, sort_keys=True).encode("utf-8"))
|
||||
except Exception:
|
||||
# Fall back to class identity if the schema can't be serialised.
|
||||
hasher.update(repr(response_schema).encode("utf-8"))
|
||||
else:
|
||||
try:
|
||||
hasher.update(json.dumps(response_schema, sort_keys=True).encode("utf-8"))
|
||||
except (TypeError, ValueError):
|
||||
hasher.update(repr(response_schema).encode("utf-8"))
|
||||
hasher.update(b"\x00")
|
||||
if tools:
|
||||
try:
|
||||
hasher.update(json.dumps(tools, sort_keys=True).encode("utf-8"))
|
||||
except (TypeError, ValueError):
|
||||
hasher.update(repr(tools).encode("utf-8"))
|
||||
else:
|
||||
hasher.update(b"no-tools")
|
||||
return hasher.hexdigest()
|
||||
|
||||
async def get_or_create(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
system_instruction: str,
|
||||
response_schema: Any | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
) -> str | None:
|
||||
"""Return a CachedContent resource name for the given prefix, or
|
||||
``None`` if Gemini rejects the create (prefix too small, model
|
||||
does not support caching, etc.).
|
||||
|
||||
``tools`` is the OpenAI-style tools list. When supplied, the tool
|
||||
definitions are baked into the CachedContent so the caller's
|
||||
``call_with_tools`` doesn't need to resend them on every
|
||||
iteration. Pass ``None`` for non-tool calls.
|
||||
|
||||
``None`` return is a normal, expected value — the caller falls
|
||||
back to an uncached call and the system continues to work.
|
||||
"""
|
||||
key = self.fingerprint(model, system_instruction, response_schema, tools)
|
||||
|
||||
async with self._lock:
|
||||
entry = self._entries.get(key)
|
||||
if entry is not None and self._is_fresh(entry):
|
||||
return entry.name
|
||||
|
||||
# Need to (re)create. Pop the stale entry first so a failed
|
||||
# create doesn't leave a name we'd return on the next call.
|
||||
self._entries.pop(key, None)
|
||||
|
||||
try:
|
||||
cache_name = await self._create_cache(
|
||||
model=model,
|
||||
system_instruction=system_instruction,
|
||||
tools=tools,
|
||||
)
|
||||
except _CacheNotEligible as e:
|
||||
logger.debug(
|
||||
"GeminiCacheManager: prefix not eligible for caching (model=%s, reason=%s) — caller will fall back",
|
||||
model,
|
||||
e,
|
||||
)
|
||||
return None
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"GeminiCacheManager: failed to create cached content "
|
||||
"(model=%s); caller will fall back to uncached call",
|
||||
model,
|
||||
)
|
||||
return None
|
||||
|
||||
if cache_name is None:
|
||||
return None
|
||||
|
||||
self._entries[key] = _CacheEntry(
|
||||
name=cache_name,
|
||||
created_at=time.monotonic(),
|
||||
ttl_seconds=self._ttl_seconds,
|
||||
)
|
||||
return cache_name
|
||||
|
||||
def _is_fresh(self, entry: _CacheEntry) -> bool:
|
||||
"""An entry is fresh if it's young enough that the next request
|
||||
won't race against the TTL expiry."""
|
||||
age = time.monotonic() - entry.created_at
|
||||
return age < (entry.ttl_seconds - self._refresh_margin_seconds)
|
||||
|
||||
def invalidate(self, name: str) -> None:
|
||||
"""Forget a cache name that the server rejected (expired/deleted/invalid).
|
||||
|
||||
Called by the provider when a generate request using this CachedContent
|
||||
fails, so the next ``get_or_create`` recreates it instead of handing back
|
||||
the dead name again. Best-effort and sync — drops the matching entry from
|
||||
the in-process map; the orphaned server-side cache (if any) ages out on
|
||||
its own TTL.
|
||||
"""
|
||||
for key, entry in list(self._entries.items()):
|
||||
if entry.name == name:
|
||||
self._entries.pop(key, None)
|
||||
|
||||
async def _create_cache(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
system_instruction: str,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
) -> str | None:
|
||||
"""Wrap ``client.aio.caches.create`` with the config we want.
|
||||
|
||||
The SDK surface differs slightly across google-genai versions;
|
||||
this implementation targets the >=1.0.0 line where caches live
|
||||
under ``client.aio.caches``.
|
||||
"""
|
||||
# Lazy import so this module doesn't require the SDK at import time.
|
||||
from google.genai import types as genai_types
|
||||
|
||||
# A CachedContent only holds reusable *input* — system_instruction,
|
||||
# contents, tools, ttl. ``response_schema``/``response_mime_type`` are
|
||||
# generation-time output constraints and the SDK rejects them here
|
||||
# (``CreateCachedContentConfig`` forbids those fields). They are applied
|
||||
# per-request on the GenerateContentConfig instead — see the call sites,
|
||||
# which set them alongside ``cached_content``. ``response_schema`` is
|
||||
# still part of the fingerprint so a schema change keys a fresh cache.
|
||||
config_kwargs: dict[str, Any] = {
|
||||
"system_instruction": system_instruction,
|
||||
"ttl": f"{self._ttl_seconds}s",
|
||||
}
|
||||
if tools:
|
||||
# OpenAI-style {"function": {...}} entries must be converted to
|
||||
# Gemini's Tool/FunctionDeclaration shape before caching.
|
||||
gemini_tools = []
|
||||
for tool in tools:
|
||||
func = tool.get("function", {})
|
||||
gemini_tools.append(
|
||||
genai_types.Tool(
|
||||
function_declarations=[
|
||||
genai_types.FunctionDeclaration(
|
||||
name=func.get("name", ""),
|
||||
description=func.get("description", ""),
|
||||
parameters=func.get("parameters"),
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
config_kwargs["tools"] = gemini_tools
|
||||
|
||||
try:
|
||||
cached = await asyncio.wait_for(
|
||||
self._client.aio.caches.create(
|
||||
model=model,
|
||||
config=genai_types.CreateCachedContentConfig(**config_kwargs),
|
||||
),
|
||||
timeout=self._create_timeout_seconds,
|
||||
)
|
||||
except Exception as e:
|
||||
# Gemini returns a 400 with a "minimum token count" message
|
||||
# when the prefix is too small. We treat this as a soft
|
||||
# "not eligible" signal rather than a real error so callers
|
||||
# silently fall back to non-cached.
|
||||
msg = str(e).lower()
|
||||
if "minimum" in msg or "too small" in msg or "too short" in msg:
|
||||
raise _CacheNotEligible(str(e)) from e
|
||||
raise
|
||||
|
||||
return getattr(cached, "name", None)
|
||||
|
||||
|
||||
class _CacheNotEligible(Exception):
|
||||
"""Raised when Gemini rejects the cache create because the prefix
|
||||
is below the model's minimum cacheable size. Treated as a soft
|
||||
fallback by the caller, not an error."""
|
||||
|
||||
|
||||
def _strip_titles(node: Any) -> None:
|
||||
"""Recursively remove auto-generated ``"title"`` keys from a JSON
|
||||
Schema-like dict tree, in place. Pydantic seeds these from the
|
||||
Python class name, which means structurally-identical schemas built
|
||||
from differently-named classes look distinct to a naive hash."""
|
||||
if isinstance(node, dict):
|
||||
node.pop("title", None)
|
||||
for v in node.values():
|
||||
_strip_titles(v)
|
||||
elif isinstance(node, list):
|
||||
for item in node:
|
||||
_strip_titles(item)
|
||||
@@ -70,6 +70,22 @@ class GeminiLLM(LLMInterface):
|
||||
# Safety settings: None means use Gemini's defaults
|
||||
self._safety_settings: list | None = kwargs.get("gemini_safety_settings")
|
||||
|
||||
# User-configured extra params merged into the GenerateContentConfig of
|
||||
# every call. Gemini's request body nests generation params, so we expose
|
||||
# them in the SDK's native config space rather than as a raw body merge:
|
||||
# keys must be GenerateContentConfig fields (e.g. temperature, top_p,
|
||||
# top_k, max_output_tokens, seed). Sourced from llm_extra_body
|
||||
# (env: HINDSIGHT_API_LLM_EXTRA_BODY).
|
||||
self._extra_body: dict[str, Any] = kwargs.get("extra_body") or {}
|
||||
|
||||
# Context-cache manager. Lazy-initialized on first cache lookup so
|
||||
# nothing happens for models/workloads that never reach it. The instance
|
||||
# default here is off (a directly-constructed GeminiLLM doesn't cache); the
|
||||
# server-level default is on and flows in via the prompt_cache_enabled kwarg
|
||||
# resolved from config in LLMProvider.
|
||||
self._cache_manager: Any | None = None
|
||||
self._prompt_cache_enabled: bool = bool(kwargs.get("prompt_cache_enabled", False))
|
||||
|
||||
if self._is_vertexai:
|
||||
self._init_vertexai(**kwargs)
|
||||
else:
|
||||
@@ -168,6 +184,7 @@ class GeminiLLM(LLMInterface):
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
cached_prefix: str | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Make a Gemini/VertexAI API call with retry logic.
|
||||
@@ -182,8 +199,17 @@ class GeminiLLM(LLMInterface):
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Use strict JSON schema enforcement (not supported by Gemini).
|
||||
strict_schema: Ignored — Gemini always grammar-enforces structured output via its
|
||||
native response_schema, so it is strict regardless of this flag.
|
||||
return_usage: If True, return tuple (result, TokenUsage).
|
||||
cached_prefix: Optional CachedContent resource name (from
|
||||
``GeminiCacheManager.get_or_create``). When set, the
|
||||
system_instruction is assumed to live in the cache; this call
|
||||
skips resending it and the cached prefix is billed at the
|
||||
cached-input rate instead of the standard input rate. The
|
||||
response_schema is still sent per-request (it is not cacheable).
|
||||
Pass ``None`` to use the
|
||||
normal uncached path.
|
||||
|
||||
Returns:
|
||||
If return_usage=False: Parsed response if response_format provided, else text.
|
||||
@@ -191,9 +217,14 @@ class GeminiLLM(LLMInterface):
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Convert OpenAI-style messages to Gemini format
|
||||
# Convert OpenAI-style messages to Gemini format. We ALWAYS build
|
||||
# system_instruction (even when a cache is in use): the config builder
|
||||
# below omits it from the request while the cache carries the prefix, but
|
||||
# it must be available so the cached-call-failed safety net can re-send it
|
||||
# inline. Whether it's actually sent is decided in _build_generation_config.
|
||||
system_instruction = None
|
||||
gemini_contents = []
|
||||
using_cache = cached_prefix is not None
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
@@ -209,7 +240,9 @@ class GeminiLLM(LLMInterface):
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
|
||||
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
# Add the JSON schema as a textual hint in the system_instruction (matching
|
||||
# the normal uncached path). Structured output is still enforced via
|
||||
# response_schema regardless; this is just guidance text.
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
@@ -218,32 +251,44 @@ class GeminiLLM(LLMInterface):
|
||||
else:
|
||||
system_instruction = schema_msg
|
||||
|
||||
# Build generation config
|
||||
config_kwargs: dict[str, Any] = {}
|
||||
if system_instruction:
|
||||
config_kwargs["system_instruction"] = system_instruction
|
||||
if response_format is not None:
|
||||
config_kwargs["response_mime_type"] = "application/json"
|
||||
config_kwargs["response_schema"] = response_format
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
# Gemini's equivalent of OpenAI-style max_completion_tokens is max_output_tokens.
|
||||
# Without it the model can produce arbitrarily long responses, ignoring the
|
||||
# caller's intended cap (e.g. mental_models max_tokens during refresh).
|
||||
if max_completion_tokens is not None:
|
||||
config_kwargs["max_output_tokens"] = max_completion_tokens
|
||||
|
||||
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
|
||||
effective_safety_settings = _safety_settings_ctx.get()
|
||||
if effective_safety_settings is None:
|
||||
effective_safety_settings = self._safety_settings
|
||||
if effective_safety_settings is not None:
|
||||
config_kwargs["safety_settings"] = [
|
||||
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
|
||||
for s in effective_safety_settings
|
||||
]
|
||||
|
||||
generation_config = genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
|
||||
# Build generation config. ``cached_content`` and ``system_instruction``
|
||||
# are mutually exclusive (the cache IS the prefix; the SDK rejects
|
||||
# re-sending it). ``response_schema``/``response_mime_type`` are
|
||||
# request-level output constraints — NOT cacheable — so they're set on
|
||||
# every structured call, including cached ones where they ride alongside
|
||||
# ``cached_content``. Built as a closure so we can rebuild it WITHOUT the
|
||||
# cache and retry inline if a stale/invalid CachedContent makes the call fail.
|
||||
def _build_generation_config(use_cache: bool) -> "genai_types.GenerateContentConfig | None":
|
||||
# Seed with user-configured extra params; explicit settings below win.
|
||||
config_kwargs: dict[str, Any] = dict(self._extra_body)
|
||||
if use_cache:
|
||||
config_kwargs["cached_content"] = cached_prefix
|
||||
elif system_instruction:
|
||||
config_kwargs["system_instruction"] = system_instruction
|
||||
if response_format is not None:
|
||||
config_kwargs["response_mime_type"] = "application/json"
|
||||
config_kwargs["response_schema"] = response_format
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
# Gemini's equivalent of OpenAI-style max_completion_tokens is max_output_tokens.
|
||||
# Without it the model can produce arbitrarily long responses, ignoring the
|
||||
# caller's intended cap (e.g. mental_models max_tokens during refresh).
|
||||
if max_completion_tokens is not None:
|
||||
config_kwargs["max_output_tokens"] = max_completion_tokens
|
||||
if effective_safety_settings is not None:
|
||||
config_kwargs["safety_settings"] = [
|
||||
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
|
||||
for s in effective_safety_settings
|
||||
]
|
||||
return genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
|
||||
|
||||
cache_active = using_cache
|
||||
generation_config = _build_generation_config(cache_active)
|
||||
|
||||
last_exception = None
|
||||
|
||||
@@ -288,13 +333,24 @@ class GeminiLLM(LLMInterface):
|
||||
else:
|
||||
result = content
|
||||
|
||||
# Extract token usage
|
||||
# Extract token usage. ``cached_content_token_count`` and
|
||||
# ``thoughts_token_count`` are populated on the Gemini 2.5+
|
||||
# family; treat missing fields as 0 so older models still
|
||||
# record sensible metrics.
|
||||
input_tokens = 0
|
||||
output_tokens = 0
|
||||
cached_input_tokens = 0
|
||||
thoughts_tokens = 0
|
||||
cached_tokens = 0
|
||||
if hasattr(response, "usage_metadata") and response.usage_metadata:
|
||||
usage = response.usage_metadata
|
||||
input_tokens = usage.prompt_token_count or 0
|
||||
output_tokens = usage.candidates_token_count or 0
|
||||
cached_input_tokens = getattr(usage, "cached_content_token_count", 0) or 0
|
||||
thoughts_tokens = getattr(usage, "thoughts_token_count", 0) or 0
|
||||
# Tracing/TokenUsage consume ``cached_tokens``; metrics consume
|
||||
# ``cached_input_tokens`` — same value, two downstream names.
|
||||
cached_tokens = cached_input_tokens
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
@@ -307,6 +363,8 @@ class GeminiLLM(LLMInterface):
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
cached_input_tokens=cached_input_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
@@ -330,6 +388,7 @@ class GeminiLLM(LLMInterface):
|
||||
duration=duration,
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
@@ -345,6 +404,7 @@ class GeminiLLM(LLMInterface):
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=input_tokens + output_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
@@ -366,6 +426,20 @@ class GeminiLLM(LLMInterface):
|
||||
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
|
||||
raise
|
||||
|
||||
# Cached-request safety net: a stale/invalid/expired CachedContent
|
||||
# (or an incompatibility like cache + tool_config) surfaces as a 400.
|
||||
# Retrying the same cached request can't recover, so on the first
|
||||
# such failure drop the cache, invalidate it so later operations
|
||||
# recreate it, and retry THIS call inline with the prefix inlined.
|
||||
# Caching must never break a request.
|
||||
if cache_active and e.code == 400:
|
||||
logger.warning(f"Gemini cached call failed (400); retrying uncached. Reason: {str(e)}")
|
||||
if self._cache_manager is not None and cached_prefix is not None:
|
||||
self._cache_manager.invalidate(cached_prefix)
|
||||
cache_active = False
|
||||
generation_config = _build_generation_config(cache_active)
|
||||
continue
|
||||
|
||||
# Retry on retryable errors (rate limits, server errors, client errors)
|
||||
if e.code in (400, 429, 500, 502, 503, 504) or (e.code and e.code >= 500):
|
||||
last_exception = e
|
||||
@@ -399,6 +473,7 @@ class GeminiLLM(LLMInterface):
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
cached_prefix: str | None = None,
|
||||
) -> LLMToolCallResult:
|
||||
"""
|
||||
Make a Gemini/VertexAI API call with tool/function calling support.
|
||||
@@ -413,27 +488,39 @@ class GeminiLLM(LLMInterface):
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools (Gemini uses "auto" only).
|
||||
cached_prefix: Optional CachedContent resource name (from
|
||||
``GeminiCacheManager.get_or_create`` with ``tools=...``). When
|
||||
set, the system_instruction and tool definitions are assumed
|
||||
to live in the cache; this call will skip resending them and
|
||||
the cached prefix is billed at the cached-input rate. The
|
||||
``tools`` argument is still required (the caller may pass
|
||||
an empty list when the cache holds them) so existing call
|
||||
sites don't break.
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
"""
|
||||
start_time = time.time()
|
||||
using_cache = cached_prefix is not None
|
||||
|
||||
# Convert tools to Gemini format
|
||||
# Convert tools to Gemini format. When the cache is in use, the
|
||||
# tool definitions are baked into the CachedContent at create time
|
||||
# and the SDK rejects re-sending them alongside ``cached_content``.
|
||||
gemini_tools = []
|
||||
for tool in tools:
|
||||
func = tool.get("function", {})
|
||||
gemini_tools.append(
|
||||
genai_types.Tool(
|
||||
function_declarations=[
|
||||
genai_types.FunctionDeclaration(
|
||||
name=func.get("name", ""),
|
||||
description=func.get("description", ""),
|
||||
parameters=func.get("parameters"),
|
||||
)
|
||||
]
|
||||
if not using_cache:
|
||||
for tool in tools:
|
||||
func = tool.get("function", {})
|
||||
gemini_tools.append(
|
||||
genai_types.Tool(
|
||||
function_declarations=[
|
||||
genai_types.FunctionDeclaration(
|
||||
name=func.get("name", ""),
|
||||
description=func.get("description", ""),
|
||||
parameters=func.get("parameters"),
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# Convert messages
|
||||
system_instruction = None
|
||||
@@ -446,6 +533,10 @@ class GeminiLLM(LLMInterface):
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
# Always capture system_instruction. _build_tools_config omits it
|
||||
# (and tools) from the request while the cache carries the prefix,
|
||||
# but it must be available so the cached-call-failed safety net can
|
||||
# re-send the prefix + tools inline.
|
||||
system_instruction = (system_instruction + "\n\n" + content) if system_instruction else content
|
||||
i += 1
|
||||
elif role == "tool":
|
||||
@@ -493,49 +584,63 @@ class GeminiLLM(LLMInterface):
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
|
||||
i += 1
|
||||
|
||||
config_kwargs: dict[str, Any] = {"tools": gemini_tools}
|
||||
if system_instruction:
|
||||
config_kwargs["system_instruction"] = system_instruction
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
# See note in `call`: Gemini's max_output_tokens is the equivalent of
|
||||
# OpenAI-style max_completion_tokens.
|
||||
if max_completion_tokens is not None:
|
||||
config_kwargs["max_output_tokens"] = max_completion_tokens
|
||||
|
||||
# Map OpenAI-style tool_choice to Gemini FunctionCallingConfig
|
||||
if tool_choice == "required":
|
||||
config_kwargs["tool_config"] = genai_types.ToolConfig(
|
||||
function_calling_config=genai_types.FunctionCallingConfig(
|
||||
mode="ANY",
|
||||
)
|
||||
)
|
||||
elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
|
||||
fn_name = tool_choice.get("function", {}).get("name")
|
||||
if fn_name:
|
||||
config_kwargs["tool_config"] = genai_types.ToolConfig(
|
||||
function_calling_config=genai_types.FunctionCallingConfig(
|
||||
mode="ANY",
|
||||
allowed_function_names=[fn_name],
|
||||
)
|
||||
)
|
||||
elif tool_choice == "none":
|
||||
config_kwargs["tool_config"] = genai_types.ToolConfig(
|
||||
function_calling_config=genai_types.FunctionCallingConfig(mode="NONE")
|
||||
)
|
||||
# "auto" is the default (no tool_config needed)
|
||||
|
||||
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
|
||||
effective_safety_settings = _safety_settings_ctx.get()
|
||||
if effective_safety_settings is None:
|
||||
effective_safety_settings = self._safety_settings
|
||||
if effective_safety_settings is not None:
|
||||
config_kwargs["safety_settings"] = [
|
||||
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
|
||||
for s in effective_safety_settings
|
||||
]
|
||||
|
||||
config = genai_types.GenerateContentConfig(**config_kwargs)
|
||||
# When using a cached prefix, the SDK rejects re-sending system_instruction
|
||||
# or tools alongside ``cached_content`` — the cache IS the prefix.
|
||||
# tool_config (mode / allowed_function_names) is a per-request decision and
|
||||
# stays out of the cache. Built as a closure so we can rebuild it WITHOUT
|
||||
# the cache and retry inline if a stale/invalid cache makes the call fail.
|
||||
def _build_tools_config(use_cache: bool) -> "genai_types.GenerateContentConfig":
|
||||
# Seed with user-configured extra params; explicit settings below win.
|
||||
config_kwargs: dict[str, Any] = dict(self._extra_body)
|
||||
if use_cache:
|
||||
config_kwargs["cached_content"] = cached_prefix
|
||||
else:
|
||||
config_kwargs["tools"] = gemini_tools
|
||||
if system_instruction:
|
||||
config_kwargs["system_instruction"] = system_instruction
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
# See note in `call`: Gemini's max_output_tokens is the equivalent of
|
||||
# OpenAI-style max_completion_tokens.
|
||||
if max_completion_tokens is not None:
|
||||
config_kwargs["max_output_tokens"] = max_completion_tokens
|
||||
|
||||
# Map OpenAI-style tool_choice to Gemini FunctionCallingConfig
|
||||
if tool_choice == "required":
|
||||
config_kwargs["tool_config"] = genai_types.ToolConfig(
|
||||
function_calling_config=genai_types.FunctionCallingConfig(
|
||||
mode="ANY",
|
||||
)
|
||||
)
|
||||
elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
|
||||
fn_name = tool_choice.get("function", {}).get("name")
|
||||
if fn_name:
|
||||
config_kwargs["tool_config"] = genai_types.ToolConfig(
|
||||
function_calling_config=genai_types.FunctionCallingConfig(
|
||||
mode="ANY",
|
||||
allowed_function_names=[fn_name],
|
||||
)
|
||||
)
|
||||
elif tool_choice == "none":
|
||||
config_kwargs["tool_config"] = genai_types.ToolConfig(
|
||||
function_calling_config=genai_types.FunctionCallingConfig(mode="NONE")
|
||||
)
|
||||
# "auto" is the default (no tool_config needed)
|
||||
|
||||
if effective_safety_settings is not None:
|
||||
config_kwargs["safety_settings"] = [
|
||||
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
|
||||
for s in effective_safety_settings
|
||||
]
|
||||
return genai_types.GenerateContentConfig(**config_kwargs)
|
||||
|
||||
cache_active = using_cache
|
||||
config = _build_tools_config(cache_active)
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
@@ -578,12 +683,18 @@ class GeminiLLM(LLMInterface):
|
||||
|
||||
finish_reason = "tool_calls" if tool_calls else "stop"
|
||||
|
||||
# Extract token usage
|
||||
# Extract token usage. ``cached_content_token_count`` and
|
||||
# ``thoughts_token_count`` are populated on the Gemini 2.5+
|
||||
# family; absent fields are treated as 0.
|
||||
input_tokens = 0
|
||||
output_tokens = 0
|
||||
cached_input_tokens = 0
|
||||
thoughts_tokens = 0
|
||||
if response.usage_metadata:
|
||||
input_tokens = response.usage_metadata.prompt_token_count or 0
|
||||
output_tokens = response.usage_metadata.candidates_token_count or 0
|
||||
cached_input_tokens = getattr(response.usage_metadata, "cached_content_token_count", 0) or 0
|
||||
thoughts_tokens = getattr(response.usage_metadata, "thoughts_token_count", 0) or 0
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
@@ -596,6 +707,8 @@ class GeminiLLM(LLMInterface):
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
cached_input_tokens=cached_input_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
|
||||
# Record OpenTelemetry span
|
||||
@@ -620,6 +733,7 @@ class GeminiLLM(LLMInterface):
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
tool_calls=tool_calls_dict,
|
||||
cached_tokens=cached_input_tokens,
|
||||
)
|
||||
|
||||
return LLMToolCallResult(
|
||||
@@ -636,6 +750,18 @@ class GeminiLLM(LLMInterface):
|
||||
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
|
||||
raise
|
||||
|
||||
# Cached-request safety net (see ``call``): a stale/invalid cache or
|
||||
# a cache+tool_config conflict surfaces as a 400. Drop the cache,
|
||||
# invalidate it for later operations, and retry THIS call inline
|
||||
# with the prefix + tools re-sent. Caching must never break a call.
|
||||
if cache_active and e.code == 400:
|
||||
logger.warning(f"Gemini cached tool call failed (400); retrying uncached. Reason: {str(e)}")
|
||||
if self._cache_manager is not None and cached_prefix is not None:
|
||||
self._cache_manager.invalidate(cached_prefix)
|
||||
cache_active = False
|
||||
config = _build_tools_config(cache_active)
|
||||
continue
|
||||
|
||||
# Retry on retryable errors
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
@@ -652,6 +778,54 @@ class GeminiLLM(LLMInterface):
|
||||
raise last_exception
|
||||
raise RuntimeError("Gemini tool call failed")
|
||||
|
||||
def supports_prompt_caching(self) -> bool:
|
||||
"""True when explicit Gemini context caching is enabled for this instance.
|
||||
|
||||
Reflects the opt-in flag so callers skip the cache lookup entirely when
|
||||
it's off; ``get_or_create_cached_prefix`` also returns None in that case.
|
||||
"""
|
||||
return self._prompt_cache_enabled
|
||||
|
||||
async def get_or_create_cached_prefix(
|
||||
self,
|
||||
*,
|
||||
system_instruction: str,
|
||||
response_schema: Any | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
) -> str | None:
|
||||
"""Return a CachedContent resource name for the given prefix, or
|
||||
``None`` if context caching is disabled, the provider doesn't
|
||||
support it, or Gemini rejects the create (prefix too small, etc.).
|
||||
|
||||
``tools`` is the OpenAI-style tools list; pass it when caching a
|
||||
prefix that will be used by ``call_with_tools()``. The fingerprint
|
||||
includes the tool definitions so a loop that swaps a tool gets a
|
||||
fresh cache automatically.
|
||||
|
||||
Callers pass the returned name to ``call(cached_prefix=...)``
|
||||
or ``call_with_tools(cached_prefix=...)`` and treat ``None``
|
||||
as "cache unavailable — use the normal path". That fallback is
|
||||
essential: the system must continue to work if caching is disabled,
|
||||
if Gemini's caching API has an outage, or if the prefix is below
|
||||
the model's minimum cacheable size.
|
||||
"""
|
||||
if not self._prompt_cache_enabled:
|
||||
return None
|
||||
if self._client is None:
|
||||
return None
|
||||
if self._cache_manager is None:
|
||||
# Lazy import so the cache module is only loaded when caching
|
||||
# is actually used.
|
||||
from hindsight_api.engine.providers.gemini_cache import GeminiCacheManager
|
||||
|
||||
self._cache_manager = GeminiCacheManager(self._client)
|
||||
return await self._cache_manager.get_or_create(
|
||||
model=self.model,
|
||||
system_instruction=system_instruction,
|
||||
response_schema=response_schema,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (close connections, etc.)."""
|
||||
# Gemini client doesn't require explicit cleanup
|
||||
|
||||
@@ -48,11 +48,18 @@ class LiteLLMLLM(LLMInterface):
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float = 300.0,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
self.timeout = timeout
|
||||
self._litellm: Any = None
|
||||
# User-configured extra params merged as top-level kwargs into every
|
||||
# completion call so LiteLLM normalizes them per-provider (e.g. maps
|
||||
# temperature/top_p/max_tokens across OpenAI, Anthropic, Bedrock, …) and
|
||||
# drops any the target model rejects (litellm.drop_params=True below).
|
||||
# Sourced from llm_extra_body (env: HINDSIGHT_API_LLM_EXTRA_BODY).
|
||||
self._extra_body: dict[str, Any] = extra_body or {}
|
||||
|
||||
try:
|
||||
import litellm
|
||||
@@ -107,6 +114,11 @@ class LiteLLMLLM(LLMInterface):
|
||||
if temperature is not None:
|
||||
kwargs["temperature"] = temperature
|
||||
|
||||
# User-configured extras fill in only where the caller didn't set a value,
|
||||
# so explicit per-call params (model, messages, temperature, …) always win.
|
||||
for key, value in self._extra_body.items():
|
||||
kwargs.setdefault(key, value)
|
||||
|
||||
return kwargs
|
||||
|
||||
# ── per-model output-tokens cap (shared with Router subclass) ────────────
|
||||
|
||||
@@ -153,11 +153,30 @@ class MockLLM(LLMInterface):
|
||||
result = self._response_callback(messages, scope)
|
||||
elif self._mock_response is not None:
|
||||
result = self._mock_response
|
||||
elif scope == "retain_extract_facts" and skip_validation:
|
||||
# Fact extraction: return canned facts derived from user message text.
|
||||
# This allows tests using a mock LLM to get real facts into the DB
|
||||
# so retain → recall → reflect pipelines work end-to-end.
|
||||
result = self._build_mock_facts(messages)
|
||||
elif scope == "consolidation" and response_format is not None:
|
||||
# Consolidation: produce a single observation from the input facts
|
||||
# so the full pipeline (retain → consolidation → observation → recall) works.
|
||||
result = self._build_mock_consolidation(messages, response_format)
|
||||
elif scope == "consolidation_dedup" and response_format is not None:
|
||||
# Observation dedup adjudication. Default to "keep" so mock-LLM consolidation never
|
||||
# spuriously merges observations — this preserves the pre-dedup behaviour that
|
||||
# deterministic consolidation tests assert (the generic branch below can't construct
|
||||
# the model because its "action" field is required and has no default).
|
||||
result = response_format(action="keep", reason="mock")
|
||||
elif scope == "memory_think":
|
||||
# Reflect: return a plausible text answer
|
||||
result = "Based on the available information, the answer is related to the context provided."
|
||||
elif response_format is not None:
|
||||
# Try to create a minimal valid instance of the response format
|
||||
# Structured output: try to return a valid empty instance of the model
|
||||
# so that callers expecting e.g. response_format with defaults
|
||||
# get a valid instance rather than a crash on {"mock": True}.
|
||||
try:
|
||||
# For Pydantic models, try to create with minimal valid data
|
||||
result = {"mock": True}
|
||||
result = response_format()
|
||||
except Exception:
|
||||
result = {"mock": True}
|
||||
else:
|
||||
@@ -243,6 +262,12 @@ class MockLLM(LLMInterface):
|
||||
else:
|
||||
result = LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
|
||||
# Set mock token usage on result if not already set
|
||||
if result.input_tokens == 0:
|
||||
result.input_tokens = 10
|
||||
if result.output_tokens == 0:
|
||||
result.output_tokens = 5
|
||||
|
||||
# Record span with mock values
|
||||
# Convert LLMToolCall objects to dicts for span recording
|
||||
tool_calls_dict = (
|
||||
@@ -266,6 +291,92 @@ class MockLLM(LLMInterface):
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _build_mock_facts(messages: list[dict]) -> dict:
|
||||
"""Build a canned fact extraction response from the user message text.
|
||||
|
||||
Splits the input into sentence-like chunks and returns each as a separate
|
||||
world fact with a simple entity extracted from the first noun-like word.
|
||||
This is intentionally simplistic — it just needs to produce structurally
|
||||
valid facts so the rest of the pipeline (embedding, storage, recall) works.
|
||||
"""
|
||||
import re
|
||||
|
||||
user_text = ""
|
||||
for m in messages:
|
||||
if m.get("role") == "user":
|
||||
user_text = m.get("content", "")
|
||||
break
|
||||
|
||||
# Split on sentence boundaries: period followed by space/EOL (not mid-number), or newlines
|
||||
sentences = [s.strip() for s in re.split(r"(?<=\.)\s+|\n+", user_text) if s.strip() and len(s.strip()) > 10]
|
||||
|
||||
if not sentences:
|
||||
sentences = [user_text[:200] if user_text else "mock fact"]
|
||||
|
||||
facts = []
|
||||
for sentence in sentences[:10]: # Cap at 10 facts per chunk
|
||||
# Extract simple entities: capitalized words that aren't common words
|
||||
words = re.findall(r"\b[A-Z][a-z]+\b", sentence)
|
||||
entities = [{"text": w} for w in dict.fromkeys(words)][:5] # Dedupe, cap at 5
|
||||
|
||||
facts.append(
|
||||
{
|
||||
"what": sentence,
|
||||
"when": "N/A",
|
||||
"where": "N/A",
|
||||
"who": "N/A",
|
||||
"why": "N/A",
|
||||
"fact_kind": "conversation",
|
||||
"fact_type": "world",
|
||||
"entities": entities,
|
||||
}
|
||||
)
|
||||
|
||||
return {"facts": facts}
|
||||
|
||||
@staticmethod
|
||||
def _build_mock_consolidation(messages: list[dict], response_format: Any) -> Any:
|
||||
"""Build a mock consolidation response that creates one observation per fact.
|
||||
|
||||
Parses fact IDs from the consolidation prompt and creates one observation
|
||||
per fact, each referencing its source fact ID. This mimics real LLM behavior
|
||||
where distinct facts produce separate observations, preserving entity
|
||||
separation so pipeline tests (graph filtering, entity linking) work correctly.
|
||||
"""
|
||||
import re
|
||||
|
||||
user_text = ""
|
||||
for m in messages:
|
||||
if m.get("role") == "user":
|
||||
user_text = m.get("content", "")
|
||||
break
|
||||
|
||||
# Extract fact UUIDs from the prompt (format: "[<uuid>] <text>")
|
||||
fact_entries = re.findall(r"\[([0-9a-f-]{36})\]\s*(.+?)(?:\n|$)", user_text)
|
||||
|
||||
if not fact_entries:
|
||||
# No facts to consolidate — return empty response
|
||||
try:
|
||||
return response_format()
|
||||
except Exception:
|
||||
return {"creates": [], "updates": [], "deletes": []}
|
||||
|
||||
# Create one observation per fact to preserve entity separation
|
||||
creates = []
|
||||
for fact_id, fact_text in fact_entries:
|
||||
creates.append({"text": fact_text.strip(), "source_fact_ids": [fact_id]})
|
||||
|
||||
try:
|
||||
return response_format(
|
||||
creates=creates,
|
||||
updates=[],
|
||||
deletes=[],
|
||||
)
|
||||
except Exception:
|
||||
# Fallback if response_format constructor doesn't accept these args
|
||||
return {"creates": creates, "updates": [], "deletes": []}
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (no-op for mock provider)."""
|
||||
pass
|
||||
@@ -318,6 +429,8 @@ class MockLLM(LLMInterface):
|
||||
return self._mock_calls
|
||||
|
||||
def clear_mock_calls(self) -> None:
|
||||
"""Clear the recorded mock calls and any set exception."""
|
||||
"""Clear all recorded calls and any configured response/exception state."""
|
||||
self._mock_calls = []
|
||||
self._mock_exception = None
|
||||
self._mock_response = None
|
||||
self._response_callback = None
|
||||
|
||||
@@ -7,7 +7,7 @@ This provider handles all OpenAI API-compatible models including:
|
||||
- Groq: Fast inference with seed control and service tiers
|
||||
- Ollama: Local models with native streaming API support
|
||||
- LMStudio: Local models with OpenAI-compatible API
|
||||
- MiniMax: MiniMax-M2.7 models with 1M context window
|
||||
- MiniMax: MiniMax-M3 / MiniMax-M2.7 models with 1M context window
|
||||
- DeepSeek: deepseek-v4-flash / deepseek-v4-pro / deepseek-chat / deepseek-reasoner via api.deepseek.com
|
||||
- Opencode Go: deepseek-v4-flash via https://opencode.ai/zen/go/v1
|
||||
|
||||
@@ -44,6 +44,16 @@ logger = logging.getLogger(__name__)
|
||||
DEFAULT_LLM_SEED = 4242
|
||||
JSON_MODE_USER_HINT = "Return valid json only."
|
||||
|
||||
# Self-hosted OpenAI-compatible servers that advertise tool_choice="required"
|
||||
# but silently ignore it: instead of forcing a tool call they return
|
||||
# finish_reason "stop"/"tool_calls" with an EMPTY tool_calls array and no error.
|
||||
# Reflect's agent loop then sees no tool call, runs synthesis with no retrieval,
|
||||
# and answers "I don't have information" even when the bank holds the answer.
|
||||
# See issues #1563 (LM Studio), #1179 (LM Studio + Qwen), #1877 (vLLM with
|
||||
# --enable-auto-tool-choice). llama-server (the "llamacpp" provider) honors
|
||||
# "required" correctly and is intentionally excluded (#1179).
|
||||
_TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS = frozenset({"lmstudio", "ollama"})
|
||||
|
||||
|
||||
class ProviderResponseError(RuntimeError):
|
||||
"""Raised when a provider returns a success response without usable content."""
|
||||
@@ -232,7 +242,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
- Groq: Fast inference with seed control and service tiers
|
||||
- Ollama: Local models with native streaming API for better structured output
|
||||
- LMStudio: Local models with OpenAI-compatible API
|
||||
- MiniMax: MiniMax-M2.7 models via OpenAI-compatible API (https://api.minimax.io/v1)
|
||||
- MiniMax: MiniMax-M3 / MiniMax-M2.7 models via OpenAI-compatible API (https://api.minimax.io/v1)
|
||||
- DeepSeek: deepseek-v4-flash / deepseek-v4-pro / deepseek-chat / deepseek-reasoner via https://api.deepseek.com
|
||||
- opencode-go: deepseek-v4-flash via https://opencode.ai/zen/go/v1
|
||||
"""
|
||||
@@ -270,6 +280,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"openai",
|
||||
"groq",
|
||||
"ollama",
|
||||
"ollama-cloud",
|
||||
"lmstudio",
|
||||
"llamacpp",
|
||||
"minimax",
|
||||
@@ -278,6 +289,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"openrouter",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"fireworks",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
raise ValueError(f"OpenAICompatibleLLM only supports: {', '.join(valid_providers)}. Got: {self.provider}")
|
||||
@@ -288,6 +300,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
self.base_url = "https://api.groq.com/openai/v1"
|
||||
elif self.provider == "ollama":
|
||||
self.base_url = "http://localhost:11434/v1"
|
||||
elif self.provider == "ollama-cloud":
|
||||
self.base_url = "https://ollama.com/v1"
|
||||
elif self.provider == "lmstudio":
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
elif self.provider == "minimax":
|
||||
@@ -300,6 +314,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
self.base_url = "https://api.z.ai/api/coding/paas/v4"
|
||||
elif self.provider == "opencode-go":
|
||||
self.base_url = "https://opencode.ai/zen/go/v1"
|
||||
elif self.provider == "fireworks":
|
||||
# OpenAI-compatible inference host (online path). The batch API
|
||||
# lives on a separate control-plane host — see FireworksLLM.
|
||||
self.base_url = "https://api.fireworks.ai/inference/v1"
|
||||
|
||||
# For ollama/lmstudio, use dummy key if not provided
|
||||
if self.provider in ("ollama", "lmstudio") and not self.api_key:
|
||||
@@ -316,6 +334,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"openrouter",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"ollama-cloud",
|
||||
)
|
||||
and not self.api_key
|
||||
):
|
||||
@@ -351,6 +370,21 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
f"base_url={self.base_url or 'default'}"
|
||||
)
|
||||
|
||||
def _drops_tool_choice_required(self) -> bool:
|
||||
"""Whether this endpoint silently ignores ``tool_choice="required"``.
|
||||
|
||||
True for self-hosted OpenAI-compatible servers known to return an empty
|
||||
tool_calls array for "required" instead of forcing a call (#1563/#1179/
|
||||
#1877). Covers LM Studio / Ollama directly, plus any server reached via
|
||||
the generic "openai" provider with a custom ``base_url`` (e.g. a local
|
||||
vLLM endpoint). The real OpenAI API (no base_url override) honors
|
||||
"required", and cloud providers keep their own default base_urls, so both
|
||||
are left untouched.
|
||||
"""
|
||||
if self.provider in _TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS:
|
||||
return True
|
||||
return self.provider == "openai" and bool(self.base_url)
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""
|
||||
Verify that the provider is configured correctly by making a simple test call.
|
||||
@@ -451,7 +485,9 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Use strict JSON schema enforcement (OpenAI only).
|
||||
strict_schema: Use strict json_schema (grammar-enforced) response_format instead of
|
||||
the soft json_object path. Supported by OpenAI and schema-capable self-hosted
|
||||
backends (llama.cpp, vLLM). Server-wide via HINDSIGHT_API_LLM_STRICT_SCHEMA.
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
|
||||
Returns:
|
||||
@@ -641,6 +677,9 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
input_tokens = usage.prompt_tokens or 0 if usage else 0
|
||||
output_tokens = usage.completion_tokens or 0 if usage else 0
|
||||
total_tokens = usage.total_tokens or 0 if usage else 0
|
||||
cached_tokens = 0
|
||||
if usage and getattr(usage, "prompt_tokens_details", None):
|
||||
cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
|
||||
|
||||
# Record LLM metrics
|
||||
metrics = get_metrics_collector()
|
||||
@@ -670,14 +709,12 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
duration=duration,
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0 and usage:
|
||||
ratio = max(1, output_tokens) / max(1, input_tokens)
|
||||
cached_tokens = 0
|
||||
if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
|
||||
cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
|
||||
cache_info = f", cached_tokens={cached_tokens}" if cached_tokens > 0 else ""
|
||||
logger.info(
|
||||
f"slow llm call: scope={scope}, model={self.provider}/{self.model}, "
|
||||
@@ -690,6 +727,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=total_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
@@ -858,6 +896,16 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
if request_tool_choice == "auto":
|
||||
request_tool_choice = None
|
||||
|
||||
# vLLM (--enable-auto-tool-choice), LM Studio, Ollama and similar
|
||||
# self-hosted servers silently drop tool_choice="required", returning an
|
||||
# empty tool_calls array instead of forcing a call (#1563/#1179/#1877).
|
||||
# Downgrade to auto (None) so the model still gets to call a tool. Named
|
||||
# tool_choice dicts were already normalized to "required" + a single
|
||||
# filtered tool above, so the call stays practically forced even under
|
||||
# auto. The real OpenAI API honors "required" and is left untouched.
|
||||
if request_tool_choice == "required" and self._drops_tool_choice_required():
|
||||
request_tool_choice = None
|
||||
|
||||
# DeepSeek tool-call replies can carry provider-specific reasoning_content.
|
||||
# The normalized tool result does not retain it, but replaying assistant
|
||||
# tool_calls without the field can trigger a 400. DeepSeek accepts an
|
||||
@@ -1073,12 +1121,17 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
last_exception = None
|
||||
|
||||
# Pass API key as Bearer token for cloud Ollama endpoints
|
||||
headers: dict[str, str] = {}
|
||||
if self.api_key and self.api_key != "local":
|
||||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||||
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await client.post(native_url, json=payload)
|
||||
response = await client.post(native_url, json=payload, headers=headers)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
|
||||
@@ -302,6 +302,24 @@ def _is_context_overflow_error(exc: Exception) -> bool:
|
||||
)
|
||||
|
||||
|
||||
def _all_mental_models_are_usable_and_fresh(tool_output: dict[str, Any]) -> bool:
|
||||
"""Return whether every retrieved mental model is explicitly fresh and has answerable content.
|
||||
|
||||
Used to decide — without an extra LLM call — whether a forced
|
||||
``search_mental_models`` result is trustworthy enough to hand control back
|
||||
to the agent. A model is usable only when it is explicitly ``is_stale ==
|
||||
False`` (an unknown/missing staleness flag is treated as unsafe) and has
|
||||
non-empty content.
|
||||
"""
|
||||
models = tool_output.get("mental_models") or []
|
||||
for model in models:
|
||||
if model.get("is_stale") is not False:
|
||||
return False
|
||||
if not str(model.get("content") or "").strip():
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
async def run_reflect_agent(
|
||||
llm_config: "LLMProvider",
|
||||
bank_id: str,
|
||||
@@ -321,6 +339,7 @@ async def run_reflect_agent(
|
||||
include_recall: bool = True,
|
||||
budget: str | None = None,
|
||||
max_context_tokens: int = 100_000,
|
||||
llm_output_language: str | None = None,
|
||||
) -> ReflectAgentResult:
|
||||
"""
|
||||
Execute the reflect agent loop using native tool calling.
|
||||
@@ -369,13 +388,40 @@ async def run_reflect_agent(
|
||||
|
||||
# Build initial messages (directives are injected into system prompt at START and END)
|
||||
system_prompt = build_system_prompt_for_tools(
|
||||
bank_profile, context, directives=directives, has_mental_models=has_mental_models, budget=budget
|
||||
bank_profile,
|
||||
context,
|
||||
directives=directives,
|
||||
has_mental_models=has_mental_models,
|
||||
include_observations=include_observations,
|
||||
budget=budget,
|
||||
)
|
||||
messages: list[dict[str, Any]] = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": query},
|
||||
]
|
||||
|
||||
# Opt into context caching for the agentic tool loop. The system
|
||||
# prompt and tool definitions are stable for the duration of this
|
||||
# reflect call (and across reflects against the same bank), so
|
||||
# caching them once and reusing across every iteration of the loop
|
||||
# collapses the dominant input cost — the prefix repeated on every
|
||||
# turn. ``get_or_create_cached_prefix`` returns None when caching is
|
||||
# disabled, unsupported, or the prefix is too small; the
|
||||
# ``call_with_tools`` invocation below transparently falls back to
|
||||
# the uncached path in that case.
|
||||
cached_prefix_name: str | None = None
|
||||
provider_impl = getattr(llm_config, "_provider_impl", None)
|
||||
if provider_impl is not None and provider_impl.supports_prompt_caching():
|
||||
try:
|
||||
cached_prefix_name = await provider_impl.get_or_create_cached_prefix(
|
||||
system_instruction=system_prompt,
|
||||
tools=tools,
|
||||
)
|
||||
except Exception:
|
||||
# Caching is a soft optimisation; never let a cache-side
|
||||
# error block a reflect.
|
||||
cached_prefix_name = None
|
||||
|
||||
# Tracking
|
||||
total_tools_called = 0
|
||||
tool_trace: list[ToolCall] = []
|
||||
@@ -436,6 +482,11 @@ async def run_reflect_agent(
|
||||
)
|
||||
|
||||
consecutive_errors = 0
|
||||
# When a forced ``search_mental_models`` returns fresh, usable models on a
|
||||
# low/mid-budget call, we stop forcing the lower retrieval layers from this
|
||||
# iteration onward and let the agent answer (or retrieve deeper itself)
|
||||
# under ``auto`` tool choice. None means the full forced path still applies.
|
||||
stop_forcing_from_iteration: int | None = None
|
||||
for iteration in range(max_iterations):
|
||||
is_last = iteration == max_iterations - 1
|
||||
|
||||
@@ -447,7 +498,10 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -504,7 +558,10 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -558,18 +615,31 @@ async def run_reflect_agent(
|
||||
if include_recall:
|
||||
forced_sequence.append("recall")
|
||||
|
||||
if iteration < len(forced_sequence):
|
||||
iter_tool_choice: str | dict = {"type": "function", "function": {"name": forced_sequence[iteration]}}
|
||||
if stop_forcing_from_iteration is not None and iteration >= stop_forcing_from_iteration:
|
||||
# A fresh mental model already short-circuited the forced path.
|
||||
iter_tool_choice: str | dict = "auto"
|
||||
elif iteration < len(forced_sequence):
|
||||
iter_tool_choice = {"type": "function", "function": {"name": forced_sequence[iteration]}}
|
||||
else:
|
||||
iter_tool_choice = "auto"
|
||||
|
||||
try:
|
||||
result = await llm_config.call_with_tools(
|
||||
ct_kwargs: dict[str, Any] = dict(
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
scope="reflect_tool_call",
|
||||
tool_choice=iter_tool_choice,
|
||||
)
|
||||
# Gemini rejects ``cached_content`` alongside a per-request
|
||||
# ``tool_config`` (forced tool choice): "CachedContent can not be used
|
||||
# with GenerateContent request setting system_instruction, tools or
|
||||
# tool_config." The forced-sequence iterations set tool_config, so only
|
||||
# the ``auto`` iterations can reference the cache; forced iterations send
|
||||
# the prefix inline. The cache (tools + system prompt) is identical
|
||||
# either way, so this just limits *which* iterations are billed cached.
|
||||
if cached_prefix_name is not None and iter_tool_choice == "auto":
|
||||
ct_kwargs["cached_prefix"] = cached_prefix_name
|
||||
result = await llm_config.call_with_tools(**ct_kwargs)
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
consecutive_errors = 0
|
||||
total_input_tokens += result.input_tokens
|
||||
@@ -607,7 +677,10 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -728,7 +801,10 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -906,6 +982,25 @@ async def run_reflect_agent(
|
||||
for mm in output["mental_models"]:
|
||||
if "id" in mm:
|
||||
available_mental_model_ids.add(mm["id"])
|
||||
# Deterministic short-circuit (no extra LLM call): on a
|
||||
# low/mid-budget call, if every retrieved mental model is
|
||||
# fresh and has usable content, stop forcing the lower
|
||||
# retrieval layers. The next iteration runs under ``auto``
|
||||
# tool choice, so the agent can answer directly when the
|
||||
# mental model suffices, or — having just read it — issue a
|
||||
# targeted ``search_observations``/``recall`` itself. Stale,
|
||||
# empty, or missing mental models keep the full forced path.
|
||||
if (
|
||||
stop_forcing_from_iteration is None
|
||||
and (budget or "low").lower() != "high"
|
||||
and output.get("mental_models")
|
||||
and _all_mental_models_are_usable_and_fresh(output)
|
||||
):
|
||||
stop_forcing_from_iteration = iteration + 1
|
||||
logger.info(
|
||||
f"[REFLECT {reflect_id}] Fresh mental models sufficient on iteration {iteration + 1}; "
|
||||
"releasing forced lower-level retrieval to auto."
|
||||
)
|
||||
|
||||
if (
|
||||
normalized_tool_name == "search_observations"
|
||||
|
||||
@@ -98,6 +98,7 @@ def build_system_prompt_for_tools(
|
||||
context: str | None = None,
|
||||
directives: list[dict[str, Any]] | None = None,
|
||||
has_mental_models: bool = False,
|
||||
include_observations: bool = True,
|
||||
budget: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -108,11 +109,17 @@ def build_system_prompt_for_tools(
|
||||
2. search_observations - Consolidated knowledge with freshness
|
||||
3. recall - Raw facts as ground truth
|
||||
|
||||
The retrieval-strategy and workflow sections are built to match the tools
|
||||
actually exposed to the LLM — mentioning a tool the agent has disabled
|
||||
causes weaker LLMs to either hallucinate the call (rejected by the agent)
|
||||
or give up with "I cannot find any information…" (see #1724).
|
||||
|
||||
Args:
|
||||
bank_profile: Bank profile with name and mission
|
||||
context: Optional additional context
|
||||
directives: Optional list of directive mental models to inject as hard rules
|
||||
has_mental_models: Whether the bank has any mental models (skip if not)
|
||||
include_observations: Whether search_observations is in the tool list.
|
||||
budget: Search depth budget - "low", "mid", or "high". Controls exploration thoroughness.
|
||||
"""
|
||||
name = bank_profile.get("name", "Assistant")
|
||||
@@ -158,56 +165,137 @@ def build_system_prompt_for_tools(
|
||||
"- If memories mention someone did an activity, you can infer they likely enjoyed it",
|
||||
"- Synthesize a coherent narrative from related memories",
|
||||
"- Be a thoughtful interpreter, not just a literal repeater",
|
||||
"- When the exact answer isn't stated, use what IS stated to give the best answer",
|
||||
"- When the exact answer isn't stated, use what IS stated to give a best-effort answer AND surface any uncertainty — never invent confidence the data doesn't support.",
|
||||
"",
|
||||
"## Temporal Reasoning",
|
||||
"Every memory and observation carries temporal fields in the JSON tool result:",
|
||||
"- `mentioned_at` — when the user retained the fact (always set).",
|
||||
"- `occurred_start` / `occurred_end` — when the underlying event happened (optional, set for dated events).",
|
||||
"",
|
||||
"When facts about the SAME facet conflict — counts, statuses, ownership, location, presence, etc. — the fact with the LATEST `mentioned_at` is authoritative. Later statements SUPERSEDE earlier ones. Do NOT average, sum, or favor an explicitly-dated fact over a more recent one.",
|
||||
"",
|
||||
"Example: three count facts come back from recall:",
|
||||
" - 'Team has 2 engineers' (mentioned_at=T1)",
|
||||
" - 'Team now has 1 engineer' (mentioned_at=T2, occurred_start=2026-05-25)",
|
||||
" - 'Team has 5 engineers' (mentioned_at=T3)",
|
||||
"with T1 < T2 < T3. The current size is 5, not 1. Then apply later events (e.g. someone leaving after T3) on top of that.",
|
||||
"",
|
||||
"For reconstructing a TIMELINE of events, order by `occurred_start` / `occurred_end` (when things happened), not `mentioned_at` (when they were retained).",
|
||||
"",
|
||||
"## Conflicts and Ambiguity",
|
||||
"Not every retrieval converges on a single answer. Distinguish two cases:",
|
||||
"",
|
||||
"- RESOLVABLE conflict — the temporal rule above (latest `mentioned_at` wins) cleanly picks a winner. Apply it and move on.",
|
||||
"- UNRESOLVABLE ambiguity — the data is internally inconsistent in a way the temporal rule does NOT settle. Examples: a recent aggregate (count, total) is incompatible with the individual entities you can enumerate; two equally-recent facts disagree and no later fact resolves them; events are described but their relative order is unclear; the user's own statements contradict each other and nothing later reconciles them.",
|
||||
"",
|
||||
"When the data is genuinely ambiguous: SAY SO in your answer. Name the conflicting facts. Explain why they can't be reconciled. Give a range or a best-effort interpretation with explicit uncertainty (e.g. 'between X and Y, depending on [unresolved condition]'; or 'the most recent statement says A, but B was stated earlier and the gap isn't accounted for in any later fact').",
|
||||
"",
|
||||
"An honest 'the data is inconsistent about X' beats a confident wrong answer. Do NOT pick a value arbitrarily, average conflicting values, or smooth over gaps in confident prose. Acknowledging ambiguity is a successful answer, not a failure mode.",
|
||||
"",
|
||||
"## Showing Your Reasoning",
|
||||
"For any answer that resolves a conflict between facts, applies events on top of a count or status, or settles an ambiguity — show your work in the answer text so a reader can audit it.",
|
||||
"",
|
||||
"Walk through these steps explicitly:",
|
||||
"1. **List the relevant facts in `mentioned_at` order (oldest → newest)**, each with the value it asserts. Use a short bulleted list.",
|
||||
"2. **Identify the authoritative fact** under the temporal rule (latest `mentioned_at` for the contested facet). Write its date down.",
|
||||
"3. **List candidate events to apply on top** — anything that changes the count, status, or state being asked about. Write each event's date down next to it.",
|
||||
"4. **Sanity-check each candidate event against the authoritative date** — for EVERY event from step 3, write a one-line check in the form `<event> (<event_date>) vs authoritative (<authoritative_date>) → BEFORE/AFTER → KEEP/DROP`. If the event is BEFORE or EQUAL to the authoritative date, DROP it: it is already reflected in the authoritative fact, and applying it again is double-counting. This is the single most common mistake — do not skip this step even if you feel confident.",
|
||||
"5. **Show the arithmetic or derivation explicitly** using only the KEEP events from step 4 — e.g. 'authoritative count = 5 (at 2025-02-12); kept events: Shadow died (2025-03-12, AFTER); 5 − 1 = 4'.",
|
||||
"6. If step 2 or 3 cannot be done cleanly (no clear winner, overlapping timestamps, unclear event order), STOP and surface this as an UNRESOLVABLE ambiguity per the section above — do not fabricate a derivation.",
|
||||
"",
|
||||
"For simple factual lookups that don't involve conflict or arithmetic, you can answer directly without this scaffolding.",
|
||||
"",
|
||||
"## HIERARCHICAL RETRIEVAL STRATEGY",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
# Build retrieval levels based on what's available
|
||||
# Assemble the retrieval-level blocks for whatever tools are exposed.
|
||||
# MM and Observations bodies are unconditional; recall's fallback wording
|
||||
# adapts to which upstream tools precede it (telling the LLM to fall back
|
||||
# to a tool that isn't in its list is the bug at the root of #1724).
|
||||
levels: list[tuple[str, list[str]]] = []
|
||||
if has_mental_models:
|
||||
parts.extend(
|
||||
levels.append(
|
||||
(
|
||||
"MENTAL MODELS (search_mental_models)",
|
||||
[
|
||||
"- User-curated summaries about specific topics",
|
||||
"- HIGHEST quality - manually created and maintained",
|
||||
"- If a relevant mental model exists and is FRESH, it may fully answer the question",
|
||||
"- Check `is_stale` field - if stale, also verify with lower levels",
|
||||
],
|
||||
)
|
||||
)
|
||||
if include_observations:
|
||||
levels.append(
|
||||
(
|
||||
"OBSERVATIONS (search_observations)",
|
||||
[
|
||||
"- Auto-consolidated knowledge from memories",
|
||||
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
|
||||
"- Good for understanding patterns and summaries",
|
||||
],
|
||||
)
|
||||
)
|
||||
recall_body = ["- Individual memories (world facts and experiences)"]
|
||||
if has_mental_models and include_observations:
|
||||
recall_body.extend(
|
||||
[
|
||||
"You have access to THREE levels of knowledge. Use them in this order:",
|
||||
"",
|
||||
"### 1. MENTAL MODELS (search_mental_models) - Try First",
|
||||
"- User-curated summaries about specific topics",
|
||||
"- HIGHEST quality - manually created and maintained",
|
||||
"- If a relevant mental model exists and is FRESH, it may fully answer the question",
|
||||
"- Check `is_stale` field - if stale, also verify with lower levels",
|
||||
"",
|
||||
"### 2. OBSERVATIONS (search_observations) - Second Priority",
|
||||
"- Auto-consolidated knowledge from memories",
|
||||
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
|
||||
"- Good for understanding patterns and summaries",
|
||||
"",
|
||||
"### 3. RAW FACTS (recall) - Ground Truth",
|
||||
"- Individual memories (world facts and experiences)",
|
||||
"- Use when: no mental models/observations exist, they're stale, or you need specific details",
|
||||
"- MANDATORY: If search_mental_models and search_observations both return 0 results, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that other levels are built from",
|
||||
"",
|
||||
"**Tool result ordering:** `recall()` and `search_observations()` return their `memories` / `observations` arrays sorted by SEMANTIC RELEVANCE to the query, NOT by time. The POSITION of an entry tells you nothing about when it was retained. For any temporal reasoning — recency, supersession, applying events on top of a state — IGNORE the position and read the per-entry `mentioned_at` field (and `occurred_start` / `occurred_end` for events).",
|
||||
"",
|
||||
]
|
||||
)
|
||||
else:
|
||||
parts.extend(
|
||||
elif has_mental_models:
|
||||
recall_body.extend(
|
||||
[
|
||||
"- Use when: no mental model exists, it's stale, or you need specific details",
|
||||
"- MANDATORY: If search_mental_models returns 0 results, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that mental models are built from",
|
||||
]
|
||||
)
|
||||
elif include_observations:
|
||||
recall_body.extend(
|
||||
[
|
||||
"You have access to TWO levels of knowledge. Use them in this order:",
|
||||
"",
|
||||
"### 1. OBSERVATIONS (search_observations) - Try First",
|
||||
"- Auto-consolidated knowledge from memories",
|
||||
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
|
||||
"- Good for understanding patterns and summaries",
|
||||
"",
|
||||
"### 2. RAW FACTS (recall) - Ground Truth",
|
||||
"- Individual memories (world facts and experiences)",
|
||||
"- Use when: no observations exist, they're stale, or you need specific details",
|
||||
"- MANDATORY: If search_observations returns 0 results or count=0, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that observations are built from",
|
||||
"",
|
||||
"**Tool result ordering:** `recall()` and `search_observations()` return their `memories` / `observations` arrays sorted by SEMANTIC RELEVANCE to the query, NOT by time. The POSITION of an entry tells you nothing about when it was retained. For any temporal reasoning — recency, supersession, applying events on top of a state — IGNORE the position and read the per-entry `mentioned_at` field (and `occurred_start` / `occurred_end` for events).",
|
||||
"",
|
||||
]
|
||||
)
|
||||
else:
|
||||
recall_body.extend(
|
||||
[
|
||||
"- MANDATORY: Call recall() to gather facts before giving up",
|
||||
"- This is the source of truth.",
|
||||
]
|
||||
)
|
||||
levels.append(("RAW FACTS (recall) - Ground Truth", recall_body))
|
||||
|
||||
# Position-dependent suffix for upstream tools; recall already carries its
|
||||
# fixed "- Ground Truth" suffix in the header text.
|
||||
suffixes = [""] * len(levels)
|
||||
if len(levels) >= 2:
|
||||
suffixes[0] = " - Try First"
|
||||
if len(levels) == 3:
|
||||
suffixes[1] = " - Second Priority"
|
||||
|
||||
if len(levels) == 1:
|
||||
parts.append("You have access to ONE level of knowledge:")
|
||||
else:
|
||||
word = "TWO" if len(levels) == 2 else "THREE"
|
||||
parts.append(f"You have access to {word} levels of knowledge. Use them in this order:")
|
||||
parts.append("")
|
||||
for idx, ((header, body), suffix) in enumerate(zip(levels, suffixes), 1):
|
||||
parts.append(f"### {idx}. {header}{suffix}")
|
||||
parts.extend(body)
|
||||
parts.append("")
|
||||
|
||||
parts.extend(
|
||||
[
|
||||
@@ -267,25 +355,28 @@ def build_system_prompt_for_tools(
|
||||
|
||||
parts.append("## Workflow")
|
||||
|
||||
steps: list[str] = []
|
||||
if has_mental_models:
|
||||
parts.extend(
|
||||
[
|
||||
"1. First, try search_mental_models() - check if a curated summary exists",
|
||||
"2. If no mental model or it's stale, try search_observations() for consolidated knowledge",
|
||||
"3. If observations are stale OR you need specific details, use recall() for raw facts",
|
||||
"4. Use expand() if you need more context on specific memories",
|
||||
"5. When ready, call done() with your answer and supporting IDs",
|
||||
]
|
||||
steps.append("First, try search_mental_models() - check if a curated summary exists")
|
||||
if include_observations:
|
||||
if has_mental_models:
|
||||
steps.append("If no mental model or it's stale, try search_observations() for consolidated knowledge")
|
||||
else:
|
||||
steps.append("First, try search_observations() - check for consolidated knowledge")
|
||||
# Recall step phrasing varies with whichever upstream tool(s) precede it.
|
||||
if include_observations:
|
||||
steps.append(
|
||||
"If observations are stale OR you need specific details, use recall() for raw facts"
|
||||
if has_mental_models
|
||||
else "If search_observations returns 0 results OR observations are stale, you MUST call recall() for raw facts"
|
||||
)
|
||||
elif has_mental_models:
|
||||
steps.append("If no mental model or it's stale, use recall() for raw facts")
|
||||
else:
|
||||
parts.extend(
|
||||
[
|
||||
"1. First, try search_observations() - check for consolidated knowledge",
|
||||
"2. If search_observations returns 0 results OR observations are stale, you MUST call recall() for raw facts",
|
||||
"3. Use expand() if you need more context on specific memories",
|
||||
"4. When ready, call done() with your answer and supporting IDs",
|
||||
]
|
||||
)
|
||||
steps.append("Call recall() to gather raw facts")
|
||||
steps.append("Use expand() if you need more context on specific memories")
|
||||
steps.append("When ready, call done() with your answer and supporting IDs")
|
||||
parts.extend(f"{idx}. {step}" for idx, step in enumerate(steps, 1))
|
||||
|
||||
parts.extend(
|
||||
[
|
||||
@@ -513,10 +604,16 @@ Just provide the direct answer with proper markdown formatting.
|
||||
CRITICAL: This is a NON-CONVERSATIONAL system. NEVER ask follow-up questions, offer to search again, suggest alternatives, or end with anything like "Would you like me to..." or "Let me know if...". The user cannot reply. Your answer must be complete and self-contained."""
|
||||
|
||||
|
||||
def build_final_system_prompt(mission: str | None = None) -> str:
|
||||
"""Build the final synthesis system prompt, using mission as role when set."""
|
||||
role_section = mission.strip() if mission else _DEFAULT_FINAL_ROLE
|
||||
return _FINAL_SYSTEM_PROMPT_BASE.format(role_section=role_section)
|
||||
def build_final_system_prompt(mission: str | None = None, llm_output_language: str | None = None) -> str:
|
||||
"""Build the final synthesis system prompt, using mission as role when set.
|
||||
|
||||
When ``llm_output_language`` is set, the response is forced into that
|
||||
language regardless of the query/source language.
|
||||
"""
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt, output_language_directive
|
||||
|
||||
role_section = escape_for_prompt(mission.strip()) if mission else _DEFAULT_FINAL_ROLE
|
||||
return _FINAL_SYSTEM_PROMPT_BASE.format(role_section=role_section) + output_language_directive(llm_output_language)
|
||||
|
||||
|
||||
# Backward-compatible constant for non-identity missions
|
||||
|
||||
@@ -1,17 +1,8 @@
|
||||
"""Token counting helpers for reflect prompts and agent control flow."""
|
||||
|
||||
from functools import lru_cache
|
||||
|
||||
import tiktoken
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _get_cl100k_base_encoding() -> tiktoken.Encoding:
|
||||
# tiktoken downloads this encoding on first lookup when it is not cached.
|
||||
# Keep the lookup lazy so importing hindsight_api does not depend on network access.
|
||||
return tiktoken.get_encoding("cl100k_base")
|
||||
from ..token_encoding import count_tokens as _count_tokens
|
||||
|
||||
|
||||
def count_cl100k_tokens(text: str) -> int:
|
||||
"""Return the number of cl100k_base tokens in text."""
|
||||
return len(_get_cl100k_base_encoding().encode(text))
|
||||
return _count_tokens(text)
|
||||
|
||||
@@ -23,6 +23,24 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _prune_nulls(d: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Drop keys whose value is None or an empty collection (``""``, ``[]``, ``{}``).
|
||||
|
||||
Reflect tools dump ``MemoryFact`` / ``ObservationResult`` via ``model_dump()``,
|
||||
which emits every field including the many that are typically null or empty
|
||||
(``context``, ``occurred_start``, ``metadata``, ``tags``, etc.). Stripping
|
||||
these before serializing to JSON for the LLM cuts token cost and removes
|
||||
fields that aren't telling the model anything.
|
||||
|
||||
Callers that need the *presence* of a specific field as a signal (e.g.
|
||||
``source_fact_ids`` for drill-down) must ensure the value is non-empty —
|
||||
pass the upstream flag that populates it (e.g. ``source_facts_max_tokens``
|
||||
> 0 on ``tool_search_observations``) rather than relying on Pydantic
|
||||
emitting ``None``.
|
||||
"""
|
||||
return {k: v for k, v in d.items() if v is not None and v != "" and v != [] and v != {}}
|
||||
|
||||
|
||||
def _document_metadata_from_retain_params(retain_params: Any) -> dict[str, Any] | None:
|
||||
"""Return document metadata stored under retain_params.metadata."""
|
||||
if isinstance(retain_params, str):
|
||||
@@ -214,8 +232,8 @@ async def tool_search_observations(
|
||||
return {
|
||||
"query": query,
|
||||
"count": len(result.results),
|
||||
"observations": [m.model_dump() for m in result.results],
|
||||
"source_facts": {k: v.model_dump() for k, v in (result.source_facts or {}).items()},
|
||||
"observations": [_prune_nulls(m.model_dump()) for m in result.results],
|
||||
"source_facts": {k: _prune_nulls(v.model_dump()) for k, v in (result.source_facts or {}).items()},
|
||||
"is_stale": is_stale,
|
||||
"freshness": freshness,
|
||||
}
|
||||
@@ -282,8 +300,8 @@ async def tool_recall(
|
||||
|
||||
return {
|
||||
"query": query,
|
||||
"memories": [m.model_dump() for m in result.results],
|
||||
"chunks": {k: v.model_dump() for k, v in (result.chunks or {}).items()},
|
||||
"memories": [_prune_nulls(m.model_dump()) for m in result.results],
|
||||
"chunks": {k: _prune_nulls(v.model_dump()) for k, v in (result.chunks or {}).items()},
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -93,6 +93,7 @@ class TokenUsage(BaseModel):
|
||||
input_tokens: int = Field(default=0, description="Number of input/prompt tokens consumed")
|
||||
output_tokens: int = Field(default=0, description="Number of output/completion tokens generated")
|
||||
total_tokens: int = Field(default=0, description="Total tokens (input + output)")
|
||||
cached_tokens: int = Field(default=0, description="Cached/cache-read prompt tokens, when reported by the provider")
|
||||
|
||||
def __add__(self, other: "TokenUsage") -> "TokenUsage":
|
||||
"""Allow aggregating token usage from multiple calls."""
|
||||
@@ -100,6 +101,7 @@ class TokenUsage(BaseModel):
|
||||
input_tokens=self.input_tokens + other.input_tokens,
|
||||
output_tokens=self.output_tokens + other.output_tokens,
|
||||
total_tokens=self.total_tokens + other.total_tokens,
|
||||
cached_tokens=self.cached_tokens + other.cached_tokens,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ import json
|
||||
import logging
|
||||
import re
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import TypedDict
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
@@ -105,6 +106,18 @@ class BankProfile(TypedDict):
|
||||
mission: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class BankProfileResult:
|
||||
"""Result of a get-or-create bank lookup.
|
||||
|
||||
``created`` is True when the bank row was freshly inserted on this call,
|
||||
which callers use to drive the one-time HINDSIGHT_API_DEFAULT_BANK_TEMPLATE hook.
|
||||
"""
|
||||
|
||||
profile: BankProfile
|
||||
created: bool
|
||||
|
||||
|
||||
class MissionMergeResponse(BaseModel):
|
||||
"""LLM response for mission merge."""
|
||||
|
||||
@@ -123,8 +136,8 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
|
||||
Returns:
|
||||
BankProfile with name, typed DispositionTraits, and mission
|
||||
"""
|
||||
profile, _ = await get_or_create_bank_profile(pool, bank_id)
|
||||
return profile
|
||||
result = await get_or_create_bank_profile(pool, bank_id)
|
||||
return result.profile
|
||||
|
||||
|
||||
async def get_bank_profile_if_exists(pool, bank_id: str) -> BankProfile | None:
|
||||
@@ -162,70 +175,89 @@ async def get_bank_profile_if_exists(pool, bank_id: str) -> BankProfile | None:
|
||||
)
|
||||
|
||||
|
||||
async def get_or_create_bank_profile(pool, bank_id: str) -> tuple[BankProfile, bool]:
|
||||
async def get_or_create_bank_profile(pool, bank_id: str) -> BankProfileResult:
|
||||
"""
|
||||
Get bank profile, auto-creating with defaults if it doesn't exist.
|
||||
|
||||
Same as get_bank_profile, but also returns a flag indicating whether the
|
||||
bank was freshly created on this call. Used by the memory engine to apply
|
||||
the HINDSIGHT_API_DEFAULT_BANK_TEMPLATE hook on first bank creation.
|
||||
Same as get_bank_profile, but also reports whether the bank was freshly
|
||||
created on this call (``BankProfileResult.created``). Used by the memory
|
||||
engine to apply the HINDSIGHT_API_DEFAULT_BANK_TEMPLATE hook on first bank
|
||||
creation.
|
||||
|
||||
Returns:
|
||||
Tuple of (BankProfile, created) where created is True if the bank
|
||||
did not exist before this call.
|
||||
Acquires its own connection. When the caller already holds a connection and
|
||||
wants the bank row to share its transaction (so the lazy bank-create commits
|
||||
or rolls back atomically with the caller's write), use
|
||||
``get_or_create_bank_profile_on_conn`` instead.
|
||||
"""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Try to get existing bank
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT name, disposition, mission
|
||||
FROM {fq_table("banks")} WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
return await get_or_create_bank_profile_on_conn(conn, bank_id, ops=pool.ops)
|
||||
|
||||
|
||||
async def get_or_create_bank_profile_on_conn(conn, bank_id: str, *, ops) -> BankProfileResult:
|
||||
"""
|
||||
Connection-bound variant of ``get_or_create_bank_profile``.
|
||||
|
||||
Runs the SELECT, the ``INSERT ... ON CONFLICT DO NOTHING`` and the per-bank
|
||||
vector index creation on the caller-supplied ``conn``. When ``conn`` is
|
||||
inside an open transaction, the lazy bank-create therefore commits (or rolls
|
||||
back) atomically with whatever bank-scoped write the caller performs on the
|
||||
same connection — closing the window where a freshly-created bank could
|
||||
outlive a write that ultimately failed.
|
||||
|
||||
``ops`` is the backend's dialect ops object (``backend.ops``), needed for
|
||||
per-bank vector index DDL.
|
||||
"""
|
||||
# Try to get existing bank
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT name, disposition, mission
|
||||
FROM {fq_table("banks")} WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if row:
|
||||
# asyncpg returns JSONB as a string, so parse it
|
||||
disposition_data = row["disposition"]
|
||||
if isinstance(disposition_data, str):
|
||||
disposition_data = json.loads(disposition_data)
|
||||
|
||||
return BankProfileResult(
|
||||
profile=BankProfile(
|
||||
name=row["name"],
|
||||
disposition=DispositionTraits(**disposition_data),
|
||||
mission=row["mission"] or "",
|
||||
),
|
||||
created=False,
|
||||
)
|
||||
|
||||
if row:
|
||||
# asyncpg returns JSONB as a string, so parse it
|
||||
disposition_data = row["disposition"]
|
||||
if isinstance(disposition_data, str):
|
||||
disposition_data = json.loads(disposition_data)
|
||||
# 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()
|
||||
inserted = await conn.fetchval(
|
||||
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
|
||||
json.dumps(DEFAULT_DISPOSITION),
|
||||
"",
|
||||
internal_id,
|
||||
)
|
||||
|
||||
return (
|
||||
BankProfile(
|
||||
name=row["name"],
|
||||
disposition=DispositionTraits(**disposition_data),
|
||||
mission=row["mission"] or "",
|
||||
),
|
||||
False,
|
||||
)
|
||||
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)
|
||||
|
||||
# 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()
|
||||
inserted = await conn.fetchval(
|
||||
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
|
||||
json.dumps(DEFAULT_DISPOSITION),
|
||||
"",
|
||||
internal_id,
|
||||
)
|
||||
|
||||
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=pool.ops)
|
||||
|
||||
return (
|
||||
BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission=""),
|
||||
created,
|
||||
)
|
||||
return BankProfileResult(
|
||||
profile=BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission=""),
|
||||
created=created,
|
||||
)
|
||||
|
||||
|
||||
async def update_bank_disposition(pool, bank_id: str, disposition: dict[str, int]) -> None:
|
||||
|
||||
@@ -4,29 +4,75 @@ Embedding generation utilities for memory units.
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Literal, Protocol
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
EmbeddingInputType = Literal["document", "query"]
|
||||
|
||||
def generate_embedding(embeddings_backend, text: str) -> list[float]:
|
||||
|
||||
class EmbeddingsBackend(Protocol):
|
||||
"""Minimal duck-typed surface used by retain/recall — the concrete `Embeddings`
|
||||
ABC supplies default implementations that delegate to `encode()`."""
|
||||
|
||||
@property
|
||||
def dimension(self) -> int: ...
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]: ...
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]: ...
|
||||
|
||||
|
||||
def _validate_embedding_vector(vector: list[float], *, index: int, expected_dimension: int) -> list[float]:
|
||||
actual_dimension = len(vector)
|
||||
if actual_dimension == 0:
|
||||
raise RuntimeError(f"embedding {index} has dimension 0; expected {expected_dimension}")
|
||||
if actual_dimension != expected_dimension:
|
||||
raise RuntimeError(f"embedding {index} has dimension {actual_dimension}; expected {expected_dimension}")
|
||||
return vector
|
||||
|
||||
|
||||
def generate_embedding(
|
||||
embeddings_backend: EmbeddingsBackend, text: str, input_type: EmbeddingInputType = "document"
|
||||
) -> list[float]:
|
||||
"""
|
||||
Generate embedding for text using the provided embeddings backend.
|
||||
|
||||
Args:
|
||||
embeddings_backend: Embeddings instance to use for encoding
|
||||
text: Text to embed
|
||||
input_type: Whether text is retained document text or recall/search query text.
|
||||
|
||||
Returns:
|
||||
Embedding vector (dimension depends on embeddings backend)
|
||||
"""
|
||||
try:
|
||||
embeddings = embeddings_backend.encode([text])
|
||||
return embeddings[0]
|
||||
embeddings = _encode_with_input_type(embeddings_backend, [text], input_type)
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to generate embedding: {str(e)}")
|
||||
|
||||
if len(embeddings) != 1:
|
||||
raise RuntimeError(
|
||||
f"Embeddings backend returned {len(embeddings)} vectors for 1 input text; expected exact 1:1 alignment"
|
||||
)
|
||||
return _validate_embedding_vector(
|
||||
embeddings[0],
|
||||
index=0,
|
||||
expected_dimension=embeddings_backend.dimension,
|
||||
)
|
||||
|
||||
async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> list[list[float]]:
|
||||
|
||||
def _encode_with_input_type(
|
||||
embeddings_backend: EmbeddingsBackend, texts: list[str], input_type: EmbeddingInputType
|
||||
) -> list[list[float]]:
|
||||
if input_type == "query":
|
||||
return embeddings_backend.encode_query(texts)
|
||||
return embeddings_backend.encode_documents(texts)
|
||||
|
||||
|
||||
async def generate_embeddings_batch(
|
||||
embeddings_backend: EmbeddingsBackend, texts: list[str], input_type: EmbeddingInputType = "document"
|
||||
) -> list[list[float]]:
|
||||
"""
|
||||
Generate embeddings for multiple texts using the provided embeddings backend.
|
||||
|
||||
@@ -36,17 +82,14 @@ async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> lis
|
||||
Args:
|
||||
embeddings_backend: Embeddings instance to use for encoding
|
||||
texts: List of texts to embed
|
||||
input_type: Whether texts are retained documents or recall/search queries.
|
||||
|
||||
Returns:
|
||||
List of embeddings in same order as input texts
|
||||
"""
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
embeddings = await loop.run_in_executor(
|
||||
None,
|
||||
embeddings_backend.encode,
|
||||
texts,
|
||||
)
|
||||
embeddings = await loop.run_in_executor(None, _encode_with_input_type, embeddings_backend, texts, input_type)
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to generate batch embeddings: {str(e)}")
|
||||
|
||||
@@ -59,4 +102,7 @@ async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> lis
|
||||
"expected exact 1:1 alignment"
|
||||
)
|
||||
|
||||
return embeddings
|
||||
return [
|
||||
_validate_embedding_vector(embedding, index=index, expected_dimension=embeddings_backend.dimension)
|
||||
for index, embedding in enumerate(embeddings)
|
||||
]
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
"""
|
||||
Entity processing for retain pipeline.
|
||||
|
||||
Handles entity extraction, resolution, and link creation for stored facts.
|
||||
Handles entity extraction and resolution for stored facts.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from . import link_utils
|
||||
from .types import EntityLink, ProcessedFact
|
||||
from .types import ProcessedFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -76,8 +76,7 @@ async def resolve_entities(
|
||||
entity_labels: Optional entity label taxonomy
|
||||
|
||||
Returns:
|
||||
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids)
|
||||
to pass to build_entity_links().
|
||||
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids).
|
||||
"""
|
||||
if not unit_ids or not facts:
|
||||
return [], [], {}
|
||||
@@ -99,68 +98,3 @@ async def resolve_entities(
|
||||
log_buffer,
|
||||
entity_labels=entity_labels,
|
||||
)
|
||||
|
||||
|
||||
async def build_entity_links(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
resolved_entity_ids: list[str],
|
||||
entity_to_unit: list[tuple],
|
||||
unit_to_entity_ids: dict[str, list[str]],
|
||||
log_buffer: list[str] = None,
|
||||
skip_unit_entities_insert: bool = False,
|
||||
ops=None,
|
||||
) -> list[EntityLink]:
|
||||
"""
|
||||
Build entity links for UI graph visualization.
|
||||
|
||||
Queries unit_entities to find shared entities between new and existing units,
|
||||
then generates EntityLink objects. When called from Phase 3 (post-transaction),
|
||||
set skip_unit_entities_insert=True since unit_entities were already inserted
|
||||
in Phase 2.
|
||||
|
||||
Args:
|
||||
entity_resolver: EntityResolver instance
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
unit_ids: Actual unit IDs (must already be inserted in the DB)
|
||||
resolved_entity_ids: From resolve_entities()
|
||||
entity_to_unit: From resolve_entities()
|
||||
unit_to_entity_ids: From resolve_entities()
|
||||
log_buffer: Optional buffer for detailed logging
|
||||
skip_unit_entities_insert: Skip unit_entities INSERT (already done in Phase 2)
|
||||
ops: DataAccessOps instance (from backend.ops)
|
||||
|
||||
Returns:
|
||||
List of EntityLink objects for batch insertion
|
||||
"""
|
||||
return await link_utils.build_entity_links_from_resolved(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id,
|
||||
unit_ids,
|
||||
resolved_entity_ids,
|
||||
entity_to_unit,
|
||||
unit_to_entity_ids,
|
||||
log_buffer,
|
||||
skip_unit_entities_insert=skip_unit_entities_insert,
|
||||
ops=ops,
|
||||
)
|
||||
|
||||
|
||||
async def insert_entity_links_batch(conn, entity_links: list[EntityLink], bank_id: str, ops=None) -> None:
|
||||
"""
|
||||
Insert entity links in batch.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
entity_links: List of EntityLink objects
|
||||
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
|
||||
ops: DataAccessOps instance (from backend.ops)
|
||||
"""
|
||||
if not entity_links:
|
||||
return
|
||||
|
||||
await link_utils.insert_entity_links_batch(conn, entity_links, bank_id, ops=ops)
|
||||
|
||||
@@ -10,12 +10,13 @@ import json
|
||||
import logging
|
||||
import re
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Literal, cast
|
||||
from typing import Any, Literal, cast
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, create_model, field_validator
|
||||
|
||||
from ...config import get_config
|
||||
from ..llm_wrapper import LLMConfig, OutputTooLongError, sanitize_llm_output
|
||||
from ..operation_metadata import RetainExtractionErrors
|
||||
from ..response_models import TokenUsage
|
||||
from .entity_labels import (
|
||||
EntityLabelsConfig,
|
||||
@@ -510,11 +511,11 @@ LANGUAGE: MANDATORY — Detect the language of the input text and produce ALL ou
|
||||
FACT FORMAT - BE CONCISE
|
||||
══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
1. **what**: Core fact - concise but complete (1-2 sentences max)
|
||||
2. **when**: Temporal info if mentioned. "N/A" if none. Use day name when known.
|
||||
3. **where**: Location if relevant. "N/A" if none.
|
||||
4. **who**: People involved with relationships. "N/A" if just general info.
|
||||
5. **why**: Context/significance ONLY if important. "N/A" if obvious.
|
||||
1. "what": Core fact - concise but complete (1-2 sentences max)
|
||||
2. "when": Temporal info if mentioned. "N/A" if none. Use day name when known.
|
||||
3. "where": Location if relevant. "N/A" if none.
|
||||
4. "who": People involved with relationships. "N/A" if just general info.
|
||||
5. "why": Context/significance ONLY if important. "N/A" if obvious.
|
||||
|
||||
CONCISENESS: Capture the essence, not every word. One good sentence beats three mediocre ones.
|
||||
|
||||
@@ -887,17 +888,15 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
extraction_mode = config.retain_extraction_mode
|
||||
extract_causal_links = config.retain_extract_causal_links
|
||||
|
||||
# Build retain_mission section if set - injected before the mode-specific guidelines
|
||||
retain_mission = getattr(config, "retain_mission", None)
|
||||
if retain_mission:
|
||||
retain_mission_section = (
|
||||
f"══════════════════════════════════════════════════════════════════════════\n"
|
||||
f"FOCUS — What to retain for this bank\n"
|
||||
f"══════════════════════════════════════════════════════════════════════════\n\n"
|
||||
f"{retain_mission}\n\n"
|
||||
)
|
||||
else:
|
||||
retain_mission_section = ""
|
||||
# The per-bank retain mission is NOT baked into this system prompt: it would
|
||||
# make the prompt bank-specific and force a separate Gemini context cache per
|
||||
# mission (one per bank). Instead the prompt is bank-agnostic so a single
|
||||
# CachedContent serves every bank, and the mission rides in the per-request
|
||||
# user message via _retain_mission_preamble(). The {retain_mission_section}
|
||||
# placeholder is kept (templates still reference it) but always empty here.
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt
|
||||
|
||||
retain_mission_section = ""
|
||||
|
||||
# Select base prompt based on extraction mode
|
||||
if extraction_mode == "custom":
|
||||
@@ -910,7 +909,7 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
base_prompt = CUSTOM_FACT_EXTRACTION_PROMPT
|
||||
prompt = base_prompt.format(
|
||||
retain_mission_section=retain_mission_section,
|
||||
custom_instructions=config.retain_custom_instructions,
|
||||
custom_instructions=escape_for_prompt(config.retain_custom_instructions),
|
||||
)
|
||||
elif extraction_mode == "verbose":
|
||||
prompt = VERBOSE_FACT_EXTRACTION_PROMPT.format(
|
||||
@@ -947,6 +946,16 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
if labels_section:
|
||||
prompt = prompt + labels_section
|
||||
|
||||
# Force the LLM to emit fact text in the configured language, regardless of
|
||||
# the source content's language. Same directive is applied to consolidation
|
||||
# and reflect so HINDSIGHT_API_LLM_OUTPUT_LANGUAGE has a uniform effect
|
||||
# across the pipeline. This is independent of the BM25 indexing language
|
||||
# (HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE) by design — search
|
||||
# tokenization and LLM output language are separate concerns.
|
||||
from ..prompt_utils import output_language_directive
|
||||
|
||||
prompt = prompt + output_language_directive(getattr(config, "llm_output_language", None))
|
||||
|
||||
response_schema = base_response_class
|
||||
|
||||
if labels_cfg and labels_cfg.attributes:
|
||||
@@ -984,6 +993,26 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
return prompt, response_schema
|
||||
|
||||
|
||||
def _retain_mission_preamble(config) -> str:
|
||||
"""The bank's retain mission, formatted for the per-request user message.
|
||||
|
||||
Kept OUT of the cached system prompt (which must stay bank-agnostic so one
|
||||
CachedContent serves every bank — otherwise each distinct mission spawns its
|
||||
own cache) and prepended to the user message instead. Returns "" when unset.
|
||||
No brace-escaping needed: unlike the system template, the user message is
|
||||
used verbatim, not passed through str.format().
|
||||
"""
|
||||
retain_mission = getattr(config, "retain_mission", None)
|
||||
if not retain_mission:
|
||||
return ""
|
||||
return (
|
||||
"══════════════════════════════════════════════════════════════════════════\n"
|
||||
"FOCUS — What to retain for this bank (takes priority over the general guidelines)\n"
|
||||
"══════════════════════════════════════════════════════════════════════════\n\n"
|
||||
f"{retain_mission}\n\n"
|
||||
)
|
||||
|
||||
|
||||
def _build_user_message(
|
||||
chunk: str,
|
||||
chunk_index: int,
|
||||
@@ -992,8 +1021,14 @@ def _build_user_message(
|
||||
context: str,
|
||||
metadata: dict[str, str] | None = None,
|
||||
agent_name: str | None = None,
|
||||
mission_preamble: str = "",
|
||||
) -> str:
|
||||
"""Build user message for fact extraction."""
|
||||
"""Build user message for fact extraction.
|
||||
|
||||
``mission_preamble`` (the bank's retain mission, possibly empty) is prepended
|
||||
so the bank-specific focus lives in the variable user turn rather than the
|
||||
cached, bank-agnostic system prompt.
|
||||
"""
|
||||
from .orchestrator import parse_datetime_flexible
|
||||
|
||||
sanitized_chunk = _sanitize_text(chunk)
|
||||
@@ -1012,9 +1047,21 @@ def _build_user_message(
|
||||
|
||||
narrator_section = ""
|
||||
if agent_name:
|
||||
narrator_section = f'\nNarrator: {agent_name} (AI agent — first-person statements like "I did X" are the agent\'s own actions; classify as "assistant")'
|
||||
narrator_section = (
|
||||
f"\nNarrator: {agent_name} (the AI agent whose memory this is). By default, "
|
||||
f'first-person statements like "I did X" are {agent_name}\'s own actions → classify as '
|
||||
f'"assistant".'
|
||||
)
|
||||
# Only defer to the Context when one was actually provided — otherwise this
|
||||
# clause points at a "Context: none" line and just adds noise.
|
||||
if context:
|
||||
narrator_section += (
|
||||
" BUT the Context above takes precedence: if it identifies a different "
|
||||
"first-person speaker (e.g. a user or customer in a transcript), attribute those "
|
||||
'statements to that speaker and classify them as "world", not "assistant".'
|
||||
)
|
||||
|
||||
return f"""Extract facts from the following text chunk.
|
||||
return f"""{mission_preamble}Extract facts from the following text chunk.
|
||||
|
||||
Chunk: {chunk_index + 1}/{total_chunks}
|
||||
Event Date: {event_date_str}
|
||||
@@ -1040,12 +1087,15 @@ def _build_request_body(llm_config, config, prompt: str, user_message: str, resp
|
||||
if llm_config.provider == "openai" and llm_config._provider_impl.openai_service_tier:
|
||||
request_body["service_tier"] = llm_config._provider_impl.openai_service_tier
|
||||
|
||||
# Add response_format (JSON schema)
|
||||
# Add response_format (JSON schema). The batch path builds the request body
|
||||
# directly instead of going through LLMProvider.call(), so honour
|
||||
# HINDSIGHT_API_LLM_STRICT_SCHEMA here too: strict=True grammar-enforces the
|
||||
# output on capable backends rather than relying on the model to emit clean JSON.
|
||||
if hasattr(response_schema, "model_json_schema"):
|
||||
schema = response_schema.model_json_schema()
|
||||
request_body["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {"name": "facts", "schema": schema},
|
||||
"json_schema": {"name": "facts", "schema": schema, "strict": config.llm_strict_schema},
|
||||
}
|
||||
|
||||
return request_body
|
||||
@@ -1081,8 +1131,38 @@ async def _extract_facts_from_chunk(
|
||||
extraction_mode = config.retain_extraction_mode
|
||||
extract_causal_links = config.retain_extract_causal_links
|
||||
|
||||
# Build user message using helper function
|
||||
user_message = _build_user_message(chunk, chunk_index, total_chunks, event_date, context, metadata, agent_name)
|
||||
# Build user message — the bank mission rides here (not in the cached prefix).
|
||||
user_message = _build_user_message(
|
||||
chunk,
|
||||
chunk_index,
|
||||
total_chunks,
|
||||
event_date,
|
||||
context,
|
||||
metadata,
|
||||
agent_name,
|
||||
mission_preamble=_retain_mission_preamble(config),
|
||||
)
|
||||
|
||||
# Opt into context caching when the provider supports it. The prompt and
|
||||
# response_schema are bank-agnostic (the mission lives in the user message),
|
||||
# so one cached prefix serves every bank; reusing it across many small-payload
|
||||
# retain calls dramatically lowers per-call input
|
||||
# cost. ``get_or_create_cached_prefix`` returns None when caching is
|
||||
# disabled, unsupported, or the prefix is too small; the LLM call
|
||||
# transparently falls back to the uncached path in that case.
|
||||
cached_prefix_name: str | None = None
|
||||
provider_impl = getattr(llm_config, "_provider_impl", None)
|
||||
if provider_impl is not None and provider_impl.supports_prompt_caching():
|
||||
try:
|
||||
cached_prefix_name = await provider_impl.get_or_create_cached_prefix(
|
||||
system_instruction=prompt,
|
||||
response_schema=response_schema,
|
||||
)
|
||||
except Exception:
|
||||
# Caching is a soft optimisation — never let a cache-side
|
||||
# error block a retain operation.
|
||||
logger.exception("Cache prefix lookup failed; falling back to uncached call")
|
||||
cached_prefix_name = None
|
||||
|
||||
# Retry logic for JSON validation errors
|
||||
# Use retain-specific overrides if set, otherwise fall back to global LLM config
|
||||
@@ -1103,7 +1183,7 @@ async def _extract_facts_from_chunk(
|
||||
config.retain_llm_max_backoff if config.retain_llm_max_backoff is not None else config.llm_max_backoff
|
||||
)
|
||||
|
||||
extraction_response_json, call_usage = await llm_config.call(
|
||||
call_kwargs: dict[str, Any] = dict(
|
||||
messages=[{"role": "system", "content": prompt}, {"role": "user", "content": user_message}],
|
||||
response_format=response_schema,
|
||||
scope="retain_extract_facts",
|
||||
@@ -1115,6 +1195,10 @@ async def _extract_facts_from_chunk(
|
||||
skip_validation=True, # Get raw JSON, we'll validate leniently
|
||||
return_usage=True,
|
||||
)
|
||||
if cached_prefix_name is not None:
|
||||
call_kwargs["cached_prefix"] = cached_prefix_name
|
||||
|
||||
extraction_response_json, call_usage = await llm_config.call(**call_kwargs)
|
||||
usage = usage + call_usage # Aggregate usage across retries
|
||||
|
||||
# Lenient parsing of facts from raw JSON
|
||||
@@ -1129,11 +1213,14 @@ async def _extract_facts_from_chunk(
|
||||
)
|
||||
continue
|
||||
else:
|
||||
logger.warning(
|
||||
f"LLM returned non-dict JSON after {llm_max_retries} attempts: {type(extraction_response_json).__name__}. "
|
||||
f"Raw: {str(extraction_response_json)[:500]}"
|
||||
# A non-dict response is malformed (the schema is {"facts": [...]}).
|
||||
# Raise instead of returning [] so the failure propagates to the
|
||||
# worker's retry machinery and ultimately fails loudly — never
|
||||
# silently commit the document with 0 facts. See issue #1833.
|
||||
raise RuntimeError(
|
||||
f"Fact extraction failed: LLM returned non-dict JSON after {llm_max_retries} attempts "
|
||||
f"({type(extraction_response_json).__name__}). Raw: {str(extraction_response_json)[:500]}"
|
||||
)
|
||||
return [], usage
|
||||
|
||||
raw_facts = extraction_response_json.get("facts", [])
|
||||
|
||||
@@ -1624,6 +1711,39 @@ logger = logging.getLogger(__name__)
|
||||
SECONDS_PER_FACT = 0.01
|
||||
|
||||
|
||||
async def _write_batch_extraction_errors(
|
||||
pool: Any,
|
||||
operation_id: str | None,
|
||||
schema: str | None,
|
||||
errors: RetainExtractionErrors,
|
||||
) -> None:
|
||||
"""Persist non-fatal Batch API extraction errors into operation result_metadata."""
|
||||
if not pool or not operation_id or errors.count == 0:
|
||||
return
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..task_backend import fq_table
|
||||
|
||||
# `errors` is the complete set for this extraction run, so overwrite the
|
||||
# extraction_errors_* keys rather than folding in what's already stored. On
|
||||
# batch crash recovery the resumed batch reprocesses every result and
|
||||
# recomputes `errors` from scratch; reading + merging the prior run's
|
||||
# counters here would double-count them. The SQL `||` merge still preserves
|
||||
# unrelated keys (e.g. batch_id) already on result_metadata.
|
||||
table = fq_table("async_operations", schema)
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET result_metadata = COALESCE(result_metadata, '{{}}'::jsonb) || $2::jsonb,
|
||||
updated_at = now()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
operation_id,
|
||||
json.dumps(errors.to_dict()),
|
||||
)
|
||||
|
||||
|
||||
async def extract_facts_from_contents_batch_api(
|
||||
contents: list[RetainContent],
|
||||
llm_config,
|
||||
@@ -1662,8 +1782,11 @@ async def extract_facts_from_contents_batch_api(
|
||||
|
||||
# Check if provider supports batch API
|
||||
if not await llm_config._provider_impl.supports_batch_api():
|
||||
logger.warning(f"Batch API not supported for provider {llm_config.provider}, falling back to sync mode")
|
||||
return await extract_facts_from_contents(contents, llm_config, agent_name, config, pool, operation_id, schema)
|
||||
raise RuntimeError(
|
||||
f"retain_batch_enabled=True but provider '{llm_config.provider}' does not "
|
||||
f"support the batch API. This should have been caught at startup — check "
|
||||
f"HINDSIGHT_API_RETAIN_BATCH_ENABLED and your LLM provider configuration."
|
||||
)
|
||||
|
||||
# Check if we're resuming an existing batch (crash recovery)
|
||||
batch_id = None
|
||||
@@ -1712,6 +1835,7 @@ async def extract_facts_from_contents_batch_api(
|
||||
item.context,
|
||||
item.metadata or None,
|
||||
agent_name,
|
||||
mission_preamble=_retain_mission_preamble(config),
|
||||
)
|
||||
|
||||
# Build request body using helper function
|
||||
@@ -1797,6 +1921,7 @@ async def extract_facts_from_contents_batch_api(
|
||||
all_facts_from_llm = []
|
||||
chunks_metadata = []
|
||||
total_usage = TokenUsage()
|
||||
extraction_errors = RetainExtractionErrors()
|
||||
|
||||
for chunk_idx, (chunk_content, content_index, chunk_index_in_content, event_date, context) in enumerate(
|
||||
all_chunks_info
|
||||
@@ -1805,7 +1930,9 @@ async def extract_facts_from_contents_batch_api(
|
||||
result = results_by_id.get(custom_id)
|
||||
|
||||
if not result:
|
||||
logger.warning(f"Missing result for {custom_id}, skipping")
|
||||
message = f"{custom_id}: missing batch result"
|
||||
logger.warning(message)
|
||||
extraction_errors.add(message)
|
||||
chunks_metadata.append(
|
||||
ChunkMetadata(
|
||||
chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
|
||||
@@ -1815,7 +1942,9 @@ async def extract_facts_from_contents_batch_api(
|
||||
|
||||
# Check for errors
|
||||
if result.get("error"):
|
||||
logger.error(f"Error in {custom_id}: {result['error']}")
|
||||
message = f"{custom_id}: {result['error']}"
|
||||
logger.error(f"Error in {message}")
|
||||
extraction_errors.add(message)
|
||||
chunks_metadata.append(
|
||||
ChunkMetadata(
|
||||
chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
|
||||
@@ -1828,7 +1957,9 @@ async def extract_facts_from_contents_batch_api(
|
||||
choices = response_body.get("choices", [])
|
||||
|
||||
if not choices:
|
||||
logger.warning(f"No choices in response for {custom_id}")
|
||||
message = f"{custom_id}: no choices in response"
|
||||
logger.warning(message)
|
||||
extraction_errors.add(message)
|
||||
chunks_metadata.append(
|
||||
ChunkMetadata(
|
||||
chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
|
||||
@@ -1843,7 +1974,9 @@ async def extract_facts_from_contents_batch_api(
|
||||
try:
|
||||
extraction_response_json = json.loads(content_str)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"Failed to parse JSON for {custom_id}: {e}")
|
||||
message = f"{custom_id}: failed to parse JSON: {e}"
|
||||
logger.error(message)
|
||||
extraction_errors.add(message)
|
||||
chunks_metadata.append(
|
||||
ChunkMetadata(
|
||||
chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
|
||||
@@ -2016,7 +2149,9 @@ async def extract_facts_from_contents_batch_api(
|
||||
fact = Fact(fact=combined_text, fact_type=fact_type, **fact_data)
|
||||
chunk_facts.append(fact)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create Fact model for fact {i}: {e}")
|
||||
message = f"{custom_id}: failed to create Fact model for fact {i}: {e}"
|
||||
logger.error(message)
|
||||
extraction_errors.add(message)
|
||||
continue
|
||||
|
||||
all_facts_from_llm.extend(chunk_facts)
|
||||
@@ -2081,6 +2216,8 @@ async def extract_facts_from_contents_batch_api(
|
||||
# Step 8: Auto-tag facts from label groups with tag=True
|
||||
_inject_label_tags(extracted_facts, config)
|
||||
|
||||
await _write_batch_extraction_errors(pool, operation_id, schema, extraction_errors)
|
||||
|
||||
logger.info(f"Batch API extracted {len(extracted_facts)} facts from {len(all_chunks_info)} chunks")
|
||||
|
||||
return extracted_facts, chunks_metadata, total_usage
|
||||
@@ -2195,7 +2332,9 @@ async def extract_facts_from_contents(
|
||||
fact_extraction_tasks.append(task)
|
||||
|
||||
# Step 2: Wait for all fact extractions to complete.
|
||||
# Use return_exceptions=True so one content item failure doesn't discard the rest.
|
||||
# return_exceptions=True so a failing item doesn't cancel its still-running
|
||||
# siblings (which would leave orphaned LLM calls / partial work); we await
|
||||
# them all, then propagate.
|
||||
all_fact_results = await asyncio.gather(*fact_extraction_tasks, return_exceptions=True)
|
||||
|
||||
# Step 3: Flatten and convert to typed objects
|
||||
@@ -2206,14 +2345,18 @@ async def extract_facts_from_contents(
|
||||
global_chunk_idx = 0
|
||||
global_fact_idx = 0
|
||||
|
||||
# Filter out failed content items
|
||||
# Never silently drop a document's memory. Any extraction failure (provider
|
||||
# rate-limit / timeout / 5xx, malformed response, token-limit, etc.)
|
||||
# propagates so the streaming producer surfaces it and the worker's
|
||||
# RetryTaskAt machinery retries the task — and ultimately fails it *loudly*
|
||||
# if the problem persists. Swallowing the error and substituting an empty
|
||||
# result here used to commit the document with 0 facts and mark the
|
||||
# operation `completed`, losing the memory with no signal. See issue #1833.
|
||||
valid_results = []
|
||||
for content, result in zip(contents, all_fact_results):
|
||||
if isinstance(result, Exception):
|
||||
logger.warning(f"Content extraction failed (skipping): {type(result).__name__}: {result}")
|
||||
valid_results.append((content, ([], [], TokenUsage())))
|
||||
else:
|
||||
valid_results.append((content, result))
|
||||
raise result
|
||||
valid_results.append((content, result))
|
||||
|
||||
for content_index, (content, (facts_from_llm, chunks_from_llm, content_usage)) in enumerate(valid_results):
|
||||
total_usage = total_usage + content_usage
|
||||
|
||||
@@ -147,12 +147,13 @@ async def ensure_bank_exists(conn, bank_id: str, ops=None) -> None:
|
||||
internal_id = uuid.uuid4()
|
||||
inserted = await conn.fetchval(
|
||||
f"""
|
||||
INSERT INTO {fq_table("banks")} (bank_id, disposition, mission, internal_id)
|
||||
VALUES ($1, $2::jsonb, $3, $4)
|
||||
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,
|
||||
@@ -321,6 +322,15 @@ async def handle_document_tracking(
|
||||
f"[RETAIN] Document {document_id} re-ingested: invalidated "
|
||||
f"{invalidated} observation(s) derived from {len(existing_unit_ids)} outgoing memory_units"
|
||||
)
|
||||
# Capture link-recompute victims BEFORE the cascade. Same staleness
|
||||
# applies on upsert as on explicit delete: surviving units in OTHER
|
||||
# documents that linked to these doomed units are about to lose
|
||||
# those links. ``ops`` may be None for older callers that haven't
|
||||
# been wired up — skip enqueue in that case rather than crash.
|
||||
if ops is not None:
|
||||
from ..graph_maintenance import enqueue_relink_victims
|
||||
|
||||
await enqueue_relink_victims(conn, bank_id, [str(uid) for uid in existing_unit_ids], ops=ops)
|
||||
# Explicitly delete memory_units by document_id BEFORE deleting the
|
||||
# document row. The CASCADE from documents→chunks→memory_units only
|
||||
# catches units that have a non-NULL chunk_id FK. Units with chunk_id=NULL
|
||||
|
||||
@@ -5,10 +5,9 @@ Link creation utilities for temporal, semantic, and entity links.
|
||||
import logging
|
||||
import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from uuid import UUID
|
||||
|
||||
from ..._vector_index import ann_search_tuning_settings, configured_vector_extension
|
||||
from ..memory_engine import fq_table
|
||||
from .types import EntityLink
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -366,136 +365,6 @@ async def resolve_entities_only(
|
||||
return resolved_entity_ids, entity_to_unit, unit_to_entity_ids
|
||||
|
||||
|
||||
async def build_entity_links_from_resolved(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
resolved_entity_ids: list[str],
|
||||
entity_to_unit: list[tuple],
|
||||
unit_to_entity_ids: dict[str, list[str]],
|
||||
log_buffer: list[str] = None,
|
||||
skip_unit_entities_insert: bool = False,
|
||||
ops=None,
|
||||
) -> list["EntityLink"]:
|
||||
"""
|
||||
Build entity links between units that share entities.
|
||||
|
||||
Queries unit_entities to find which existing units share entities with the
|
||||
new units, then generates EntityLink objects for UI graph visualization.
|
||||
|
||||
Args:
|
||||
entity_resolver: EntityResolver instance
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
unit_ids: Actual unit IDs (must already be inserted in the DB)
|
||||
resolved_entity_ids: Entity IDs from resolve_entities_only
|
||||
entity_to_unit: Mapping from resolve_entities_only
|
||||
unit_to_entity_ids: Mapping from resolve_entities_only
|
||||
log_buffer: Optional logging buffer
|
||||
skip_unit_entities_insert: If True, skip unit_entities INSERT (already done in Phase 2)
|
||||
|
||||
Returns:
|
||||
List of EntityLink objects for batch insertion
|
||||
"""
|
||||
if not resolved_entity_ids:
|
||||
return []
|
||||
|
||||
if not skip_unit_entities_insert:
|
||||
# Insert unit-entity links (used in fallback path where Phase 2 didn't do this)
|
||||
substep_start = time.time()
|
||||
unit_entity_pairs = []
|
||||
for idx, (unit_id, _local_idx, fact_date) in enumerate(entity_to_unit):
|
||||
# Propagate the unit's fact_date so entity_cooccurrences.last_cooccurred
|
||||
# reflects the event timeline, not the ingest moment.
|
||||
unit_entity_pairs.append((unit_id, resolved_entity_ids[idx], fact_date))
|
||||
|
||||
await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_start:.3f}s",
|
||||
level="debug",
|
||||
)
|
||||
|
||||
# Build entity links between units that share entities
|
||||
substep_start = time.time()
|
||||
all_entity_ids = set()
|
||||
for entity_ids_list in unit_to_entity_ids.values():
|
||||
all_entity_ids.update(entity_ids_list)
|
||||
|
||||
_log(log_buffer, f" [6.3] Creating entity links for {len(all_entity_ids)} unique entities...", level="debug")
|
||||
|
||||
MAX_LINKS_PER_ENTITY = 10
|
||||
|
||||
entity_to_units = {}
|
||||
if all_entity_ids:
|
||||
query_start = time.time()
|
||||
import uuid
|
||||
|
||||
entity_id_list = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in all_entity_ids]
|
||||
limit_per_entity = MAX_LINKS_PER_ENTITY + len(unit_ids) # room for new units + existing cap
|
||||
|
||||
rows = await ops.fetch_entity_unit_fanout(
|
||||
conn,
|
||||
fq_table("unit_entities"),
|
||||
entity_id_list,
|
||||
limit_per_entity,
|
||||
)
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [6.3.1] Query unit_entities (LATERAL): {len(rows)} rows in {time.time() - query_start:.3f}s",
|
||||
level="debug",
|
||||
)
|
||||
|
||||
group_start = time.time()
|
||||
for row in rows:
|
||||
entity_id = row["entity_id"]
|
||||
if entity_id not in entity_to_units:
|
||||
entity_to_units[entity_id] = []
|
||||
entity_to_units[entity_id].append(row["unit_id"])
|
||||
_log(log_buffer, f" [6.3.2] Group by entity_id: {time.time() - group_start:.3f}s", level="debug")
|
||||
link_gen_start = time.time()
|
||||
links: list[EntityLink] = []
|
||||
new_unit_set = set(unit_ids)
|
||||
|
||||
def to_uuid(val) -> UUID:
|
||||
return UUID(val) if isinstance(val, str) else val
|
||||
|
||||
for entity_id, units_with_entity in entity_to_units.items():
|
||||
entity_uuid = to_uuid(entity_id)
|
||||
new_units = [u for u in units_with_entity if str(u) in new_unit_set or u in new_unit_set]
|
||||
existing_units = [u for u in units_with_entity if str(u) not in new_unit_set and u not in new_unit_set]
|
||||
|
||||
new_units_to_link = new_units[-MAX_LINKS_PER_ENTITY:] if len(new_units) > MAX_LINKS_PER_ENTITY else new_units
|
||||
for i, unit_id_1 in enumerate(new_units_to_link):
|
||||
for unit_id_2 in new_units_to_link[i + 1 :]:
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(unit_id_1), to_unit_id=to_uuid(unit_id_2), entity_id=entity_uuid)
|
||||
)
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(unit_id_2), to_unit_id=to_uuid(unit_id_1), entity_id=entity_uuid)
|
||||
)
|
||||
|
||||
existing_to_link = existing_units[-MAX_LINKS_PER_ENTITY:]
|
||||
for new_unit in new_units:
|
||||
for existing_unit in existing_to_link:
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(new_unit), to_unit_id=to_uuid(existing_unit), entity_id=entity_uuid)
|
||||
)
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(existing_unit), to_unit_id=to_uuid(new_unit), entity_id=entity_uuid)
|
||||
)
|
||||
|
||||
_log(log_buffer, f" [6.3.3] Generate {len(links)} links: {time.time() - link_gen_start:.3f}s", level="debug")
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [6.3] Entity link creation: {len(links)} links for {len(all_entity_ids)} unique entities in {time.time() - substep_start:.3f}s",
|
||||
level="debug",
|
||||
)
|
||||
|
||||
return links
|
||||
|
||||
|
||||
async def create_temporal_links_batch_per_fact(
|
||||
conn,
|
||||
bank_id: str,
|
||||
@@ -701,16 +570,18 @@ async def compute_semantic_links_ann(
|
||||
# `relation "_ann_seeds" does not exist` on the second statement.
|
||||
#
|
||||
# Using ON COMMIT DROP + SET LOCAL also means we don't have to remember to
|
||||
# manually drop the temp table or reset hnsw.ef_search — the transaction
|
||||
# end handles both.
|
||||
# manually drop the temp table or reset the per-backend ANN tuning GUC —
|
||||
# the transaction end handles both.
|
||||
rows: list = []
|
||||
async with conn.transaction():
|
||||
# Transaction-local ef_search. Default 400 is tuned for recall precision
|
||||
# but at 164k units each HNSW probe takes 94ms. ef_search=60 gives 2.7ms
|
||||
# per probe (35x faster) with sufficient accuracy for top-50 semantic
|
||||
# link creation. SET LOCAL auto-reverts at commit, so we don't pollute
|
||||
# the pool for subsequent recall queries.
|
||||
await conn.execute("SET LOCAL hnsw.ef_search = 60")
|
||||
# Transaction-local ANN tuning. Each supported backend exposes its own
|
||||
# GUC (hnsw.ef_search on pgvector, vchordrq.probes on vchord); the
|
||||
# dispatcher returns the right knob for the configured backend with a
|
||||
# value tuned for top-50 semantic link creation (lower recall but much
|
||||
# lower latency than the recall-side default). SET LOCAL auto-reverts
|
||||
# at commit, so we don't pollute the pool for subsequent queries.
|
||||
for guc, value in ann_search_tuning_settings(configured_vector_extension(), kind="low_latency"):
|
||||
await conn.execute(f"SET LOCAL {guc} = {value}")
|
||||
|
||||
t_setup = time_mod.time()
|
||||
await conn.execute("CREATE TEMP TABLE _ann_seeds (unit_id text, emb_text text, fact_type text) ON COMMIT DROP")
|
||||
@@ -728,23 +599,35 @@ async def compute_semantic_links_ann(
|
||||
t_query = time_mod.time()
|
||||
seed_count = sum(1 for ft in fact_types if ft == fact_type)
|
||||
logger.debug(f"[ANN] Querying fact_type={fact_type}: {seed_count} seeds")
|
||||
# Cast each seed's text embedding to `vector` exactly once in a
|
||||
# MATERIALIZED CTE. Casting inside the LATERAL (s.emb_text::vector)
|
||||
# re-parses the ~5KB embedding string for every candidate row the
|
||||
# probe touches — seeds × bank_units text-parses per batch, which
|
||||
# dominated the whole job on small banks (see #1919: ~50 seeds over
|
||||
# ~1k units took 1.5-3.7s, ~25-48x slower than casting once). The
|
||||
# stable `vector` column also lets the planner consider an HNSW
|
||||
# index scan, which a cast expression inhibits.
|
||||
ft_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH seeds AS MATERIALIZED (
|
||||
SELECT unit_id, emb_text::vector AS emb
|
||||
FROM _ann_seeds
|
||||
WHERE fact_type = $2
|
||||
)
|
||||
SELECT s.unit_id AS from_id,
|
||||
n.id::text AS to_id,
|
||||
n.similarity
|
||||
FROM _ann_seeds s
|
||||
FROM seeds s
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT mu.id,
|
||||
1 - (mu.embedding <=> s.emb_text::vector) AS similarity
|
||||
1 - (mu.embedding <=> s.emb) AS similarity
|
||||
FROM {fq_table("memory_units")} mu
|
||||
WHERE mu.bank_id = $1
|
||||
AND mu.fact_type = $2
|
||||
AND mu.embedding IS NOT NULL
|
||||
ORDER BY mu.embedding <=> s.emb_text::vector
|
||||
ORDER BY mu.embedding <=> s.emb
|
||||
LIMIT $3
|
||||
) n
|
||||
WHERE s.fact_type = $2
|
||||
""",
|
||||
bank_id,
|
||||
fact_type,
|
||||
@@ -889,29 +772,6 @@ async def create_semantic_links_batch(
|
||||
raise
|
||||
|
||||
|
||||
async def insert_entity_links_batch(conn, links: list[EntityLink], bank_id: str, chunk_size: int = 5000, ops=None):
|
||||
"""
|
||||
Bulk-insert entity links via sorted INSERT FROM unnest().
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
links: List of EntityLink objects
|
||||
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
|
||||
chunk_size: Number of rows per INSERT chunk (default 5000)
|
||||
"""
|
||||
if not links:
|
||||
return
|
||||
|
||||
import time as time_mod
|
||||
|
||||
total_start = time_mod.time()
|
||||
tuples = [(link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id) for link in links]
|
||||
await _bulk_insert_links(conn, tuples, bank_id=bank_id, chunk_size=chunk_size, ops=ops)
|
||||
logger.debug(
|
||||
f" [9.TOTAL] Entity links batch insert ({len(tuples)} rows): {time_mod.time() - total_start:.3f}s"
|
||||
)
|
||||
|
||||
|
||||
async def create_causal_links_batch(
|
||||
conn,
|
||||
bank_id: str,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user