Compare commits
197
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a5d6b7cc90 | ||
|
|
32e421d3a5 | ||
|
|
b83f621611 | ||
|
|
19e607b287 | ||
|
|
bc83eacc7b | ||
|
|
4e7780e593 | ||
|
|
e0221cae6e | ||
|
|
07c85da988 | ||
|
|
0afa046fc5 | ||
|
|
a1228ec3a7 | ||
|
|
f83fafa45b | ||
|
|
0a74ce3f07 | ||
|
|
5b5188af82 | ||
|
|
621ab7e66b | ||
|
|
72d9881a6a | ||
|
|
de22b606e7 | ||
|
|
51d25d84a3 | ||
|
|
c06e85b64f | ||
|
|
1d1f718ce5 | ||
|
|
a814b97197 | ||
|
|
61f980f5a3 | ||
|
|
fe404efb20 | ||
|
|
fb16fc4fdd | ||
|
|
82c7df7266 | ||
|
|
a0e6bedcf1 | ||
|
|
fd848a18c1 | ||
|
|
a0d91408cf | ||
|
|
0d55a9b78d | ||
|
|
90ee101bec | ||
|
|
27cb1c6843 | ||
|
|
a942b1c817 | ||
|
|
0280b3486f | ||
|
|
96910071c9 | ||
|
|
68d947b8db | ||
|
|
c815329c14 | ||
|
|
41b6e5746a | ||
|
|
85402035ed | ||
|
|
3f025c1fc3 | ||
|
|
37a20ec524 | ||
|
|
a67a8f774d | ||
|
|
91d767cdcb | ||
|
|
22ae72a907 | ||
|
|
4cb78173ba | ||
|
|
36e31c675c | ||
|
|
1c133dbd6c | ||
|
|
c6dd089445 | ||
|
|
ba44e0205d | ||
|
|
d5df9ad083 | ||
|
|
d7ff44b984 | ||
|
|
3069bb41af | ||
|
|
4dc149a1ac | ||
|
|
1296e9fc12 | ||
|
|
109e1bd955 | ||
|
|
c0f0c3a769 | ||
|
|
6ba4aeaf03 | ||
|
|
9ea1ef164a | ||
|
|
b0f86f9c0d | ||
|
|
9ca6617813 | ||
|
|
6dc56498ce | ||
|
|
9a9aef6225 | ||
|
|
66e58a23af | ||
|
|
e76021add3 | ||
|
|
ccf0dc8268 | ||
|
|
394d66e607 | ||
|
|
9891f53177 | ||
|
|
8170fe880e | ||
|
|
95d77233bf | ||
|
|
4a0a599473 | ||
|
|
bfdc1c5e65 | ||
|
|
37e28fac09 | ||
|
|
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 | ||
|
|
b708302187 | ||
|
|
18c45c9d01 | ||
|
|
06f36b8b25 | ||
|
|
a74e5e6b5a | ||
|
|
c01fc12f7e | ||
|
|
df7f45e698 | ||
|
|
dfe74b1de9 | ||
|
|
a933a417cd | ||
|
|
05602730e8 | ||
|
|
b67e813a83 | ||
|
|
2e011d279c | ||
|
|
c94935bfa2 | ||
|
|
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 |
@@ -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
|
||||
@@ -173,6 +192,18 @@ If a migration adds a new PostgreSQL table (look for `CREATE TABLE` / `op.create
|
||||
- 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.
|
||||
|
||||
### 11b. Check new config flags update the env template
|
||||
|
||||
If the diff adds a new configuration field (a new `ENV_*` / `HINDSIGHT_*` env var
|
||||
in `hindsight-api-slim/hindsight_api/config.py`):
|
||||
- **`.env.example`** (repo root) — must add the variable (commented if optional)
|
||||
alongside the docs entry in `hindsight-docs/docs/developer/configuration.md`.
|
||||
A flag added to `config.py` but absent from `.env.example` is a **should fix**.
|
||||
- **`hindsight-embed/hindsight_embed/env.example`** — the bundled copy must stay
|
||||
byte-identical to the repo-root `.env.example` (it seeds embed/profile configs).
|
||||
The `test_bundled_template_matches_repo_root` sync test fails on drift; if the
|
||||
root file changed without re-copying, flag it as a **must fix**.
|
||||
|
||||
### 12. Review against other coding standards
|
||||
|
||||
Check the diff for violations of the standards listed above:
|
||||
@@ -196,7 +227,10 @@ 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:
|
||||
|
||||
@@ -0,0 +1,116 @@
|
||||
---
|
||||
name: hs-release
|
||||
description: Cut a core Hindsight release (vX.Y.Z) and open the changelog + blog PR. Use when asked to cut/start a release, bump the version, or publish a new Hindsight version.
|
||||
user_invocable: true
|
||||
---
|
||||
|
||||
# Hindsight Release
|
||||
|
||||
Cut a **core** Hindsight release and open the accompanying changelog/blog PR. This is for the core
|
||||
product version (API, clients, CLI, control plane, Helm). **Integrations are versioned
|
||||
independently** — use `scripts/release-integration.sh` for those, not this skill.
|
||||
|
||||
The release is **irreversible and outward-facing**: it tags a version and pushes it straight to
|
||||
`main`, which triggers CI that publishes packages to PyPI / npm / Helm. Confirm the version number
|
||||
and that the intended fixes are already merged to `main` before you start.
|
||||
|
||||
## Step 0 — Pre-flight
|
||||
|
||||
1. **Decide the base.** A release is cut from the latest `origin/main`, never from a feature
|
||||
branch. `git fetch origin --tags` first. Confirm the "couple of fixes" the user means are
|
||||
actually merged to `main` (`git log v<prev>..origin/main --oneline`).
|
||||
2. **Find where `main` is checked out.** `main` is often already checked out in a sibling worktree
|
||||
(`git worktree list`). You **cannot** check out `main` in a second worktree — run the release in
|
||||
the worktree that already holds it. If that worktree is dirty with throwaway cruft
|
||||
(`.next-*` tsconfig paths, screenshots), `git stash push -u`, fast-forward to `origin/main`,
|
||||
run the release, then `git stash pop`.
|
||||
3. **Pitfall:** never pipe the checkout in an `&&` chain like
|
||||
`git checkout main 2>&1 | tail && git reset --hard ...` — the pipe's exit status is `tail`'s
|
||||
(always 0), so a failed checkout won't stop the chain and the `reset` fires on the **wrong
|
||||
branch**. Check out as its own command and verify `git branch --show-current` before resetting.
|
||||
|
||||
## Step 1 — Cut the release
|
||||
|
||||
Run from the worktree on a clean `main`:
|
||||
|
||||
```bash
|
||||
./scripts/release.sh <version> # e.g. 0.8.1 (no leading v)
|
||||
```
|
||||
|
||||
`release.sh` bumps the version in every component, regenerates the OpenAPI spec + all client SDKs,
|
||||
updates docs versioning, commits `Release v<version>`, tags `v<version>`, and **pushes the commit
|
||||
and tag directly to `main`**. The push triggers the `Release` GitHub Actions workflow that builds
|
||||
and publishes the packages. It is **not** a PR.
|
||||
|
||||
Verify after: `gh run list --limit 5` should show the `Release v<version>` workflow running, and
|
||||
`git ls-remote --tags origin v<version>` should return the tag.
|
||||
|
||||
## Step 2 — Changelog + blog PR (separate)
|
||||
|
||||
Done **after** the tag exists, as its own PR (precedent: v0.8.0 = #2053, v0.8.1 = #2080). Work on a
|
||||
branch off the new `main`:
|
||||
|
||||
```bash
|
||||
git checkout -b docs-changelog-<version> origin/main
|
||||
```
|
||||
|
||||
Only spin up a separate worktree (`git worktree add ../hindsight-changelog-<version> -b
|
||||
docs-changelog-<version> origin/main`) if you can't get a clean checkout otherwise — e.g. `main` is
|
||||
held in another worktree and the current one has work you don't want to disturb.
|
||||
|
||||
**Branch naming:** use the `docs-` (hyphen) convention, e.g. `docs-changelog-0.8.1`. A remote
|
||||
branch literally named `docs` exists, so any `docs/...` branch is rejected on push with
|
||||
`directory file conflict`.
|
||||
|
||||
### Changelog
|
||||
|
||||
```bash
|
||||
uv run --directory hindsight-dev generate-changelog <version>
|
||||
```
|
||||
|
||||
LLM-summarizes the commits between the previous tag and `v<version>` and prepends an entry to
|
||||
`hindsight-docs/src/pages/changelog/index.md`. Requires `OPENAI_API_KEY` (already in the repo
|
||||
`.env`). It excludes `hindsight-integrations/` source, but new integrations whose commits also
|
||||
touched docs will still appear — that matches precedent, leave them in the **changelog**.
|
||||
|
||||
### Blog post
|
||||
|
||||
Hand-write `hindsight-docs/blog/YYYY-MM-DD-version-X-Y-Z.md` (mirror an existing one; patch
|
||||
releases are short — see `2026-06-02-version-0-7-2.md`). Guidance:
|
||||
|
||||
- **Explain user impact, not internals/mechanism.** Lead with what the user can now do and what to
|
||||
set. Config/env-var names are fine (developer-facing), code symbols and internals are not.
|
||||
- **Do not list integrations in the release blog.** The core blog covers core engine / API /
|
||||
ops changes; each integration ships its own changelog. (Integrations may still appear in the
|
||||
generated `changelog/index.md` — that's fine; just keep them out of the blog.)
|
||||
- Call out an upgrade recommendation when there are operational/data-integrity fixes.
|
||||
- Validate formatting: `npx prettier --check <blog file>`.
|
||||
|
||||
### Sync the docs skill
|
||||
|
||||
```bash
|
||||
./scripts/generate-docs-skill.sh
|
||||
```
|
||||
|
||||
Refreshes `skills/hindsight-docs/references/changelog/index.md`. It will also bump
|
||||
`skills/hindsight-docs/references/openapi.json` by one version — `release.sh` regenerates the skill
|
||||
*before* bumping OpenAPI, so the skill copy lags a version in the release commit; this step syncs
|
||||
it. Expect a one-line `version` diff there; keep it.
|
||||
|
||||
### Commit, push, PR
|
||||
|
||||
```bash
|
||||
git add -A
|
||||
git commit --no-verify -m "docs: changelog and blog post for v<version>"
|
||||
git push -u origin docs-changelog-<version>
|
||||
gh pr create --base main --title "docs: changelog and blog post for v<version>" --body "..."
|
||||
```
|
||||
|
||||
Expected files in the PR: the changelog entry, the new blog post, the regenerated skill changelog
|
||||
mirror, and the skill `openapi.json` version sync.
|
||||
|
||||
## Cleanup
|
||||
|
||||
If you created a temporary worktree, remove it once the PR is up
|
||||
(`git worktree remove ../hindsight-changelog-<version>`; the branch stays on origin). Restore any
|
||||
stash you popped in Step 0.
|
||||
+15
-2
@@ -25,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
|
||||
@@ -80,10 +80,23 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER=
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default), "tei", "openai", "cohere", "google", "openrouter", "zeroentropy", "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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -9,7 +9,11 @@ jobs:
|
||||
publish:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
id-token: write # for PyPI trusted publishing
|
||||
id-token: write # for PyPI trusted publishing + build-provenance attestations
|
||||
attestations: write # for actions/attest-build-provenance (Obsidian assets)
|
||||
# No `contents: write`: we never create releases in this repo. The Obsidian
|
||||
# plugin's distribution release is pushed to its dedicated repo using
|
||||
# OBSIDIAN_DIST_TOKEN (see the "Mirror Obsidian plugin" step below).
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -112,6 +116,71 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
|
||||
run: npm run build
|
||||
|
||||
# Build-provenance attestations for the Obsidian release assets (community-store
|
||||
# recommendation). Runs after the build so main.js exists. The assets are
|
||||
# released in the dedicated repo while the build runs here, so users verify at
|
||||
# owner scope: `gh attestation verify main.js --owner vectorize-io`.
|
||||
- name: Attest Obsidian plugin build provenance
|
||||
if: steps.type.outputs.type == 'typescript' && steps.info.outputs.integration == 'obsidian'
|
||||
uses: actions/attest-build-provenance@v2
|
||||
with:
|
||||
subject-path: |
|
||||
hindsight-integrations/obsidian/main.js
|
||||
hindsight-integrations/obsidian/styles.css
|
||||
|
||||
# ── Obsidian plugin — mirror to its dedicated repo + cut the BRAT release ──
|
||||
# We do NOT create a GitHub Release in this monorepo: per-integration
|
||||
# releases pollute the repo's release list (it's for the core product) and
|
||||
# steal the "Latest" badge, and BRAT / the community store read a repo's
|
||||
# *latest* release — not a tag — so they can't target a tag in a monorepo.
|
||||
#
|
||||
# Instead this monorepo stays the source of truth, and on each obsidian
|
||||
# release we mirror hindsight-integrations/obsidian/ → the *root* of
|
||||
# github.com/vectorize-io/hindsight-obsidian (git subtree, history
|
||||
# preserved) and cut the BRAT / community-store release *there*.
|
||||
#
|
||||
# Requires secret OBSIDIAN_DIST_TOKEN — a token with `contents: write` on
|
||||
# vectorize-io/hindsight-obsidian (fine-grained PAT or app installation
|
||||
# token). The dedicated repo is generated; never edit it directly.
|
||||
- name: Mirror Obsidian plugin to its dedicated repo
|
||||
if: steps.type.outputs.type == 'typescript' && steps.info.outputs.integration == 'obsidian'
|
||||
env:
|
||||
DIST_TOKEN: ${{ secrets.OBSIDIAN_DIST_TOKEN }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
VERSION="${{ steps.info.outputs.version }}"
|
||||
DIST_REPO="vectorize-io/hindsight-obsidian"
|
||||
OBS_DIR="hindsight-integrations/obsidian"
|
||||
|
||||
# `git subtree split` needs full history; the default checkout is shallow.
|
||||
git fetch --unshallow 2>/dev/null || true
|
||||
git config user.name "github-actions[bot]"
|
||||
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
|
||||
|
||||
# The runner injects the default GITHUB_TOKEN as an http.extraheader via
|
||||
# an *included* config file (/home/runner/work/_temp/git-credentials-*.config),
|
||||
# so `git config --local --unset-all` can't remove it and it authenticates
|
||||
# the push as github-actions[bot] (no access to the dedicated repo → 403).
|
||||
# The documented way to drop an inherited extraheader is to RESET the list
|
||||
# with an empty value: since command-line `-c` is read last, the empty
|
||||
# value clears the accumulated headers (including the included one) at
|
||||
# request-build time. The dist token then comes from the push URL → a
|
||||
# single Authorization header.
|
||||
git subtree split --prefix="$OBS_DIR" -b _obs_dist
|
||||
git -c "http.https://github.com/.extraheader=" \
|
||||
push "https://x-access-token:${DIST_TOKEN}@github.com/${DIST_REPO}.git" _obs_dist:main
|
||||
|
||||
# Cut the BRAT / community-store release. Bare version tag (e.g. 0.1.0)
|
||||
# to match manifest.json — idempotent so re-runs just refresh the assets.
|
||||
export GH_TOKEN="$DIST_TOKEN"
|
||||
ASSETS="$OBS_DIR/main.js $OBS_DIR/manifest.json $OBS_DIR/styles.css"
|
||||
NOTES="Hindsight for Obsidian v${VERSION}. Install via BRAT (add ${DIST_REPO}) or copy main.js/manifest.json/styles.css into <vault>/.obsidian/plugins/hindsight/."
|
||||
if gh release view "$VERSION" --repo "$DIST_REPO" >/dev/null 2>&1; then
|
||||
gh release upload "$VERSION" $ASSETS --repo "$DIST_REPO" --clobber
|
||||
else
|
||||
gh release create "$VERSION" $ASSETS --repo "$DIST_REPO" --title "$VERSION" --notes "$NOTES"
|
||||
fi
|
||||
|
||||
- name: Publish TypeScript package to npm
|
||||
if: steps.type.outputs.type == 'typescript'
|
||||
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
|
||||
@@ -121,7 +190,12 @@ jobs:
|
||||
EXIT_CODE=$?
|
||||
echo "$OUTPUT"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
if echo "$OUTPUT" | grep -q "cannot publish over"; then
|
||||
# Treat "already published" as success so re-pointed-tag re-runs stay green.
|
||||
# "cannot publish over" = the version exists. TLOG_CREATE_ENTRY_ERROR / 409
|
||||
# "equivalent entry already exists in the transparency log" = the identical
|
||||
# --provenance artifact was already logged on a prior run (Sigstore tlog is
|
||||
# idempotent); the package is published, so this is benign.
|
||||
if echo "$OUTPUT" | grep -qE "cannot publish over|TLOG_CREATE_ENTRY_ERROR|already exists in the transparency log"; then
|
||||
echo "Package version already published, skipping..."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
+471
-7
@@ -34,25 +34,36 @@ jobs:
|
||||
integrations-ai-sdk: ${{ steps.filter.outputs.integrations-ai-sdk }}
|
||||
integrations-chat: ${{ steps.filter.outputs.integrations-chat }}
|
||||
integrations-claude-code: ${{ steps.filter.outputs.integrations-claude-code }}
|
||||
integrations-cline: ${{ steps.filter.outputs.integrations-cline }}
|
||||
integrations-codex: ${{ steps.filter.outputs.integrations-codex }}
|
||||
integrations-cursor-cli: ${{ steps.filter.outputs.integrations-cursor-cli }}
|
||||
integrations-crewai: ${{ steps.filter.outputs.integrations-crewai }}
|
||||
integrations-litellm: ${{ steps.filter.outputs.integrations-litellm }}
|
||||
integrations-pydantic-ai: ${{ steps.filter.outputs.integrations-pydantic-ai }}
|
||||
integrations-ag2: ${{ steps.filter.outputs.integrations-ag2 }}
|
||||
integrations-autogen: ${{ steps.filter.outputs.integrations-autogen }}
|
||||
integrations-langgraph: ${{ steps.filter.outputs.integrations-langgraph }}
|
||||
integrations-llamaindex: ${{ steps.filter.outputs.integrations-llamaindex }}
|
||||
integrations-paperclip: ${{ steps.filter.outputs.integrations-paperclip }}
|
||||
integrations-opencode: ${{ steps.filter.outputs.integrations-opencode }}
|
||||
integrations-cursor: ${{ steps.filter.outputs.integrations-cursor }}
|
||||
integrations-n8n: ${{ steps.filter.outputs.integrations-n8n }}
|
||||
integrations-cloudflare-oauth-proxy: ${{ steps.filter.outputs.integrations-cloudflare-oauth-proxy }}
|
||||
integrations-superagent: ${{ steps.filter.outputs.integrations-superagent }}
|
||||
integrations-lockfiles: ${{ steps.filter.outputs.integrations-lockfiles }}
|
||||
integrations-openai-agents: ${{ steps.filter.outputs.integrations-openai-agents }}
|
||||
integrations-pipecat: ${{ steps.filter.outputs.integrations-pipecat }}
|
||||
integrations-agentcore: ${{ steps.filter.outputs.integrations-agentcore }}
|
||||
integrations-smolagents: ${{ steps.filter.outputs.integrations-smolagents }}
|
||||
integrations-claude-agent-sdk: ${{ steps.filter.outputs.integrations-claude-agent-sdk }}
|
||||
integrations-dify: ${{ steps.filter.outputs.integrations-dify }}
|
||||
integrations-gemini-spark: ${{ steps.filter.outputs.integrations-gemini-spark }}
|
||||
integrations-vapi: ${{ steps.filter.outputs.integrations-vapi }}
|
||||
integrations-flowise: ${{ steps.filter.outputs.integrations-flowise }}
|
||||
integrations-google-adk: ${{ steps.filter.outputs.integrations-google-adk }}
|
||||
integrations-obsidian: ${{ steps.filter.outputs.integrations-obsidian }}
|
||||
integrations-omo: ${{ steps.filter.outputs.integrations-omo }}
|
||||
integrations-haystack: ${{ steps.filter.outputs.integrations-haystack }}
|
||||
tools-agent-sdk: ${{ steps.filter.outputs.tools-agent-sdk }}
|
||||
integrations-roo-code: ${{ steps.filter.outputs.integrations-roo-code }}
|
||||
dev: ${{ steps.filter.outputs.dev }}
|
||||
@@ -96,6 +107,9 @@ jobs:
|
||||
docs:
|
||||
- 'hindsight-docs/**'
|
||||
- '*.md'
|
||||
# Integration changes can add/rename integrations, which the docs
|
||||
# build's integrations check validates against integrations.json.
|
||||
- 'hindsight-integrations/**'
|
||||
embed:
|
||||
- 'hindsight-embed/**'
|
||||
all-npm:
|
||||
@@ -114,8 +128,12 @@ jobs:
|
||||
- 'hindsight-integrations/chat/**'
|
||||
integrations-claude-code:
|
||||
- 'hindsight-integrations/claude-code/**'
|
||||
integrations-cline:
|
||||
- 'hindsight-integrations/cline/**'
|
||||
integrations-codex:
|
||||
- 'hindsight-integrations/codex/**'
|
||||
integrations-cursor-cli:
|
||||
- 'hindsight-integrations/cursor-cli/**'
|
||||
integrations-crewai:
|
||||
- 'hindsight-integrations/crewai/**'
|
||||
integrations-litellm:
|
||||
@@ -124,16 +142,26 @@ jobs:
|
||||
- 'hindsight-integrations/pydantic-ai/**'
|
||||
integrations-ag2:
|
||||
- 'hindsight-integrations/ag2/**'
|
||||
integrations-autogen:
|
||||
- 'hindsight-integrations/autogen/**'
|
||||
integrations-langgraph:
|
||||
- 'hindsight-integrations/langgraph/**'
|
||||
integrations-llamaindex:
|
||||
- 'hindsight-integrations/llamaindex/**'
|
||||
integrations-haystack:
|
||||
- 'hindsight-integrations/haystack/**'
|
||||
integrations-paperclip:
|
||||
- 'hindsight-integrations/paperclip/**'
|
||||
integrations-opencode:
|
||||
- 'hindsight-integrations/opencode/**'
|
||||
integrations-cursor:
|
||||
- 'hindsight-integrations/cursor/**'
|
||||
integrations-n8n:
|
||||
- 'hindsight-integrations/n8n/**'
|
||||
integrations-cloudflare-oauth-proxy:
|
||||
- 'hindsight-integrations/cloudflare-oauth-proxy/**'
|
||||
integrations-superagent:
|
||||
- 'hindsight-integrations/superagent/**'
|
||||
integrations-lockfiles:
|
||||
- 'hindsight-integrations/*/package-lock.json'
|
||||
- 'hindsight-integrations/*/package.json'
|
||||
@@ -146,6 +174,8 @@ jobs:
|
||||
- 'hindsight-integrations/agentcore/**'
|
||||
integrations-smolagents:
|
||||
- 'hindsight-integrations/smolagents/**'
|
||||
integrations-claude-agent-sdk:
|
||||
- 'hindsight-integrations/claude-agent-sdk/**'
|
||||
integrations-dify:
|
||||
- 'hindsight-integrations/dify/**'
|
||||
integrations-gemini-spark:
|
||||
@@ -154,6 +184,12 @@ jobs:
|
||||
- 'hindsight-integrations/vapi/**'
|
||||
integrations-flowise:
|
||||
- 'hindsight-integrations/flowise/**'
|
||||
integrations-google-adk:
|
||||
- 'hindsight-integrations/google-adk/**'
|
||||
integrations-obsidian:
|
||||
- 'hindsight-integrations/obsidian/**'
|
||||
integrations-omo:
|
||||
- 'hindsight-integrations/omo/**'
|
||||
tools-agent-sdk:
|
||||
- 'hindsight-tools/hindsight-agent-sdk/**'
|
||||
integrations-roo-code:
|
||||
@@ -414,6 +450,95 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/claude-code
|
||||
run: python -m pytest tests/ -v
|
||||
|
||||
test-cursor-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-cursor == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.11'
|
||||
|
||||
- name: Install pytest
|
||||
run: pip install pytest
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/cursor
|
||||
run: python -m pytest tests/ -v
|
||||
|
||||
test-omo-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-omo == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.11'
|
||||
|
||||
- name: Install pytest
|
||||
run: pip install pytest
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/omo
|
||||
run: python -m pytest tests/ -v
|
||||
|
||||
test-cline-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-cline == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- 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: Build cline integration
|
||||
working-directory: ./hindsight-integrations/cline
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/cline
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/cline
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-codex-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -440,6 +565,43 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/codex
|
||||
run: python -m pytest tests/ -v
|
||||
|
||||
test-cursor-cli-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-cursor-cli == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- 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: Build cursor-cli integration
|
||||
working-directory: ./hindsight-integrations/cursor-cli
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/cursor-cli
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/cursor-cli
|
||||
run: uv run pytest tests -v
|
||||
|
||||
build-ai-sdk-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -728,6 +890,43 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/pipecat
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-google-adk-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-google-adk == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- 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: Build google-adk integration
|
||||
working-directory: ./hindsight-integrations/google-adk
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/google-adk
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/google-adk
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-gemini-spark-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -776,17 +975,28 @@ jobs:
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- 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: '3.11'
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Install pytest
|
||||
run: pip install pytest
|
||||
- name: Build roo-code integration
|
||||
working-directory: ./hindsight-integrations/roo-code
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/roo-code
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/roo-code
|
||||
run: python -m pytest tests/ -v
|
||||
run: uv run pytest tests -v
|
||||
|
||||
build-control-plane:
|
||||
needs: [detect-changes]
|
||||
@@ -868,6 +1078,7 @@ jobs:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
fetch-depth: 0 # fetch tags so check-released-integrations can see them
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
@@ -876,6 +1087,12 @@ jobs:
|
||||
cache: 'npm'
|
||||
cache-dependency-path: package-lock.json
|
||||
|
||||
# Fail fast before the (slow) build: every integrations.json entry must have a
|
||||
# doc page, and every released integration tag must be in integrations.json.
|
||||
# Needs no npm install (pure Node) and uses the tags fetched above.
|
||||
- name: Check integrations (single source of truth)
|
||||
run: node hindsight-docs/scripts/check-integrations.mjs
|
||||
|
||||
- name: Install dependencies
|
||||
run: npm ci --workspace=hindsight-docs
|
||||
|
||||
@@ -2761,6 +2978,45 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/ag2
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-autogen-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-autogen == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- 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: Build autogen integration
|
||||
working-directory: ./hindsight-integrations/autogen
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/autogen
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/autogen
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-smolagents-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -2864,6 +3120,41 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/flowise
|
||||
run: npm test
|
||||
|
||||
test-obsidian-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-obsidian == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/obsidian
|
||||
run: npm install --no-audit --no-fund
|
||||
|
||||
- name: Type check
|
||||
working-directory: ./hindsight-integrations/obsidian
|
||||
run: npx tsc --noEmit
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/obsidian
|
||||
run: npm run build
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/obsidian
|
||||
run: npm test
|
||||
|
||||
test-crewai-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -2939,6 +3230,45 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/vapi
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-superagent-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-superagent == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- 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: Build superagent integration
|
||||
working-directory: ./hindsight-integrations/superagent
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/superagent
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/superagent
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs live Hindsight + provider keys and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-litellm-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -2974,7 +3304,9 @@ jobs:
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/litellm
|
||||
run: uv run pytest tests -v
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs live Hindsight + provider keys and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-pydantic-ai-integration:
|
||||
needs: [detect-changes]
|
||||
@@ -3013,6 +3345,45 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/pydantic-ai
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-langgraph-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-langgraph == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- 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: Build langgraph integration
|
||||
working-directory: ./hindsight-integrations/langgraph
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/langgraph
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/langgraph
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-llamaindex-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -3048,7 +3419,48 @@ jobs:
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/llamaindex
|
||||
run: uv run pytest tests -v
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-haystack-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-haystack == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- 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: Build haystack integration
|
||||
working-directory: ./hindsight-integrations/haystack
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/haystack
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/haystack
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-openai-agents-integration:
|
||||
needs: [detect-changes]
|
||||
@@ -3085,7 +3497,47 @@ jobs:
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/openai-agents
|
||||
run: uv run pytest tests -v
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-claude-agent-sdk-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-claude-agent-sdk == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- 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: Build claude-agent-sdk integration
|
||||
working-directory: ./hindsight-integrations/claude-agent-sdk
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/claude-agent-sdk
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/claude-agent-sdk
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-agentcore-integration:
|
||||
needs: [detect-changes]
|
||||
@@ -3950,16 +4402,21 @@ jobs:
|
||||
- build-openclaw-integration
|
||||
- smoke-openclaw-install
|
||||
- test-claude-code-integration
|
||||
- test-cursor-integration
|
||||
- test-cline-integration
|
||||
- test-codex-integration
|
||||
- test-cursor-cli-integration
|
||||
- build-ai-sdk-integration
|
||||
- test-ai-sdk-integration-deno
|
||||
- test-opencode-integration
|
||||
- test-omo-integration
|
||||
- test-cloudflare-oauth-proxy-integration
|
||||
- build-chat-integration
|
||||
- test-paperclip-integration
|
||||
- test-pipecat-integration
|
||||
- test-gemini-spark-integration
|
||||
- test-vapi-integration
|
||||
- test-google-adk-integration
|
||||
- test-roo-code-integration
|
||||
- build-control-plane
|
||||
- build-docs
|
||||
@@ -3980,19 +4437,26 @@ jobs:
|
||||
- test-openclaw-integration
|
||||
- test-integration
|
||||
- test-ag2-integration
|
||||
- test-autogen-integration
|
||||
- test-smolagents-integration
|
||||
- test-dify-integration
|
||||
- test-flowise-integration
|
||||
- test-obsidian-integration
|
||||
- test-crewai-integration
|
||||
- test-langgraph-integration
|
||||
- test-superagent-integration
|
||||
- test-litellm-integration
|
||||
- test-pydantic-ai-integration
|
||||
- test-llamaindex-integration
|
||||
- test-openai-agents-integration
|
||||
- test-agentcore-integration
|
||||
- test-haystack-integration
|
||||
- test-pip-slim
|
||||
- test-embed
|
||||
- test-embed-windows
|
||||
- test-hindsight-all
|
||||
- test-hindsight-agent-sdk
|
||||
- test-claude-agent-sdk-integration
|
||||
- test-doc-examples
|
||||
- test-upgrade
|
||||
- verify-generated-files
|
||||
|
||||
+4
-1
@@ -15,6 +15,8 @@ node_modules/
|
||||
|
||||
# Environment variables and local config
|
||||
.env
|
||||
.env.bak*
|
||||
.env.*.bak
|
||||
docker-compose.yml
|
||||
docker-compose.override.yml
|
||||
|
||||
@@ -59,4 +61,5 @@ 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
|
||||
@@ -311,6 +335,16 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
|
||||
- Add to appropriate section table with Variable, Description, Default
|
||||
- Mark if it's hierarchical (can be overridden per-bank)
|
||||
|
||||
6. **Env template** (`.env.example`):
|
||||
- Add the variable to the appropriate section, commented if optional, with a
|
||||
short inline comment describing it (mirror the documentation entry).
|
||||
- This file is the single source of truth for the env template:
|
||||
`scripts/dev/setup.sh` copies it to `.env`, and `hindsight-embed` ships a
|
||||
bundled copy (`hindsight-embed/hindsight_embed/env.example`) that seeds
|
||||
embed/profile configs. After editing `.env.example`, re-copy it to the
|
||||
embed package (`cp .env.example hindsight-embed/hindsight_embed/env.example`)
|
||||
or the `test_bundled_template_matches_repo_root` sync test will fail.
|
||||
|
||||
#### Hierarchical vs Static Guidelines
|
||||
|
||||
**Hierarchical** (per-bank overridable):
|
||||
|
||||
+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
|
||||
```
|
||||
|
||||
@@ -143,6 +143,8 @@ main();
|
||||
pip install hindsight-all -U
|
||||
```
|
||||
|
||||
On Intel (x86_64) Macs, install `hindsight-all-slim` instead — see [Supported Platforms](#supported-platforms).
|
||||
|
||||
```python
|
||||
import os
|
||||
from hindsight import HindsightServer, HindsightClient
|
||||
@@ -300,6 +302,19 @@ client.reflect(bank_id="my-bank", query="What should I know about Alice?")
|
||||
[](https://www.star-history.com/#vectorize-io/hindsight&type=date&legend=top-left)
|
||||
---
|
||||
|
||||
## Supported Platforms
|
||||
|
||||
| Platform | Docker | Bare Metal (pip) | Embedded DB (pg0) |
|
||||
|----------|--------|------------------|--------------------|
|
||||
| **Linux** (x86_64, ARM64) | ✅ | ✅ | ✅ |
|
||||
| **macOS** (Apple Silicon / arm64) | ✅ | ✅ | ✅ |
|
||||
| **macOS** (Intel / x86_64) | ✅ | ⚠️ | ✅ |
|
||||
| **Windows** (x86_64) | ✅ | ✅ | ✅ |
|
||||
|
||||
⚠️ Intel Macs: use `hindsight-all-slim` — see the [installation guide](https://docs.hindsight.vectorize.io/docs/developer/installation#supported-platforms) for details.
|
||||
|
||||
---
|
||||
|
||||
## Contributing
|
||||
|
||||
See [CONTRIBUTING.md](./CONTRIBUTING.md).
|
||||
|
||||
@@ -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 \
|
||||
|
||||
@@ -43,11 +43,56 @@ check_pg0_data_integrity() {
|
||||
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}"
|
||||
|
||||
@@ -8,7 +8,7 @@ source "$SCRIPT_DIR/start-all.sh"
|
||||
unset HINDSIGHT_START_ALL_SOURCE_ONLY
|
||||
|
||||
TMP_DIR="$(mktemp -d)"
|
||||
trap 'rm -rf "$TMP_DIR"' EXIT
|
||||
trap 'chmod -R u+rwx "$TMP_DIR" 2>/dev/null || true; rm -rf "$TMP_DIR"' EXIT
|
||||
|
||||
assert_contains() {
|
||||
local output="$1"
|
||||
@@ -71,3 +71,51 @@ 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
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.7.1
|
||||
appVersion: "0.7.1"
|
||||
version: 0.8.1
|
||||
appVersion: "0.8.1"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.7.1",
|
||||
"version": "0.8.1",
|
||||
"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.7.1"
|
||||
version = "0.8.1"
|
||||
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.7.1",
|
||||
"hindsight-api-slim==0.8.1",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.7.1"
|
||||
version = "0.8.1"
|
||||
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.7.1",
|
||||
"hindsight-api-slim[all]==0.8.1",
|
||||
"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.7.1",
|
||||
"hindsight-api-slim[local-llm]==0.8.1",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -53,4 +53,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.7.1"
|
||||
__version__ = "0.8.1"
|
||||
|
||||
@@ -54,23 +54,20 @@ _INDEX_TYPE_KEYWORDS = {
|
||||
# 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.
|
||||
# - vchord exposes vchordrq.probes, but its shape must match the index's
|
||||
# build.internal.lists hierarchy. VectorChord 1.1 added per-index fallback
|
||||
# parameters for this reason: a session GUC overrides every vchordrq index,
|
||||
# and a single value can be invalid for listless or mixed-layout indexes.
|
||||
# Hindsight's built-in vchord clause does not set lists, so the safe default
|
||||
# is no session-level probe override; deployments that partition vchordrq
|
||||
# indexes should attach probes to the index storage parameters instead.
|
||||
# - pgvectorscale / pg_diskann / scann do not expose an equivalent per-statement
|
||||
# knob in the engine today, so the dispatcher returns no statements for them.
|
||||
_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 = {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -47,21 +49,42 @@ BACKUP_TABLES = [
|
||||
"entities",
|
||||
"chunks",
|
||||
"memory_units",
|
||||
"invalidated_memory_units",
|
||||
"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)
|
||||
@@ -330,6 +353,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)
|
||||
|
||||
+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 TEXT 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 TEXT 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)
|
||||
+75
@@ -0,0 +1,75 @@
|
||||
"""Repair: widen ``*_history.bank_id`` from VARCHAR(64) to TEXT on PostgreSQL.
|
||||
|
||||
The original split-history migration (``a7b8c9d0e1f2``) declared
|
||||
``observation_history.bank_id`` and ``mental_model_history.bank_id`` as
|
||||
``VARCHAR(64)`` on PostgreSQL. But ``memory_units.bank_id`` — the backfill
|
||||
source for observations — is ``TEXT`` (unbounded), as are ``banks``,
|
||||
``documents`` and ``entities``. Any deployment whose ``bank_id`` exceeds 64
|
||||
characters aborts the backfill ``INSERT`` with::
|
||||
|
||||
psycopg2.errors.StringDataRightTruncation: value too long for type
|
||||
character varying(64)
|
||||
|
||||
Because the migration runs in ``lifespan`` startup inside a transaction, the
|
||||
whole migration rolls back and the API never comes up — unrecoverable from the
|
||||
running container. See https://github.com/vectorize-io/hindsight/issues/2106.
|
||||
|
||||
``a7b8c9d0e1f2`` itself has been corrected to create the column as ``TEXT``,
|
||||
which unblocks deployments that *failed* (the migration rolled back, so it
|
||||
re-runs the fixed DDL). This forward migration covers deployments that already
|
||||
*succeeded* with the narrow ``VARCHAR(64)`` column — where editing
|
||||
``a7b8c9d0e1f2`` has no effect because it will not re-run — by widening the
|
||||
column in place. ``ALTER COLUMN ... TYPE TEXT`` is a no-op on a column that is
|
||||
already ``TEXT`` (fresh installs and re-run failures), so every upgrade path
|
||||
converges on ``TEXT``.
|
||||
|
||||
The history tables are per-tenant (they live in each tenant schema, not
|
||||
``public``), so this runs for every migrated schema via the search-path-aware
|
||||
prefix — unlike the shared-``public`` routines repaired in ``b2d4f6a8c1e3``.
|
||||
|
||||
PostgreSQL only. On Oracle both ``memory_units.bank_id`` and the history
|
||||
``bank_id`` columns are already ``VARCHAR2(256)`` (consistent, never
|
||||
truncates), so the Oracle slot is intentionally absent.
|
||||
|
||||
Revision ID: c3e5a7b9d1f4
|
||||
Revises: c9a1b2d3e4f5
|
||||
Create Date: 2026-06-10
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c3e5a7b9d1f4"
|
||||
down_revision: str | Sequence[str] | None = "c9a1b2d3e4f5"
|
||||
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 ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}observation_history ALTER COLUMN bank_id TYPE TEXT")
|
||||
op.execute(f"ALTER TABLE {schema}mental_model_history ALTER COLUMN bank_id TYPE TEXT")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# No-op: narrowing back to VARCHAR(64) could truncate real data and would
|
||||
# re-introduce the bug this migration repairs. The column type is owned by
|
||||
# ``a7b8c9d0e1f2``'s lifecycle.
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+100
@@ -0,0 +1,100 @@
|
||||
"""Add invalidated_memory_units table for curation (edit/invalidate).
|
||||
|
||||
Curation keeps the recall hot-path (``memory_units``) clean by *moving*
|
||||
invalidated facts into a sibling archive table rather than flagging them in
|
||||
place. If a row is in ``memory_units`` it is live; if it is in
|
||||
``invalidated_memory_units`` it has been retired. Recall/consolidation/graph
|
||||
queries never need a state predicate — the rows simply aren't there.
|
||||
|
||||
The archive mirrors ``memory_units`` column-for-column (so a row round-trips
|
||||
losslessly on revert) plus:
|
||||
- ``invalidation_reason`` optional free text recorded on invalidate
|
||||
- ``invalidated_at`` when it was retired
|
||||
- ``entity_ids`` snapshot of the unit's entity associations, so revert
|
||||
can restore them (``unit_entities`` is cascade-deleted
|
||||
when the live row is removed)
|
||||
|
||||
This migration also adds ``edited_at`` to ``memory_units``: set whenever a user
|
||||
edits a memory's fields (text, context, dates, fact_type, entities) via curation.
|
||||
NULL means never manually modified; a non-NULL value answers "has the user ever
|
||||
changed this?" with the time of the last edit (distinct from ``updated_at``,
|
||||
which background operations also bump). It is added to ``memory_units`` *before*
|
||||
the archive is cloned below, so the archive inherits the column and the marker
|
||||
travels with a fact when it is invalidated.
|
||||
|
||||
Revision ID: c9a1b2d3e4f5
|
||||
Revises: b2d4f6a8c1e3
|
||||
Create Date: 2026-06-03
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c9a1b2d3e4f5"
|
||||
down_revision: str | Sequence[str] | None = "b2d4f6a8c1e3"
|
||||
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()
|
||||
# Add edited_at to the live table FIRST so the archive's LIKE clone below
|
||||
# inherits it (keeps the two tables column-for-column identical for round-trip).
|
||||
op.execute(f"ALTER TABLE {schema}memory_units ADD COLUMN IF NOT EXISTS edited_at TIMESTAMPTZ")
|
||||
# LIKE ... INCLUDING DEFAULTS clones every memory_units column (incl. the
|
||||
# embedding vector and edited_at) so an invalidated row can move back verbatim.
|
||||
# We deliberately omit indexes/constraints — the archive is cold storage, not a
|
||||
# recall surface; only the lookups below need indexing.
|
||||
op.execute(
|
||||
f"CREATE TABLE IF NOT EXISTS {schema}invalidated_memory_units (LIKE {schema}memory_units INCLUDING DEFAULTS)"
|
||||
)
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}invalidated_memory_units "
|
||||
f"ADD COLUMN IF NOT EXISTS invalidation_reason TEXT, "
|
||||
f"ADD COLUMN IF NOT EXISTS invalidated_at TIMESTAMPTZ DEFAULT now(), "
|
||||
f"ADD COLUMN IF NOT EXISTS entity_ids UUID[]"
|
||||
)
|
||||
op.execute(f"CREATE UNIQUE INDEX IF NOT EXISTS idx_invalidated_mu_id ON {schema}invalidated_memory_units (id)")
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_invalidated_mu_bank "
|
||||
f"ON {schema}invalidated_memory_units (bank_id, invalidated_at)"
|
||||
)
|
||||
# Deleting a document (or bank) should clear its archived facts too, mirroring
|
||||
# the memory_units → documents cascade.
|
||||
op.execute(
|
||||
f"""
|
||||
DO $$ BEGIN
|
||||
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'invalidated_mu_document_fkey') THEN
|
||||
ALTER TABLE {schema}invalidated_memory_units
|
||||
ADD CONSTRAINT invalidated_mu_document_fkey
|
||||
FOREIGN KEY (document_id, bank_id)
|
||||
REFERENCES {schema}documents(id, bank_id) ON DELETE CASCADE;
|
||||
END IF; END $$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# Drops the archive (and its inherited edited_at) wholesale, then removes
|
||||
# edited_at from the live table.
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}invalidated_memory_units")
|
||||
op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS edited_at")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# PG-only: Oracle gets the table from the baseline snapshot, matching the
|
||||
# convention used by sibling column/index migrations in this tree.
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+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)
|
||||
+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)
|
||||
@@ -122,6 +122,7 @@ _TABLES: tuple[str, ...] = (
|
||||
text_signals CLOB,
|
||||
consolidation_failed_at TIMESTAMP WITH TIME ZONE,
|
||||
search_vector CLOB,
|
||||
edited_at TIMESTAMP WITH TIME ZONE,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_memory_units PRIMARY KEY (id),
|
||||
@@ -138,6 +139,48 @@ _TABLES: tuple[str, ...] = (
|
||||
PARTITION BY LIST (bank_id) AUTOMATIC
|
||||
(PARTITION p_default VALUES ('__default__'))
|
||||
""",
|
||||
# Cold archive for curation: invalidated facts are MOVED here out of
|
||||
# memory_units so the recall hot-path never sees them. Mirrors memory_units
|
||||
# plus invalidation bookkeeping and an entity-id snapshot for lossless revert.
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS invalidated_memory_units (
|
||||
id RAW(16) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
document_id VARCHAR2(512),
|
||||
chunk_id VARCHAR2(512),
|
||||
text CLOB NOT NULL,
|
||||
embedding VECTOR(384, FLOAT32),
|
||||
context CLOB,
|
||||
event_date TIMESTAMP WITH TIME ZONE NOT NULL,
|
||||
occurred_start TIMESTAMP WITH TIME ZONE,
|
||||
occurred_end TIMESTAMP WITH TIME ZONE,
|
||||
mentioned_at TIMESTAMP WITH TIME ZONE,
|
||||
fact_type VARCHAR2(64) DEFAULT 'world' NOT NULL,
|
||||
confidence_score BINARY_DOUBLE,
|
||||
access_count NUMBER(10) DEFAULT 0 NOT NULL,
|
||||
consolidated_at TIMESTAMP WITH TIME ZONE,
|
||||
observation_scopes CLOB CONSTRAINT imu_obs_scopes_json CHECK (observation_scopes IS JSON OR observation_scopes IS NULL),
|
||||
tags CLOB DEFAULT '[]' NOT NULL,
|
||||
metadata CLOB DEFAULT '{}' NOT NULL
|
||||
CONSTRAINT imu_metadata_json CHECK (metadata IS JSON),
|
||||
proof_count NUMBER(10) DEFAULT 1,
|
||||
source_memory_ids CLOB,
|
||||
history CLOB DEFAULT '[]'
|
||||
CONSTRAINT imu_history_json CHECK (history IS JSON OR history IS NULL),
|
||||
text_signals CLOB,
|
||||
consolidation_failed_at TIMESTAMP WITH TIME ZONE,
|
||||
search_vector CLOB,
|
||||
edited_at TIMESTAMP WITH TIME ZONE,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
invalidation_reason CLOB,
|
||||
invalidated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP,
|
||||
entity_ids CLOB CONSTRAINT imu_entity_ids_json CHECK (entity_ids IS JSON OR entity_ids IS NULL),
|
||||
CONSTRAINT pk_invalidated_memory_units PRIMARY KEY (id),
|
||||
CONSTRAINT fk_imu_document FOREIGN KEY (document_id, bank_id)
|
||||
REFERENCES documents(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
""",
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS entities (
|
||||
id RAW(16) DEFAULT SYS_GUID() NOT NULL,
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
"""Client-disconnect detection that works behind ``BaseHTTPMiddleware``.
|
||||
|
||||
``Request.is_disconnected()`` is the obvious way to notice an abandoned HTTP
|
||||
request, but it is silently broken once any ``@app.middleware("http")``
|
||||
(Starlette ``BaseHTTPMiddleware``) is installed: that middleware runs the route
|
||||
in a child task behind anyio memory streams, so the ``http.disconnect`` ASGI
|
||||
event never reaches the route's ``Request``. This app has such middlewares, so
|
||||
the recall/reflect cancellation in #2122/#2127 never actually fired in
|
||||
production — the disconnect was never observed.
|
||||
|
||||
This pure-ASGI middleware sits *outside* the ``BaseHTTPMiddleware`` layer, where
|
||||
it still owns the real ``receive`` channel. For the recall and reflect routes it
|
||||
drains ``receive`` in a background task and trips a :class:`CancellationToken`
|
||||
the moment ``http.disconnect`` arrives, stashing the token on the ASGI ``scope``.
|
||||
The route copies that token onto its ``RequestContext`` and the engine checks it
|
||||
at stage boundaries — so abandoned work stops instead of running to completion.
|
||||
|
||||
It only wraps recall/reflect (small JSON bodies); every other request — uploads,
|
||||
MCP streams, etc. — passes straight through untouched, so there is no buffering
|
||||
or latency cost elsewhere.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
from collections.abc import Awaitable, Callable, MutableMapping
|
||||
from typing import Any
|
||||
|
||||
from ..cancellation import CancellationToken
|
||||
|
||||
# Key under which the per-request CancellationToken is stored on the ASGI scope.
|
||||
# A dedicated top-level scope key (not scope["state"]) avoids any interaction
|
||||
# with Starlette's per-request state copying.
|
||||
SCOPE_CANCELLATION_TOKEN = "hindsight.cancellation_token"
|
||||
|
||||
_CLIENT_DISCONNECTED_REASON = "client disconnected"
|
||||
|
||||
Scope = MutableMapping[str, Any]
|
||||
Receive = Callable[[], Awaitable[MutableMapping[str, Any]]]
|
||||
Send = Callable[[MutableMapping[str, Any]], Awaitable[None]]
|
||||
|
||||
|
||||
def _should_monitor(path: str) -> bool:
|
||||
"""Only the two long-running, abandon-prone read endpoints need monitoring."""
|
||||
return path.endswith("/memories/recall") or path.endswith("/reflect")
|
||||
|
||||
|
||||
class ClientDisconnectCancellationMiddleware:
|
||||
"""Trip a scope-level CancellationToken when the client disconnects.
|
||||
|
||||
Must be installed *outside* any ``BaseHTTPMiddleware`` so it owns the real
|
||||
ASGI ``receive`` channel.
|
||||
"""
|
||||
|
||||
def __init__(self, app: Callable) -> None:
|
||||
self.app = app
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["type"] != "http" or not _should_monitor(scope.get("path", "")):
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
token = CancellationToken()
|
||||
scope[SCOPE_CANCELLATION_TOKEN] = token
|
||||
|
||||
# The downstream app still needs to read the request body, so we cannot
|
||||
# simply consume `receive` ourselves. Instead a single pump task drains
|
||||
# the real channel, forwards every message to a queue the app reads from,
|
||||
# and trips the token the instant `http.disconnect` shows up — which the
|
||||
# app would otherwise never pull once it has finished reading the body.
|
||||
queue: asyncio.Queue = asyncio.Queue()
|
||||
|
||||
async def pump() -> None:
|
||||
while True:
|
||||
message = await receive()
|
||||
if message["type"] == "http.disconnect":
|
||||
token.cancel(_CLIENT_DISCONNECTED_REASON)
|
||||
await queue.put(message)
|
||||
return
|
||||
await queue.put(message)
|
||||
|
||||
async def proxied_receive() -> MutableMapping[str, Any]:
|
||||
return await queue.get()
|
||||
|
||||
pump_task = asyncio.create_task(pump())
|
||||
try:
|
||||
await self.app(scope, proxied_receive, send)
|
||||
finally:
|
||||
pump_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await pump_task
|
||||
|
||||
|
||||
def get_scope_cancellation_token(scope: Scope) -> CancellationToken | None:
|
||||
"""Return the CancellationToken the middleware attached, if any."""
|
||||
return scope.get(SCOPE_CANCELLATION_TOKEN)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -113,6 +113,8 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"delete_directive",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"update_memory",
|
||||
"invalidate_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"delete_document",
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
"""Cooperative cancellation for long-running engine operations.
|
||||
|
||||
Recall runs as a staged pipeline whose heavy stages — graph expansion and
|
||||
cross-encoder reranking — execute in worker threads (``run_in_executor``) that
|
||||
asyncio task cancellation cannot interrupt once they have started. Cancelling
|
||||
the awaiting task only unblocks the ``await``; the thread keeps burning CPU to
|
||||
completion. So rather than rely on task cancellation, callers thread a
|
||||
``CancellationToken`` through ``RequestContext`` and the engine checks it at
|
||||
stage boundaries (``raise_if_cancelled``), bailing out *before* dispatching the
|
||||
next expensive stage.
|
||||
|
||||
This is cooperative by design: it cannot stop a computation already inside a
|
||||
worker thread, but it does stop an abandoned recall from progressing into — or
|
||||
past — that work, which is what starves the instance in issue #2122. The token
|
||||
lives on ``RequestContext``, so any operation that receives one (recall today;
|
||||
reflect/consolidation/MCP later) can adopt the same checkpoints, and any driver
|
||||
(client disconnect today; a deadline tomorrow) can fire it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
|
||||
class OperationCancelledError(Exception):
|
||||
"""Raised at a checkpoint when the operation has been cancelled.
|
||||
|
||||
Carries the ``reason`` set by whoever cancelled (e.g. "client disconnected")
|
||||
so the HTTP layer can translate it into the appropriate status code instead
|
||||
of a generic 500.
|
||||
|
||||
NOTE: this is a plain ``Exception`` on purpose, NOT ``BaseException``. The
|
||||
recall/reflect pipelines have broad ``except Exception`` handlers that would
|
||||
otherwise swallow it — those handlers re-raise ``OperationCancelledError``
|
||||
explicitly (see ``_search_with_retries``) so cancellation propagates to the
|
||||
HTTP layer. A ``BaseException`` would dodge those handlers but also slip past
|
||||
legitimate ``isinstance(result, Exception)`` checks (e.g. the reflect agent's
|
||||
``asyncio.gather(..., return_exceptions=True)`` tool-result handling), which
|
||||
expect every non-tuple result to be an ``Exception``.
|
||||
"""
|
||||
|
||||
def __init__(self, reason: str = "operation cancelled") -> None:
|
||||
super().__init__(reason)
|
||||
self.reason = reason
|
||||
|
||||
|
||||
class CancellationToken:
|
||||
"""A one-shot, cooperative cancellation signal.
|
||||
|
||||
Cheap to poll (``raise_if_cancelled``) at stage boundaries and awaitable
|
||||
(``wait``) so a driver task can block until cancellation. Safe to share
|
||||
across an engine call tree; polling is a no-op until something cancels, and
|
||||
cancellation is idempotent (the first reason wins).
|
||||
"""
|
||||
|
||||
__slots__ = ("_event", "_reason")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._event = asyncio.Event()
|
||||
self._reason = "operation cancelled"
|
||||
|
||||
def cancel(self, reason: str = "operation cancelled") -> None:
|
||||
"""Signal cancellation. Idempotent; the first reason recorded wins."""
|
||||
if not self._event.is_set():
|
||||
self._reason = reason
|
||||
self._event.set()
|
||||
|
||||
@property
|
||||
def cancelled(self) -> bool:
|
||||
"""Whether cancellation has been signalled."""
|
||||
return self._event.is_set()
|
||||
|
||||
@property
|
||||
def reason(self) -> str:
|
||||
"""The reason recorded by the first ``cancel`` call."""
|
||||
return self._reason
|
||||
|
||||
def raise_if_cancelled(self) -> None:
|
||||
"""Raise ``OperationCancelledError`` if cancellation has been signalled."""
|
||||
if self._event.is_set():
|
||||
raise OperationCancelledError(self._reason)
|
||||
|
||||
async def wait(self) -> None:
|
||||
"""Block until cancellation is signalled."""
|
||||
await self._event.wait()
|
||||
@@ -141,8 +141,11 @@ ENV_LLM_TIMEOUT = "HINDSIGHT_API_LLM_TIMEOUT"
|
||||
ENV_LLM_REASONING_EFFORT = "HINDSIGHT_API_LLM_REASONING_EFFORT"
|
||||
ENV_LLM_GROQ_SERVICE_TIER = "HINDSIGHT_API_LLM_GROQ_SERVICE_TIER"
|
||||
ENV_LLM_OPENAI_SERVICE_TIER = "HINDSIGHT_API_LLM_OPENAI_SERVICE_TIER"
|
||||
ENV_LLM_BEDROCK_SERVICE_TIER = "HINDSIGHT_API_LLM_BEDROCK_SERVICE_TIER"
|
||||
ENV_LLM_EXTRA_BODY = "HINDSIGHT_API_LLM_EXTRA_BODY"
|
||||
ENV_LLM_DEFAULT_HEADERS = "HINDSIGHT_API_LLM_DEFAULT_HEADERS"
|
||||
ENV_LLM_STRICT_SCHEMA = "HINDSIGHT_API_LLM_STRICT_SCHEMA"
|
||||
ENV_LLM_SEND_BANK_AS_USER = "HINDSIGHT_API_LLM_SEND_BANK_AS_USER"
|
||||
|
||||
# LiteLLM Router chain — provider-specific config consumed by the "litellmrouter"
|
||||
# provider. Each entry is a deployment; the Router tries them in declared order and
|
||||
@@ -155,6 +158,7 @@ ENV_LLM_LITELLMROUTER_CONFIG = "HINDSIGHT_API_LLM_LITELLMROUTER_CONFIG"
|
||||
# Defaults for service tiers
|
||||
DEFAULT_LLM_GROQ_SERVICE_TIER = "auto" # "on_demand", "flex", or "auto"
|
||||
DEFAULT_LLM_OPENAI_SERVICE_TIER = None # None (default) or "flex" (50% cheaper)
|
||||
DEFAULT_LLM_BEDROCK_SERVICE_TIER = None # None (default), "flex", "priority", or "reserved"
|
||||
DEFAULT_LLM_EXTRA_BODY = None # None = no extra body params; JSON dict merged into OpenAI extra_body
|
||||
DEFAULT_LLM_DEFAULT_HEADERS = (
|
||||
None # None = no extra headers; JSON dict passed as default_headers to provider SDK clients
|
||||
@@ -209,6 +213,17 @@ ENV_EMBEDDINGS_PROVIDER = "HINDSIGHT_API_EMBEDDINGS_PROVIDER"
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL"
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU = "HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU"
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE = "HINDSIGHT_API_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE"
|
||||
ENV_EMBEDDINGS_ONNX_MODEL_ID = "HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_ID"
|
||||
ENV_EMBEDDINGS_ONNX_MODEL_PATH = "HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_PATH"
|
||||
ENV_EMBEDDINGS_ONNX_TOKENIZER_NAME_OR_PATH = "HINDSIGHT_API_EMBEDDINGS_ONNX_TOKENIZER_NAME_OR_PATH"
|
||||
ENV_EMBEDDINGS_ONNX_FILE = "HINDSIGHT_API_EMBEDDINGS_ONNX_FILE"
|
||||
ENV_EMBEDDINGS_ONNX_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_ONNX_DIMENSIONS"
|
||||
ENV_EMBEDDINGS_ONNX_MAX_TOKENS = "HINDSIGHT_API_EMBEDDINGS_ONNX_MAX_TOKENS"
|
||||
ENV_EMBEDDINGS_ONNX_POOLING = "HINDSIGHT_API_EMBEDDINGS_ONNX_POOLING"
|
||||
ENV_EMBEDDINGS_ONNX_NORMALIZE = "HINDSIGHT_API_EMBEDDINGS_ONNX_NORMALIZE"
|
||||
ENV_EMBEDDINGS_ONNX_QUERY_PREFIX = "HINDSIGHT_API_EMBEDDINGS_ONNX_QUERY_PREFIX"
|
||||
ENV_EMBEDDINGS_ONNX_PASSAGE_PREFIX = "HINDSIGHT_API_EMBEDDINGS_ONNX_PASSAGE_PREFIX"
|
||||
ENV_EMBEDDINGS_ONNX_OUTPUT_NAME = "HINDSIGHT_API_EMBEDDINGS_ONNX_OUTPUT_NAME"
|
||||
ENV_EMBEDDINGS_TEI_URL = "HINDSIGHT_API_EMBEDDINGS_TEI_URL"
|
||||
ENV_EMBEDDINGS_OPENAI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY"
|
||||
ENV_EMBEDDINGS_OPENAI_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL"
|
||||
@@ -240,6 +255,7 @@ ENV_EMBEDDINGS_OPENROUTER_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENROUTER_API_KEY
|
||||
ENV_EMBEDDINGS_OPENROUTER_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENROUTER_MODEL"
|
||||
ENV_RERANKER_OPENROUTER_API_KEY = "HINDSIGHT_API_RERANKER_OPENROUTER_API_KEY"
|
||||
ENV_RERANKER_OPENROUTER_MODEL = "HINDSIGHT_API_RERANKER_OPENROUTER_MODEL"
|
||||
ENV_RERANKER_OPENROUTER_BASE_URL = "HINDSIGHT_API_RERANKER_OPENROUTER_BASE_URL"
|
||||
|
||||
# ZeroEntropy configuration (embeddings)
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_API_KEY = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_API_KEY"
|
||||
@@ -297,6 +313,7 @@ ENV_RERANKER_LITELLM_TIMEOUT = "HINDSIGHT_API_RERANKER_LITELLM_TIMEOUT"
|
||||
ENV_RERANKER_LITELLM_SDK_TIMEOUT = "HINDSIGHT_API_RERANKER_LITELLM_SDK_TIMEOUT"
|
||||
ENV_RERANKER_GOOGLE_TIMEOUT = "HINDSIGHT_API_RERANKER_GOOGLE_TIMEOUT"
|
||||
ENV_RERANKER_MAX_CANDIDATES = "HINDSIGHT_API_RERANKER_MAX_CANDIDATES"
|
||||
ENV_SEMANTIC_MIN_SIMILARITY = "HINDSIGHT_API_SEMANTIC_MIN_SIMILARITY"
|
||||
ENV_RERANKER_FLASHRANK_MODEL = "HINDSIGHT_API_RERANKER_FLASHRANK_MODEL"
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR"
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA = "HINDSIGHT_API_RERANKER_FLASHRANK_CPU_MEM_ARENA"
|
||||
@@ -338,6 +355,7 @@ ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
|
||||
ENV_MCP_ENABLED_TOOLS = "HINDSIGHT_API_MCP_ENABLED_TOOLS"
|
||||
ENV_MCP_STATELESS = "HINDSIGHT_API_MCP_STATELESS"
|
||||
ENV_ENABLE_BANK_CONFIG_API = "HINDSIGHT_API_ENABLE_BANK_CONFIG_API"
|
||||
ENV_ENABLE_BANK_LLM_HEALTH = "HINDSIGHT_API_ENABLE_BANK_LLM_HEALTH"
|
||||
ENV_DEFAULT_BANK_TEMPLATE = "HINDSIGHT_API_DEFAULT_BANK_TEMPLATE"
|
||||
ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER"
|
||||
ENV_RECALL_MAX_CONCURRENT = "HINDSIGHT_API_RECALL_MAX_CONCURRENT"
|
||||
@@ -346,6 +364,8 @@ ENV_RECALL_MAX_QUERY_TOKENS = "HINDSIGHT_API_RECALL_MAX_QUERY_TOKENS"
|
||||
ENV_MENTAL_MODEL_REFRESH_CONCURRENCY = "HINDSIGHT_API_MENTAL_MODEL_REFRESH_CONCURRENCY"
|
||||
ENV_LINK_EXPANSION_PER_ENTITY_LIMIT = "HINDSIGHT_API_LINK_EXPANSION_PER_ENTITY_LIMIT"
|
||||
ENV_LINK_EXPANSION_TIMEOUT = "HINDSIGHT_API_LINK_EXPANSION_TIMEOUT"
|
||||
ENV_BANK_STATS_CACHE_TTL_SECONDS = "HINDSIGHT_API_BANK_STATS_CACHE_TTL_SECONDS"
|
||||
ENV_BANK_STATS_CACHE_MAX_ENTRIES = "HINDSIGHT_API_BANK_STATS_CACHE_MAX_ENTRIES"
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
ENV_OTEL_TRACES_ENABLED = "HINDSIGHT_API_OTEL_TRACES_ENABLED"
|
||||
@@ -363,6 +383,16 @@ ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOU
|
||||
# Gemini safety settings
|
||||
ENV_LLM_GEMINI_SAFETY_SETTINGS = "HINDSIGHT_API_LLM_GEMINI_SAFETY_SETTINGS"
|
||||
|
||||
# Gemini prompt caching. When enabled, retain fact-extraction reuses a
|
||||
# CachedContent prefix for the static system_instruction + response_schema,
|
||||
# cutting per-call input cost on workloads with many small documents.
|
||||
# Provider-agnostic prompt-prefix caching. Providers that support it (currently
|
||||
# Gemini/Vertex via CachedContent) reuse the large, fixed, bank-agnostic system
|
||||
# prefix at the cached-input rate; providers that don't simply ignore it. On by
|
||||
# default — the prefix is bank-agnostic so a single cache is shared across all
|
||||
# banks, and creation soft-fails to an uncached call, so it never breaks a request.
|
||||
ENV_LLM_PROMPT_CACHE_ENABLED = "HINDSIGHT_API_LLM_PROMPT_CACHE_ENABLED"
|
||||
|
||||
# Retain settings
|
||||
ENV_RETAIN_MAX_COMPLETION_TOKENS = "HINDSIGHT_API_RETAIN_MAX_COMPLETION_TOKENS"
|
||||
ENV_RETAIN_CHUNK_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_SIZE"
|
||||
@@ -399,6 +429,11 @@ ENV_FILE_CONVERSION_MAX_BATCH_SIZE_MB = "HINDSIGHT_API_FILE_CONVERSION_MAX_BATCH
|
||||
ENV_FILE_CONVERSION_MAX_BATCH_SIZE = "HINDSIGHT_API_FILE_CONVERSION_MAX_BATCH_SIZE"
|
||||
ENV_ENABLE_FILE_UPLOAD_API = "HINDSIGHT_API_ENABLE_FILE_UPLOAD_API"
|
||||
ENV_FILE_DELETE_AFTER_RETAIN = "HINDSIGHT_API_FILE_DELETE_AFTER_RETAIN"
|
||||
ENV_STORE_DOCUMENT_TEXT = "HINDSIGHT_API_STORE_DOCUMENT_TEXT"
|
||||
|
||||
# Document transfer (export/import documents between banks without re-running the LLM)
|
||||
ENV_ENABLE_DOCUMENT_EXPORT_API = "HINDSIGHT_API_ENABLE_DOCUMENT_EXPORT_API"
|
||||
ENV_ENABLE_DOCUMENT_IMPORT_API = "HINDSIGHT_API_ENABLE_DOCUMENT_IMPORT_API"
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
ENV_ENABLE_OBSERVATIONS = "HINDSIGHT_API_ENABLE_OBSERVATIONS"
|
||||
@@ -406,8 +441,10 @@ ENV_ENABLE_AUTO_CONSOLIDATION = "HINDSIGHT_API_ENABLE_AUTO_CONSOLIDATION"
|
||||
ENV_CONSOLIDATION_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_BATCH_SIZE"
|
||||
ENV_CONSOLIDATION_MAX_MEMORIES_PER_ROUND = "HINDSIGHT_API_CONSOLIDATION_MAX_MEMORIES_PER_ROUND"
|
||||
ENV_CONSOLIDATION_LLM_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_LLM_BATCH_SIZE"
|
||||
ENV_CONSOLIDATION_DEDUP_THRESHOLD = "HINDSIGHT_API_CONSOLIDATION_DEDUP_THRESHOLD"
|
||||
ENV_CONSOLIDATION_LLM_PARALLELISM = "HINDSIGHT_API_CONSOLIDATION_LLM_PARALLELISM"
|
||||
ENV_CONSOLIDATION_MAX_TOKENS = "HINDSIGHT_API_CONSOLIDATION_MAX_TOKENS"
|
||||
ENV_CONSOLIDATION_MAX_COMPLETION_TOKENS = "HINDSIGHT_API_CONSOLIDATION_MAX_COMPLETION_TOKENS"
|
||||
ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS = "HINDSIGHT_API_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS"
|
||||
ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION = (
|
||||
"HINDSIGHT_API_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION"
|
||||
@@ -417,6 +454,7 @@ ENV_CONSOLIDATION_MAX_ATTEMPTS = "HINDSIGHT_API_CONSOLIDATION_MAX_ATTEMPTS"
|
||||
ENV_OBSERVATIONS_MISSION = "HINDSIGHT_API_OBSERVATIONS_MISSION"
|
||||
ENV_MAX_OBSERVATIONS_PER_SCOPE = "HINDSIGHT_API_MAX_OBSERVATIONS_PER_SCOPE"
|
||||
ENV_ENABLE_OBSERVATION_HISTORY = "HINDSIGHT_API_ENABLE_OBSERVATION_HISTORY"
|
||||
ENV_OBSERVATION_HISTORY_MAX_ENTRIES = "HINDSIGHT_API_OBSERVATION_HISTORY_MAX_ENTRIES"
|
||||
ENV_ENABLE_MENTAL_MODEL_HISTORY = "HINDSIGHT_API_ENABLE_MENTAL_MODEL_HISTORY"
|
||||
ENV_MENTAL_MODEL_HISTORY_MAX_ENTRIES = "HINDSIGHT_API_MENTAL_MODEL_HISTORY_MAX_ENTRIES"
|
||||
|
||||
@@ -448,6 +486,11 @@ ENV_DB_COMMAND_TIMEOUT = "HINDSIGHT_API_DB_COMMAND_TIMEOUT"
|
||||
ENV_DB_ACQUIRE_TIMEOUT = "HINDSIGHT_API_DB_ACQUIRE_TIMEOUT"
|
||||
ENV_DB_STATEMENT_TIMEOUT = "HINDSIGHT_API_DB_STATEMENT_TIMEOUT"
|
||||
|
||||
# Wall-clock cap on model/connection initialization at startup. If embeddings,
|
||||
# cross-encoder, or LLM verification hang (e.g. an offline HuggingFace download
|
||||
# or an unreachable provider), the daemon fails fast instead of hanging forever.
|
||||
ENV_MODEL_INIT_TIMEOUT = "HINDSIGHT_API_MODEL_INIT_TIMEOUT"
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
ENV_WORKER_ENABLED = "HINDSIGHT_API_WORKER_ENABLED"
|
||||
ENV_WORKER_ID = "HINDSIGHT_API_WORKER_ID"
|
||||
@@ -468,6 +511,7 @@ WORKER_SLOT_RESERVATION_TYPES: dict[str, tuple[str, int]] = {
|
||||
"file_convert_retain": ("HINDSIGHT_API_WORKER_FILE_CONVERT_RETAIN_MAX_SLOTS", 0),
|
||||
"refresh_mental_model": ("HINDSIGHT_API_WORKER_REFRESH_MENTAL_MODEL_MAX_SLOTS", 0),
|
||||
"graph_maintenance": ("HINDSIGHT_API_WORKER_GRAPH_MAINTENANCE_MAX_SLOTS", 0),
|
||||
"import_documents": ("HINDSIGHT_API_WORKER_IMPORT_DOCUMENTS_MAX_SLOTS", 0),
|
||||
}
|
||||
ENV_WORKER_CONSOLIDATION_BANK_PRIORITY = "HINDSIGHT_API_WORKER_CONSOLIDATION_BANK_PRIORITY"
|
||||
ENV_RETAIN_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_MAX_CONCURRENT"
|
||||
@@ -493,11 +537,32 @@ ENV_RECALL_BUDGET_ADAPTIVE_HIGH = "HINDSIGHT_API_RECALL_BUDGET_ADAPTIVE_HIGH"
|
||||
ENV_RECALL_BUDGET_MIN = "HINDSIGHT_API_RECALL_BUDGET_MIN"
|
||||
ENV_RECALL_BUDGET_MAX = "HINDSIGHT_API_RECALL_BUDGET_MAX"
|
||||
|
||||
# Recall candidate gating (per-source cap + BM25 score floor)
|
||||
ENV_BM25_MIN_SCORE = "HINDSIGHT_API_BM25_MIN_SCORE"
|
||||
ENV_RECALL_MAX_CANDIDATES_PER_SOURCE = "HINDSIGHT_API_RECALL_MAX_CANDIDATES_PER_SOURCE"
|
||||
# Per-strategy recall boost. Prioritises specific retrieval arms (semantic,
|
||||
# bm25, graph, temporal) on recall via a human priority level — e.g.
|
||||
# "graph:high" to strongly favour graph hits, or "graph:high,semantic:low".
|
||||
# Valid levels: low | medium | high. The level (not a raw number) is the knob
|
||||
# because the boost is applied on two different score scales — see
|
||||
# engine/search/recall_boost.py for the level -> magnitude mapping and rationale.
|
||||
# Empty disables the feature.
|
||||
ENV_RECALL_STRATEGY_BOOSTS = "HINDSIGHT_API_RECALL_STRATEGY_BOOSTS"
|
||||
|
||||
# Audit log settings
|
||||
ENV_AUDIT_LOG_ENABLED = "HINDSIGHT_API_AUDIT_LOG_ENABLED"
|
||||
ENV_AUDIT_LOG_ACTIONS = "HINDSIGHT_API_AUDIT_LOG_ACTIONS"
|
||||
ENV_AUDIT_LOG_RETENTION_DAYS = "HINDSIGHT_API_AUDIT_LOG_RETENTION_DAYS"
|
||||
|
||||
# LLM request tracing settings
|
||||
ENV_LLM_TRACE_ENABLED = "HINDSIGHT_API_LLM_TRACE_ENABLED"
|
||||
ENV_LLM_TRACE_SCOPES = "HINDSIGHT_API_LLM_TRACE_SCOPES"
|
||||
ENV_LLM_TRACE_RETENTION_DAYS = "HINDSIGHT_API_LLM_TRACE_RETENTION_DAYS"
|
||||
ENV_LLM_TRACE_MAX_CHARS = "HINDSIGHT_API_LLM_TRACE_MAX_CHARS"
|
||||
|
||||
# Background maintenance settings
|
||||
ENV_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS = "HINDSIGHT_API_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS"
|
||||
|
||||
# Disposition settings
|
||||
ENV_DISPOSITION_SKEPTICISM = "HINDSIGHT_API_DISPOSITION_SKEPTICISM"
|
||||
ENV_DISPOSITION_LITERALISM = "HINDSIGHT_API_DISPOSITION_LITERALISM"
|
||||
@@ -513,9 +578,9 @@ DEFAULT_LLM_PROVIDER = "openai"
|
||||
PROVIDER_DEFAULT_MODELS = {
|
||||
"openai": "gpt-4o-mini",
|
||||
"anthropic": "claude-haiku-4-5",
|
||||
"gemini": "gemini-2.5-flash",
|
||||
"gemini": "gemini-3.5-flash",
|
||||
"groq": "openai/gpt-oss-120b",
|
||||
"minimax": "MiniMax-M2.7",
|
||||
"minimax": "MiniMax-M3",
|
||||
"deepseek": "deepseek-v4-flash",
|
||||
"zai": "glm-4.5-flash",
|
||||
"opencode-go": "deepseek-v4-flash",
|
||||
@@ -523,7 +588,7 @@ PROVIDER_DEFAULT_MODELS = {
|
||||
"ollama-cloud": "gemma3:12b",
|
||||
"llamacpp": "gemma-4-e2b-it",
|
||||
"lmstudio": "local-model",
|
||||
"vertexai": "google/gemini-2.5-flash-lite",
|
||||
"vertexai": "google/gemini-3.1-flash-lite",
|
||||
"openai-codex": "gpt-5.4-mini",
|
||||
"claude-code": "claude-sonnet-4-5-20250929",
|
||||
"mock": "mock-model",
|
||||
@@ -533,6 +598,7 @@ PROVIDER_DEFAULT_MODELS = {
|
||||
"volcano": "doubao-pro-32k",
|
||||
"openrouter": "qwen/qwen3.5-9b",
|
||||
"fireworks": "accounts/fireworks/models/llama-v3p1-8b-instruct",
|
||||
"nous": "deepseek/deepseek-v4-flash",
|
||||
}
|
||||
DEFAULT_LLM_MODEL = "gpt-4o-mini" # Fallback if provider not in table
|
||||
# Built-in llama.cpp defaults
|
||||
@@ -542,12 +608,21 @@ DEFAULT_LLAMACPP_CHAT_FORMAT = None # None = auto-detect from GGUF metadata
|
||||
DEFAULT_LLAMACPP_NO_GRAMMAR = False # True = disable JSON grammar enforcement (faster but less reliable)
|
||||
DEFAULT_LLAMACPP_EXTRA_ARGS = None # Space-separated extra CLI args for llama.cpp server
|
||||
|
||||
# True = ask schema-capable backends to grammar-enforce structured output via
|
||||
# json_schema strict (OpenAI-compatible, LiteLLM; Gemini already enforces its
|
||||
# native response_schema). Default False keeps the soft "schema-in-prompt +
|
||||
# json_object" path, which weaker self-hosted instruction-followers can violate
|
||||
# (prose preambles, markdown fences, invalid JSON) — wedging retain/consolidation
|
||||
# on parse retries.
|
||||
DEFAULT_LLM_STRICT_SCHEMA = False
|
||||
|
||||
DEFAULT_LLM_MAX_CONCURRENT = 32
|
||||
DEFAULT_LLM_MAX_RETRIES = 3 # Max retry attempts for LLM API calls
|
||||
DEFAULT_LLM_INITIAL_BACKOFF = 1.0 # Initial backoff in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_MAX_BACKOFF = 60.0 # Max backoff cap in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_TIMEOUT = 120.0 # seconds
|
||||
DEFAULT_LLM_REASONING_EFFORT = "low"
|
||||
DEFAULT_LLM_SEND_BANK_AS_USER = False # Opt-in: tag provider calls with user=<bank_id>
|
||||
|
||||
# Vertex AI defaults
|
||||
DEFAULT_LLM_VERTEXAI_PROJECT_ID = None # Required for Vertex AI
|
||||
@@ -561,6 +636,13 @@ DEFAULT_EMBEDDINGS_PROVIDER = "local"
|
||||
DEFAULT_EMBEDDINGS_LOCAL_MODEL = "BAAI/bge-small-en-v1.5"
|
||||
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU = False # Force CPU mode for local embeddings
|
||||
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE = False # Security: disabled by default, required for some models
|
||||
DEFAULT_EMBEDDINGS_ONNX_MODEL_ID = "intfloat/multilingual-e5-small"
|
||||
DEFAULT_EMBEDDINGS_ONNX_FILE = "onnx/model.onnx"
|
||||
DEFAULT_EMBEDDINGS_ONNX_MAX_TOKENS = 512
|
||||
DEFAULT_EMBEDDINGS_ONNX_POOLING = "mean"
|
||||
DEFAULT_EMBEDDINGS_ONNX_NORMALIZE = True
|
||||
DEFAULT_EMBEDDINGS_ONNX_QUERY_PREFIX = "query: "
|
||||
DEFAULT_EMBEDDINGS_ONNX_PASSAGE_PREFIX = "passage: "
|
||||
DEFAULT_EMBEDDINGS_OPENAI_MODEL = "text-embedding-3-small"
|
||||
DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE = 100
|
||||
DEFAULT_EMBEDDINGS_GEMINI_MODEL = "gemini-embedding-001"
|
||||
@@ -592,6 +674,65 @@ DEFAULT_RERANKER_LITELLM_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_LITELLM_SDK_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_GOOGLE_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_MAX_CANDIDATES = 300
|
||||
DEFAULT_SEMANTIC_MIN_SIMILARITY = 0.3
|
||||
# Minimum BM25 score a row must exceed to enter fusion. 0.0 gates out
|
||||
# zero-score (non-matching) rows on backends — notably VectorChord — whose
|
||||
# operator ranks every document rather than pre-filtering to term matches.
|
||||
DEFAULT_BM25_MIN_SCORE = 0.0
|
||||
# Per-source candidate cap applied to each retrieval arm (semantic, BM25, graph,
|
||||
# temporal) before RRF, so a single over-expanding backend cannot fill the
|
||||
# reranker's global candidate budget on its own. 0 disables the cap.
|
||||
DEFAULT_RECALL_MAX_CANDIDATES_PER_SOURCE = 0
|
||||
# Per-strategy recall boost, as a comma-separated "strategy:level" list (e.g.
|
||||
# "graph:high,semantic:low"). Empty disables the feature. See
|
||||
# ENV_RECALL_STRATEGY_BOOSTS for the full rationale.
|
||||
DEFAULT_RECALL_STRATEGY_BOOSTS = ""
|
||||
# Retrieval arms that can be boosted; mirrors fusion.py source_names.
|
||||
RECALL_STRATEGY_NAMES = ("semantic", "bm25", "graph", "temporal")
|
||||
# User-facing priority levels. Kept in sync with recall_boost.BOOST_LEVELS by a
|
||||
# guard test; defined here (not imported) so config stays free of the heavy
|
||||
# engine.search import graph.
|
||||
RECALL_BOOST_LEVELS = ("low", "medium", "high")
|
||||
# Level applied when a strategy is listed without one (e.g. "graph" or "graph:").
|
||||
DEFAULT_RECALL_BOOST_LEVEL = "medium"
|
||||
|
||||
|
||||
def _parse_strategy_boosts(raw: str | None) -> dict[str, str]:
|
||||
"""Parse a "strategy:level,strategy:level" string into a boost map.
|
||||
|
||||
A strategy listed without a level (``"graph"`` or ``"graph:"``) defaults to
|
||||
``medium``. Only the strategies you list are boosted; any strategy you omit
|
||||
keeps its normal, unboosted weight. Unknown strategy names, unknown levels,
|
||||
and malformed entries are skipped with a warning so a typo degrades to a
|
||||
no-op boost rather than breaking recall.
|
||||
"""
|
||||
if not raw or not raw.strip():
|
||||
return {}
|
||||
boosts: dict[str, str] = {}
|
||||
for entry in raw.split(","):
|
||||
entry = entry.strip()
|
||||
if not entry:
|
||||
continue
|
||||
name, _sep, level = entry.partition(":")
|
||||
name = name.strip().lower()
|
||||
level = level.strip().lower() or DEFAULT_RECALL_BOOST_LEVEL
|
||||
if name not in RECALL_STRATEGY_NAMES:
|
||||
logger.warning(
|
||||
"Ignoring unknown recall strategy %r in boost (valid: %s)", name, ", ".join(RECALL_STRATEGY_NAMES)
|
||||
)
|
||||
continue
|
||||
if level not in RECALL_BOOST_LEVELS:
|
||||
logger.warning(
|
||||
"Ignoring unknown recall boost level %r for %r (valid: %s)",
|
||||
level,
|
||||
name,
|
||||
", ".join(RECALL_BOOST_LEVELS),
|
||||
)
|
||||
continue
|
||||
boosts[name] = level
|
||||
return boosts
|
||||
|
||||
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL = "ms-marco-MiniLM-L-12-v2" # Best balance of speed and quality
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR = None # Use default cache directory
|
||||
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA = False # Disable ONNX CPU memory arena to bound RSS
|
||||
@@ -602,6 +743,7 @@ DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
|
||||
# OpenRouter defaults
|
||||
DEFAULT_EMBEDDINGS_OPENROUTER_MODEL = "perplexity/pplx-embed-v1-0.6b"
|
||||
DEFAULT_RERANKER_OPENROUTER_MODEL = "cohere/rerank-v3.5"
|
||||
DEFAULT_RERANKER_OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1/rerank"
|
||||
|
||||
# ZeroEntropy defaults
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL = "zembed-1"
|
||||
@@ -659,6 +801,9 @@ DEFAULT_MCP_ENABLED = True
|
||||
DEFAULT_MCP_ENABLED_TOOLS: list[str] | None = None # None = all tools enabled
|
||||
DEFAULT_MCP_STATELESS = False # False = stateful (supports SSE/GET); True = stateless (POST-only)
|
||||
DEFAULT_ENABLE_BANK_CONFIG_API = True
|
||||
# The per-bank LLM connectivity probe makes a real provider call, so it's OFF by
|
||||
# default (cost/abuse concerns) and must be explicitly enabled to expose the endpoint.
|
||||
DEFAULT_ENABLE_BANK_LLM_HEALTH = False
|
||||
DEFAULT_DEFAULT_BANK_TEMPLATE: dict | None = None # BankTemplateManifest dict applied to newly-created banks
|
||||
DEFAULT_GRAPH_RETRIEVER = "link_expansion"
|
||||
DEFAULT_RECALL_MAX_CONCURRENT = 32 # Max concurrent recall operations per worker
|
||||
@@ -667,6 +812,8 @@ DEFAULT_RECALL_MAX_QUERY_TOKENS = 500 # Maximum tokens allowed in recall query
|
||||
DEFAULT_MENTAL_MODEL_REFRESH_CONCURRENCY = 8 # Max concurrent mental model refreshes
|
||||
DEFAULT_LINK_EXPANSION_PER_ENTITY_LIMIT = 200 # Max target units per entity in graph expansion
|
||||
DEFAULT_LINK_EXPANSION_TIMEOUT = 10.0 # Timeout (seconds) for entity expansion query
|
||||
DEFAULT_BANK_STATS_CACHE_TTL_SECONDS = 60.0 # TTL for get_bank_stats result cache; 0 disables
|
||||
DEFAULT_BANK_STATS_CACHE_MAX_ENTRIES = 1024 # LRU bound across (schema, bank) keys
|
||||
|
||||
# Retain settings
|
||||
DEFAULT_RETAIN_MAX_COMPLETION_TOKENS = 64000 # Max tokens for fact extraction LLM call
|
||||
@@ -685,6 +832,7 @@ DEFAULT_RETAIN_BATCH_TOKENS = 10_000 # ~40KB of text # Max chars per sub-batch
|
||||
DEFAULT_RETAIN_ENTITY_LOOKUP = "trigram" # "full" or "trigram"
|
||||
DEFAULT_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE = 100 # Unique entity names per pg_trgm candidate lookup query
|
||||
DEFAULT_RETAIN_BATCH_ENABLED = False # Use LLM Batch API for fact extraction (only when async=True)
|
||||
DEFAULT_LLM_PROMPT_CACHE_ENABLED = True # Reuse the fixed system prefix via provider prompt caching
|
||||
DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60 # Batch API polling interval in seconds
|
||||
|
||||
# File storage defaults
|
||||
@@ -695,29 +843,47 @@ DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE_MB = 100 # Max total batch size in MB (a
|
||||
DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE = 10 # Max files per batch upload
|
||||
DEFAULT_ENABLE_FILE_UPLOAD_API = True # Enable file upload endpoint
|
||||
DEFAULT_FILE_DELETE_AFTER_RETAIN = True # Delete file bytes after retain (saves storage)
|
||||
DEFAULT_STORE_DOCUMENT_TEXT = True # Persist raw source text in documents.original_text / chunks.chunk_text
|
||||
|
||||
# Document transfer defaults (export/import enabled by default; gated independently)
|
||||
DEFAULT_ENABLE_DOCUMENT_EXPORT_API = True
|
||||
DEFAULT_ENABLE_DOCUMENT_IMPORT_API = True
|
||||
|
||||
# Observations defaults (consolidated knowledge from facts)
|
||||
DEFAULT_ENABLE_OBSERVATIONS = True # Observations enabled by default
|
||||
DEFAULT_ENABLE_AUTO_CONSOLIDATION = True # Auto-consolidation after retain enabled by default
|
||||
DEFAULT_ENABLE_OBSERVATION_HISTORY = True # Observation history tracking enabled by default
|
||||
DEFAULT_ENABLE_MENTAL_MODEL_HISTORY = True # Mental model history tracking enabled by default
|
||||
# Each history entry snapshots previous_content + previous_reflect_response. Without
|
||||
# a cap, sustained mental-model refresh load grows the jsonb array unboundedly until
|
||||
# it crosses Postgres's hard 256MB jsonb limit and subsequent UPDATEs fail with
|
||||
# SQLSTATE 54000. 50 keeps the array well under 100MB even with large reflect
|
||||
# responses, while preserving enough recent history for meaningful audit / rollback.
|
||||
# History (mental-model refresh snapshots and observation update snapshots) lives in
|
||||
# the dedicated mental_model_history / observation_history tables, one row per change.
|
||||
# On every write we insert the new entry and delete the oldest rows beyond the cap,
|
||||
# so the per-item history can never grow unboundedly (the old single-JSONB-column
|
||||
# design hit Postgres's hard 256MB jsonb limit -> SQLSTATE 54000 and stuck rows).
|
||||
# 50 preserves enough recent history for meaningful audit / rollback per item.
|
||||
# A cap <= 0 removes the trim (unbounded growth) — to turn history OFF use the
|
||||
# enable_* flag, not a zero cap.
|
||||
DEFAULT_MENTAL_MODEL_HISTORY_MAX_ENTRIES = 50
|
||||
DEFAULT_OBSERVATION_HISTORY_MAX_ENTRIES = 50
|
||||
DEFAULT_CONSOLIDATION_MAX_ATTEMPTS = 3 # Outer retry attempts for consolidation LLM batch calls
|
||||
DEFAULT_CONSOLIDATION_BATCH_SIZE = 50 # Memories to load per batch (internal memory optimization)
|
||||
DEFAULT_CONSOLIDATION_MAX_MEMORIES_PER_ROUND = (
|
||||
100 # Max memories per consolidation round (0 = unlimited). Limits how long one bank holds a worker slot.
|
||||
)
|
||||
DEFAULT_CONSOLIDATION_LLM_BATCH_SIZE = 8 # Facts per LLM call (1 = no batching; >1 = batch mode)
|
||||
# Cosine >= this between a newly-created or freshly-updated observation and an existing one
|
||||
# triggers a focused 1-by-1 LLM "merge or keep" pass (the LLM reads both, so numbers/negation/
|
||||
# entities are respected). Enabled by default; set to 1.0 to disable. Postgres only — the merge
|
||||
# path uses Postgres-only SQL, so consolidation skips it on Oracle regardless of this value.
|
||||
DEFAULT_CONSOLIDATION_DEDUP_THRESHOLD = 0.97
|
||||
DEFAULT_CONSOLIDATION_LLM_PARALLELISM = (
|
||||
4 # Max tag groups consolidated concurrently per op. Locks on overlapping write
|
||||
# scopes degrade to sequential automatically; matches retain_max_concurrent.
|
||||
)
|
||||
DEFAULT_CONSOLIDATION_MAX_TOKENS = 512 # Max tokens for recall when finding related observations
|
||||
# Unset by default: the key is omitted from the LLM call so every provider keeps its current implicit output
|
||||
# budget — 100% backwards compatible. Operators on providers with a low hidden default (notably Bedrock imported
|
||||
# models, which cap at 4096 and truncate structured consolidation JSON) set this explicitly to fix #1939.
|
||||
DEFAULT_CONSOLIDATION_MAX_COMPLETION_TOKENS = None
|
||||
DEFAULT_CONSOLIDATION_RECALL_BUDGET = "low" # Budget level for consolidation recall (low/mid/high)
|
||||
DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS = (
|
||||
4096 # Total token budget for source facts in consolidation recall (-1 = unlimited)
|
||||
@@ -737,6 +903,7 @@ DEFAULT_DB_POOL_MAX_SIZE = 100
|
||||
DEFAULT_DB_COMMAND_TIMEOUT = 60 # seconds
|
||||
DEFAULT_DB_ACQUIRE_TIMEOUT = 30 # seconds
|
||||
DEFAULT_DB_STATEMENT_TIMEOUT = 600 # seconds (Postgres statement_timeout applied on every pool connection; 0 disables)
|
||||
DEFAULT_MODEL_INIT_TIMEOUT = 300 # seconds (cap on startup model/connection init; covers first-time downloads)
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
DEFAULT_WORKER_ENABLED = True # API runs worker by default (standalone mode)
|
||||
@@ -789,6 +956,18 @@ DEFAULT_AUDIT_LOG_ENABLED = False # Disabled by default
|
||||
DEFAULT_AUDIT_LOG_ACTIONS = "" # Empty = audit all eligible actions
|
||||
DEFAULT_AUDIT_LOG_RETENTION_DAYS = -1 # -1 = keep forever
|
||||
|
||||
# LLM request tracing defaults
|
||||
DEFAULT_LLM_TRACE_ENABLED = True # Enabled by default
|
||||
DEFAULT_LLM_TRACE_SCOPES = "" # Empty = trace all call scopes
|
||||
DEFAULT_LLM_TRACE_RETENTION_DAYS = 1 # Retain trace rows for 1 day by default
|
||||
DEFAULT_LLM_TRACE_MAX_CHARS = 50000 # Truncate stored input/output beyond this many chars
|
||||
|
||||
# Background maintenance defaults
|
||||
# Periodic reconcile that re-schedules consolidation for banks with eligible-but-unscheduled
|
||||
# facts (e.g. after a consolidation operation failed terminally and left them unscheduled).
|
||||
# 0 disables the reconcile sweep.
|
||||
DEFAULT_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS = 300
|
||||
|
||||
# Default MCP tool descriptions (can be customized via env vars)
|
||||
DEFAULT_MCP_RETAIN_DESCRIPTION = """Store important information to long-term memory.
|
||||
|
||||
@@ -1046,12 +1225,19 @@ class HindsightConfig:
|
||||
llm_reasoning_effort: str
|
||||
llm_groq_service_tier: str # Groq: "on_demand", "flex", or "auto"
|
||||
llm_openai_service_tier: str | None # OpenAI: None (default) or "flex" (50% cheaper)
|
||||
llm_bedrock_service_tier: str | None # Bedrock: None (default), "flex", "priority", or "reserved"
|
||||
llm_extra_body: (
|
||||
dict | None
|
||||
) # Extra body params merged into OpenAI-compatible API calls (e.g. {"chat_template_kwargs": {"enable_thinking": true}})
|
||||
llm_default_headers: (
|
||||
dict | None
|
||||
) # Custom headers passed as default_headers to provider SDK clients (e.g. {"X-Component-Id": "hindsight"} for proxies / request tracing)
|
||||
llm_strict_schema: bool # Grammar-enforce structured output via the provider's strongest schema mode (see DEFAULT_LLM_STRICT_SCHEMA)
|
||||
# Tags outbound OpenAI-compatible LLM + embedding calls with `user=<bank_id>` for
|
||||
# per-bank cost attribution. Downstream cost gateways (OpenRouter usage accounting,
|
||||
# LiteLLM, Helicone) key attribution on the OpenAI `user` field. Opt-in; never
|
||||
# overrides a `user` the caller already set.
|
||||
llm_send_bank_as_user: bool
|
||||
|
||||
# LiteLLM Router chain (provider-specific; consumed by the "litellmrouter" provider).
|
||||
# List of deployment dicts evaluated in order with fallback on transient errors.
|
||||
@@ -1067,6 +1253,10 @@ class HindsightConfig:
|
||||
# Gemini safety settings (None = use Gemini defaults; list of dicts with category/threshold)
|
||||
llm_gemini_safety_settings: list | None
|
||||
|
||||
# Gemini prompt caching toggle. When True, retain extraction reuses a
|
||||
# CachedContent prefix for its system prompt + response schema.
|
||||
llm_prompt_cache_enabled: bool
|
||||
|
||||
# Built-in llama.cpp configuration (for provider=llamacpp)
|
||||
llamacpp_model_path: str | None # Path to GGUF file (None = auto-download default)
|
||||
llamacpp_gpu_layers: int # -1 = all layers on GPU, 0 = CPU only
|
||||
@@ -1119,6 +1309,17 @@ class HindsightConfig:
|
||||
embeddings_local_model: str
|
||||
embeddings_local_force_cpu: bool
|
||||
embeddings_local_trust_remote_code: bool
|
||||
embeddings_onnx_model_id: str
|
||||
embeddings_onnx_model_path: str | None
|
||||
embeddings_onnx_tokenizer_name_or_path: str | None
|
||||
embeddings_onnx_file: str
|
||||
embeddings_onnx_dimensions: int | None
|
||||
embeddings_onnx_max_tokens: int
|
||||
embeddings_onnx_pooling: str
|
||||
embeddings_onnx_normalize: bool
|
||||
embeddings_onnx_query_prefix: str
|
||||
embeddings_onnx_passage_prefix: str
|
||||
embeddings_onnx_output_name: str | None
|
||||
embeddings_tei_url: str | None
|
||||
embeddings_openai_base_url: str | None
|
||||
embeddings_cohere_api_key: str | None
|
||||
@@ -1158,12 +1359,17 @@ class HindsightConfig:
|
||||
reranker_tei_max_concurrent: int
|
||||
reranker_tei_http_timeout: float
|
||||
reranker_max_candidates: int
|
||||
semantic_min_similarity: float
|
||||
bm25_min_score: float
|
||||
recall_max_candidates_per_source: int
|
||||
recall_strategy_boosts: dict[str, str]
|
||||
reranker_cohere_api_key: str | None
|
||||
reranker_cohere_model: str
|
||||
reranker_cohere_base_url: str | None
|
||||
reranker_cohere_timeout: float
|
||||
reranker_openrouter_api_key: str | None
|
||||
reranker_openrouter_model: str
|
||||
reranker_openrouter_base_url: str
|
||||
reranker_openrouter_timeout: float
|
||||
reranker_litellm_api_base: str
|
||||
reranker_litellm_api_key: str | None
|
||||
@@ -1201,6 +1407,7 @@ class HindsightConfig:
|
||||
mcp_enabled_tools: list[str] | None # None = all tools; explicit list = allowlist
|
||||
mcp_stateless: bool # True = stateless HTTP (POST-only); False = stateful (supports GET/SSE)
|
||||
enable_bank_config_api: bool
|
||||
enable_bank_llm_health: bool
|
||||
# Default bank template (static, server-level only). When set, the manifest is applied
|
||||
# to every newly-created bank, overriding the env/config defaults for any fields it sets.
|
||||
default_bank_template: dict | None
|
||||
@@ -1213,6 +1420,8 @@ class HindsightConfig:
|
||||
mental_model_refresh_concurrency: int
|
||||
link_expansion_per_entity_limit: int
|
||||
link_expansion_timeout: float
|
||||
bank_stats_cache_ttl_seconds: float
|
||||
bank_stats_cache_max_entries: int
|
||||
|
||||
# Retain settings
|
||||
retain_max_completion_tokens: int
|
||||
@@ -1251,18 +1460,24 @@ class HindsightConfig:
|
||||
file_conversion_max_batch_size: int # Max files per request
|
||||
enable_file_upload_api: bool
|
||||
file_delete_after_retain: bool
|
||||
store_document_text: bool # When False, store NULL original_text / empty chunk_text
|
||||
enable_document_export_api: bool
|
||||
enable_document_import_api: bool
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations: bool
|
||||
enable_auto_consolidation: bool
|
||||
enable_observation_history: bool
|
||||
observation_history_max_entries: int
|
||||
enable_mental_model_history: bool
|
||||
mental_model_history_max_entries: int
|
||||
consolidation_batch_size: int
|
||||
consolidation_dedup_threshold: float
|
||||
consolidation_max_memories_per_round: int
|
||||
consolidation_llm_batch_size: int
|
||||
consolidation_llm_parallelism: int
|
||||
consolidation_max_tokens: int
|
||||
consolidation_max_completion_tokens: int | None
|
||||
consolidation_recall_budget: str
|
||||
consolidation_source_facts_max_tokens: int
|
||||
consolidation_source_facts_max_tokens_per_observation: int
|
||||
@@ -1277,6 +1492,10 @@ class HindsightConfig:
|
||||
# When False: only label entities are extracted (or no entities at all if no labels configured)
|
||||
entities_allow_free_form: bool
|
||||
|
||||
# Memory Defense policy (dict matching DefensePolicy schema — validated on write)
|
||||
# None = Memory Defense disabled / not configured for this bank
|
||||
memory_defense: dict | None
|
||||
|
||||
# Reflect agent settings
|
||||
reflect_mission: str | None
|
||||
reflect_source_facts_max_tokens: int
|
||||
@@ -1318,6 +1537,7 @@ class HindsightConfig:
|
||||
db_command_timeout: int
|
||||
db_acquire_timeout: int
|
||||
db_statement_timeout: int
|
||||
model_init_timeout: float
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
worker_enabled: bool
|
||||
@@ -1349,6 +1569,17 @@ class HindsightConfig:
|
||||
audit_log_actions: list[str] # Allowlist of action types (empty = all)
|
||||
audit_log_retention_days: int # -1 = keep forever, >0 = delete after N days
|
||||
|
||||
# LLM request tracing configuration (static - server-level only)
|
||||
llm_trace_enabled: bool # Master switch for per-bank LLM request tracing
|
||||
llm_trace_scopes: list[str] # Allowlist of call scopes to trace (empty = all)
|
||||
llm_trace_retention_days: int # -1 = keep forever, >0 = delete after N days
|
||||
llm_trace_max_chars: int # Truncate stored input/output beyond this many chars
|
||||
|
||||
# Background maintenance configuration (static - server-level only)
|
||||
# Interval for the periodic sweep that re-schedules consolidation for banks with
|
||||
# eligible-but-unscheduled facts. 0 = disabled.
|
||||
consolidation_reconcile_interval_seconds: int
|
||||
|
||||
# Webhook configuration (static - server-level only, not per-bank)
|
||||
webhook_url: str | None # Global webhook URL (None = disabled)
|
||||
webhook_secret: str | None # HMAC signing secret (None = unsigned)
|
||||
@@ -1389,6 +1620,7 @@ class HindsightConfig:
|
||||
"embeddings_tei_base_url",
|
||||
"reranker_tei_base_url",
|
||||
"reranker_cohere_base_url",
|
||||
"reranker_openrouter_base_url",
|
||||
"embeddings_zeroentropy_base_url",
|
||||
"reranker_zeroentropy_base_url",
|
||||
"reranker_siliconflow_base_url",
|
||||
@@ -1459,6 +1691,8 @@ class HindsightConfig:
|
||||
"disposition_empathy",
|
||||
# Gemini safety settings (controls content filtering for Gemini/VertexAI providers)
|
||||
"llm_gemini_safety_settings",
|
||||
# Memory Defense policy (validated against DefensePolicy schema on write)
|
||||
"memory_defense",
|
||||
}
|
||||
|
||||
@property
|
||||
@@ -1549,6 +1783,21 @@ class HindsightConfig:
|
||||
self.text_search_extension_pg_search_tokenizer
|
||||
)
|
||||
|
||||
if not 0.0 <= self.semantic_min_similarity <= 1.0:
|
||||
raise ValueError(
|
||||
f"Invalid semantic_min_similarity: {self.semantic_min_similarity}. Must be between 0.0 and 1.0"
|
||||
)
|
||||
|
||||
# Validate bedrock_service_tier
|
||||
valid_bedrock_tiers = (None, "flex", "priority", "reserved")
|
||||
if self.llm_bedrock_service_tier not in valid_bedrock_tiers:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_LLM_BEDROCK_SERVICE_TIER: "
|
||||
f"{self.llm_bedrock_service_tier!r}. Must be one of: "
|
||||
f"{', '.join(t for t in valid_bedrock_tiers if t is not None)}. "
|
||||
f"Note: 'standard' is not a valid Bedrock service tier -- use unset for default tier."
|
||||
)
|
||||
|
||||
# When LLM provider is "none", force chunks-only mode and disable LLM-dependent features
|
||||
if self.llm_provider == "none":
|
||||
self.retain_extraction_mode = "chunks"
|
||||
@@ -1597,6 +1846,21 @@ class HindsightConfig:
|
||||
" and ".join(missing),
|
||||
)
|
||||
|
||||
if self.embeddings_provider == "onnx":
|
||||
try:
|
||||
import importlib
|
||||
|
||||
importlib.import_module("onnxruntime")
|
||||
importlib.import_module("transformers")
|
||||
except ImportError:
|
||||
logger.warning(
|
||||
"ONNX embeddings provider configured, but 'onnxruntime' and/or "
|
||||
"'transformers' is not installed. The API will fail at model init time. Either:\n"
|
||||
" 1. Install ONNX deps: pip install hindsight-api-slim[local-onnx]\n"
|
||||
" 2. Use a different embeddings provider, e.g. HINDSIGHT_API_EMBEDDINGS_PROVIDER=local "
|
||||
"or openai"
|
||||
)
|
||||
|
||||
# Validate that sum of per-operation slot reservations does not exceed max_slots
|
||||
total_reserved = sum(self.worker_slot_reservations.values())
|
||||
if total_reserved > self.worker_max_slots:
|
||||
@@ -1647,8 +1911,12 @@ class HindsightConfig:
|
||||
llm_reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
|
||||
llm_groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
|
||||
llm_openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
|
||||
llm_bedrock_service_tier=os.getenv(ENV_LLM_BEDROCK_SERVICE_TIER) or None,
|
||||
llm_extra_body=json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null")),
|
||||
llm_default_headers=json.loads(os.getenv(ENV_LLM_DEFAULT_HEADERS, "null")),
|
||||
llm_strict_schema=os.getenv(ENV_LLM_STRICT_SCHEMA, str(DEFAULT_LLM_STRICT_SCHEMA)).lower() in ("true", "1"),
|
||||
llm_send_bank_as_user=os.getenv(ENV_LLM_SEND_BANK_AS_USER, str(DEFAULT_LLM_SEND_BANK_AS_USER)).lower()
|
||||
in ("true", "1"),
|
||||
llm_litellmrouter_config=_parse_llm_router_config(ENV_LLM_LITELLMROUTER_CONFIG),
|
||||
# Vertex AI
|
||||
llm_vertexai_project_id=os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID) or DEFAULT_LLM_VERTEXAI_PROJECT_ID,
|
||||
@@ -1657,6 +1925,10 @@ class HindsightConfig:
|
||||
or DEFAULT_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY,
|
||||
# Gemini safety settings (JSON-encoded list of {category, threshold} dicts)
|
||||
llm_gemini_safety_settings=json.loads(os.getenv(ENV_LLM_GEMINI_SAFETY_SETTINGS, "null")),
|
||||
llm_prompt_cache_enabled=os.getenv(
|
||||
ENV_LLM_PROMPT_CACHE_ENABLED, str(DEFAULT_LLM_PROMPT_CACHE_ENABLED)
|
||||
).lower()
|
||||
in ("1", "true", "yes", "on"),
|
||||
# Built-in llama.cpp configuration
|
||||
llamacpp_model_path=os.getenv(ENV_LLAMACPP_MODEL_PATH) or None,
|
||||
llamacpp_gpu_layers=int(os.getenv(ENV_LLAMACPP_GPU_LAYERS, str(DEFAULT_LLAMACPP_GPU_LAYERS))),
|
||||
@@ -1755,6 +2027,36 @@ class HindsightConfig:
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE, str(DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE)
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
embeddings_onnx_model_id=os.getenv(ENV_EMBEDDINGS_ONNX_MODEL_ID, DEFAULT_EMBEDDINGS_ONNX_MODEL_ID),
|
||||
embeddings_onnx_model_path=os.getenv(ENV_EMBEDDINGS_ONNX_MODEL_PATH) or None,
|
||||
embeddings_onnx_tokenizer_name_or_path=os.getenv(ENV_EMBEDDINGS_ONNX_TOKENIZER_NAME_OR_PATH) or None,
|
||||
embeddings_onnx_file=os.getenv(ENV_EMBEDDINGS_ONNX_FILE, DEFAULT_EMBEDDINGS_ONNX_FILE),
|
||||
embeddings_onnx_dimensions=_parse_optional_positive_int(
|
||||
ENV_EMBEDDINGS_ONNX_DIMENSIONS,
|
||||
os.getenv(ENV_EMBEDDINGS_ONNX_DIMENSIONS),
|
||||
),
|
||||
embeddings_onnx_max_tokens=_parse_positive_int(
|
||||
ENV_EMBEDDINGS_ONNX_MAX_TOKENS,
|
||||
os.getenv(ENV_EMBEDDINGS_ONNX_MAX_TOKENS),
|
||||
DEFAULT_EMBEDDINGS_ONNX_MAX_TOKENS,
|
||||
),
|
||||
embeddings_onnx_pooling=_parse_optional_choice(
|
||||
ENV_EMBEDDINGS_ONNX_POOLING,
|
||||
os.getenv(ENV_EMBEDDINGS_ONNX_POOLING),
|
||||
frozenset({"mean", "cls"}),
|
||||
)
|
||||
or DEFAULT_EMBEDDINGS_ONNX_POOLING,
|
||||
embeddings_onnx_normalize=os.getenv(
|
||||
ENV_EMBEDDINGS_ONNX_NORMALIZE, str(DEFAULT_EMBEDDINGS_ONNX_NORMALIZE)
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
embeddings_onnx_query_prefix=os.getenv(
|
||||
ENV_EMBEDDINGS_ONNX_QUERY_PREFIX, DEFAULT_EMBEDDINGS_ONNX_QUERY_PREFIX
|
||||
),
|
||||
embeddings_onnx_passage_prefix=os.getenv(
|
||||
ENV_EMBEDDINGS_ONNX_PASSAGE_PREFIX, DEFAULT_EMBEDDINGS_ONNX_PASSAGE_PREFIX
|
||||
),
|
||||
embeddings_onnx_output_name=os.getenv(ENV_EMBEDDINGS_ONNX_OUTPUT_NAME) or None,
|
||||
embeddings_tei_url=os.getenv(ENV_EMBEDDINGS_TEI_URL),
|
||||
embeddings_openai_base_url=os.getenv(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None,
|
||||
embeddings_openai_batch_size=_parse_positive_int(
|
||||
@@ -1876,6 +2178,14 @@ class HindsightConfig:
|
||||
os.getenv(ENV_RERANKER_TEI_HTTP_TIMEOUT, str(DEFAULT_RERANKER_TEI_HTTP_TIMEOUT))
|
||||
),
|
||||
reranker_max_candidates=int(os.getenv(ENV_RERANKER_MAX_CANDIDATES, str(DEFAULT_RERANKER_MAX_CANDIDATES))),
|
||||
semantic_min_similarity=float(os.getenv(ENV_SEMANTIC_MIN_SIMILARITY, str(DEFAULT_SEMANTIC_MIN_SIMILARITY))),
|
||||
bm25_min_score=float(os.getenv(ENV_BM25_MIN_SCORE, str(DEFAULT_BM25_MIN_SCORE))),
|
||||
recall_max_candidates_per_source=int(
|
||||
os.getenv(ENV_RECALL_MAX_CANDIDATES_PER_SOURCE, str(DEFAULT_RECALL_MAX_CANDIDATES_PER_SOURCE))
|
||||
),
|
||||
recall_strategy_boosts=_parse_strategy_boosts(
|
||||
os.getenv(ENV_RECALL_STRATEGY_BOOSTS, DEFAULT_RECALL_STRATEGY_BOOSTS)
|
||||
),
|
||||
# Cohere reranker (with backward-compatible fallback to shared API key)
|
||||
reranker_cohere_api_key=os.getenv(ENV_RERANKER_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
|
||||
reranker_cohere_model=os.getenv(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL),
|
||||
@@ -1886,6 +2196,9 @@ class HindsightConfig:
|
||||
or os.getenv(ENV_OPENROUTER_API_KEY)
|
||||
or os.getenv(ENV_LLM_API_KEY),
|
||||
reranker_openrouter_model=os.getenv(ENV_RERANKER_OPENROUTER_MODEL, DEFAULT_RERANKER_OPENROUTER_MODEL),
|
||||
reranker_openrouter_base_url=os.getenv(
|
||||
ENV_RERANKER_OPENROUTER_BASE_URL, DEFAULT_RERANKER_OPENROUTER_BASE_URL
|
||||
),
|
||||
reranker_openrouter_timeout=float(
|
||||
os.getenv(ENV_RERANKER_OPENROUTER_TIMEOUT, str(DEFAULT_RERANKER_OPENROUTER_TIMEOUT))
|
||||
),
|
||||
@@ -1948,6 +2261,8 @@ class HindsightConfig:
|
||||
if os.getenv(ENV_MCP_ENABLED_TOOLS)
|
||||
else DEFAULT_MCP_ENABLED_TOOLS,
|
||||
mcp_stateless=os.getenv(ENV_MCP_STATELESS, str(DEFAULT_MCP_STATELESS)).lower() == "true",
|
||||
enable_bank_llm_health=os.getenv(ENV_ENABLE_BANK_LLM_HEALTH, str(DEFAULT_ENABLE_BANK_LLM_HEALTH)).lower()
|
||||
== "true",
|
||||
enable_bank_config_api=os.getenv(ENV_ENABLE_BANK_CONFIG_API, str(DEFAULT_ENABLE_BANK_CONFIG_API)).lower()
|
||||
== "true",
|
||||
default_bank_template=_parse_default_bank_template(os.getenv(ENV_DEFAULT_BANK_TEMPLATE)),
|
||||
@@ -1965,6 +2280,12 @@ class HindsightConfig:
|
||||
os.getenv(ENV_LINK_EXPANSION_PER_ENTITY_LIMIT, str(DEFAULT_LINK_EXPANSION_PER_ENTITY_LIMIT))
|
||||
),
|
||||
link_expansion_timeout=float(os.getenv(ENV_LINK_EXPANSION_TIMEOUT, str(DEFAULT_LINK_EXPANSION_TIMEOUT))),
|
||||
bank_stats_cache_ttl_seconds=float(
|
||||
os.getenv(ENV_BANK_STATS_CACHE_TTL_SECONDS, str(DEFAULT_BANK_STATS_CACHE_TTL_SECONDS))
|
||||
),
|
||||
bank_stats_cache_max_entries=int(
|
||||
os.getenv(ENV_BANK_STATS_CACHE_MAX_ENTRIES, str(DEFAULT_BANK_STATS_CACHE_MAX_ENTRIES))
|
||||
),
|
||||
# Optimization flags
|
||||
skip_llm_verification=os.getenv(ENV_SKIP_LLM_VERIFICATION, "false").lower() == "true",
|
||||
lazy_reranker=os.getenv(ENV_LAZY_RERANKER, "false").lower() == "true",
|
||||
@@ -2028,6 +2349,15 @@ class HindsightConfig:
|
||||
ENV_FILE_DELETE_AFTER_RETAIN, str(DEFAULT_FILE_DELETE_AFTER_RETAIN)
|
||||
).lower()
|
||||
== "true",
|
||||
store_document_text=os.getenv(ENV_STORE_DOCUMENT_TEXT, str(DEFAULT_STORE_DOCUMENT_TEXT)).lower() == "true",
|
||||
enable_document_export_api=os.getenv(
|
||||
ENV_ENABLE_DOCUMENT_EXPORT_API, str(DEFAULT_ENABLE_DOCUMENT_EXPORT_API)
|
||||
).lower()
|
||||
== "true",
|
||||
enable_document_import_api=os.getenv(
|
||||
ENV_ENABLE_DOCUMENT_IMPORT_API, str(DEFAULT_ENABLE_DOCUMENT_IMPORT_API)
|
||||
).lower()
|
||||
== "true",
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations=os.getenv(ENV_ENABLE_OBSERVATIONS, str(DEFAULT_ENABLE_OBSERVATIONS)).lower() == "true",
|
||||
enable_auto_consolidation=os.getenv(
|
||||
@@ -2038,6 +2368,12 @@ class HindsightConfig:
|
||||
ENV_ENABLE_OBSERVATION_HISTORY, str(DEFAULT_ENABLE_OBSERVATION_HISTORY)
|
||||
).lower()
|
||||
== "true",
|
||||
observation_history_max_entries=int(
|
||||
os.getenv(
|
||||
ENV_OBSERVATION_HISTORY_MAX_ENTRIES,
|
||||
str(DEFAULT_OBSERVATION_HISTORY_MAX_ENTRIES),
|
||||
)
|
||||
),
|
||||
enable_mental_model_history=os.getenv(
|
||||
ENV_ENABLE_MENTAL_MODEL_HISTORY, str(DEFAULT_ENABLE_MENTAL_MODEL_HISTORY)
|
||||
).lower()
|
||||
@@ -2057,6 +2393,9 @@ class HindsightConfig:
|
||||
str(DEFAULT_CONSOLIDATION_MAX_MEMORIES_PER_ROUND),
|
||||
)
|
||||
),
|
||||
consolidation_dedup_threshold=float(
|
||||
os.getenv(ENV_CONSOLIDATION_DEDUP_THRESHOLD, str(DEFAULT_CONSOLIDATION_DEDUP_THRESHOLD))
|
||||
),
|
||||
consolidation_llm_batch_size=int(
|
||||
os.getenv(ENV_CONSOLIDATION_LLM_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_LLM_BATCH_SIZE))
|
||||
),
|
||||
@@ -2072,6 +2411,11 @@ class HindsightConfig:
|
||||
consolidation_max_tokens=int(
|
||||
os.getenv(ENV_CONSOLIDATION_MAX_TOKENS, str(DEFAULT_CONSOLIDATION_MAX_TOKENS))
|
||||
),
|
||||
consolidation_max_completion_tokens=(
|
||||
int(os.getenv(ENV_CONSOLIDATION_MAX_COMPLETION_TOKENS))
|
||||
if os.getenv(ENV_CONSOLIDATION_MAX_COMPLETION_TOKENS)
|
||||
else DEFAULT_CONSOLIDATION_MAX_COMPLETION_TOKENS
|
||||
),
|
||||
consolidation_recall_budget=os.getenv(ENV_CONSOLIDATION_RECALL_BUDGET, DEFAULT_CONSOLIDATION_RECALL_BUDGET),
|
||||
consolidation_source_facts_max_tokens=int(
|
||||
os.getenv(ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS, str(DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS))
|
||||
@@ -2091,6 +2435,7 @@ class HindsightConfig:
|
||||
),
|
||||
entity_labels=None,
|
||||
entities_allow_free_form=True,
|
||||
memory_defense=None,
|
||||
# Database migrations
|
||||
run_migrations_on_startup=os.getenv(ENV_RUN_MIGRATIONS_ON_STARTUP, "true").lower() == "true",
|
||||
# Database connection pool
|
||||
@@ -2099,6 +2444,7 @@ class HindsightConfig:
|
||||
db_command_timeout=int(os.getenv(ENV_DB_COMMAND_TIMEOUT, str(DEFAULT_DB_COMMAND_TIMEOUT))),
|
||||
db_acquire_timeout=int(os.getenv(ENV_DB_ACQUIRE_TIMEOUT, str(DEFAULT_DB_ACQUIRE_TIMEOUT))),
|
||||
db_statement_timeout=int(os.getenv(ENV_DB_STATEMENT_TIMEOUT, str(DEFAULT_DB_STATEMENT_TIMEOUT))),
|
||||
model_init_timeout=float(os.getenv(ENV_MODEL_INIT_TIMEOUT, str(DEFAULT_MODEL_INIT_TIMEOUT))),
|
||||
# Worker configuration
|
||||
worker_enabled=os.getenv(ENV_WORKER_ENABLED, str(DEFAULT_WORKER_ENABLED)).lower() == "true",
|
||||
worker_id=os.getenv(ENV_WORKER_ID) or DEFAULT_WORKER_ID,
|
||||
@@ -2183,6 +2529,22 @@ class HindsightConfig:
|
||||
audit_log_retention_days=int(
|
||||
os.getenv(ENV_AUDIT_LOG_RETENTION_DAYS, str(DEFAULT_AUDIT_LOG_RETENTION_DAYS))
|
||||
),
|
||||
# LLM request tracing configuration (static, server-level only)
|
||||
llm_trace_enabled=os.getenv(ENV_LLM_TRACE_ENABLED, str(DEFAULT_LLM_TRACE_ENABLED)).lower() == "true",
|
||||
llm_trace_scopes=[
|
||||
s.strip() for s in os.getenv(ENV_LLM_TRACE_SCOPES, DEFAULT_LLM_TRACE_SCOPES).split(",") if s.strip()
|
||||
],
|
||||
llm_trace_retention_days=int(
|
||||
os.getenv(ENV_LLM_TRACE_RETENTION_DAYS, str(DEFAULT_LLM_TRACE_RETENTION_DAYS))
|
||||
),
|
||||
llm_trace_max_chars=int(os.getenv(ENV_LLM_TRACE_MAX_CHARS, str(DEFAULT_LLM_TRACE_MAX_CHARS))),
|
||||
# Background maintenance configuration (static, server-level only)
|
||||
consolidation_reconcile_interval_seconds=int(
|
||||
os.getenv(
|
||||
ENV_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS,
|
||||
str(DEFAULT_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS),
|
||||
)
|
||||
),
|
||||
# Webhook configuration (static, server-level only)
|
||||
webhook_url=os.getenv(ENV_WEBHOOK_URL) or DEFAULT_WEBHOOK_URL,
|
||||
webhook_secret=os.getenv(ENV_WEBHOOK_SECRET) or DEFAULT_WEBHOOK_SECRET,
|
||||
|
||||
@@ -266,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,34 @@
|
||||
"""Per-bank provider cost attribution via the OpenAI ``user`` field.
|
||||
|
||||
Shared by the OpenAI-compatible LLM path and the OpenAI embeddings path so both
|
||||
tag outbound requests identically. Opt-in via ``HINDSIGHT_API_LLM_SEND_BANK_AS_USER``;
|
||||
downstream cost gateways (OpenRouter usage accounting, LiteLLM, Helicone) key spend
|
||||
on the OpenAI ``user`` field.
|
||||
|
||||
Note: when enabled, the bank id is transmitted to the upstream provider as the
|
||||
end-user identifier. Banks that are themselves end-user identifiers are therefore
|
||||
forwarded to the provider — which is exactly what the OpenAI ``user`` field is for,
|
||||
but operators should opt in with that in mind.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def apply_bank_attribution(request: dict[str, Any]) -> None:
|
||||
"""Tag ``request`` with ``user=<bank_id>`` for per-bank cost attribution.
|
||||
|
||||
Mutates ``request`` in place. No-op when the flag is off, no bank is in context,
|
||||
or the caller already set ``user`` — we never override an explicit value.
|
||||
"""
|
||||
if "user" in request:
|
||||
return
|
||||
# Lazy imports: memory_engine imports the embeddings/provider modules that call
|
||||
# this, so a top-level import of memory_engine here would be circular.
|
||||
from ..config import get_config
|
||||
from .memory_engine import get_current_bank_id
|
||||
|
||||
if not get_config().llm_send_bank_as_user:
|
||||
return
|
||||
bank_id = get_current_bank_id()
|
||||
if bank_id:
|
||||
request["user"] = bank_id
|
||||
@@ -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()
|
||||
@@ -22,19 +22,30 @@ import time
|
||||
import uuid
|
||||
from collections import defaultdict
|
||||
from contextlib import AsyncExitStack
|
||||
from dataclasses import dataclass, field
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from itertools import combinations
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import TYPE_CHECKING, Any, Literal
|
||||
|
||||
from pydantic import BaseModel, field_validator
|
||||
|
||||
from ...config import get_config
|
||||
from ...worker.stage import set_stage
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..llm_trace import (
|
||||
record_created_memory_ids,
|
||||
record_source_memory_ids,
|
||||
reset_trace_context,
|
||||
set_trace_context,
|
||||
trace_context_of,
|
||||
)
|
||||
from ..llm_wrapper import sanitize_llm_output
|
||||
from ..memory_engine import Budget, fq_table
|
||||
from ..retain import embedding_utils
|
||||
from .prompts import build_batch_consolidation_prompt
|
||||
from .prompts import (
|
||||
build_consolidation_input,
|
||||
build_consolidation_system_prompt,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from asyncpg import Connection
|
||||
@@ -46,6 +57,254 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _norm_obs_text(text: str) -> str:
|
||||
"""Whitespace-normalised observation text for exact-duplicate matching.
|
||||
|
||||
Collapses runs of whitespace only; case is preserved. The reconciliation guard
|
||||
drops a CREATE on the premise that an exact-text match loses no information — but
|
||||
case-folding would also drop a create differing only in case (e.g. "TLS" vs "tls"),
|
||||
which *does* lose information, so we match case-sensitively.
|
||||
"""
|
||||
return " ".join((text or "").split()).strip()
|
||||
|
||||
|
||||
def _duplicate_create_target(
|
||||
create_text: str,
|
||||
shown_obs_by_text: "dict[str, MemoryFact]",
|
||||
update_texts: set[str],
|
||||
) -> str | None:
|
||||
"""Return a human label for what ``create_text`` duplicates, or None if novel.
|
||||
|
||||
A CREATE is a duplicate when its normalised text matches an observation that was
|
||||
already shown to the LLM, or the text of an UPDATE issued in the same response
|
||||
(the model occasionally UPDATEs the twin to text X and also CREATEs X). Exact-text
|
||||
match means no information is lost by dropping the CREATE.
|
||||
"""
|
||||
norm = _norm_obs_text(create_text)
|
||||
matched = shown_obs_by_text.get(norm)
|
||||
if matched is not None:
|
||||
return f"shown observation {str(matched.id)[:8]}"
|
||||
if norm in update_texts:
|
||||
return "an UPDATE in this response"
|
||||
return None
|
||||
|
||||
|
||||
# Top-K existing observations probed (by the new observation's own embedding) when
|
||||
# semantic dedup is enabled. Small: we only need the nearest few candidates.
|
||||
_DEDUP_TOP_K = 5
|
||||
|
||||
|
||||
class _DedupDecision(BaseModel):
|
||||
"""Focused 1-by-1 verdict for whether a new observation duplicates an existing one."""
|
||||
|
||||
action: Literal["merge", "keep"]
|
||||
text: str = "" # the synthesized merged observation (when action == "merge")
|
||||
reason: str = ""
|
||||
|
||||
|
||||
_DEDUP_PROMPT = """You reconcile long-term memory observations. A NEW observation is about to be \
|
||||
stored, and it is highly similar to an EXISTING one:
|
||||
|
||||
[NEW] {new}
|
||||
[EXISTING] {existing}
|
||||
|
||||
If they assert the SAME fact (wording aside), respond action="merge" and provide `text`: a single \
|
||||
observation that preserves EVERY detail from both. If they differ in ANY important detail — a \
|
||||
number/quantity, a named entity or language, a negation, or a condition — respond action="keep"."""
|
||||
|
||||
|
||||
def _dedup_active(config: Any) -> bool:
|
||||
"""Whether create/update semantic dedup runs for this consolidation.
|
||||
|
||||
Enabled when the resolved threshold is < 1.0, EXCEPT on Oracle: the merge path uses
|
||||
Postgres-only SQL (``unnest``/``array_agg``, ``UPDATE ... FROM``), so on Oracle dedup is
|
||||
skipped — it behaves exactly as it did before this feature, regardless of the configured
|
||||
threshold. This is why the feature can ship enabled-by-default without breaking Oracle.
|
||||
"""
|
||||
if config is None or getattr(config, "consolidation_dedup_threshold", 1.0) >= 1.0:
|
||||
return False
|
||||
return get_config().database_backend != "oracle"
|
||||
|
||||
|
||||
@dataclass
|
||||
class _DedupOutcome:
|
||||
"""Result of probing one observation against its in-scope neighbours.
|
||||
|
||||
``best_id`` is the nearest observation at/above the threshold (None if none),
|
||||
``merged_text`` is the LLM-synthesized union text (set only when ``should_merge``).
|
||||
"""
|
||||
|
||||
best_id: str | None
|
||||
merged_text: str
|
||||
should_merge: bool
|
||||
|
||||
|
||||
async def _dedup_adjudicate(
|
||||
conn: "Connection",
|
||||
memory_engine: "MemoryEngine",
|
||||
bank_id: str,
|
||||
config: Any,
|
||||
dedup_llm_config: Any,
|
||||
anchor_text: str,
|
||||
anchor_emb_str: str | None,
|
||||
tags: list[str] | None,
|
||||
exclude_id: str | None,
|
||||
) -> _DedupOutcome:
|
||||
"""Probe one observation's embedding against in-scope observations and adjudicate a merge.
|
||||
|
||||
Anchored on the observation text — the correct obs<->obs comparison, unlike consolidation
|
||||
recall which is anchored on the raw fact. Returns the nearest observation at/above
|
||||
``consolidation_dedup_threshold`` and, when found, the LLM's focused 1-by-1 merge-or-keep
|
||||
verdict (scope ``consolidation_dedup``): the LLM reads both texts, so a word-level difference
|
||||
(number / negation / entity) is respected. ``exclude_id`` skips the anchor observation itself
|
||||
(used by the UPDATE path, where the anchor row already exists and would self-match at 1.0).
|
||||
``anchor_emb_str`` reuses an already-computed embedding (the UPDATE path just embedded it);
|
||||
pass None to embed ``anchor_text`` here (the CREATE path).
|
||||
"""
|
||||
from ..search.retrieval import retrieve_semantic_bm25_combined
|
||||
|
||||
threshold = config.consolidation_dedup_threshold
|
||||
if anchor_emb_str is None:
|
||||
embs = await embedding_utils.generate_embeddings_batch(memory_engine.embeddings, [anchor_text])
|
||||
if not embs:
|
||||
return _DedupOutcome(best_id=None, merged_text="", should_merge=False)
|
||||
anchor_emb_str = str(embs[0])
|
||||
tags_match = "all_strict" if tags else "any"
|
||||
grouped = await retrieve_semantic_bm25_combined(
|
||||
conn, anchor_emb_str, anchor_text, bank_id, ["observation"], _DEDUP_TOP_K, tags=tags, tags_match=tags_match
|
||||
)
|
||||
results = grouped.get("observation", ([], []))[0]
|
||||
best_id: str | None = None
|
||||
best_text = ""
|
||||
best_sim = threshold # only candidates at/above the threshold are considered
|
||||
for r in results:
|
||||
rid = str(r.id)
|
||||
if exclude_id is not None and rid == exclude_id:
|
||||
continue # never match the anchor observation against itself
|
||||
sim = r.similarity or 0.0
|
||||
if sim >= best_sim:
|
||||
best_id, best_text, best_sim = rid, r.text, sim
|
||||
|
||||
if best_id is None:
|
||||
return _DedupOutcome(best_id=None, merged_text="", should_merge=False)
|
||||
|
||||
decision: _DedupDecision = await dedup_llm_config.call(
|
||||
messages=[{"role": "user", "content": _DEDUP_PROMPT.format(new=anchor_text, existing=best_text)}],
|
||||
response_format=_DedupDecision,
|
||||
scope="consolidation_dedup",
|
||||
)
|
||||
if decision.action != "merge":
|
||||
return _DedupOutcome(best_id=best_id, merged_text="", should_merge=False)
|
||||
return _DedupOutcome(best_id=best_id, merged_text=decision.text.strip() or best_text, should_merge=True)
|
||||
|
||||
|
||||
async def _dedup_reconcile_create(
|
||||
conn: "Connection",
|
||||
memory_engine: "MemoryEngine",
|
||||
bank_id: str,
|
||||
config: Any,
|
||||
dedup_llm_config: Any,
|
||||
create_text: str,
|
||||
create_source_ids: list[uuid.UUID],
|
||||
tags: list[str] | None,
|
||||
) -> str | None:
|
||||
"""Semantic dedup for a single CREATE (create-time, focused 1-by-1).
|
||||
|
||||
On "merge", folds the new source facts + the synthesized text into the existing
|
||||
observation and returns its id (caller skips the CREATE). Returns None when there is
|
||||
no near twin or the LLM keeps them distinct.
|
||||
"""
|
||||
outcome = await _dedup_adjudicate(
|
||||
conn, memory_engine, bank_id, config, dedup_llm_config, create_text, None, tags, exclude_id=None
|
||||
)
|
||||
if not outcome.should_merge or outcome.best_id is None:
|
||||
return None
|
||||
|
||||
# Fold the new source facts into the twin and persist the merged text. We keep the twin's
|
||||
# existing embedding: the merged text is >= threshold similar, so the stored vector stays
|
||||
# representative and we avoid a re-embed + a dialect-specific vector UPDATE.
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
SET text = $1,
|
||||
source_memory_ids = (SELECT array_agg(DISTINCT e) FROM unnest(source_memory_ids || $2::uuid[]) e),
|
||||
proof_count = (SELECT count(DISTINCT e) FROM unnest(source_memory_ids || $2::uuid[]) e),
|
||||
updated_at = now()
|
||||
WHERE id = $3::uuid
|
||||
""",
|
||||
outcome.merged_text,
|
||||
create_source_ids,
|
||||
uuid.UUID(outcome.best_id),
|
||||
)
|
||||
return outcome.best_id
|
||||
|
||||
|
||||
async def _dedup_reconcile_update(
|
||||
conn: "Connection",
|
||||
memory_engine: "MemoryEngine",
|
||||
bank_id: str,
|
||||
config: Any,
|
||||
dedup_llm_config: Any,
|
||||
updated_id: str,
|
||||
updated_text: str,
|
||||
updated_emb_str: str | None,
|
||||
tags: list[str] | None,
|
||||
) -> None:
|
||||
"""Semantic dedup for an UPDATE (after the observation was rewritten + re-embedded).
|
||||
|
||||
An UPDATE rewrites an observation's text and re-embeds it, so its vector can drift to
|
||||
within threshold of a DIFFERENT existing observation. The create-time guard never sees
|
||||
this (it only runs on CREATE), so without this the two persist as a near-duplicate pair —
|
||||
the measured residual-duplicate source. Probe the updated observation's new embedding
|
||||
against the others (excluding itself); on "merge", fold the just-updated observation's
|
||||
sources into the twin, persist the merged text, and DELETE the updated row. Unlike the
|
||||
CREATE path the row already exists, so reconciliation is a fold-and-delete, not a skip.
|
||||
"""
|
||||
outcome = await _dedup_adjudicate(
|
||||
conn,
|
||||
memory_engine,
|
||||
bank_id,
|
||||
config,
|
||||
dedup_llm_config,
|
||||
updated_text,
|
||||
updated_emb_str,
|
||||
tags,
|
||||
exclude_id=updated_id,
|
||||
)
|
||||
if not outcome.should_merge or outcome.best_id is None:
|
||||
return
|
||||
|
||||
# Fold the updated observation's sources into the twin (keeping the twin's embedding, as in
|
||||
# the create path) then delete the now-redundant updated row. The all_strict/any tag match
|
||||
# guarantees twin and updated share scope, so dropping the updated row's tags loses no
|
||||
# visibility. Temporal fields follow the surviving twin (minimal scope; matches create).
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")} t
|
||||
SET text = $1,
|
||||
source_memory_ids = (
|
||||
SELECT array_agg(DISTINCT e) FROM unnest(t.source_memory_ids || u.source_memory_ids) e
|
||||
),
|
||||
proof_count = (
|
||||
SELECT count(DISTINCT e) FROM unnest(t.source_memory_ids || u.source_memory_ids) e
|
||||
),
|
||||
updated_at = now()
|
||||
FROM {fq_table("memory_units")} u
|
||||
WHERE t.id = $2::uuid AND u.id = $3::uuid
|
||||
""",
|
||||
outcome.merged_text,
|
||||
uuid.UUID(outcome.best_id),
|
||||
uuid.UUID(updated_id),
|
||||
)
|
||||
await _execute_delete_action(conn, bank_id, updated_id)
|
||||
logger.info(
|
||||
"[CONSOLIDATION] dedup-merged updated observation %s into %s (cosine>=%.2f)",
|
||||
updated_id[:8],
|
||||
outcome.best_id[:8],
|
||||
config.consolidation_dedup_threshold,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _BatchDeltas:
|
||||
"""Per-LLM-batch deltas, merged into the job's running stats after dispatch.
|
||||
@@ -167,6 +426,9 @@ async def _filter_live_source_memories(
|
||||
class _CreateAction(BaseModel):
|
||||
text: str
|
||||
source_fact_ids: list[str] # memory UUIDs from the NEW FACTS list
|
||||
# One-sentence justification from the LLM (why CREATE vs UPDATE). Diagnostic
|
||||
# only — surfaced in the consolidation trace to explain duplicate creates.
|
||||
reason: str = ""
|
||||
|
||||
@field_validator("text", mode="before")
|
||||
@classmethod
|
||||
@@ -178,6 +440,7 @@ class _UpdateAction(BaseModel):
|
||||
text: str
|
||||
observation_id: str # UUID of the existing observation to update
|
||||
source_fact_ids: list[str] # memory UUIDs from the NEW FACTS list
|
||||
reason: str = "" # LLM's one-sentence justification (diagnostic only)
|
||||
|
||||
@field_validator("text", mode="before")
|
||||
@classmethod
|
||||
@@ -187,6 +450,7 @@ class _UpdateAction(BaseModel):
|
||||
|
||||
class _DeleteAction(BaseModel):
|
||||
observation_id: str # UUID of the observation to remove
|
||||
reason: str = "" # LLM's one-sentence justification (diagnostic only)
|
||||
|
||||
|
||||
class _ConsolidationBatchResponse(BaseModel):
|
||||
@@ -361,8 +625,36 @@ async def run_consolidation_job(
|
||||
|
||||
# Build a configured LLM wrapper that applies per-bank settings (e.g. safety settings)
|
||||
# to every call without leaking across operations.
|
||||
llm_config = memory_engine._consolidation_llm_config.with_config(config)
|
||||
llm_config = memory_engine._consolidation_llm_config.with_config(config, bank_id=bank_id, operation="consolidation")
|
||||
|
||||
# Bind the operation trace context for the whole run so the create/update DB
|
||||
# sites (deep inside _process_memory_batch) can accumulate the observations
|
||||
# this consolidation produced and the source memories it consumed onto the
|
||||
# trace — flushed onto every trace row on exit by attach_memory_ids.
|
||||
trace_ctx = trace_context_of(llm_config)
|
||||
trace_token = set_trace_context(trace_ctx) if trace_ctx is not None else None
|
||||
try:
|
||||
return await _run_consolidation_job(
|
||||
memory_engine, bank_id, request_context, config, llm_config, operation_id, observation_scopes
|
||||
)
|
||||
finally:
|
||||
if trace_token is not None:
|
||||
reset_trace_context(trace_token)
|
||||
# Fire-and-forget: patched on a background task, off the consolidation
|
||||
# critical path.
|
||||
memory_engine._llm_recorder.attach_memory_ids(trace_ctx)
|
||||
|
||||
|
||||
async def _run_consolidation_job(
|
||||
memory_engine: "MemoryEngine",
|
||||
bank_id: str,
|
||||
request_context: "RequestContext",
|
||||
config: Any,
|
||||
llm_config: Any,
|
||||
operation_id: str | None = None,
|
||||
observation_scopes: list[list[str]] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Core consolidation flow. See ``run_consolidation_job`` for the public entrypoint."""
|
||||
perf = ConsolidationPerfLog(bank_id)
|
||||
max_memories_per_batch = config.consolidation_batch_size
|
||||
max_memories_per_round = config.consolidation_max_memories_per_round
|
||||
@@ -426,6 +718,44 @@ async def run_consolidation_job(
|
||||
logger.info(f"[CONSOLIDATION] bank={bank_id} total_unconsolidated={total_count}")
|
||||
perf.log(f"[1] Found {total_count} pending memories to consolidate")
|
||||
|
||||
# Initial durable progress snapshot so an operator polling the operation status
|
||||
# API sees the job has started and how much work it found, before the first batch
|
||||
# of LLM work completes (which can take minutes on a dense bank). Uses the same
|
||||
# "consolidating" stage as the per-batch heartbeat so the operator sees a single
|
||||
# phase advancing 0/N -> N/N rather than an opaque "scanning" -> "processing" hop.
|
||||
set_stage("consolidation.consolidating")
|
||||
await memory_engine._write_operation_progress(operation_id, stage="consolidating", processed=0, total=total_count)
|
||||
|
||||
async def _count_unconsolidated() -> int:
|
||||
"""Re-count memories still pending consolidation in this job's scope.
|
||||
|
||||
``total_count`` is a point-in-time estimate from job start; memories retained
|
||||
while consolidation runs get picked up by later fetches, so processed can pass
|
||||
it. When that happens we re-count to report a real total (processed + remaining)
|
||||
instead of pinning the bar at 100%."""
|
||||
async with acquire_with_retry(pool) as count_conn:
|
||||
pending = await count_conn.fetchval(
|
||||
f"""
|
||||
SELECT COUNT(*)
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
{scope_clause}
|
||||
""",
|
||||
*scope_params,
|
||||
)
|
||||
return pending or 0
|
||||
|
||||
async def _progress_total(processed: int) -> int:
|
||||
# Cheap path: while we're still within the start-of-job estimate it's exact, so
|
||||
# no extra query. Only re-count once the estimate is exhausted (≈the final batch
|
||||
# normally, or repeatedly only if memories keep arriving mid-run).
|
||||
if processed < total_count:
|
||||
return total_count
|
||||
return processed + await _count_unconsolidated()
|
||||
|
||||
# Process each memory with individual commits for crash recovery
|
||||
stats: dict[str, int] = {
|
||||
"memories_processed": 0,
|
||||
@@ -446,10 +776,18 @@ async def run_consolidation_job(
|
||||
hit_round_limit = False
|
||||
|
||||
llm_batch_num = 0
|
||||
# Cumulative count of memories processed across the whole job, shared by
|
||||
# the per-batch log so it can still report processed/total under parallelism.
|
||||
# Mutable container so the inner closure can update without a `nonlocal`.
|
||||
cumulative_progress = {"processed": 0}
|
||||
# Cumulative counters across the whole job, shared by the per-batch log and the
|
||||
# durable progress snapshot so both report processed/total (and observation
|
||||
# tallies) under parallelism. Mutable container so the inner closure can update
|
||||
# without a `nonlocal`.
|
||||
cumulative_progress = {
|
||||
"processed": 0,
|
||||
"observations_created": 0,
|
||||
"observations_updated": 0,
|
||||
"observations_merged": 0,
|
||||
"observations_deleted": 0,
|
||||
"memories_failed": 0,
|
||||
}
|
||||
while True:
|
||||
# Cap fetch size by remaining round budget
|
||||
fetch_limit = (
|
||||
@@ -670,12 +1008,18 @@ async def run_consolidation_job(
|
||||
local_stats["memories_failed"] += 1
|
||||
|
||||
# Maintain the cumulative-progress indicator under parallelism:
|
||||
# increment a shared counter and snapshot under the same statement
|
||||
# so the snapshot includes this batch. No await between the read
|
||||
# and write, so single-threaded asyncio gives us atomicity for free
|
||||
# — no lock needed.
|
||||
# increment shared counters and snapshot under the same statements so
|
||||
# the snapshot includes this batch. No await between the reads and
|
||||
# writes, so single-threaded asyncio gives us atomicity for free —
|
||||
# no lock needed.
|
||||
cumulative_progress["processed"] += local_stats["memories_processed"]
|
||||
cumulative_progress["observations_created"] += local_stats["observations_created"]
|
||||
cumulative_progress["observations_updated"] += local_stats["observations_updated"]
|
||||
cumulative_progress["observations_merged"] += local_stats["observations_merged"]
|
||||
cumulative_progress["observations_deleted"] += local_stats["observations_deleted"]
|
||||
cumulative_progress["memories_failed"] += local_stats["memories_failed"]
|
||||
cum_processed = cumulative_progress["processed"]
|
||||
cum_snapshot = dict(cumulative_progress)
|
||||
|
||||
# Per-batch log uses batch_perf so timings/llm-calls/tokens reflect
|
||||
# only this batch's own work, even when other batches are running
|
||||
@@ -703,6 +1047,27 @@ async def run_consolidation_job(
|
||||
f" | avg={llm_batch_time / max(1, len(llm_batch_local)):.3f}s/memory"
|
||||
)
|
||||
|
||||
# Durable progress snapshot per LLM batch — this is the heartbeat an
|
||||
# operator polls. The whole fetched batch is processed inside one outer
|
||||
# round, so a round-boundary write would sit at the pre-round count for
|
||||
# the entire (often minutes-long) LLM phase; writing here advances
|
||||
# processed/total as each batch commits. set_stage mirrors it for the
|
||||
# live worker log.
|
||||
set_stage(f"consolidation.llm_batch.{batch_num_local}")
|
||||
await memory_engine._write_operation_progress(
|
||||
operation_id,
|
||||
stage="consolidating",
|
||||
processed=cum_processed,
|
||||
total=await _progress_total(cum_processed),
|
||||
detail={
|
||||
"observations_created": cum_snapshot["observations_created"],
|
||||
"observations_updated": cum_snapshot["observations_updated"],
|
||||
"observations_merged": cum_snapshot["observations_merged"],
|
||||
"observations_deleted": cum_snapshot["observations_deleted"],
|
||||
"memories_failed": cum_snapshot["memories_failed"],
|
||||
},
|
||||
)
|
||||
|
||||
# Fold batch counters into the job-level perf so the final summary
|
||||
# (perf.flush) totals every batch correctly. Safe without a lock —
|
||||
# ConsolidationPerfLog.merge_from is a series of += on Python ints
|
||||
@@ -841,6 +1206,13 @@ async def run_consolidation_job(
|
||||
stats["mental_models_refreshed"] = 0
|
||||
logger.info(f"[CONSOLIDATION] bank={bank_id} skipping mental model refresh (round limit hit, re-queued)")
|
||||
else:
|
||||
set_stage("consolidation.refreshing_mental_models")
|
||||
await memory_engine._write_operation_progress(
|
||||
operation_id,
|
||||
stage="refreshing_mental_models",
|
||||
processed=stats["memories_processed"],
|
||||
total=await _progress_total(stats["memories_processed"]),
|
||||
)
|
||||
# SECURITY: Only refresh mental models with matching tags (or all if no tags were consolidated)
|
||||
mental_models_refreshed = await _trigger_mental_model_refreshes(
|
||||
memory_engine=memory_engine,
|
||||
@@ -984,6 +1356,9 @@ async def _process_memory_batch(
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# Map the source memories this batch consumes onto the consolidation trace.
|
||||
record_source_memory_ids([str(m["id"]) for m in memories])
|
||||
|
||||
# 1. Parallel recalls — one per fact
|
||||
# When obs_tags_override is set, use it as the observation scope for all facts.
|
||||
t0 = time.time()
|
||||
@@ -1063,6 +1438,20 @@ async def _process_memory_batch(
|
||||
|
||||
mem_by_id = {str(m["id"]): m for m in memories}
|
||||
|
||||
# Semantic dedup: when enabled, an observation that is >= the threshold cosine to a DIFFERENT
|
||||
# existing observation is reconciled by a focused 1-by-1 LLM merge (anchored on the observation
|
||||
# text, not the source fact). It runs on both CREATE (a near-dup emitted despite the twin being
|
||||
# in context — weak-model failure mode) and UPDATE (a rewrite+re-embed that drifts an existing
|
||||
# observation into a twin — the create-time guard can't see this). The trace operation/scope is
|
||||
# "consolidation_dedup" (routes through the consolidation concurrency bucket via llm_wrapper's
|
||||
# "consolidation" prefix; recorded distinctly in llm_requests).
|
||||
dedup_enabled = _dedup_active(config)
|
||||
dedup_llm_config = (
|
||||
memory_engine._consolidation_llm_config.with_config(config, bank_id=bank_id, operation="consolidation_dedup")
|
||||
if dedup_enabled
|
||||
else None
|
||||
)
|
||||
|
||||
# Execute deletes first to free observation slots before creates consume them
|
||||
deleted_count = 0
|
||||
for delete in llm_result.deletes:
|
||||
@@ -1087,7 +1476,7 @@ async def _process_memory_batch(
|
||||
)
|
||||
continue
|
||||
agg = _aggregate_source_fields(source_mems, tags=fact_tags)
|
||||
await _execute_update_action(
|
||||
updated_emb_str = await _execute_update_action(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
@@ -1103,17 +1492,76 @@ async def _process_memory_batch(
|
||||
)
|
||||
for m in source_mems:
|
||||
per_memory_updated.add(str(m["id"]))
|
||||
# Reconcile the rewritten observation against its neighbours: the re-embed may have
|
||||
# drifted it into a near-twin of another existing observation (the residual-duplicate
|
||||
# source). updated_emb_str is None when the update was skipped — nothing to reconcile.
|
||||
if dedup_enabled and updated_emb_str is not None:
|
||||
await _dedup_reconcile_update(
|
||||
conn,
|
||||
memory_engine,
|
||||
bank_id,
|
||||
config,
|
||||
dedup_llm_config,
|
||||
update.observation_id,
|
||||
update.text,
|
||||
updated_emb_str,
|
||||
agg.tags,
|
||||
)
|
||||
|
||||
# Deterministic dedup guard: map the observations the LLM was SHOWN by their
|
||||
# normalised text. The model intermittently emits a CREATE whose text is identical
|
||||
# to an observation already in its context (over-aggregation / incoherence — it even
|
||||
# UPDATEs the twin and creates a sibling). When that happens we drop the duplicate
|
||||
# CREATE instead of inserting a redundant row. No extra LLM/embedding cost — the
|
||||
# match is exact text against the in-memory set.
|
||||
shown_obs_by_text = {_norm_obs_text(o.text): o for o in union_observations}
|
||||
# Also collapse a CREATE that reproduces the text of an UPDATE issued in the SAME
|
||||
# response (the model occasionally UPDATEs the twin to text X and also CREATEs X).
|
||||
update_texts = {_norm_obs_text(u.text) for u in llm_result.updates if u.text}
|
||||
|
||||
for create in llm_result.creates:
|
||||
source_mems = [mem_by_id[fid] for fid in create.source_fact_ids if fid in mem_by_id]
|
||||
if not source_mems:
|
||||
continue
|
||||
agg = _aggregate_source_fields(source_mems, tags=fact_tags)
|
||||
create_source_ids = [m["id"] for m in source_mems]
|
||||
|
||||
# Reconcile against observations shown to the LLM: an exact-text match means
|
||||
# this CREATE reproduces verbatim an observation the model already had in context.
|
||||
# Since that observation already carries this exact text, drop the duplicate CREATE
|
||||
# — no row is inserted, nothing is lost. We deliberately do NOT also UPDATE the twin
|
||||
# here: the LLM frequently UPDATEd it earlier in this same batch, and a second update
|
||||
# would run off the pre-LLM snapshot and clobber that change (see _dedupe_updates).
|
||||
duplicate_of = _duplicate_create_target(create.text, shown_obs_by_text, update_texts)
|
||||
if duplicate_of is not None:
|
||||
logger.warning(
|
||||
"[CONSOLIDATION] dropped duplicate observation CREATE — verbatim match of %s; llm_reason=%r",
|
||||
duplicate_of,
|
||||
create.reason or "(none given)",
|
||||
)
|
||||
continue
|
||||
|
||||
# Semantic near-duplicate reconciliation: merge this CREATE into an existing
|
||||
# near-identical observation (LLM-adjudicated, 1-by-1) instead of inserting a dup.
|
||||
if dedup_enabled:
|
||||
merged_into = await _dedup_reconcile_create(
|
||||
conn, memory_engine, bank_id, config, dedup_llm_config, create.text, create_source_ids, agg.tags
|
||||
)
|
||||
if merged_into is not None:
|
||||
logger.info(
|
||||
"[CONSOLIDATION] dedup-merged observation CREATE into %s (cosine>=%.2f)",
|
||||
merged_into[:8],
|
||||
config.consolidation_dedup_threshold,
|
||||
)
|
||||
for m in source_mems:
|
||||
per_memory_created.add(str(m["id"]))
|
||||
continue
|
||||
|
||||
await _execute_create_action(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[m["id"] for m in source_mems],
|
||||
source_memory_ids=create_source_ids,
|
||||
text=create.text,
|
||||
source_fact_tags=agg.tags,
|
||||
event_date=agg.event_date,
|
||||
@@ -1153,6 +1601,64 @@ def _max_date(dates: "Any") -> "datetime | None":
|
||||
return max((d for d in dates if d is not None), default=None)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ObservationHistorySnapshot:
|
||||
"""Pre-update state of an observation, persisted as the ``content`` JSON blob
|
||||
of one observation_history row.
|
||||
|
||||
Temporal fields are the ISO strings carried on MemoryFact; new_source_memory_ids
|
||||
are the ids added by the update.
|
||||
"""
|
||||
|
||||
previous_text: str | None
|
||||
previous_tags: list[str]
|
||||
previous_occurred_start: str | None
|
||||
previous_occurred_end: str | None
|
||||
previous_mentioned_at: str | None
|
||||
new_source_memory_ids: list[str]
|
||||
|
||||
|
||||
async def _append_observation_history(
|
||||
conn: "Connection",
|
||||
bank_id: str,
|
||||
observation_id: str,
|
||||
snapshot: _ObservationHistorySnapshot,
|
||||
max_entries: int,
|
||||
) -> None:
|
||||
"""Insert one pre-update snapshot into ``observation_history``, then delete the
|
||||
oldest rows beyond ``max_entries`` for this observation.
|
||||
|
||||
The snapshot is stored as a single JSONB ``content`` blob (per-row, so it stays
|
||||
small). Bounding by row count keeps a frequently-reinforced observation's
|
||||
history from growing without bound.
|
||||
"""
|
||||
obs_uuid = uuid.UUID(observation_id)
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("observation_history")} (observation_id, bank_id, content, changed_at)
|
||||
VALUES ($1, $2, $3::jsonb, now())
|
||||
""",
|
||||
obs_uuid,
|
||||
bank_id,
|
||||
json.dumps(asdict(snapshot)),
|
||||
)
|
||||
if max_entries and max_entries > 0:
|
||||
await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {fq_table("observation_history")}
|
||||
WHERE observation_id = $1
|
||||
AND id NOT IN (
|
||||
SELECT id FROM {fq_table("observation_history")}
|
||||
WHERE observation_id = $1
|
||||
ORDER BY changed_at DESC, id DESC
|
||||
LIMIT $2
|
||||
)
|
||||
""",
|
||||
obs_uuid,
|
||||
max_entries,
|
||||
)
|
||||
|
||||
|
||||
async def _execute_update_action(
|
||||
conn: "Connection",
|
||||
memory_engine: "MemoryEngine",
|
||||
@@ -1166,12 +1672,15 @@ async def _execute_update_action(
|
||||
source_occurred_end: datetime | None = None,
|
||||
source_mentioned_at: datetime | None = None,
|
||||
perf: ConsolidationPerfLog | None = None,
|
||||
) -> None:
|
||||
) -> str | None:
|
||||
"""
|
||||
Update an existing observation.
|
||||
|
||||
Extends source_memory_ids with all contributing memories, updates temporal fields
|
||||
(LEAST for occurred_start, GREATEST for occurred_end / mentioned_at), and merges tags.
|
||||
|
||||
Returns the observation's freshly-computed embedding (pgvector literal) so the caller can
|
||||
run UPDATE-path dedup without re-embedding, or None when the update was skipped.
|
||||
"""
|
||||
model = next((m for m in observations if str(m.id) == observation_id), None)
|
||||
if not model:
|
||||
@@ -1189,15 +1698,14 @@ async def _execute_update_action(
|
||||
|
||||
from ...config import get_config
|
||||
|
||||
history_entry = {
|
||||
"previous_text": model.text,
|
||||
"previous_tags": list(model.tags or []),
|
||||
"previous_occurred_start": model.occurred_start,
|
||||
"previous_occurred_end": model.occurred_end,
|
||||
"previous_mentioned_at": model.mentioned_at,
|
||||
"changed_at": datetime.now(timezone.utc).isoformat(),
|
||||
"new_source_memory_ids": [str(mid) for mid in source_memory_ids],
|
||||
}
|
||||
history_entry = _ObservationHistorySnapshot(
|
||||
previous_text=model.text,
|
||||
previous_tags=list(model.tags or []),
|
||||
previous_occurred_start=model.occurred_start,
|
||||
previous_occurred_end=model.occurred_end,
|
||||
previous_mentioned_at=model.mentioned_at,
|
||||
new_source_memory_ids=[str(mid) for mid in source_memory_ids],
|
||||
)
|
||||
|
||||
source_ids = list(model.source_fact_ids or []) + source_memory_ids
|
||||
|
||||
@@ -1213,9 +1721,6 @@ async def _execute_update_action(
|
||||
perf.record_timing("embedding", time.time() - t0)
|
||||
|
||||
config = get_config()
|
||||
history_clause = (
|
||||
"history = COALESCE(history, '[]'::jsonb) || $3::jsonb," if config.enable_observation_history else ""
|
||||
)
|
||||
|
||||
t0 = time.time()
|
||||
await conn.execute(
|
||||
@@ -1223,19 +1728,17 @@ async def _execute_update_action(
|
||||
UPDATE {fq_table("memory_units")}
|
||||
SET text = $1,
|
||||
embedding = $2::vector,
|
||||
{history_clause}
|
||||
source_memory_ids = $4,
|
||||
proof_count = $5,
|
||||
tags = $10,
|
||||
source_memory_ids = $3,
|
||||
proof_count = $4,
|
||||
tags = $9,
|
||||
updated_at = now(),
|
||||
occurred_start = LEAST(occurred_start, COALESCE($7, occurred_start)),
|
||||
occurred_end = GREATEST(occurred_end, COALESCE($8, occurred_end)),
|
||||
mentioned_at = GREATEST(mentioned_at, COALESCE($9, mentioned_at))
|
||||
WHERE id = $6
|
||||
occurred_start = LEAST(occurred_start, COALESCE($6, occurred_start)),
|
||||
occurred_end = GREATEST(occurred_end, COALESCE($7, occurred_end)),
|
||||
mentioned_at = GREATEST(mentioned_at, COALESCE($8, mentioned_at))
|
||||
WHERE id = $5
|
||||
""",
|
||||
new_text,
|
||||
embedding_str,
|
||||
json.dumps([history_entry]),
|
||||
source_ids,
|
||||
len(source_ids),
|
||||
uuid.UUID(observation_id),
|
||||
@@ -1245,6 +1748,15 @@ async def _execute_update_action(
|
||||
merged_tags,
|
||||
)
|
||||
|
||||
# Record the pre-update snapshot in the dedicated observation_history table
|
||||
# (one row per change), then trim to the configured cap. History lived in a
|
||||
# single unbounded JSONB column before; an often-reinforced observation grew
|
||||
# it until it crossed Postgres's 256MB jsonb limit and got stuck.
|
||||
if config.enable_observation_history:
|
||||
await _append_observation_history(
|
||||
conn, bank_id, observation_id, history_entry, config.observation_history_max_entries
|
||||
)
|
||||
|
||||
# Sync observation_sources junction table (Oracle only — PG uses native array ops).
|
||||
if memory_engine._backend.ops.uses_observation_sources_table:
|
||||
obs_uuid = uuid.UUID(observation_id)
|
||||
@@ -1265,7 +1777,10 @@ async def _execute_update_action(
|
||||
if perf:
|
||||
perf.record_timing("db_write", time.time() - t0)
|
||||
|
||||
# Map the updated observation onto the consolidation trace as a produced memory.
|
||||
record_created_memory_ids([observation_id])
|
||||
logger.debug(f"Updated observation {observation_id} from {len(source_memory_ids)} source memories")
|
||||
return embedding_str
|
||||
|
||||
|
||||
async def _execute_create_action(
|
||||
@@ -1287,7 +1802,7 @@ async def _execute_create_action(
|
||||
Tags are inherited from the source facts (determined algorithmically, not by LLM)
|
||||
to maintain visibility scope.
|
||||
"""
|
||||
await _create_observation_directly(
|
||||
created = await _create_observation_directly(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
@@ -1300,6 +1815,10 @@ async def _execute_create_action(
|
||||
mentioned_at=mentioned_at,
|
||||
perf=perf,
|
||||
)
|
||||
# Map the new observation onto the consolidation trace as a produced memory.
|
||||
new_id = created.get("observation_id")
|
||||
if new_id:
|
||||
record_created_memory_ids([new_id])
|
||||
logger.debug(f"Created observation from {len(source_memory_ids)} source memories")
|
||||
|
||||
|
||||
@@ -1398,6 +1917,13 @@ async def _find_related_observations(
|
||||
include_source_facts=True, # Embed source facts so we avoid a separate DB fetch
|
||||
max_source_facts_tokens=config.consolidation_source_facts_max_tokens,
|
||||
max_source_facts_tokens_per_observation=config.consolidation_source_facts_max_tokens_per_observation,
|
||||
# Round-robin interleave fusion (no cross-encoder): consolidation is looking
|
||||
# for an existing near-identical observation to merge into. Both the
|
||||
# cross-encoder (semantic #1 -> reranked #37) and RRF (semantic #1 -> outside
|
||||
# the 512-token budget) were measured to bury that twin; interleave guarantees
|
||||
# each retrieval arm's top hits a slot, so the semantic-#1 twin is always shown
|
||||
# to the LLM, which then UPDATEs instead of creating a duplicate.
|
||||
reranking="interleave",
|
||||
_quiet=True, # Suppress logging
|
||||
)
|
||||
finally:
|
||||
@@ -1532,16 +2058,38 @@ async def _consolidate_batch_with_llm(
|
||||
f"(out of {max_observations_per_scope}). Prefer UPDATE over CREATE when possible."
|
||||
)
|
||||
|
||||
prompt_template = build_batch_consolidation_prompt(
|
||||
config.observations_mission,
|
||||
observation_capacity_note,
|
||||
# Split the prompt: a bank-agnostic system instruction (rules + input format +
|
||||
# decision guide + output format) that is byte-identical across batches AND
|
||||
# across banks, and a per-batch user message (mission + capacity note + facts +
|
||||
# existing observations). The split lets the system prefix be served from a
|
||||
# single Gemini context cache shared by every bank — the bank mission, capacity
|
||||
# note, and response_schema (all bank/batch-variable) are kept OUT of the
|
||||
# cached prefix so one cache serves all and it never busts within a run.
|
||||
system_prompt = build_consolidation_system_prompt(
|
||||
llm_output_language=getattr(config, "llm_output_language", None),
|
||||
)
|
||||
prompt = prompt_template.format(
|
||||
user_content = build_consolidation_input(
|
||||
facts_text=facts_lines,
|
||||
observations_text=observations_text,
|
||||
observations_mission=config.observations_mission,
|
||||
observation_capacity_note=observation_capacity_note,
|
||||
)
|
||||
|
||||
# Opt into context caching of the stable system prefix when the provider
|
||||
# supports it (gemini/vertexai with the flag on). response_schema is NOT
|
||||
# passed to the fingerprint: it varies per batch (max_creates) but is not
|
||||
# part of the cached prefix, so keying on it would needlessly bust the cache.
|
||||
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,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Consolidation cache prefix lookup failed; falling back to uncached call")
|
||||
cached_prefix_name = None
|
||||
|
||||
# Use a constrained response model when observation limit is active
|
||||
response_model = _build_response_model(max_creates=remaining_observation_slots)
|
||||
|
||||
@@ -1561,12 +2109,23 @@ async def _consolidate_batch_with_llm(
|
||||
for attempt in range(1, max_attempts + 1):
|
||||
try:
|
||||
call_kwargs: dict[str, Any] = {
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"messages": [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_content},
|
||||
],
|
||||
"response_format": response_model,
|
||||
"scope": "consolidation",
|
||||
}
|
||||
# Only request an explicit output budget when configured. Left unset by default the key is
|
||||
# omitted, so each provider keeps its implicit default (backwards compatible). Operators on
|
||||
# providers with a low hidden cap (notably Bedrock imported models, which truncate structured
|
||||
# consolidation JSON) set HINDSIGHT_API_CONSOLIDATION_MAX_COMPLETION_TOKENS to fix it.
|
||||
if config.consolidation_max_completion_tokens is not None:
|
||||
call_kwargs["max_completion_tokens"] = config.consolidation_max_completion_tokens
|
||||
if inner_max_retries is not None:
|
||||
call_kwargs["max_retries"] = inner_max_retries
|
||||
if cached_prefix_name is not None:
|
||||
call_kwargs["cached_prefix"] = cached_prefix_name
|
||||
response: _ConsolidationBatchResponse = await llm_config.call(**call_kwargs)
|
||||
# Defensive truncation: some LLM providers may not enforce JSON schema max_length
|
||||
creates = response.creates
|
||||
@@ -1583,7 +2142,7 @@ async def _consolidate_batch_with_llm(
|
||||
updates=updates,
|
||||
deletes=response.deletes,
|
||||
obs_count=len(union_observations),
|
||||
prompt_chars=len(prompt),
|
||||
prompt_chars=len(system_prompt) + len(user_content),
|
||||
)
|
||||
except Exception as exc:
|
||||
last_exc = exc
|
||||
@@ -1595,7 +2154,9 @@ async def _consolidate_batch_with_llm(
|
||||
f"[CONSOLIDATION] LLM batch call failed after {max_attempts} attempts for {batch_label}, "
|
||||
f"skipping batch. Last error: {last_exc}"
|
||||
)
|
||||
return _BatchLLMResult(obs_count=len(union_observations), prompt_chars=len(prompt), failed=True)
|
||||
return _BatchLLMResult(
|
||||
obs_count=len(union_observations), prompt_chars=len(system_prompt) + len(user_content), failed=True
|
||||
)
|
||||
|
||||
|
||||
async def _create_observation_directly(
|
||||
@@ -1642,10 +2203,10 @@ async def _create_observation_directly(
|
||||
# VectorChord: manually tokenize and insert search_vector
|
||||
query = f"""
|
||||
INSERT INTO {fq_table("memory_units")} (
|
||||
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history,
|
||||
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids,
|
||||
tags, event_date, occurred_start, occurred_end, mentioned_at, search_vector
|
||||
)
|
||||
VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8, $9, $10,
|
||||
VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, $6, $7, $8, $9, $10,
|
||||
tokenize($3, 'llmlingua2')::bm25_catalog.bm25vector)
|
||||
RETURNING id
|
||||
"""
|
||||
@@ -1661,10 +2222,10 @@ async def _create_observation_directly(
|
||||
# re-ingested. Tracking a separate fix for that gap.
|
||||
query = f"""
|
||||
INSERT INTO {fq_table("memory_units")} (
|
||||
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history,
|
||||
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids,
|
||||
tags, event_date, occurred_start, occurred_end, mentioned_at
|
||||
)
|
||||
VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8, $9, $10)
|
||||
VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, $6, $7, $8, $9, $10)
|
||||
RETURNING id
|
||||
"""
|
||||
|
||||
|
||||
@@ -37,6 +37,33 @@ _PROCESSING_RULES = """## PROCESSING RULES
|
||||
|
||||
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
|
||||
|
||||
{facts_text}
|
||||
|
||||
### Existing observations
|
||||
|
||||
{observations_text}"""
|
||||
|
||||
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
|
||||
_INPUT_SECTION = """## INPUT
|
||||
|
||||
@@ -65,7 +92,7 @@ _DECISION_GUIDE = """## DECISION GUIDE
|
||||
# Output format — JSON braces escaped as {{ }} so .format() leaves them literal
|
||||
_OUTPUT_SECTION = """## OUTPUT FORMAT
|
||||
|
||||
Return a JSON object with three arrays: `creates`, `updates`, `deletes`.
|
||||
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
|
||||
|
||||
@@ -79,7 +106,7 @@ Existing observation:
|
||||
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"]}}],
|
||||
"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
|
||||
@@ -93,8 +120,8 @@ Existing observation:
|
||||
|
||||
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"]}}],
|
||||
"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"]}}],
|
||||
{{"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
|
||||
@@ -110,6 +137,7 @@ Expected output (UPDATE for the state change; CREATE for the unrelated work-hour
|
||||
- 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."""
|
||||
|
||||
@@ -145,3 +173,55 @@ def build_batch_consolidation_prompt(
|
||||
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)
|
||||
|
||||
@@ -46,7 +46,6 @@ from ..config import (
|
||||
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,
|
||||
@@ -1199,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,
|
||||
@@ -1209,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)
|
||||
@@ -1284,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
|
||||
|
||||
@@ -1678,7 +1679,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
return CohereCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_openrouter_model,
|
||||
base_url="https://openrouter.ai/api/v1/rerank",
|
||||
base_url=config.reranker_openrouter_base_url,
|
||||
timeout=config.reranker_openrouter_timeout,
|
||||
)
|
||||
elif provider == "flashrank":
|
||||
@@ -1697,13 +1698,8 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -155,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:
|
||||
@@ -685,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)
|
||||
|
||||
@@ -862,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):
|
||||
@@ -872,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()
|
||||
@@ -1059,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 []]
|
||||
@@ -1097,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 []]
|
||||
@@ -1130,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())
|
||||
|
||||
@@ -43,6 +43,10 @@ from ..config import (
|
||||
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,
|
||||
@@ -53,6 +57,7 @@ from ..config import (
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
ENV_LLM_API_KEY,
|
||||
)
|
||||
from .bank_attribution import apply_bank_attribution
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -252,6 +257,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.
|
||||
@@ -535,6 +706,7 @@ class OpenAIEmbeddings(Embeddings):
|
||||
}
|
||||
if self.dimensions is not None:
|
||||
request["dimensions"] = self.dimensions
|
||||
apply_bank_attribution(request)
|
||||
|
||||
response = self._client.embeddings.create(**request)
|
||||
|
||||
@@ -1177,6 +1349,21 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
# Gemini Embedding 2+ multimodal models return a SINGLE aggregated embedding
|
||||
# for a multi-input request instead of one vector per input (see
|
||||
# https://ai.google.dev/gemini-api/docs/embeddings#embedding-aggregation). For
|
||||
# these models we must embed one input per call to preserve the 1:1 input→vector
|
||||
# alignment the rest of the pipeline relies on. The marker matches preview and GA
|
||||
# names (e.g. "gemini-embedding-2-preview", "gemini-embedding-2"), with or
|
||||
# without a "google/" or "models/" prefix.
|
||||
_GEMINI_AGGREGATING_MODEL_MARKER = "gemini-embedding-2"
|
||||
|
||||
|
||||
def _gemini_model_aggregates_inputs(model: str) -> bool:
|
||||
"""Whether the model aggregates a multi-input request into one embedding."""
|
||||
return _GEMINI_AGGREGATING_MODEL_MARKER in model.lower()
|
||||
|
||||
|
||||
class GeminiEmbeddings(Embeddings):
|
||||
"""
|
||||
Google embeddings via the google.genai SDK.
|
||||
@@ -1186,6 +1373,10 @@ class GeminiEmbeddings(Embeddings):
|
||||
2. Vertex AI with service account or Application Default Credentials (ADC)
|
||||
|
||||
Uses the embed_content API: client.models.embed_content(model, contents)
|
||||
|
||||
Gemini Embedding 2+ multimodal models aggregate a multi-input request into a
|
||||
single embedding, so for those the batch size is forced to 1 (one input per
|
||||
call) to keep one vector per input.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -1340,9 +1531,13 @@ class GeminiEmbeddings(Embeddings):
|
||||
|
||||
all_embeddings = []
|
||||
|
||||
# Gemini Embedding 2+ multimodal models return one aggregated vector for a
|
||||
# multi-input request, so embed one input per call to keep 1:1 alignment.
|
||||
batch_size = 1 if _gemini_model_aggregates_inputs(self.model) else self.batch_size
|
||||
|
||||
# Process in batches
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
|
||||
embed_kwargs = {"model": self.model, "contents": batch}
|
||||
if self._embed_config is not None:
|
||||
@@ -1350,7 +1545,13 @@ class GeminiEmbeddings(Embeddings):
|
||||
|
||||
result = self._client.models.embed_content(**embed_kwargs)
|
||||
|
||||
all_embeddings.extend([emb.values for emb in result.embeddings])
|
||||
embeddings = result.embeddings or []
|
||||
if len(embeddings) != len(batch):
|
||||
raise RuntimeError(
|
||||
f"Gemini embeddings backend returned {len(embeddings)} vectors for "
|
||||
f"{len(batch)} input texts (model {self.model}); expected exact 1:1 alignment"
|
||||
)
|
||||
all_embeddings.extend([emb.values for emb in embeddings])
|
||||
|
||||
# L2-normalize when output_dimensionality is set — Gemini only returns
|
||||
# normalized vectors at full 3072 dims; truncated dims need re-normalization
|
||||
@@ -1391,6 +1592,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)
|
||||
@@ -1492,6 +1707,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown embeddings provider: {provider}. "
|
||||
f"Supported: 'local', 'tei', 'openai', 'openai-codex', 'openrouter', 'cohere', 'google', "
|
||||
f"Supported: 'local', 'onnx', 'tei', 'openai', 'openai-codex', 'openrouter', 'cohere', 'google', "
|
||||
f"'zeroentropy', 'litellm', 'litellm-sdk'"
|
||||
)
|
||||
|
||||
@@ -10,7 +10,7 @@ from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.engine.memory_engine import BankLlmHealthInfo, Budget
|
||||
from hindsight_api.engine.response_models import RecallResult, ReflectResult
|
||||
from hindsight_api.engine.search.tags import TagsMatch
|
||||
from hindsight_api.models import RequestContext
|
||||
@@ -458,8 +458,42 @@ 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).
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def check_bank_llm(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> "BankLlmHealthInfo":
|
||||
"""
|
||||
Probe the LLM consolidation would use for this bank. Deliberate connectivity
|
||||
test (one real minimal call); never returns the API key. See
|
||||
MemoryEngine.check_bank_llm.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@@ -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}")
|
||||
@@ -114,6 +114,32 @@ def _semaphores_for_scope(scope: str) -> list[asyncio.Semaphore]:
|
||||
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.
|
||||
@@ -206,6 +232,7 @@ _PROVIDERS_WITHOUT_API_KEY = frozenset(
|
||||
"litellm",
|
||||
"litellmrouter",
|
||||
"bedrock",
|
||||
"nous",
|
||||
}
|
||||
)
|
||||
|
||||
@@ -223,12 +250,14 @@ def create_llm_provider(
|
||||
reasoning_effort: str,
|
||||
groq_service_tier: str | None = None,
|
||||
openai_service_tier: str | None = None,
|
||||
bedrock_service_tier: str | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
default_headers: dict[str, str] | None = None,
|
||||
vertexai_project_id: str | None = None,
|
||||
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
|
||||
"""
|
||||
@@ -242,7 +271,12 @@ 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.
|
||||
bedrock_service_tier: Bedrock service tier (for Bedrock provider) - None (default), "flex", "priority", or "reserved".
|
||||
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.
|
||||
@@ -317,6 +351,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":
|
||||
@@ -327,6 +363,7 @@ def create_llm_provider(
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
default_headers=default_headers,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower == "litellm":
|
||||
@@ -336,6 +373,7 @@ def create_llm_provider(
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower == "litellmrouter":
|
||||
@@ -353,6 +391,7 @@ def create_llm_provider(
|
||||
model=model,
|
||||
config=litellmrouter_config,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower == "bedrock":
|
||||
@@ -364,6 +403,8 @@ def create_llm_provider(
|
||||
base_url=base_url,
|
||||
model=bedrock_model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
bedrock_service_tier=bedrock_service_tier,
|
||||
)
|
||||
|
||||
elif provider_lower == "llamacpp":
|
||||
@@ -397,6 +438,21 @@ def create_llm_provider(
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower == "nous":
|
||||
# Nous Portal is OpenAI-compatible on the wire; NousLLM adds rotating
|
||||
# inference:invoke JWT auth read natively from ~/.hermes/auth.json
|
||||
# (no static api_key, no hermes_cli dependency — same shape as Codex).
|
||||
from hindsight_api.engine.providers.nous_llm import NousLLM
|
||||
|
||||
return NousLLM(
|
||||
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",
|
||||
@@ -441,7 +497,9 @@ class LLMProvider:
|
||||
reasoning_effort: str = "low",
|
||||
groq_service_tier: str | None = None,
|
||||
openai_service_tier: str | None = None,
|
||||
bedrock_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,
|
||||
@@ -457,8 +515,10 @@ class LLMProvider:
|
||||
reasoning_effort: Reasoning effort level for supported providers.
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto") - from config.
|
||||
openai_service_tier: OpenAI service tier (None or "flex") - from config.
|
||||
bedrock_service_tier: Bedrock service tier (None, "flex", "priority", "reserved") - 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``)
|
||||
@@ -478,8 +538,14 @@ class LLMProvider:
|
||||
# Service tiers from hierarchical config (not env vars)
|
||||
self.groq_service_tier = groq_service_tier
|
||||
self.openai_service_tier = openai_service_tier
|
||||
self.bedrock_service_tier = bedrock_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).
|
||||
@@ -519,6 +585,7 @@ class LLMProvider:
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"fireworks",
|
||||
"nous",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
raise ValueError(f"Invalid LLM provider: {self.provider}. Must be one of: {', '.join(valid_providers)}")
|
||||
@@ -543,6 +610,8 @@ class LLMProvider:
|
||||
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 == "nous":
|
||||
self.base_url = "https://inference-api.nousresearch.com/v1"
|
||||
|
||||
# Prepare Vertex AI config (if applicable)
|
||||
vertexai_project_id = None
|
||||
@@ -598,6 +667,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
|
||||
@@ -620,12 +704,14 @@ class LLMProvider:
|
||||
reasoning_effort=self.reasoning_effort,
|
||||
groq_service_tier=self.groq_service_tier,
|
||||
openai_service_tier=self.openai_service_tier,
|
||||
bedrock_service_tier=self.bedrock_service_tier,
|
||||
extra_body=self.extra_body,
|
||||
default_headers=self.default_headers,
|
||||
vertexai_project_id=vertexai_project_id,
|
||||
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,
|
||||
)
|
||||
|
||||
@@ -689,6 +775,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.
|
||||
@@ -703,7 +790,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:
|
||||
@@ -723,33 +813,83 @@ class LLMProvider:
|
||||
structured = "+structured" if response_format is not None else ""
|
||||
set_stage(f"llm.{self.provider}.{scope}{structured}")
|
||||
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
# 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
|
||||
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call(
|
||||
messages=messages,
|
||||
response_format=response_format,
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
@@ -764,6 +904,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.
|
||||
@@ -786,31 +927,66 @@ class LLMProvider:
|
||||
|
||||
set_stage(f"llm.{self.provider}.{scope}+tools")
|
||||
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
# 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
|
||||
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call_with_tools(
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
@@ -914,7 +1090,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.
|
||||
|
||||
@@ -924,12 +1107,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)."""
|
||||
@@ -944,6 +1146,7 @@ class LLMProvider:
|
||||
DEFAULT_LLM_REASONING_EFFORT,
|
||||
ENV_LLM_API_KEY,
|
||||
ENV_LLM_BASE_URL,
|
||||
ENV_LLM_BEDROCK_SERVICE_TIER,
|
||||
ENV_LLM_DEFAULT_HEADERS,
|
||||
ENV_LLM_EXTRA_BODY,
|
||||
ENV_LLM_MODEL,
|
||||
@@ -975,6 +1178,7 @@ class LLMProvider:
|
||||
reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
bedrock_service_tier=os.getenv(ENV_LLM_BEDROCK_SERVICE_TIER) or None,
|
||||
)
|
||||
|
||||
|
||||
@@ -993,10 +1197,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 ──────────────────────────────────────────────────
|
||||
|
||||
@@ -1009,10 +1219,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,
|
||||
@@ -1023,12 +1235,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."""
|
||||
|
||||
@@ -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
|
||||
@@ -178,6 +186,9 @@ class AnthropicLLM(LLMInterface):
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if self._extra_body:
|
||||
call_params["extra_body"] = self._extra_body
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
@@ -216,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()
|
||||
@@ -245,6 +257,7 @@ class AnthropicLLM(LLMInterface):
|
||||
duration=duration,
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
@@ -260,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
|
||||
@@ -394,6 +408,9 @@ class AnthropicLLM(LLMInterface):
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if self._extra_body:
|
||||
call_params["extra_body"] = self._extra_body
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
|
||||
@@ -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)
|
||||
@@ -8,6 +8,7 @@ This provider supports both:
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
@@ -43,6 +44,14 @@ except ImportError:
|
||||
VERTEXAI_AVAILABLE = False
|
||||
|
||||
|
||||
def _to_int(value: Any) -> int:
|
||||
"""Coerce Gemini's optional/string completion counts to int, defaulting to 0."""
|
||||
try:
|
||||
return int(value)
|
||||
except (ValueError, TypeError):
|
||||
return 0
|
||||
|
||||
|
||||
class GeminiLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider for Google Gemini and Vertex AI.
|
||||
@@ -70,6 +79,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 +193,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 +208,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 +226,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 +249,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 +260,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 +342,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 +372,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 +397,7 @@ class GeminiLLM(LLMInterface):
|
||||
duration=duration,
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
@@ -345,6 +413,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 +435,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 +482,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 +497,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 +542,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 +593,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 +692,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 +716,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 +742,7 @@ class GeminiLLM(LLMInterface):
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
tool_calls=tool_calls_dict,
|
||||
cached_tokens=cached_input_tokens,
|
||||
)
|
||||
|
||||
return LLMToolCallResult(
|
||||
@@ -636,6 +759,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 +787,330 @@ 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,
|
||||
)
|
||||
|
||||
# ── Batch API (Gemini API only — not Vertex AI) ─────────────────────────
|
||||
#
|
||||
# Google's Gemini Batch API gives a flat 50% discount on input + output
|
||||
# tokens with a 24h completion SLA (https://ai.google.dev/gemini-api/docs/batch-api).
|
||||
# The retain orchestrator and ``fact_extraction`` consumer speak the
|
||||
# OpenAI-batch interface contract, so these overrides translate that shape
|
||||
# to/from Gemini's file-upload → ``batches.create`` → ``batches.get`` →
|
||||
# download flow — nothing downstream changes (same pattern as FireworksLLM).
|
||||
#
|
||||
# Interface contract preserved (see fact_extraction.py result handling)::
|
||||
# result["response"]["body"]["choices"][0]["message"]["content"]
|
||||
|
||||
async def supports_batch_api(self) -> bool:
|
||||
"""True for the Gemini API; False for Vertex AI.
|
||||
|
||||
Only ``provider="gemini"`` is supported: it exposes the file-upload
|
||||
Batch API used below. Vertex AI's batch path is GCS/BigQuery-backed (no
|
||||
file-upload analogue), so it stays unsupported here — the startup
|
||||
validation then surfaces a clear error instead of silently falling back
|
||||
to synchronous, full-price calls.
|
||||
"""
|
||||
return self.provider == "gemini"
|
||||
|
||||
async def submit_batch(
|
||||
self,
|
||||
requests: list[dict[str, Any]],
|
||||
endpoint: str = "/v1/chat/completions",
|
||||
completion_window: str = "24h",
|
||||
) -> dict[str, Any]:
|
||||
"""Submit a batch of (OpenAI-shaped) requests to the Gemini Batch API."""
|
||||
if not await self.supports_batch_api():
|
||||
raise NotImplementedError(f"Batch API not supported for provider: {self.provider}")
|
||||
|
||||
# endpoint/completion_window are part of the shared LLMInterface batch
|
||||
# contract (used by the OpenAI path) but have no analogue on Gemini: the
|
||||
# request shape is fixed (generateContent) and the SLA is server-side.
|
||||
# Kept for signature compatibility with the shared retain driver.
|
||||
logger.info(f"Submitting Gemini batch with {len(requests)} requests")
|
||||
|
||||
jsonl = self._translate_requests(requests)
|
||||
|
||||
# Upload the JSONL as a Gemini file (mime_type must be "jsonl"; a
|
||||
# BytesIO has no path for the SDK to infer it from).
|
||||
file_obj = io.BytesIO(jsonl.encode("utf-8"))
|
||||
uploaded = await self._client.aio.files.upload(
|
||||
file=file_obj,
|
||||
config=genai_types.UploadFileConfig(mime_type="jsonl", display_name="hindsight-batch-input"),
|
||||
)
|
||||
|
||||
batch = await self._client.aio.batches.create(
|
||||
model=self.model,
|
||||
src=uploaded.name,
|
||||
config=genai_types.CreateBatchJobConfig(display_name="hindsight-batch"),
|
||||
)
|
||||
|
||||
logger.info(f"Gemini batch submitted: {batch.name}, state={self._state_name(batch.state)}")
|
||||
|
||||
return {
|
||||
"batch_id": batch.name,
|
||||
"status": self._normalize_state(batch.state),
|
||||
"input_file_id": uploaded.name,
|
||||
"request_count": len(requests),
|
||||
}
|
||||
|
||||
async def get_batch_status(self, batch_id: str) -> dict[str, Any]:
|
||||
"""Get the status of a Gemini batch job, in the shared status shape."""
|
||||
if not await self.supports_batch_api():
|
||||
raise NotImplementedError(f"Batch API not supported for provider: {self.provider}")
|
||||
|
||||
batch = await self._client.aio.batches.get(name=batch_id)
|
||||
|
||||
stats = batch.completion_stats
|
||||
successful = _to_int(getattr(stats, "successful_count", None)) if stats else 0
|
||||
failed = _to_int(getattr(stats, "failed_count", None)) if stats else 0
|
||||
incomplete = _to_int(getattr(stats, "incomplete_count", None)) if stats else 0
|
||||
|
||||
result: dict[str, Any] = {
|
||||
"batch_id": batch.name,
|
||||
"status": self._normalize_state(batch.state),
|
||||
"request_counts": {
|
||||
"total": successful + failed + incomplete,
|
||||
"completed": successful,
|
||||
"failed": failed,
|
||||
},
|
||||
}
|
||||
|
||||
if batch.dest and getattr(batch.dest, "file_name", None):
|
||||
result["output_file_id"] = batch.dest.file_name
|
||||
if batch.error:
|
||||
result["errors"] = self._error_to_dict(batch.error)
|
||||
|
||||
return result
|
||||
|
||||
async def retrieve_batch_results(self, batch_id: str) -> list[dict[str, Any]]:
|
||||
"""Download and normalize completed Gemini batch results."""
|
||||
if not await self.supports_batch_api():
|
||||
raise NotImplementedError(f"Batch API not supported for provider: {self.provider}")
|
||||
|
||||
batch = await self._client.aio.batches.get(name=batch_id)
|
||||
|
||||
status = self._normalize_state(batch.state)
|
||||
if status != "completed":
|
||||
raise ValueError(f"Gemini batch {batch_id} is not completed yet (state: {self._state_name(batch.state)})")
|
||||
|
||||
dest = batch.dest
|
||||
if not dest or not getattr(dest, "file_name", None):
|
||||
raise ValueError(
|
||||
f"Gemini batch {batch_id} completed but reported no output file "
|
||||
f"(submit_batch always uses file mode, so this is unexpected)"
|
||||
)
|
||||
|
||||
content = await self._client.aio.files.download(file=dest.file_name)
|
||||
text = content.decode("utf-8") if isinstance(content, (bytes, bytearray)) else str(content)
|
||||
|
||||
# The output is a JSONL error file plus results merged into one stream;
|
||||
# error lines carry an `error` so partial failures surface per key
|
||||
# instead of vanishing (JOB_STATE_PARTIALLY_SUCCEEDED maps to completed).
|
||||
results: list[dict[str, Any]] = []
|
||||
for line in text.strip().split("\n"):
|
||||
if line.strip():
|
||||
results.append(self._normalize_output_line(json.loads(line)))
|
||||
|
||||
logger.info(f"Retrieved {len(results)} results for Gemini batch {batch_id}")
|
||||
return results
|
||||
|
||||
# ----- pure translation/normalization helpers (unit-tested) ----------
|
||||
|
||||
@staticmethod
|
||||
def _translate_requests(requests: list[dict[str, Any]]) -> str:
|
||||
"""OpenAI batch requests -> Gemini batch input JSONL.
|
||||
|
||||
Each output line is ``{"key": <custom_id>, "request": <GenerateContentRequest>}``;
|
||||
the model is supplied to ``batches.create`` so it is omitted per-line.
|
||||
"""
|
||||
lines = []
|
||||
for req in requests:
|
||||
gemini_request = GeminiLLM._openai_body_to_gemini_request(req.get("body") or {})
|
||||
lines.append(json.dumps({"key": req.get("custom_id"), "request": gemini_request}, ensure_ascii=False))
|
||||
return "\n".join(lines)
|
||||
|
||||
@staticmethod
|
||||
def _openai_body_to_gemini_request(body: dict[str, Any]) -> dict[str, Any]:
|
||||
"""OpenAI chat-completions body -> Gemini ``GenerateContentRequest`` JSON.
|
||||
|
||||
Mirrors the synchronous ``call`` path: system messages become
|
||||
``systemInstruction``; a ``response_format`` json_schema forces JSON
|
||||
output (``responseMimeType``), appends the schema as a textual hint, and
|
||||
grammar-enforces via ``responseJsonSchema`` when ``strict`` is set.
|
||||
"""
|
||||
system_texts: list[str] = []
|
||||
contents: list[dict[str, Any]] = []
|
||||
for msg in body.get("messages") or []:
|
||||
role = msg.get("role", "user")
|
||||
text = msg.get("content", "") or ""
|
||||
if role == "system":
|
||||
system_texts.append(text)
|
||||
elif role == "assistant":
|
||||
contents.append({"role": "model", "parts": [{"text": text}]})
|
||||
else:
|
||||
contents.append({"role": "user", "parts": [{"text": text}]})
|
||||
|
||||
generation_config: dict[str, Any] = {}
|
||||
if body.get("temperature") is not None:
|
||||
generation_config["temperature"] = body["temperature"]
|
||||
if body.get("max_completion_tokens") is not None:
|
||||
generation_config["maxOutputTokens"] = body["max_completion_tokens"]
|
||||
|
||||
response_format = body.get("response_format")
|
||||
if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
|
||||
json_schema = response_format.get("json_schema") or {}
|
||||
schema = json_schema.get("schema")
|
||||
generation_config["responseMimeType"] = "application/json"
|
||||
if schema:
|
||||
system_texts.append(
|
||||
"You must respond with valid JSON matching this schema:\n" + json.dumps(schema, ensure_ascii=False)
|
||||
)
|
||||
if json_schema.get("strict"):
|
||||
generation_config["responseJsonSchema"] = schema
|
||||
|
||||
request: dict[str, Any] = {"contents": contents}
|
||||
if system_texts:
|
||||
request["systemInstruction"] = {"parts": [{"text": "\n\n".join(system_texts)}]}
|
||||
if generation_config:
|
||||
request["generationConfig"] = generation_config
|
||||
return request
|
||||
|
||||
@staticmethod
|
||||
def _normalize_output_line(line: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Gemini batch output line -> OpenAI-batch-output shape.
|
||||
|
||||
Target: ``{"custom_id", "response": {"body": {"choices": [...], "usage": {...}}}, "error"}``
|
||||
so the consumer's ``result["response"]["body"]["choices"][0]...`` works and
|
||||
it can read ``body["usage"]`` for token accounting (the consumer reports
|
||||
zero usage otherwise).
|
||||
"""
|
||||
custom_id = line.get("key") if line.get("key") is not None else line.get("custom_id")
|
||||
error = line.get("error")
|
||||
if error:
|
||||
return {"custom_id": custom_id, "response": None, "error": error}
|
||||
|
||||
response = line.get("response") or {}
|
||||
body: dict[str, Any] = {"choices": [{"message": {"content": GeminiLLM._extract_text_from_response(response)}}]}
|
||||
usage = GeminiLLM._usage_from_response(response)
|
||||
if usage is not None:
|
||||
body["usage"] = usage
|
||||
return {"custom_id": custom_id, "response": {"body": body}, "error": None}
|
||||
|
||||
@staticmethod
|
||||
def _extract_text_from_response(response: dict[str, Any]) -> str:
|
||||
"""Concatenate the text parts of a (JSON) GenerateContentResponse."""
|
||||
candidates = response.get("candidates") or []
|
||||
if not candidates:
|
||||
return ""
|
||||
content = candidates[0].get("content") or {}
|
||||
parts = content.get("parts") or []
|
||||
return "".join(p.get("text", "") for p in parts if isinstance(p, dict) and p.get("text"))
|
||||
|
||||
@staticmethod
|
||||
def _usage_from_response(response: dict[str, Any]) -> dict[str, Any] | None:
|
||||
"""Gemini ``usageMetadata`` -> OpenAI-shaped ``usage`` block, or None.
|
||||
|
||||
The batch consumer accumulates token usage from ``body["usage"]`` using
|
||||
OpenAI key names, so translate here to keep the output contract uniform
|
||||
across providers. Handles both the REST camelCase (downloaded JSONL) and
|
||||
snake_case spellings defensively.
|
||||
"""
|
||||
meta = response.get("usageMetadata") or response.get("usage_metadata")
|
||||
if not isinstance(meta, dict):
|
||||
return None
|
||||
prompt = meta.get("promptTokenCount") or meta.get("prompt_token_count") or 0
|
||||
completion = meta.get("candidatesTokenCount") or meta.get("candidates_token_count") or 0
|
||||
total = meta.get("totalTokenCount") or meta.get("total_token_count") or 0
|
||||
return {"prompt_tokens": prompt, "completion_tokens": completion, "total_tokens": total}
|
||||
|
||||
@staticmethod
|
||||
def _normalize_state(state: Any) -> str:
|
||||
"""Gemini ``JobState`` -> the retain driver's status strings.
|
||||
|
||||
Unknown / in-flight states map to ``in_progress`` so the driver keeps
|
||||
polling; ``PARTIALLY_SUCCEEDED`` maps to ``completed`` (per-line errors
|
||||
surface the partial failures during retrieval).
|
||||
"""
|
||||
name = GeminiLLM._state_name(state).upper()
|
||||
if name in ("JOB_STATE_SUCCEEDED", "JOB_STATE_PARTIALLY_SUCCEEDED"):
|
||||
return "completed"
|
||||
if name == "JOB_STATE_FAILED":
|
||||
return "failed"
|
||||
if name in ("JOB_STATE_CANCELLED", "JOB_STATE_CANCELLING"):
|
||||
return "cancelled"
|
||||
if name == "JOB_STATE_EXPIRED":
|
||||
return "expired"
|
||||
return "in_progress"
|
||||
|
||||
@staticmethod
|
||||
def _state_name(state: Any) -> str:
|
||||
"""Extract the bare ``JOB_STATE_*`` name from a JobState enum or string."""
|
||||
if state is None:
|
||||
return ""
|
||||
name = getattr(state, "name", None)
|
||||
if name:
|
||||
return str(name)
|
||||
text = str(state)
|
||||
if "." in text:
|
||||
text = text.rsplit(".", 1)[-1]
|
||||
return text
|
||||
|
||||
@staticmethod
|
||||
def _error_to_dict(error: Any) -> dict[str, Any]:
|
||||
"""Coerce a Gemini JobError into a JSON-serializable dict for logging."""
|
||||
if hasattr(error, "model_dump"):
|
||||
try:
|
||||
return error.model_dump(exclude_none=True)
|
||||
except Exception:
|
||||
pass
|
||||
return {"message": str(error)}
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (close connections, etc.)."""
|
||||
# Gemini client doesn't require explicit cleanup
|
||||
|
||||
@@ -48,11 +48,20 @@ class LiteLLMLLM(LLMInterface):
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float = 300.0,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
bedrock_service_tier: str | 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 {}
|
||||
self.bedrock_service_tier = bedrock_service_tier
|
||||
|
||||
try:
|
||||
import litellm
|
||||
@@ -107,6 +116,15 @@ 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)
|
||||
|
||||
# Bedrock service tier: flex (50% cheaper), priority, or reserved
|
||||
if self.model.startswith("bedrock/") and self.bedrock_service_tier is not None:
|
||||
kwargs["service_tier"] = self.bedrock_service_tier
|
||||
|
||||
return kwargs
|
||||
|
||||
# ── per-model output-tokens cap (shared with Router subclass) ────────────
|
||||
|
||||
@@ -162,6 +162,12 @@ class MockLLM(LLMInterface):
|
||||
# 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."
|
||||
|
||||
@@ -0,0 +1,463 @@
|
||||
"""
|
||||
Native Nous Portal OAuth authentication manager.
|
||||
|
||||
The Nous Portal inference endpoint (https://inference-api.nousresearch.com/v1)
|
||||
speaks the OpenAI-compatible wire format but authenticates with a short-lived,
|
||||
inference-scoped JWT rather than a static API key. Hermes obtains that JWT once
|
||||
via an interactive browser login (``hermes portal``) and persists the resulting
|
||||
OAuth state — ``access_token`` + ``refresh_token`` — under ``providers.nous`` in
|
||||
``~/.hermes/auth.json``.
|
||||
|
||||
This manager reads that file *directly* and refreshes the access token itself,
|
||||
exactly mirroring ``codex_auth.py`` (read ``~/.codex/auth.json`` + native
|
||||
refresh). It deliberately does **not** import the Hermes ``hermes_cli`` package:
|
||||
that package is the interactive CLI, not a library Hindsight can depend on. The
|
||||
refresh request shape is mirrored from Hermes' own resolver
|
||||
(``POST {portal}/api/oauth/token`` with an ``x-nous-refresh-token`` header and a
|
||||
``grant_type=refresh_token`` form body), so server-side changes affect both
|
||||
clients identically. The inference bearer is the access token itself — in
|
||||
Hermes' state the ``agent_key`` field is literally ``= access_token``.
|
||||
|
||||
Single-use refresh tokens
|
||||
-------------------------
|
||||
Nous refresh tokens are single-use with server-side reuse-detection: if two
|
||||
processes refresh with the same ``refresh_token``, or a rotated token is not
|
||||
persisted back, the Portal revokes the whole session as a theft signal. Because
|
||||
Hindsight shares ``~/.hermes/auth.json`` with a possibly-running Hermes agent,
|
||||
every refresh here is performed while holding the **same cross-process advisory
|
||||
lock Hermes uses** (``~/.hermes/auth.lock`` via ``fcntl.flock``) and re-reads the
|
||||
latest ``refresh_token`` from disk under that lock before exchanging it. That is
|
||||
the protocol Hermes follows too, so the two coordinate safely through the file.
|
||||
|
||||
Usage
|
||||
-----
|
||||
mgr = NousAuthManager.from_file()
|
||||
token = mgr.ensure_fresh_token() # proactive; refreshes if near expiry
|
||||
... # use token as Bearer
|
||||
mgr.refresh_tokens(force=True) # reactive, on a 401
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError: # pragma: no cover - Windows
|
||||
fcntl = None # type: ignore[assignment]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Constants — mirrored from Hermes' canonical Nous resolver
|
||||
# (hermes_cli/auth.py: DEFAULT_NOUS_* and _refresh_access_token). Endpoints and
|
||||
# client id are overridable via the same env vars Hermes honours, so a staging
|
||||
# Portal or a future change can be pointed at without a code change.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_NOUS_PORTAL_BASE_URL = (
|
||||
os.environ.get("HERMES_PORTAL_BASE_URL")
|
||||
or os.environ.get("NOUS_PORTAL_BASE_URL")
|
||||
or "https://portal.nousresearch.com"
|
||||
)
|
||||
_NOUS_INFERENCE_BASE_URL = os.environ.get("NOUS_INFERENCE_BASE_URL") or "https://inference-api.nousresearch.com/v1"
|
||||
_NOUS_CLIENT_ID = "hermes-cli"
|
||||
|
||||
# Proactively refresh this many seconds before the JWT ``exp`` claim — matches
|
||||
# the 120s skew Hermes' own runtime resolver uses for Nous.
|
||||
_NOUS_TOKEN_REFRESH_SKEW_SECONDS = 120
|
||||
|
||||
# OAuth error codes the Portal returns when the refresh_token itself is no
|
||||
# longer usable. These are terminal — retrying will not succeed; the user must
|
||||
# re-run ``hermes portal``.
|
||||
_NOUS_TERMINAL_REFRESH_ERROR_CODES = frozenset(
|
||||
{"invalid_grant", "invalid_token", "refresh_token_reused", "refresh_token_expired"}
|
||||
)
|
||||
|
||||
_AUTH_LOCK_TIMEOUT_SECONDS = 20.0
|
||||
|
||||
|
||||
def _default_auth_file() -> Path:
|
||||
return Path.home() / ".hermes" / "auth.json"
|
||||
|
||||
|
||||
class NousNotLoggedInError(RuntimeError):
|
||||
"""Raised when ``~/.hermes/auth.json`` has no usable Nous OAuth state.
|
||||
|
||||
Remediation: run ``hermes portal`` to log in to Nous Portal.
|
||||
"""
|
||||
|
||||
|
||||
class NousRefreshExpiredError(RuntimeError):
|
||||
"""Raised when the Nous refresh_token itself is permanently invalid.
|
||||
|
||||
The user must re-run ``hermes portal`` to obtain new credentials. Callers
|
||||
should surface a clear remediation message and stop retrying.
|
||||
"""
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _hermes_auth_lock(auth_file: Path, timeout_seconds: float = _AUTH_LOCK_TIMEOUT_SECONDS) -> Iterator[None]:
|
||||
"""Cross-process advisory lock on the Hermes auth store.
|
||||
|
||||
Uses ``<auth_file>.lock`` (i.e. ``~/.hermes/auth.lock``) with
|
||||
``fcntl.flock(LOCK_EX)`` — the exact same lock file and primitive Hermes'
|
||||
``_auth_store_lock`` takes — so a refresh here is mutually exclusive with a
|
||||
concurrently-running Hermes agent. Degrades to a no-op (with a debug log)
|
||||
where ``fcntl`` is unavailable (Windows); the single-process in-memory lock
|
||||
still serialises this process's own refreshes.
|
||||
"""
|
||||
if fcntl is None: # pragma: no cover - Windows
|
||||
logger.debug("fcntl unavailable; Nous refresh proceeds without a cross-process lock.")
|
||||
yield
|
||||
return
|
||||
|
||||
lock_path = auth_file.with_suffix(".lock")
|
||||
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(lock_path, "a+") as lock_file:
|
||||
deadline = time.monotonic() + max(1.0, timeout_seconds)
|
||||
while True:
|
||||
try:
|
||||
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
break
|
||||
except (BlockingIOError, OSError):
|
||||
if time.monotonic() >= deadline:
|
||||
raise TimeoutError("Timed out waiting for the Hermes auth store lock") from None
|
||||
time.sleep(0.05)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
with contextlib.suppress(OSError):
|
||||
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
|
||||
|
||||
|
||||
class NousAuthManager:
|
||||
"""Sync Nous Portal OAuth credential manager.
|
||||
|
||||
Holds the access_token + refresh_token in memory and handles
|
||||
proactive/reactive refresh. A ``threading.Lock`` gives single-flight
|
||||
semantics within the process; the cross-process ``fcntl`` lock guards
|
||||
against a concurrent Hermes agent (see module docstring).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
access_token: str,
|
||||
refresh_token: str | None,
|
||||
auth_file: Path,
|
||||
*,
|
||||
portal_base_url: str = _NOUS_PORTAL_BASE_URL,
|
||||
inference_base_url: str = _NOUS_INFERENCE_BASE_URL,
|
||||
client_id: str = _NOUS_CLIENT_ID,
|
||||
) -> None:
|
||||
self.access_token = access_token
|
||||
self.refresh_token = refresh_token
|
||||
self._auth_file = auth_file
|
||||
self._portal_base_url = portal_base_url.rstrip("/")
|
||||
self._inference_base_url = inference_base_url.rstrip("/")
|
||||
self._client_id = client_id
|
||||
self._lock = threading.Lock()
|
||||
self._http_client = httpx.Client(timeout=30.0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Construction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@classmethod
|
||||
def from_file(cls, auth_file: Path | None = None) -> "NousAuthManager":
|
||||
"""Build a manager from ``providers.nous`` in the Hermes auth store.
|
||||
|
||||
Raises
|
||||
------
|
||||
NousNotLoggedInError:
|
||||
If the file is missing, unreadable, or has no Nous OAuth state with
|
||||
an ``access_token``.
|
||||
"""
|
||||
if auth_file is None:
|
||||
auth_file = _default_auth_file()
|
||||
|
||||
if not auth_file.exists():
|
||||
raise NousNotLoggedInError(
|
||||
f"Hermes auth file not found: {auth_file}. Run 'hermes portal' to log in to Nous Portal."
|
||||
)
|
||||
|
||||
try:
|
||||
with open(auth_file) as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError) as e:
|
||||
raise NousNotLoggedInError(f"Could not read Hermes auth file {auth_file}: {type(e).__name__}") from e
|
||||
|
||||
state = cls._nous_state(data)
|
||||
if not state:
|
||||
raise NousNotLoggedInError(
|
||||
"Hermes is not logged into Nous Portal (no providers.nous OAuth state). Run 'hermes portal'."
|
||||
)
|
||||
|
||||
access_token = state.get("access_token")
|
||||
if not isinstance(access_token, str) or not access_token:
|
||||
raise NousNotLoggedInError("Nous OAuth state has no access_token. Re-authenticate with 'hermes portal'.")
|
||||
|
||||
return cls(
|
||||
access_token=access_token,
|
||||
refresh_token=state.get("refresh_token"),
|
||||
auth_file=auth_file,
|
||||
portal_base_url=cls._optional_url(state.get("portal_base_url")) or _NOUS_PORTAL_BASE_URL,
|
||||
inference_base_url=cls._optional_url(state.get("inference_base_url")) or _NOUS_INFERENCE_BASE_URL,
|
||||
client_id=str(state.get("client_id") or _NOUS_CLIENT_ID),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _nous_state(data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Pull the ``providers.nous`` state dict out of a loaded auth store."""
|
||||
providers = data.get("providers")
|
||||
if not isinstance(providers, dict):
|
||||
return {}
|
||||
state = providers.get("nous")
|
||||
return state if isinstance(state, dict) else {}
|
||||
|
||||
@staticmethod
|
||||
def _optional_url(value: Any) -> str | None:
|
||||
return value.rstrip("/") if isinstance(value, str) and value.strip() else None
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
return self._inference_base_url
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Token state
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def load_refresh_token_from_file(auth_file: Path) -> str | None:
|
||||
"""Read ``providers.nous.refresh_token`` from ``auth_file``.
|
||||
|
||||
Returns ``None`` when the file is unreadable or omits the field. Does
|
||||
not raise — the caller degrades to using the in-memory token.
|
||||
"""
|
||||
try:
|
||||
with open(auth_file) as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
return NousAuthManager._nous_state(data).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 failure.
|
||||
|
||||
The signature is not verified — the server is the source of truth on
|
||||
acceptance. This only schedules proactive refresh.
|
||||
"""
|
||||
try:
|
||||
parts = token.split(".")
|
||||
if len(parts) < 2:
|
||||
return None
|
||||
payload_b64 = parts[1]
|
||||
padding = "=" * (-len(payload_b64) % 4)
|
||||
payload = json.loads(base64.urlsafe_b64decode(payload_b64 + padding).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 = _NOUS_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`` is unparseable.
|
||||
"""
|
||||
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_state_atomic(self, updated: dict[str, Any]) -> None:
|
||||
"""Patch ``providers.nous`` in ``_auth_file`` and write atomically.
|
||||
|
||||
Re-reads the on-disk store first so fields written by Hermes (other
|
||||
providers, the credential pool, rotated tokens) are never clobbered,
|
||||
then patches only the Nous OAuth fields and ``os.replace``s into place
|
||||
(atomic on POSIX within the same filesystem). Must be called while
|
||||
holding :func:`_hermes_auth_lock`.
|
||||
"""
|
||||
try:
|
||||
with open(self._auth_file) as f:
|
||||
loaded = json.load(f)
|
||||
current: dict[str, Any] = loaded if isinstance(loaded, dict) else {}
|
||||
except (OSError, json.JSONDecodeError):
|
||||
current = {}
|
||||
|
||||
providers = current.get("providers")
|
||||
if not isinstance(providers, dict):
|
||||
providers = {}
|
||||
current["providers"] = providers
|
||||
state = providers.get("nous")
|
||||
if not isinstance(state, dict):
|
||||
state = {}
|
||||
providers["nous"] = state
|
||||
|
||||
state.update(updated)
|
||||
# The inference bearer is the access token itself; keep agent_key in
|
||||
# sync so Hermes' own resolver/status sees the rotation too.
|
||||
state["agent_key"] = updated.get("access_token", state.get("access_token"))
|
||||
current["updated_at"] = 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())
|
||||
with contextlib.suppress(OSError):
|
||||
os.chmod(tmp_path, 0o600)
|
||||
os.replace(tmp_path, self._auth_file)
|
||||
except Exception:
|
||||
with contextlib.suppress(OSError):
|
||||
os.unlink(tmp_path)
|
||||
raise
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Refresh
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _extract_oauth_error_code(response: httpx.Response) -> str | None:
|
||||
"""Pull the OAuth error code out of a 4xx refresh response, if present."""
|
||||
try:
|
||||
body = response.json()
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return None
|
||||
if not isinstance(body, dict):
|
||||
return None
|
||||
err = body.get("error")
|
||||
if isinstance(err, str):
|
||||
return err
|
||||
if isinstance(err, dict) and isinstance(err.get("code"), str):
|
||||
return err["code"]
|
||||
code = body.get("error_code")
|
||||
return code if isinstance(code, str) else None
|
||||
|
||||
def refresh_tokens(self, reason: str = "", *, force: bool = False) -> None:
|
||||
"""Single-flight Nous OAuth token refresh.
|
||||
|
||||
Serialised through ``self._lock`` (in-process single-flight) and
|
||||
:func:`_hermes_auth_lock` (cross-process, vs a running Hermes agent).
|
||||
The latest ``refresh_token`` is re-read from disk under the lock before
|
||||
the exchange — single-use tokens make using a stale in-memory RT a
|
||||
session-revoking mistake.
|
||||
|
||||
Raises
|
||||
------
|
||||
NousRefreshExpiredError:
|
||||
On a terminal refresh error (expired/reused/invalid grant).
|
||||
RuntimeError:
|
||||
For other refresh failures (network, 5xx, missing refresh_token).
|
||||
"""
|
||||
token_before_lock = self.access_token
|
||||
with self._lock:
|
||||
if force:
|
||||
if self.access_token != token_before_lock:
|
||||
return # another caller already refreshed while we waited
|
||||
elif not self._token_is_stale():
|
||||
return
|
||||
|
||||
with _hermes_auth_lock(self._auth_file):
|
||||
# Re-read the freshest refresh_token persisted by whoever rotated
|
||||
# last (this process or Hermes). Using a stale RT is exactly what
|
||||
# trips the Portal's single-use reuse-detection.
|
||||
disk_rt = self.load_refresh_token_from_file(self._auth_file)
|
||||
if disk_rt:
|
||||
self.refresh_token = disk_rt
|
||||
|
||||
if not self.refresh_token:
|
||||
raise RuntimeError(
|
||||
"Nous access_token is expired but no refresh_token is available. "
|
||||
"Run 'hermes portal' to re-authenticate."
|
||||
)
|
||||
|
||||
log_reason = f" ({reason})" if reason else ""
|
||||
logger.info(f"Refreshing Nous Portal access_token{log_reason}")
|
||||
|
||||
try:
|
||||
response = self._http_client.post(
|
||||
f"{self._portal_base_url}/api/oauth/token",
|
||||
headers={"x-nous-refresh-token": self.refresh_token},
|
||||
data={"grant_type": "refresh_token", "client_id": self._client_id},
|
||||
timeout=30.0,
|
||||
)
|
||||
except httpx.RequestError as e:
|
||||
raise RuntimeError(f"Nous OAuth refresh network error: {type(e).__name__}") from e
|
||||
|
||||
if response.status_code != 200:
|
||||
code = self._extract_oauth_error_code(response)
|
||||
if code in _NOUS_TERMINAL_REFRESH_ERROR_CODES or response.status_code in (400, 401):
|
||||
raise NousRefreshExpiredError(
|
||||
f"Nous refresh_token is no longer valid (status={response.status_code}, "
|
||||
f"error={code or 'none'}). Run 'hermes portal' to re-authenticate."
|
||||
)
|
||||
raise RuntimeError(f"Nous OAuth refresh failed with HTTP {response.status_code}")
|
||||
|
||||
try:
|
||||
body = response.json()
|
||||
except (json.JSONDecodeError, ValueError) as e:
|
||||
raise RuntimeError(f"Nous OAuth refresh returned non-JSON body: {e}") from e
|
||||
|
||||
new_access = body.get("access_token")
|
||||
if not new_access:
|
||||
raise RuntimeError("Nous OAuth refresh returned no access_token")
|
||||
new_refresh = body.get("refresh_token") or self.refresh_token
|
||||
|
||||
# Update in-memory state first so waiters see fresh credentials
|
||||
# even if the 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}
|
||||
expires_in = body.get("expires_in")
|
||||
if isinstance(expires_in, (int, float)):
|
||||
persisted["expires_at"] = datetime.fromtimestamp(
|
||||
time.time() + float(expires_in), tz=timezone.utc
|
||||
).isoformat()
|
||||
try:
|
||||
self._persist_state_atomic(persisted)
|
||||
except OSError as e:
|
||||
logger.warning(
|
||||
f"Nous refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
|
||||
"In-memory credentials are current; the on-disk rotated token was not saved."
|
||||
)
|
||||
logger.info("Nous Portal access_token refreshed successfully")
|
||||
|
||||
def ensure_fresh_token(self) -> str:
|
||||
"""Refresh proactively if near/at expiry, then return the bearer token.
|
||||
|
||||
Cheap when fresh (a JWT exp decode + comparison).
|
||||
"""
|
||||
if self._token_is_stale():
|
||||
self.refresh_tokens(reason="proactive (token near expiry)")
|
||||
return self.access_token
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the underlying HTTP client."""
|
||||
self._http_client.close()
|
||||
@@ -0,0 +1,167 @@
|
||||
"""
|
||||
Nous Portal LLM provider for Hindsight.
|
||||
|
||||
Thin wrapper over :class:`OpenAICompatibleLLM`. The Nous Portal speaks the
|
||||
OpenAI chat-completions wire format, so all request/response handling is
|
||||
inherited unchanged. The only thing Nous needs on top is a rotating,
|
||||
inference-scoped JWT (there is no static API key in the Hermes login flow),
|
||||
which :class:`NousAuthManager` reads from ``~/.hermes/auth.json`` and refreshes
|
||||
natively — the same pattern as the Codex provider, with no dependency on the
|
||||
``hermes_cli`` package. See ``nous_auth.py`` for the auth mechanics.
|
||||
|
||||
Configure with::
|
||||
|
||||
llm_provider = "nous"
|
||||
llm_base_url = "https://inference-api.nousresearch.com/v1" # or omit
|
||||
llm_model = "deepseek/deepseek-v4-flash" # any Nous slug
|
||||
|
||||
No API key is set in config; the token comes from the shared Hermes auth store
|
||||
after a one-time ``hermes portal`` login.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from openai import APIStatusError, AsyncOpenAI
|
||||
|
||||
from hindsight_api.engine.providers.nous_auth import (
|
||||
NousAuthManager,
|
||||
NousNotLoggedInError,
|
||||
NousRefreshExpiredError,
|
||||
)
|
||||
from hindsight_api.engine.providers.openai_compatible_llm import OpenAICompatibleLLM
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
__all__ = ["NousLLM", "NousAuthManager", "NousNotLoggedInError", "NousRefreshExpiredError"]
|
||||
|
||||
|
||||
class NousLLM(OpenAICompatibleLLM):
|
||||
"""OpenAI-compatible provider for the Nous Portal with rotating-JWT auth."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str, # Ignored — the token is read from ~/.hermes/auth.json
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
**kwargs: Any,
|
||||
):
|
||||
try:
|
||||
self._auth = NousAuthManager.from_file()
|
||||
except NousNotLoggedInError as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to load Nous Portal credentials: {e}\n\n"
|
||||
"To set up Nous authentication:\n"
|
||||
"1. Install Hermes: https://hermes-agent.nousresearch.com\n"
|
||||
"2. Log in to Nous Portal: hermes portal\n"
|
||||
"3. Verify: hermes portal status\n\n"
|
||||
"Or use a different provider (openai, anthropic, gemini) with an API key."
|
||||
) from e
|
||||
|
||||
# Single-flight async refresh lock — concurrent coroutines racing toward
|
||||
# an expired token produce one network refresh.
|
||||
self._auth_lock = asyncio.Lock()
|
||||
|
||||
token = self._auth.access_token
|
||||
resolved_base = base_url or self._auth.base_url
|
||||
# Parent validates provider against a fixed list; present as "openai"
|
||||
# (identical wire format) while retaining the true identity for logs.
|
||||
super().__init__(
|
||||
provider="openai",
|
||||
api_key=token,
|
||||
base_url=resolved_base,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
**kwargs,
|
||||
)
|
||||
self._nous_provider_name = provider
|
||||
logger.info(
|
||||
"Nous LLM initialized: model=%s base_url=%s (rotating inference:invoke JWT)",
|
||||
self.model,
|
||||
self.base_url,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Token lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _rebuild_client(self) -> None:
|
||||
"""Rebuild the OpenAI SDK client against the current token."""
|
||||
self.api_key = self._auth.access_token
|
||||
self._client = AsyncOpenAI(
|
||||
api_key=self.api_key,
|
||||
base_url=self.base_url,
|
||||
max_retries=0,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
async def _ensure_fresh_token(self) -> None:
|
||||
"""Proactively refresh if the JWT is near expiry; rebuild on change.
|
||||
|
||||
Cheap when fresh (a JWT exp decode). The blocking refresh (network +
|
||||
cross-process file lock) is offloaded to a thread so the event loop is
|
||||
never stalled.
|
||||
"""
|
||||
if not self._auth._token_is_stale():
|
||||
return
|
||||
await self._refresh(reason="proactive (token near expiry)", force=False)
|
||||
|
||||
async def _refresh(self, *, reason: str, force: bool) -> None:
|
||||
token_before = self.api_key
|
||||
async with self._auth_lock:
|
||||
if force:
|
||||
if self.api_key != token_before:
|
||||
return # another coroutine already refreshed
|
||||
elif not self._auth._token_is_stale():
|
||||
return
|
||||
await asyncio.to_thread(lambda: self._auth.refresh_tokens(reason, force=force))
|
||||
if self._auth.access_token != self.api_key:
|
||||
self._rebuild_client()
|
||||
|
||||
async def _with_auth_retry(self, fn: Any, label: str, *args: Any, **kwargs: Any) -> Any:
|
||||
"""Run an OpenAI-compatible call, refreshing once on a 401.
|
||||
|
||||
The proactive refresh covers most expiries; a token can still be
|
||||
rejected mid-flight if Hermes rotated it out from under us or the exp
|
||||
claim was unparseable. One reactive refresh + retry is the safety net.
|
||||
"""
|
||||
await self._ensure_fresh_token()
|
||||
try:
|
||||
return await fn(*args, **kwargs)
|
||||
except APIStatusError as e:
|
||||
if getattr(e, "status_code", None) != 401:
|
||||
raise
|
||||
logger.warning("Nous 401 (%s) — forcing token refresh and retrying once.", label)
|
||||
try:
|
||||
await self._refresh(reason=f"reactive (HTTP 401 on {label})", force=True)
|
||||
except NousRefreshExpiredError as refresh_err:
|
||||
raise RuntimeError(
|
||||
"Nous authentication failed and the refresh_token is no longer valid.\n"
|
||||
"Run 'hermes portal' to re-authenticate."
|
||||
) from refresh_err
|
||||
return await fn(*args, **kwargs)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Overrides
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
await self._ensure_fresh_token()
|
||||
return await super().verify_connection()
|
||||
|
||||
async def call(self, *args: Any, **kwargs: Any) -> Any:
|
||||
return await self._with_auth_retry(super().call, "call", *args, **kwargs)
|
||||
|
||||
async def call_with_tools(self, *args: Any, **kwargs: Any) -> Any:
|
||||
return await self._with_auth_retry(super().call_with_tools, "call_with_tools", *args, **kwargs)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
self._auth.close()
|
||||
parent_cleanup = getattr(super(), "cleanup", None)
|
||||
if parent_cleanup is not None:
|
||||
await parent_cleanup()
|
||||
@@ -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
|
||||
|
||||
@@ -33,6 +33,7 @@ import httpx
|
||||
from openai import APIConnectionError, APIStatusError, AsyncOpenAI, LengthFinishReasonError
|
||||
|
||||
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
|
||||
from hindsight_api.engine.bank_attribution import apply_bank_attribution
|
||||
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
|
||||
@@ -44,6 +45,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 +243,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
|
||||
"""
|
||||
@@ -360,6 +371,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.
|
||||
@@ -460,7 +486,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:
|
||||
@@ -568,6 +596,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
call_params["messages"] = _ensure_json_word_in_user_message(call_params["messages"])
|
||||
call_params["response_format"] = {"type": "json_object"}
|
||||
|
||||
apply_bank_attribution(call_params)
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
@@ -650,6 +680,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()
|
||||
@@ -679,14 +712,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}, "
|
||||
@@ -699,6 +730,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
|
||||
@@ -867,6 +899,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
|
||||
@@ -906,6 +948,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
if extra_body:
|
||||
call_params["extra_body"] = extra_body
|
||||
|
||||
apply_bank_attribution(call_params)
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
|
||||
@@ -14,6 +14,7 @@ import re
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||
|
||||
from ...config import get_config
|
||||
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, TokenUsageSummary, ToolCall
|
||||
from .prompts import (
|
||||
_extract_directive_rules,
|
||||
@@ -302,6 +303,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,
|
||||
@@ -322,6 +341,7 @@ async def run_reflect_agent(
|
||||
budget: str | None = None,
|
||||
max_context_tokens: int = 100_000,
|
||||
llm_output_language: str | None = None,
|
||||
cancel_check: Callable[[], None] | None = None,
|
||||
) -> ReflectAgentResult:
|
||||
"""
|
||||
Execute the reflect agent loop using native tool calling.
|
||||
@@ -358,12 +378,16 @@ async def run_reflect_agent(
|
||||
# Extract directive rules for tool schema (if any)
|
||||
directive_rules = _extract_directive_rules(directives) if directives else None
|
||||
|
||||
# Get tools for this agent (with directive compliance field if directives exist)
|
||||
# Get tools for this agent (with directive compliance field if directives exist).
|
||||
# The expand tool only reads back raw source text (chunks/documents), so it is
|
||||
# useless and excluded when document text storage is disabled.
|
||||
include_expand = get_config().store_document_text
|
||||
tools = get_reflect_tools(
|
||||
directive_rules=directive_rules,
|
||||
include_mental_models=has_mental_models,
|
||||
include_observations=include_observations,
|
||||
include_recall=include_recall,
|
||||
include_expand=include_expand,
|
||||
)
|
||||
# Build set of enabled tool names to guard against LLM hallucinating disabled tool calls
|
||||
enabled_tools: frozenset[str] = frozenset(t["function"]["name"] for t in tools if t.get("type") == "function")
|
||||
@@ -382,6 +406,28 @@ async def run_reflect_agent(
|
||||
{"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] = []
|
||||
@@ -442,7 +488,19 @@ 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):
|
||||
# Cooperative cancellation checkpoint: abort the agent loop between
|
||||
# iterations if the caller (e.g. an HTTP client) has gone away, rather
|
||||
# than spending another LLM round-trip on a result nobody will read
|
||||
# (issue #2122). Raises OperationCancelledError when fired.
|
||||
if cancel_check is not None:
|
||||
cancel_check()
|
||||
|
||||
is_last = iteration == max_iterations - 1
|
||||
|
||||
if is_last:
|
||||
@@ -455,7 +513,9 @@ async def run_reflect_agent(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
"content": build_final_system_prompt(
|
||||
bank_profile.get("mission"), llm_output_language, directives
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
@@ -515,7 +575,9 @@ async def run_reflect_agent(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
"content": build_final_system_prompt(
|
||||
bank_profile.get("mission"), llm_output_language, directives
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
@@ -570,18 +632,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
|
||||
@@ -621,7 +696,9 @@ async def run_reflect_agent(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
"content": build_final_system_prompt(
|
||||
bank_profile.get("mission"), llm_output_language, directives
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
@@ -745,7 +822,9 @@ async def run_reflect_agent(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
"content": build_final_system_prompt(
|
||||
bank_profile.get("mission"), llm_output_language, directives
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
@@ -850,7 +929,9 @@ async def run_reflect_agent(
|
||||
hallucinated_tools = []
|
||||
for tc in other_tools:
|
||||
norm = _normalize_tool_name(tc.name)
|
||||
if enabled_tools is not None and norm not in enabled_tools and norm not in ("done", "expand"):
|
||||
# "done" is always available. "expand" is governed by enabled_tools
|
||||
# (it is excluded when text storage is disabled), so it is not hardcoded here.
|
||||
if enabled_tools is not None and norm not in enabled_tools and norm != "done":
|
||||
hallucinated_tools.append(tc)
|
||||
else:
|
||||
allowed_tools.append(tc)
|
||||
@@ -924,6 +1005,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"
|
||||
@@ -1159,8 +1259,10 @@ async def _execute_tool(
|
||||
# Normalize tool name for various LLM output formats
|
||||
tool_name = _normalize_tool_name(tool_name)
|
||||
|
||||
# Guard against LLMs hallucinating calls to tools that were not provided
|
||||
if enabled_tools is not None and tool_name not in enabled_tools and tool_name not in ("done", "expand"):
|
||||
# Guard against LLMs hallucinating calls to tools that were not provided.
|
||||
# "done" is always available; "expand" is governed by enabled_tools (excluded
|
||||
# when text storage is disabled), so it is not hardcoded as always-allowed here.
|
||||
if enabled_tools is not None and tool_name not in enabled_tools and tool_name != "done":
|
||||
return {"error": f"Tool '{tool_name}' is not available. Use only the tools provided to you."}
|
||||
|
||||
if tool_name == "search_mental_models":
|
||||
|
||||
@@ -604,16 +604,44 @@ 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, llm_output_language: str | None = None) -> str:
|
||||
# The final synthesis is a SEPARATE LLM call with its own system prompt — the
|
||||
# agent/reasoning system prompt (which carries directives and the language rule)
|
||||
# is NOT in scope here. So this default language rule, and the directives, must
|
||||
# be repeated for the answer-writing model. Without it, weaker models drift to
|
||||
# English even when the question/facts are in another language or a directive
|
||||
# demands a specific one (the cause of flaky multilingual reflect tests).
|
||||
_FINAL_LANGUAGE_RULE = (
|
||||
"## LANGUAGE\n"
|
||||
"- Respond in the SAME language as the user's question "
|
||||
"(e.g. a question in Chinese gets a Chinese answer; Japanese → Japanese).\n"
|
||||
"- If a directive above specifies a response language, follow the directive — "
|
||||
"it takes precedence over this default."
|
||||
)
|
||||
|
||||
|
||||
def build_final_system_prompt(
|
||||
mission: str | None = None,
|
||||
llm_output_language: str | None = None,
|
||||
directives: list[dict[str, Any]] | 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.
|
||||
``directives`` are re-injected here (they live in the agent/reasoning prompt,
|
||||
but the final answer is a separate call) so output-constraining rules — most
|
||||
visibly response language — are honoured by the model that actually writes
|
||||
the answer. When ``llm_output_language`` is set it forces that language
|
||||
regardless of the query/source/directive language (config override wins).
|
||||
"""
|
||||
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)
|
||||
|
||||
parts = [build_directives_section(directives) if directives else ""]
|
||||
parts.append(_FINAL_SYSTEM_PROMPT_BASE.format(role_section=role_section))
|
||||
parts.append(_FINAL_LANGUAGE_RULE)
|
||||
parts.append(build_directives_reminder(directives) if directives else "")
|
||||
|
||||
return "\n\n".join(p.strip() for p in parts if p.strip()) + output_language_directive(llm_output_language)
|
||||
|
||||
|
||||
# Backward-compatible constant for non-identity missions
|
||||
|
||||
@@ -232,6 +232,7 @@ def get_reflect_tools(
|
||||
include_mental_models: bool = True,
|
||||
include_observations: bool = True,
|
||||
include_recall: bool = True,
|
||||
include_expand: bool = True,
|
||||
) -> list[dict]:
|
||||
"""
|
||||
Get the list of tools for the reflect agent.
|
||||
@@ -247,6 +248,9 @@ def get_reflect_tools(
|
||||
include_mental_models: Whether to include the search_mental_models tool.
|
||||
include_observations: Whether to include the search_observations tool.
|
||||
include_recall: Whether to include the recall tool.
|
||||
include_expand: Whether to include the expand tool. Disabled when raw
|
||||
document/chunk text is not stored, since expand only reads back
|
||||
source text and would return empty results.
|
||||
|
||||
Returns:
|
||||
List of tool definitions in OpenAI format
|
||||
@@ -260,7 +264,8 @@ def get_reflect_tools(
|
||||
if include_recall:
|
||||
tools.append(TOOL_RECALL)
|
||||
|
||||
tools.append(TOOL_EXPAND)
|
||||
if include_expand:
|
||||
tools.append(TOOL_EXPAND)
|
||||
|
||||
# Use directive-aware done tool if directives are present
|
||||
if directive_rules:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -8,6 +8,7 @@ import hashlib
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
from ...config import get_config
|
||||
from ..memory_engine import fq_table
|
||||
from .types import ChunkMetadata
|
||||
|
||||
@@ -88,6 +89,11 @@ async def store_chunks_batch(
|
||||
if not chunks:
|
||||
return {}
|
||||
|
||||
# When document text storage is disabled, persist empty chunk_text (the
|
||||
# column is NOT NULL) while still computing content_hash from the real text
|
||||
# so delta-retain dedup is unaffected.
|
||||
store_text = get_config().store_document_text
|
||||
|
||||
# Prepare chunk data for batch insert
|
||||
chunk_ids = []
|
||||
chunk_texts = []
|
||||
@@ -98,7 +104,7 @@ async def store_chunks_batch(
|
||||
for chunk in chunks:
|
||||
chunk_id = f"{bank_id}_{document_id}_{chunk.chunk_index}"
|
||||
chunk_ids.append(chunk_id)
|
||||
chunk_texts.append(chunk.chunk_text)
|
||||
chunk_texts.append(chunk.chunk_text if store_text else "")
|
||||
chunk_indices.append(chunk.chunk_index)
|
||||
content_hashes.append(compute_chunk_hash(chunk.chunk_text))
|
||||
chunk_id_map[chunk.chunk_index] = chunk_id
|
||||
|
||||
@@ -3,6 +3,7 @@ Embedding generation utilities for memory units.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import logging
|
||||
from typing import Literal, Protocol
|
||||
|
||||
@@ -15,11 +16,23 @@ 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]:
|
||||
@@ -36,10 +49,19 @@ def generate_embedding(
|
||||
"""
|
||||
try:
|
||||
embeddings = _encode_with_input_type(embeddings_backend, [text], input_type)
|
||||
return embeddings[0]
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
def _encode_with_input_type(
|
||||
embeddings_backend: EmbeddingsBackend, texts: list[str], input_type: EmbeddingInputType
|
||||
@@ -68,7 +90,14 @@ async def generate_embeddings_batch(
|
||||
"""
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
embeddings = await loop.run_in_executor(None, _encode_with_input_type, embeddings_backend, texts, input_type)
|
||||
# run_in_executor runs the encode in a worker thread, which does NOT inherit
|
||||
# the caller's contextvars. Capture the current context and run the encode
|
||||
# inside it so context-dependent behavior (e.g. per-bank `user` attribution
|
||||
# read via get_current_bank_id()) survives the thread hop.
|
||||
ctx = contextvars.copy_context()
|
||||
embeddings = await loop.run_in_executor(
|
||||
None, lambda: ctx.run(_encode_with_input_type, embeddings_backend, texts, input_type)
|
||||
)
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to generate batch embeddings: {str(e)}")
|
||||
|
||||
@@ -81,4 +110,7 @@ async def generate_embeddings_batch(
|
||||
"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)
|
||||
]
|
||||
|
||||
@@ -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,
|
||||
@@ -405,22 +406,63 @@ class VerbatimFactExtractionResponse(BaseModel):
|
||||
facts: list[VerbatimExtractedFact] = Field(description="List of metadata entries (one per chunk)")
|
||||
|
||||
|
||||
# Separators for sentence-aware recursive text splitting, ordered most- to
|
||||
# least-preferred. The final "" lets the splitter break mid-word as a last
|
||||
# resort so a chunk can never exceed the size budget.
|
||||
_RECURSIVE_TEXT_SEPARATORS = [
|
||||
"\n\n", # Paragraph breaks
|
||||
"\n", # Line breaks
|
||||
". ", # Sentence endings
|
||||
"! ", # Exclamations
|
||||
"? ", # Questions
|
||||
"; ", # Semicolons
|
||||
", ", # Commas
|
||||
" ", # Words
|
||||
"", # Characters (last resort)
|
||||
]
|
||||
|
||||
# A single structured unit (a JSONL line or a conversation turn) is kept whole
|
||||
# even when it overflows the budget — but only up to this multiple. Beyond it,
|
||||
# the unit is split as text rather than handed to the LLM wildly over budget
|
||||
# (the extractor has no second re-chunk pass; an oversized chunk just errors).
|
||||
_CHUNK_OVERFLOW_FACTOR = 1.5
|
||||
|
||||
|
||||
def _split_oversized_unit(text: str, max_chars: int) -> list[str]:
|
||||
"""Sentence-aware split of a single unit that overflowed the budget.
|
||||
|
||||
Used when one JSONL line / conversation turn is so large it can't be kept
|
||||
whole within ``_CHUNK_OVERFLOW_FACTOR``. The resulting fragments are no
|
||||
longer valid JSON, but the fact extractor treats every chunk as plain text.
|
||||
"""
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
|
||||
splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=max_chars,
|
||||
chunk_overlap=0,
|
||||
length_function=len,
|
||||
is_separator_regex=False,
|
||||
separators=_RECURSIVE_TEXT_SEPARATORS,
|
||||
)
|
||||
return splitter.split_text(text)
|
||||
|
||||
|
||||
def chunk_text(text: str, max_chars: int) -> list[str]:
|
||||
"""
|
||||
Split text into chunks, preserving conversation structure when possible.
|
||||
|
||||
For JSON conversation arrays (user/assistant turns), splits at turn boundaries
|
||||
while preserving speaker context. For plain text, uses sentence-aware splitting.
|
||||
For JSON conversation arrays (user/assistant turns) and JSONL (newline-delimited
|
||||
JSON objects), splits at turn/line boundaries so no object is split across chunks.
|
||||
A single turn/line that overflows is kept whole up to ``_CHUNK_OVERFLOW_FACTOR``×
|
||||
the budget, then split as text. For plain text, uses sentence-aware splitting.
|
||||
|
||||
Args:
|
||||
text: Input text to chunk (plain text or JSON conversation)
|
||||
text: Input text to chunk (plain text, JSON conversation, or JSONL)
|
||||
max_chars: Maximum characters per chunk (default 120k ≈ 30k tokens)
|
||||
|
||||
Returns:
|
||||
List of text chunks, roughly under max_chars
|
||||
"""
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
|
||||
# If text is small enough, return as-is
|
||||
if len(text) <= max_chars:
|
||||
return [text]
|
||||
@@ -434,26 +476,13 @@ def chunk_text(text: str, max_chars: int) -> list[str]:
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
|
||||
# Fall back to sentence-aware text splitting
|
||||
splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=max_chars,
|
||||
chunk_overlap=0,
|
||||
length_function=len,
|
||||
is_separator_regex=False,
|
||||
separators=[
|
||||
"\n\n", # Paragraph breaks
|
||||
"\n", # Line breaks
|
||||
". ", # Sentence endings
|
||||
"! ", # Exclamations
|
||||
"? ", # Questions
|
||||
"; ", # Semicolons
|
||||
", ", # Commas
|
||||
" ", # Words
|
||||
"", # Characters (last resort)
|
||||
],
|
||||
)
|
||||
# Try to parse as JSONL (newline-delimited JSON objects, e.g. session logs)
|
||||
jsonl_chunks = _chunk_jsonl(text, max_chars)
|
||||
if jsonl_chunks is not None:
|
||||
return jsonl_chunks
|
||||
|
||||
return splitter.split_text(text)
|
||||
# Fall back to sentence-aware text splitting
|
||||
return _split_oversized_unit(text, max_chars)
|
||||
|
||||
|
||||
def _chunk_conversation(turns: list[dict], max_chars: int) -> list[str]:
|
||||
@@ -468,32 +497,109 @@ def _chunk_conversation(turns: list[dict], max_chars: int) -> list[str]:
|
||||
List of JSON-serialized chunks, each containing complete turns
|
||||
"""
|
||||
|
||||
overflow_limit = int(max_chars * _CHUNK_OVERFLOW_FACTOR)
|
||||
|
||||
chunks = []
|
||||
current_chunk = []
|
||||
current_size = 2 # Account for "[]"
|
||||
|
||||
def _flush() -> None:
|
||||
nonlocal current_chunk, current_size
|
||||
if current_chunk:
|
||||
chunks.append(json.dumps(current_chunk, ensure_ascii=False))
|
||||
current_chunk = []
|
||||
current_size = 2 # Reset to "[]"
|
||||
|
||||
for turn in turns:
|
||||
# Estimate size of this turn when serialized (with comma separator)
|
||||
turn_json = json.dumps(turn, ensure_ascii=False)
|
||||
turn_size = len(turn_json) + 1 # +1 for comma
|
||||
|
||||
# A turn too large to keep whole even alone: flush, then split it as
|
||||
# text so no chunk runs far over budget (the extractor won't re-chunk).
|
||||
if turn_size > overflow_limit:
|
||||
_flush()
|
||||
chunks.extend(_split_oversized_unit(turn_json, max_chars))
|
||||
continue
|
||||
|
||||
# If adding this turn would exceed limit and we have turns, save current chunk
|
||||
if current_size + turn_size > max_chars and current_chunk:
|
||||
chunks.append(json.dumps(current_chunk, ensure_ascii=False))
|
||||
current_chunk = []
|
||||
current_size = 2 # Reset to "[]"
|
||||
_flush()
|
||||
|
||||
# Add turn to current chunk
|
||||
current_chunk.append(turn)
|
||||
current_size += turn_size
|
||||
|
||||
# Add final chunk if non-empty
|
||||
if current_chunk:
|
||||
chunks.append(json.dumps(current_chunk, ensure_ascii=False))
|
||||
_flush()
|
||||
|
||||
return chunks if chunks else [json.dumps(turns, ensure_ascii=False)]
|
||||
|
||||
|
||||
def _chunk_jsonl(text: str, max_chars: int) -> list[str] | None:
|
||||
"""Chunk newline-delimited JSON (JSONL) at line boundaries.
|
||||
|
||||
Detects JSONL — two or more non-empty lines, each a complete JSON object —
|
||||
and packs whole lines into chunks so no line is split across chunks (multiple
|
||||
short lines may share a chunk). A line that overflows is kept whole up to
|
||||
``_CHUNK_OVERFLOW_FACTOR``× the budget, then split as text. Returns ``None``
|
||||
if the input is not JSONL, so the caller falls back to plain-text splitting.
|
||||
|
||||
Args:
|
||||
text: Input text to inspect/chunk.
|
||||
max_chars: Maximum characters per chunk.
|
||||
|
||||
Returns:
|
||||
List of JSONL chunks (lines joined by newline), or ``None`` if not JSONL.
|
||||
"""
|
||||
lines = [line for line in text.splitlines() if line.strip()]
|
||||
if len(lines) < 2:
|
||||
return None
|
||||
|
||||
for line in lines:
|
||||
try:
|
||||
obj = json.loads(line)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return None
|
||||
if not isinstance(obj, dict):
|
||||
return None
|
||||
|
||||
overflow_limit = int(max_chars * _CHUNK_OVERFLOW_FACTOR)
|
||||
|
||||
chunks: list[str] = []
|
||||
current_chunk: list[str] = []
|
||||
current_size = 0
|
||||
|
||||
def _flush() -> None:
|
||||
nonlocal current_chunk, current_size
|
||||
if current_chunk:
|
||||
chunks.append("\n".join(current_chunk))
|
||||
current_chunk = []
|
||||
current_size = 0
|
||||
|
||||
for line in lines:
|
||||
line_size = len(line) + 1 # +1 for the joining newline
|
||||
|
||||
# A line too large to keep whole even alone: flush, then split it as
|
||||
# text so no chunk runs far over budget (the extractor won't re-chunk).
|
||||
if line_size > overflow_limit:
|
||||
_flush()
|
||||
chunks.extend(_split_oversized_unit(line, max_chars))
|
||||
continue
|
||||
|
||||
# If adding this line would exceed the limit and we have lines, flush.
|
||||
# A line up to overflow_limit is kept whole (a small, bounded overflow).
|
||||
if current_size + line_size > max_chars and current_chunk:
|
||||
_flush()
|
||||
|
||||
current_chunk.append(line)
|
||||
current_size += line_size
|
||||
|
||||
_flush()
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# FACT EXTRACTION PROMPTS
|
||||
# =============================================================================
|
||||
@@ -510,11 +616,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,20 +993,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
|
||||
# Escape braces so user-supplied text survives str.format() on the prompt template.
|
||||
# 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 = 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"{escape_for_prompt(retain_mission)}\n\n"
|
||||
)
|
||||
else:
|
||||
retain_mission_section = ""
|
||||
retain_mission_section = ""
|
||||
|
||||
# Select base prompt based on extraction mode
|
||||
if extraction_mode == "custom":
|
||||
@@ -997,6 +1098,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,
|
||||
@@ -1005,8 +1126,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)
|
||||
@@ -1025,9 +1152,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}
|
||||
@@ -1053,12 +1192,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
|
||||
@@ -1094,8 +1236,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
|
||||
@@ -1116,7 +1288,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",
|
||||
@@ -1128,6 +1300,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
|
||||
@@ -1640,6 +1816,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,
|
||||
@@ -1731,6 +1940,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
|
||||
@@ -1816,6 +2026,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
|
||||
@@ -1824,7 +2035,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
|
||||
@@ -1834,7 +2047,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
|
||||
@@ -1847,7 +2062,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
|
||||
@@ -1862,7 +2079,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
|
||||
@@ -2035,7 +2254,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)
|
||||
@@ -2100,6 +2321,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
|
||||
|
||||
@@ -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,
|
||||
@@ -398,7 +399,12 @@ async def _upsert_document_row(
|
||||
INSERT so that re-ingesting a document (which deletes + inserts the row)
|
||||
keeps the original creation timestamp. ``updated_at`` is always set to
|
||||
``NOW()`` on both INSERT and the ON CONFLICT UPDATE branch.
|
||||
|
||||
When ``store_document_text`` is disabled, the raw source text
|
||||
is dropped and ``original_text`` is stored as NULL. The ``content_hash`` is
|
||||
still computed from the real content so delta-retain dedup is unaffected.
|
||||
"""
|
||||
original_text = combined_content if get_config().store_document_text else None
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, retain_params, tags, created_at, updated_at)
|
||||
@@ -412,7 +418,7 @@ async def _upsert_document_row(
|
||||
""",
|
||||
document_id,
|
||||
bank_id,
|
||||
combined_content,
|
||||
original_text,
|
||||
content_hash,
|
||||
json.dumps(retain_params) if retain_params else None,
|
||||
document_tags or [],
|
||||
|
||||
@@ -574,12 +574,10 @@ async def compute_semantic_links_ann(
|
||||
# the transaction end handles both.
|
||||
rows: list = []
|
||||
async with conn.transaction():
|
||||
# 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.
|
||||
# Transaction-local ANN tuning. The dispatcher only returns GUCs that
|
||||
# are safe to apply at session/transaction scope for the configured
|
||||
# backend. VectorChord probe values are index-shaped, so vchordrq uses
|
||||
# index storage fallback parameters instead of a blanket SET LOCAL.
|
||||
for guc, value in ann_search_tuning_settings(configured_vector_extension(), kind="low_latency"):
|
||||
await conn.execute(f"SET LOCAL {guc} = {value}")
|
||||
|
||||
@@ -599,23 +597,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,
|
||||
@@ -624,7 +634,7 @@ async def compute_semantic_links_ann(
|
||||
logger.debug(f"[ANN] fact_type={fact_type}: {len(ft_rows)} rows in {time_mod.time() - t_query:.3f}s")
|
||||
rows.extend(ft_rows)
|
||||
# Transaction commits here. _ann_seeds is dropped (ON COMMIT DROP).
|
||||
# hnsw.ef_search reverts (SET LOCAL).
|
||||
# Transaction-local ANN tuning reverts (SET LOCAL).
|
||||
|
||||
for row in rows:
|
||||
sim = float(min(1.0, max(0.0, row["similarity"])))
|
||||
|
||||
@@ -11,9 +11,17 @@ import logging
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from ...extensions.memory_defense import (
|
||||
DefenseAction,
|
||||
DefenseDecision,
|
||||
MemoryDefenseExtension,
|
||||
apply_redaction,
|
||||
parse_policy,
|
||||
)
|
||||
from ...worker.stage import set_stage
|
||||
from ..db.base import DatabaseBackend
|
||||
from ..db_utils import acquire_with_retry
|
||||
@@ -21,11 +29,120 @@ from ..memory_engine import count_tokens, fq_table
|
||||
from . import bank_utils
|
||||
|
||||
|
||||
@dataclass
|
||||
class BlockedViolation:
|
||||
"""One item blocked by the Memory Defense policy (surfaced in the 422 body)."""
|
||||
|
||||
index: int
|
||||
detector: str | None
|
||||
message: str
|
||||
|
||||
|
||||
class MemoryDefenseAllBlockedError(Exception):
|
||||
"""Raised when every item in a retain batch is blocked by the Memory Defense policy."""
|
||||
|
||||
def __init__(self, violations: list[BlockedViolation]) -> None:
|
||||
self.violations = violations
|
||||
super().__init__(f"all {len(violations)} items blocked by Memory Defense policy")
|
||||
|
||||
|
||||
def utcnow():
|
||||
"""Get current UTC time."""
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
def _redact_document_body(body: str, config: Any) -> str:
|
||||
"""Apply Memory Defense redaction to a document body.
|
||||
|
||||
Per-item screening only scrubs the chunked content that goes through
|
||||
`screen()`. When a sub-batch carries `document_body_override` (the full
|
||||
original text of an oversized item — see `_split_contents_into_sub_batches`),
|
||||
that override bypasses screening and would persist verbatim into
|
||||
`documents.original_text`. Apply the same redactor here so the document
|
||||
body is scrubbed regardless of which path produced it.
|
||||
"""
|
||||
try:
|
||||
policy = parse_policy(getattr(config, "memory_defense", None))
|
||||
except Exception:
|
||||
return body
|
||||
if not policy.enabled:
|
||||
return body
|
||||
if not any(r.on == "sensitive_data" for r in policy.rules):
|
||||
return body
|
||||
return apply_redaction(body).content
|
||||
|
||||
|
||||
async def _fire_memory_defense_webhook(
|
||||
webhook_manager: Any,
|
||||
*,
|
||||
conn: Any,
|
||||
schema: str | None,
|
||||
bank_id: str,
|
||||
operation_id: str | None,
|
||||
document_id: str | None,
|
||||
decision: DefenseDecision,
|
||||
) -> None:
|
||||
"""Fire a memory_defense.triggered webhook for a non-allow decision.
|
||||
|
||||
No-op when no webhook manager is wired or none is subscribed. Delivery
|
||||
failures are swallowed so screening never blocks a retain.
|
||||
"""
|
||||
if webhook_manager is None:
|
||||
return
|
||||
try:
|
||||
from ...webhooks import MemoryDefenseEventData, WebhookEvent, WebhookEventType
|
||||
|
||||
event = WebhookEvent(
|
||||
event=WebhookEventType.MEMORY_DEFENSE_TRIGGERED,
|
||||
bank_id=bank_id,
|
||||
operation_id=operation_id or "",
|
||||
status=decision.action.value,
|
||||
timestamp=utcnow(),
|
||||
data=MemoryDefenseEventData(
|
||||
action=decision.action.value,
|
||||
detector=decision.detector,
|
||||
document_id=document_id,
|
||||
matched_types=decision.matched_types or None,
|
||||
message=decision.message or None,
|
||||
),
|
||||
)
|
||||
await webhook_manager.fire_event_with_conn(event, conn, schema=schema)
|
||||
except Exception:
|
||||
logger.warning("memory_defense webhook delivery failed", exc_info=True)
|
||||
|
||||
|
||||
def _audit_memory_defense(
|
||||
audit_logger: Any,
|
||||
*,
|
||||
bank_id: str,
|
||||
document_id: str | None,
|
||||
decision: DefenseDecision,
|
||||
) -> None:
|
||||
"""Write a fire-and-forget ``memory_defense`` audit entry for a non-allow decision.
|
||||
|
||||
No-op when audit logging is disabled (the logger gates on its own config).
|
||||
The action taken (redact/block) and what matched live in the entry metadata.
|
||||
"""
|
||||
if audit_logger is None:
|
||||
return
|
||||
from ..audit import AuditEntry
|
||||
|
||||
entry = AuditEntry(
|
||||
action="memory_defense",
|
||||
transport="system",
|
||||
bank_id=bank_id,
|
||||
metadata={
|
||||
"action": decision.action.value,
|
||||
"detector": decision.detector,
|
||||
"document_id": document_id,
|
||||
"matched_types": decision.matched_types,
|
||||
"message": decision.message,
|
||||
},
|
||||
)
|
||||
entry.ended_at = entry.started_at # point-in-time policy decision (duration 0)
|
||||
audit_logger.log_fire_and_forget(entry)
|
||||
|
||||
|
||||
def _merge_processed_content_tokens(a: int | None, b: int | None) -> int | None:
|
||||
"""Combine the processed-content-tokens signal across sub-results.
|
||||
|
||||
@@ -111,6 +228,25 @@ RetainOutboxCallback = Callable[[asyncpg.Connection], Awaitable[None]]
|
||||
RetainOutboxCallbackFactory = Callable[[list[RetainContentDict]], RetainOutboxCallback | None]
|
||||
|
||||
|
||||
def _resolve_narrator(profile_name: str, bank_id: str) -> str | None:
|
||||
"""Resolve the narrator (memory owner) used to prime fact extraction.
|
||||
|
||||
The narrator is injected as a "Narrator: {name}" line in fact extraction and
|
||||
is stamped into the who-dimension of every first-person fact — and the
|
||||
observations later consolidated from those facts. That is correct for a named
|
||||
agent retaining its own logs, but harmful when ``name`` is just the bank_id:
|
||||
on auto-create the bank ``name`` defaults to ``bank_id``, which is typically a
|
||||
routing key (e.g. ``my-agent::channel-456::user-789``), not a speaker. Priming
|
||||
extraction with a routing key embeds that string into stored fact text and
|
||||
pollutes downstream observations (issue #1680). Suppress it in that case.
|
||||
|
||||
Returns the narrator name, or ``None`` to omit the Narrator line entirely.
|
||||
"""
|
||||
if profile_name == bank_id:
|
||||
return None
|
||||
return profile_name
|
||||
|
||||
|
||||
def _build_retain_params(contents_dicts, document_tags=None, doc_contents=None):
|
||||
"""Build retain_params and merged_tags from content dicts."""
|
||||
if doc_contents is not None:
|
||||
@@ -404,6 +540,10 @@ async def retain_batch(
|
||||
db_semaphore: "asyncio.Semaphore | None" = None,
|
||||
document_body_override: str | None = None,
|
||||
chunk_index_offset: int = 0,
|
||||
progress_callback: "Callable[..., Awaitable[None]] | None" = None,
|
||||
webhook_manager: Any = None,
|
||||
memory_defense_extension: "MemoryDefenseExtension | None" = None,
|
||||
audit_logger: Any = None,
|
||||
) -> tuple[list[list[str]], TokenUsage, int | None]:
|
||||
"""
|
||||
Process a batch of content through the retain pipeline.
|
||||
@@ -439,7 +579,9 @@ async def retain_batch(
|
||||
|
||||
# Get bank profile
|
||||
profile = await bank_utils.get_bank_profile(pool, bank_id)
|
||||
agent_name = profile["name"]
|
||||
# Suppress the narrator when name == bank_id (auto-create default) — see
|
||||
# _resolve_narrator for why a routing-key narrator pollutes extraction (#1680).
|
||||
agent_name = _resolve_narrator(profile["name"], bank_id)
|
||||
|
||||
# Convert dicts to RetainContent objects
|
||||
contents = _build_contents(contents_dicts, document_tags)
|
||||
@@ -492,6 +634,10 @@ async def retain_batch(
|
||||
db_semaphore=db_semaphore,
|
||||
document_body_override=document_body_override,
|
||||
chunk_index_offset=chunk_index_offset,
|
||||
progress_callback=progress_callback,
|
||||
webhook_manager=webhook_manager,
|
||||
memory_defense_extension=memory_defense_extension,
|
||||
audit_logger=audit_logger,
|
||||
)
|
||||
for group_idx, orig_idx in enumerate(original_indices[doc_key]):
|
||||
if group_idx < len(group_ids):
|
||||
@@ -500,6 +646,80 @@ async def retain_batch(
|
||||
total_processed_tokens = _merge_processed_content_tokens(total_processed_tokens, group_processed)
|
||||
return result_unit_ids, total_usage, total_processed_tokens
|
||||
|
||||
# --- Memory Defense pre-extraction screening ---
|
||||
# Delegate to the loaded extension. `config` is a resolved HindsightConfig
|
||||
# object at this point (see _retain_batch_async_internal). On a non-allow
|
||||
# decision we redact in place or drop the item, and fire a
|
||||
# memory_defense.triggered webhook when one is configured.
|
||||
_policy = parse_policy(getattr(config, "memory_defense", None))
|
||||
_blocked_violations: list[BlockedViolation] = []
|
||||
|
||||
if memory_defense_extension is not None and _policy.enabled:
|
||||
async with acquire_with_retry(pool) as _defense_conn:
|
||||
for _idx, _content in enumerate(contents):
|
||||
# Prefer the per-item document_id over the batch-level value so
|
||||
# the decision and webhook carry the document the caller
|
||||
# submitted, not whichever doc_id the batch happens to share.
|
||||
_item_doc_id = contents_dicts[_idx].get("document_id") or document_id
|
||||
|
||||
_decision = await memory_defense_extension.screen(
|
||||
policy=_policy,
|
||||
bank_id=bank_id,
|
||||
document_id=_item_doc_id,
|
||||
content=_content.content,
|
||||
tags=_content.tags,
|
||||
)
|
||||
|
||||
if _decision.action is DefenseAction.ALLOW:
|
||||
continue
|
||||
|
||||
if _decision.action is DefenseAction.REDACT:
|
||||
_redacted = _decision.redacted_content or _content.content
|
||||
_content.content = _redacted
|
||||
# Mirror the redaction into the raw dict so the document
|
||||
# body persisted further down the pipeline also stores the
|
||||
# redacted text, not the verbatim secret.
|
||||
contents_dicts[_idx]["content"] = _redacted
|
||||
elif _decision.action is DefenseAction.BLOCK:
|
||||
_blocked_violations.append(
|
||||
BlockedViolation(
|
||||
index=_idx,
|
||||
detector=_decision.detector,
|
||||
message=_decision.message,
|
||||
)
|
||||
)
|
||||
|
||||
await _fire_memory_defense_webhook(
|
||||
webhook_manager,
|
||||
conn=_defense_conn,
|
||||
schema=schema,
|
||||
bank_id=bank_id,
|
||||
operation_id=operation_id,
|
||||
document_id=_item_doc_id,
|
||||
decision=_decision,
|
||||
)
|
||||
_audit_memory_defense(
|
||||
audit_logger,
|
||||
bank_id=bank_id,
|
||||
document_id=_item_doc_id,
|
||||
decision=_decision,
|
||||
)
|
||||
|
||||
if _blocked_violations:
|
||||
# All items blocked → raise so the HTTP layer can return 422.
|
||||
if len(_blocked_violations) == len(contents):
|
||||
raise MemoryDefenseAllBlockedError(_blocked_violations)
|
||||
|
||||
# Remove blocked items from the pipeline.
|
||||
_skip_indices = {v.index for v in _blocked_violations}
|
||||
if _skip_indices:
|
||||
_surviving = [i for i in range(len(contents)) if i not in _skip_indices]
|
||||
contents = [contents[i] for i in _surviving]
|
||||
contents_dicts = [contents_dicts[i] for i in _surviving]
|
||||
# If nothing survives, return empty results immediately.
|
||||
if not contents:
|
||||
return [[] for _ in contents_dicts], TokenUsage(), 0
|
||||
|
||||
# Resolve effective document_id early so both delta and streaming paths
|
||||
# can find existing chunks from a prior attempt. On retry, a generated
|
||||
# document_id is recovered from operation result_metadata.document_ids[0].
|
||||
@@ -692,6 +912,7 @@ async def retain_batch(
|
||||
db_semaphore=db_semaphore,
|
||||
document_body_override=document_body_override,
|
||||
chunk_index_offset=chunk_index_offset,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
|
||||
@@ -831,6 +1052,7 @@ async def _streaming_retain_batch(
|
||||
db_semaphore: "asyncio.Semaphore | None" = None,
|
||||
document_body_override: str | None = None,
|
||||
chunk_index_offset: int = 0,
|
||||
progress_callback: "Callable[..., Awaitable[None]] | None" = None,
|
||||
) -> tuple[list[list[str]], TokenUsage]:
|
||||
"""
|
||||
Process a large document in streaming mini-batches to bound memory usage.
|
||||
@@ -870,7 +1092,9 @@ async def _streaming_retain_batch(
|
||||
# so documents.original_text stores the complete payload, not just this
|
||||
# slice (issue #1838).
|
||||
if document_body_override is not None:
|
||||
combined_content = document_body_override
|
||||
# The override is the unmodified original body — apply redaction so
|
||||
# secrets in oversized inputs don't bypass screening.
|
||||
combined_content = _redact_document_body(document_body_override, config)
|
||||
else:
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
# Memory: contents_dicts content strings are now captured in combined_content.
|
||||
@@ -952,19 +1176,29 @@ async def _streaming_retain_batch(
|
||||
tags=source.tags,
|
||||
observation_scopes=source.observation_scopes,
|
||||
)
|
||||
extracted, processed, chunk_meta, usage = await _extract_and_embed(
|
||||
[content],
|
||||
llm_config,
|
||||
agent_name,
|
||||
config,
|
||||
embeddings_model,
|
||||
format_date_fn,
|
||||
fact_type_override,
|
||||
log_buffer,
|
||||
pool,
|
||||
operation_id,
|
||||
schema,
|
||||
)
|
||||
# Attribute this chunk's extraction LLM call to its document, so the
|
||||
# trace row carries document_id (a document accrues one such trace
|
||||
# per retain/re-retain). Per-call: the operation-level trace context
|
||||
# is shared across a batch's documents.
|
||||
from ..llm_trace import reset_call_metadata, set_call_metadata
|
||||
|
||||
meta_token = set_call_metadata({"document_id": effective_doc_id})
|
||||
try:
|
||||
extracted, processed, chunk_meta, usage = await _extract_and_embed(
|
||||
[content],
|
||||
llm_config,
|
||||
agent_name,
|
||||
config,
|
||||
embeddings_model,
|
||||
format_date_fn,
|
||||
fact_type_override,
|
||||
log_buffer,
|
||||
pool,
|
||||
operation_id,
|
||||
schema,
|
||||
)
|
||||
finally:
|
||||
reset_call_metadata(meta_token)
|
||||
await chunk_queue.put((global_idx, content, extracted, processed, chunk_meta, usage))
|
||||
# Memory: release the chunk text from the shared list now that it's
|
||||
# been extracted and queued. The queued RetainContent holds its own copy.
|
||||
@@ -999,6 +1233,25 @@ async def _streaming_retain_batch(
|
||||
async def _db_consumer() -> None:
|
||||
batch: list[tuple] = []
|
||||
consumer_batch_idx = 0
|
||||
chunks_committed = 0
|
||||
|
||||
# Best-effort durable progress: how many chunks of this document have been
|
||||
# extracted+committed so far. Written per consumer batch so an operator polling
|
||||
# the retain operation sees "storing 200/1200 chunks" advancing instead of a
|
||||
# single opaque sub-batch tick. Never lets a heartbeat failure break retain.
|
||||
async def _emit_chunk_progress() -> None:
|
||||
if not (progress_callback and operation_id):
|
||||
return
|
||||
try:
|
||||
await progress_callback(
|
||||
operation_id,
|
||||
stage="storing",
|
||||
processed=chunks_committed,
|
||||
total=total_chunks,
|
||||
detail={"facts_committed": len(all_unit_ids)},
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("retain chunk-progress write failed", exc_info=True)
|
||||
|
||||
while True:
|
||||
item = await chunk_queue.get()
|
||||
@@ -1010,6 +1263,8 @@ async def _streaming_retain_batch(
|
||||
consumer_batch_idx,
|
||||
is_last=True,
|
||||
)
|
||||
chunks_committed += len(batch)
|
||||
await _emit_chunk_progress()
|
||||
break
|
||||
|
||||
batch.append(item)
|
||||
@@ -1029,6 +1284,8 @@ async def _streaming_retain_batch(
|
||||
is_last=False,
|
||||
)
|
||||
consumer_batch_idx += 1
|
||||
chunks_committed += len(batch)
|
||||
await _emit_chunk_progress()
|
||||
batch = []
|
||||
|
||||
async def _process_db_batch(
|
||||
@@ -1181,20 +1438,17 @@ async def _streaming_retain_batch(
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# --- Document ownership gate ---
|
||||
# Lock the document row to serialize all concurrent writers.
|
||||
# SELECT ... FOR UPDATE doesn't lock non-existent rows, so we
|
||||
# first ensure the row exists with a lightweight upsert, THEN lock it.
|
||||
# The content_hash='__pending__' placeholder is immediately overwritten
|
||||
# by handle_document_tracking or upsert_document_metadata below.
|
||||
await conn.execute(
|
||||
f"INSERT INTO {fq_table('documents')} (id, bank_id, original_text, content_hash) "
|
||||
f"VALUES ($1, $2, '', '__pending__') "
|
||||
f"ON CONFLICT (id, bank_id) DO NOTHING",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
existing_hash = await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
# Ensure the document row exists, lock it to serialize all
|
||||
# concurrent same-document writers, and read its pre-existing
|
||||
# hash. The lock prevents interleaved retains from corrupting
|
||||
# each other in handle_document_tracking; the returned hash
|
||||
# ('__pending__' for a freshly inserted row) drives the
|
||||
# takeover check for later batches below. The PG/Oracle split
|
||||
# lives in the ops layer because Oracle can't do this upsert +
|
||||
# RETURNING in a single statement.
|
||||
existing_hash = await pool.ops.lock_document_for_write(
|
||||
conn,
|
||||
fq_table("documents"),
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
@@ -1501,6 +1755,35 @@ async def _streaming_retain_batch(
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ChunkDiff:
|
||||
"""Classification of chunk indices when diffing new content vs stored chunks."""
|
||||
|
||||
unchanged: list[int]
|
||||
changed: list[int]
|
||||
new: list[int]
|
||||
removed: list[int]
|
||||
|
||||
|
||||
def _classify_chunk_diff(existing_by_index: dict[int, Any], new_hashes: dict[int, str]) -> _ChunkDiff:
|
||||
"""Classify chunk indices by comparing freshly computed ``new_hashes``
|
||||
(index -> content hash) against the currently stored chunks
|
||||
(``existing_by_index``: index -> chunk row)."""
|
||||
diff = _ChunkDiff(unchanged=[], changed=[], new=[], removed=[])
|
||||
for idx, new_hash in new_hashes.items():
|
||||
existing = existing_by_index.get(idx)
|
||||
if existing and existing.content_hash == new_hash:
|
||||
diff.unchanged.append(idx)
|
||||
elif existing:
|
||||
diff.changed.append(idx)
|
||||
else:
|
||||
diff.new.append(idx)
|
||||
for idx in existing_by_index:
|
||||
if idx not in new_hashes:
|
||||
diff.removed.append(idx)
|
||||
return diff
|
||||
|
||||
|
||||
async def _try_delta_retain(
|
||||
pool: Any,
|
||||
embeddings_model,
|
||||
@@ -1570,18 +1853,11 @@ async def _try_delta_retain(
|
||||
existing_by_index = {c.chunk_index: c for c in existing_chunks}
|
||||
new_hashes = {idx: chunk_storage.compute_chunk_hash(text) for idx, text in new_chunks_with_contents.items()}
|
||||
|
||||
unchanged_indices, changed_indices, new_indices, removed_indices = [], [], [], []
|
||||
for idx, new_hash in new_hashes.items():
|
||||
existing = existing_by_index.get(idx)
|
||||
if existing and existing.content_hash == new_hash:
|
||||
unchanged_indices.append(idx)
|
||||
elif existing:
|
||||
changed_indices.append(idx)
|
||||
else:
|
||||
new_indices.append(idx)
|
||||
for idx in existing_by_index:
|
||||
if idx not in new_hashes:
|
||||
removed_indices.append(idx)
|
||||
diff = _classify_chunk_diff(existing_by_index, new_hashes)
|
||||
unchanged_indices = diff.unchanged
|
||||
changed_indices = diff.changed
|
||||
new_indices = diff.new
|
||||
removed_indices = diff.removed
|
||||
|
||||
log_buffer.append(
|
||||
f"[delta] Chunk diff: {len(unchanged_indices)} unchanged, "
|
||||
@@ -1609,6 +1885,7 @@ async def _try_delta_retain(
|
||||
start_time,
|
||||
outbox_callback,
|
||||
document_body_override=document_body_override,
|
||||
config=config,
|
||||
)
|
||||
|
||||
# Build content items for only the changed/new chunks
|
||||
@@ -1626,22 +1903,89 @@ async def _try_delta_retain(
|
||||
start_time,
|
||||
outbox_callback,
|
||||
document_body_override=document_body_override,
|
||||
config=config,
|
||||
)
|
||||
|
||||
# Extract facts and generate embeddings (shared pipeline)
|
||||
extracted_facts, processed_facts, new_chunk_metadata, usage = await _extract_and_embed(
|
||||
delta_contents,
|
||||
llm_config,
|
||||
agent_name,
|
||||
config,
|
||||
embeddings_model,
|
||||
format_date_fn,
|
||||
fact_type_override,
|
||||
log_buffer,
|
||||
pool,
|
||||
operation_id,
|
||||
schema,
|
||||
)
|
||||
# Freshness recheck BEFORE the (expensive) LLM extraction.
|
||||
#
|
||||
# We snapshotted the document hash and chunks outside any lock. A concurrent
|
||||
# retain for the same document may have committed a new version while we were
|
||||
# chunking and diffing. Re-read the current hash; if it changed, recompute the
|
||||
# diff against the now-committed chunk state. If the concurrent writer already
|
||||
# produced content identical to ours, there is nothing left to extract — skip
|
||||
# the LLM call entirely (metadata-only). If it still differs, fall back to the
|
||||
# streaming path (which dedups per-chunk and re-locks the document).
|
||||
#
|
||||
# This narrows — but cannot fully close — the race window: a writer can still
|
||||
# commit during our extraction. The post-extraction hash gate inside the write
|
||||
# transaction remains the correctness backstop; this check exists purely to
|
||||
# avoid burning LLM tokens on work a concurrent request already did.
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
recheck_hash = await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
if recheck_hash is not None and doc_hash_at_load is not None and recheck_hash != doc_hash_at_load:
|
||||
log_buffer.append(
|
||||
f"[delta] Document {effective_doc_id} changed before extraction "
|
||||
f"(concurrent retain) — rechecking diff against current state"
|
||||
)
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
current_chunks = await chunk_storage.load_existing_chunks(conn, bank_id, effective_doc_id)
|
||||
if not current_chunks or any(c.content_hash is None for c in current_chunks):
|
||||
log_buffer.append("[delta] Recheck: current chunks unavailable — falling back to full retain")
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
return None
|
||||
current_by_index = {c.chunk_index: c for c in current_chunks}
|
||||
recheck = _classify_chunk_diff(current_by_index, new_hashes)
|
||||
if not (recheck.changed or recheck.new or recheck.removed):
|
||||
log_buffer.append(
|
||||
"[delta] Recheck: concurrent retain already stored identical content — "
|
||||
"skipping extraction, updating metadata only"
|
||||
)
|
||||
return await _delta_metadata_only(
|
||||
pool,
|
||||
bank_id,
|
||||
contents_dicts,
|
||||
contents,
|
||||
effective_doc_id,
|
||||
document_tags,
|
||||
log_buffer,
|
||||
start_time,
|
||||
outbox_callback,
|
||||
document_body_override=document_body_override,
|
||||
config=config,
|
||||
)
|
||||
log_buffer.append(
|
||||
f"[delta] Recheck: {len(recheck.changed) + len(recheck.new) + len(recheck.removed)} chunks still differ — "
|
||||
f"falling back to full retain"
|
||||
)
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
return None
|
||||
|
||||
# Extract facts and generate embeddings (shared pipeline). Attribute these
|
||||
# extraction calls to the document so the delta re-retain's trace also binds
|
||||
# to it (a document accrues one trace per full/delta retain).
|
||||
from ..llm_trace import reset_call_metadata, set_call_metadata
|
||||
|
||||
meta_token = set_call_metadata({"document_id": effective_doc_id})
|
||||
try:
|
||||
extracted_facts, processed_facts, new_chunk_metadata, usage = await _extract_and_embed(
|
||||
delta_contents,
|
||||
llm_config,
|
||||
agent_name,
|
||||
config,
|
||||
embeddings_model,
|
||||
format_date_fn,
|
||||
fact_type_override,
|
||||
log_buffer,
|
||||
pool,
|
||||
operation_id,
|
||||
schema,
|
||||
)
|
||||
finally:
|
||||
reset_call_metadata(meta_token)
|
||||
|
||||
# Database transaction
|
||||
result_unit_ids: list[list[str]] = []
|
||||
@@ -1687,9 +2031,10 @@ async def _try_delta_retain(
|
||||
step_start = time.time()
|
||||
# When this sub-batch is one slice of an oversized item
|
||||
# split across multiple sub-batches, store the full body
|
||||
# (issue #1838) instead of just the slice.
|
||||
# (issue #1838) instead of just the slice. Redact the
|
||||
# override since it bypassed per-chunk screening.
|
||||
if document_body_override is not None:
|
||||
combined_content = document_body_override
|
||||
combined_content = _redact_document_body(document_body_override, config)
|
||||
else:
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
|
||||
@@ -1818,6 +2163,7 @@ async def _delta_metadata_only(
|
||||
outbox_callback,
|
||||
*,
|
||||
document_body_override: str | None = None,
|
||||
config: Any = None,
|
||||
):
|
||||
"""Handle the case where no chunks changed — just update document metadata and tags."""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
@@ -1830,8 +2176,9 @@ async def _delta_metadata_only(
|
||||
)
|
||||
# When this sub-batch is a slice of an oversized item, write the
|
||||
# full original body (issue #1838) instead of just the slice.
|
||||
# Redact the override since it bypassed per-chunk screening.
|
||||
if document_body_override is not None:
|
||||
combined_content = document_body_override
|
||||
combined_content = _redact_document_body(document_body_override, config)
|
||||
else:
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
|
||||
|
||||
@@ -7,6 +7,27 @@ from typing import Any
|
||||
from .types import MergedCandidate, RetrievalResult
|
||||
|
||||
|
||||
def cap_per_source(results: list[RetrievalResult], cap: int) -> list[RetrievalResult]:
|
||||
"""Truncate a single retrieval arm to its top-``cap`` results.
|
||||
|
||||
Applied per source (semantic, BM25, graph, temporal) before fusion so that
|
||||
one over-expanding backend cannot crowd out the others when the merged pool
|
||||
is later trimmed to the reranker's global candidate budget. The caller is
|
||||
responsible for sorting ``results`` by relevance first; this only slices.
|
||||
|
||||
Args:
|
||||
results: Results for a single source, already sorted best-first.
|
||||
cap: Maximum results to keep. ``0`` (or negative) disables the cap.
|
||||
|
||||
Returns:
|
||||
The original list when the cap is disabled or not exceeded, otherwise a
|
||||
truncated copy of the top ``cap`` results.
|
||||
"""
|
||||
if cap <= 0 or len(results) <= cap:
|
||||
return results
|
||||
return results[:cap]
|
||||
|
||||
|
||||
def reciprocal_rank_fusion(result_lists: list[list[RetrievalResult]], k: int = 60) -> list[MergedCandidate]:
|
||||
"""
|
||||
Merge multiple ranked result lists using Reciprocal Rank Fusion.
|
||||
@@ -77,6 +98,66 @@ def reciprocal_rank_fusion(result_lists: list[list[RetrievalResult]], k: int = 6
|
||||
return merged_results
|
||||
|
||||
|
||||
def interleave_fusion(result_lists: list[list[RetrievalResult]]) -> list[MergedCandidate]:
|
||||
"""Round-robin (interleaved) fusion — an alternative to RRF for dedup-style recall.
|
||||
|
||||
RRF scores a doc by the *sum* of its reciprocal ranks across arms, so a result
|
||||
that is #1 in one arm but absent/low in the others gets averaged down. That is
|
||||
exactly the consolidation-dedup failure mode: the near-identical existing
|
||||
observation (the "twin" to merge into) is semantic rank #1, yet shares no
|
||||
source-fact graph link and little lexical overlap, so RRF drops it below the
|
||||
recall budget cutoff and the LLM never sees it → creates a duplicate.
|
||||
|
||||
Interleave instead *guarantees every arm's top hits a slot*: take each arm's
|
||||
#1, then each arm's #2, … in arm-priority order, de-duplicating, until all
|
||||
results are placed. The arm priority is the order of ``result_lists``
|
||||
(semantic, bm25, graph, temporal), so semantic #1 is always first.
|
||||
|
||||
``rrf_score`` is assigned strictly decreasing by final interleave position so
|
||||
downstream order-by-score sorts preserve the interleave order; ``source_ranks``
|
||||
mirrors the RRF bookkeeping (each doc's rank within every arm it appears in).
|
||||
"""
|
||||
source_names = ["semantic", "bm25", "graph", "temporal"]
|
||||
source_ranks: dict[str, dict[str, int]] = {}
|
||||
all_retrievals: dict[str, RetrievalResult] = {}
|
||||
|
||||
for source_idx, results in enumerate(result_lists):
|
||||
source_name = source_names[source_idx] if source_idx < len(source_names) else f"source_{source_idx}"
|
||||
for rank, retrieval in enumerate(results, start=1):
|
||||
if not isinstance(retrieval, RetrievalResult):
|
||||
raise TypeError(
|
||||
f"Expected RetrievalResult but got {type(retrieval).__name__} in {source_name} results at rank {rank}"
|
||||
)
|
||||
doc_id = retrieval.id
|
||||
all_retrievals.setdefault(doc_id, retrieval)
|
||||
source_ranks.setdefault(doc_id, {})[f"{source_name}_rank"] = rank
|
||||
|
||||
# Round-robin pick across arms in priority order: all #1s, then all #2s, ...
|
||||
ordered_ids: list[str] = []
|
||||
seen: set[str] = set()
|
||||
max_len = max((len(r) for r in result_lists), default=0)
|
||||
for r in range(max_len):
|
||||
for results in result_lists:
|
||||
if r < len(results):
|
||||
doc_id = results[r].id
|
||||
if doc_id not in seen:
|
||||
seen.add(doc_id)
|
||||
ordered_ids.append(doc_id)
|
||||
|
||||
n = len(ordered_ids)
|
||||
return [
|
||||
MergedCandidate(
|
||||
retrieval=all_retrievals[doc_id],
|
||||
# Strictly decreasing by interleave position → sorting desc by rrf_score
|
||||
# reproduces the interleave order downstream.
|
||||
rrf_score=float(n - pos),
|
||||
rrf_rank=pos + 1,
|
||||
source_ranks=source_ranks[doc_id],
|
||||
)
|
||||
for pos, doc_id in enumerate(ordered_ids)
|
||||
]
|
||||
|
||||
|
||||
def normalize_scores_on_deltas(results: list[dict[str, Any]], score_keys: list[str]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Normalize scores based on deltas (min-max normalization within result set).
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
"""Per-strategy recall boosting.
|
||||
|
||||
A deployment can prioritise one retrieval arm (semantic, bm25, graph, temporal)
|
||||
over the others via ``HINDSIGHT_API_RECALL_STRATEGY_BOOSTS``, expressed as a
|
||||
human priority *level* rather than an opaque number — e.g. ``graph:high`` to
|
||||
strongly favour graph hits.
|
||||
|
||||
A level is chosen instead of a raw weight because the boost is applied in two
|
||||
structurally different places that live on different score scales, so a single
|
||||
number could not mean the same thing in both. The level maps to a tuned
|
||||
:class:`BoostWeights` pair:
|
||||
|
||||
1. **Before the reranker cap** — :func:`boosted_rrf_score` uses ``BoostWeights.rrf``
|
||||
as a weighted-RRF multiplier on the boosted arm's rank contribution, so its
|
||||
candidates survive the global reranker candidate budget instead of being
|
||||
trimmed by raw RRF score. Rank-aware: a candidate ranked #1 in the boosted
|
||||
arm is protected more than one ranked #200.
|
||||
|
||||
2. **After the reranker** — :func:`additive_strategy_boost` uses
|
||||
``BoostWeights.additive`` as a flat bump to the final ranking weight (which
|
||||
sits in ~[0, 1] after cross-encoder + recency/temporal scoring), nudging the
|
||||
boosted arm's candidates up the final ordering.
|
||||
|
||||
Both functions are no-ops when ``boosts`` is empty, preserving current behaviour.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .types import MergedCandidate
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BoostWeights:
|
||||
"""Per-stage boost magnitudes for one priority level.
|
||||
|
||||
The two fields live on different scales on purpose (see module docstring):
|
||||
``rrf`` multiplies an arm's ``1/(k+rank)`` RRF contribution; ``additive`` is
|
||||
added directly to the post-rerank weight in ~[0, 1].
|
||||
"""
|
||||
|
||||
rrf: float
|
||||
additive: float
|
||||
|
||||
|
||||
# Priority level -> per-stage boost magnitudes. Tuned against real recall traces
|
||||
# (LoCoMo bank, 336 merged candidates → 300-cap, local ms-marco cross-encoder):
|
||||
#
|
||||
# Stage 1 (rrf, weighted-RRF multiplier on the arm's 1/(k+rank) contribution).
|
||||
# The observed 300-cap boundary RRF score was ~0.0055; a graph-only candidate
|
||||
# falls below it past graph-rank ~120. The multipliers map to that boundary:
|
||||
# low=1.0 doubles the arm's vote — rescues at-risk candidates from the cut
|
||||
# (graph-rank 150: 0.0048 → 0.0095) without reshuffling much.
|
||||
# medium=3.0 promotes them into the middle of the pool (~rank 60).
|
||||
# high=6.0 makes the boosted arm dominate the top of the candidate pool.
|
||||
#
|
||||
# Stage 2 (additive, flat bump to the post-rerank weight in [0, 1]). The local
|
||||
# cross-encoder is sharply bimodal: strong direct matches score 0.5–0.999, while
|
||||
# everything else — including graph hits the CE undervalues, which is exactly
|
||||
# what we boost — collapses near 0. So the additive lifts a ~0 candidate up the
|
||||
# weight scale. Levels are calibrated as relevance thresholds it can outrank:
|
||||
# low=0.05 nudges above the near-0 tail; loses to any real CE match.
|
||||
# medium=0.2 competes with weak/moderate matches.
|
||||
# high=0.5 wins over most semantic matches (honouring "prioritise graph over
|
||||
# semantic"); only a strong direct match (>0.5 normalized) still wins.
|
||||
#
|
||||
# The keys are the user-facing contract; config.py validates env input against
|
||||
# them (kept in sync by a guard test).
|
||||
BOOST_LEVELS: dict[str, BoostWeights] = {
|
||||
"low": BoostWeights(rrf=1.0, additive=0.05),
|
||||
"medium": BoostWeights(rrf=3.0, additive=0.2),
|
||||
"high": BoostWeights(rrf=6.0, additive=0.5),
|
||||
}
|
||||
|
||||
|
||||
def boosted_rrf_score(candidate: MergedCandidate, boosts: dict[str, str], k: int = 60) -> float:
|
||||
"""Return ``candidate``'s RRF score plus a weighted-RRF boost delta.
|
||||
|
||||
For each boosted arm the candidate appeared in, adds ``level.rrf * 1/(k+rank)``
|
||||
— i.e. scales that arm's RRF contribution by the level's multiplier. Staying
|
||||
in RRF units keeps the boost comparable to the base score and rank-aware.
|
||||
|
||||
Args:
|
||||
candidate: Merged candidate carrying ``rrf_score`` and ``source_ranks``.
|
||||
boosts: Map of strategy name -> priority level. Empty means no boost.
|
||||
k: RRF constant; must match the value used during fusion.
|
||||
|
||||
Returns:
|
||||
The (possibly) boosted score to sort by. Equal to ``rrf_score`` when no
|
||||
boosted arm surfaced this candidate.
|
||||
"""
|
||||
if not boosts:
|
||||
return candidate.rrf_score
|
||||
delta = 0.0
|
||||
for strategy, level in boosts.items():
|
||||
rank = candidate.source_ranks.get(f"{strategy}_rank")
|
||||
if rank is not None:
|
||||
delta += BOOST_LEVELS[level].rrf * (1.0 / (k + rank))
|
||||
return candidate.rrf_score + delta
|
||||
|
||||
|
||||
def additive_strategy_boost(source_ranks: dict[str, int], boosts: dict[str, str]) -> float:
|
||||
"""Return the flat additive boost for a candidate given its source ranks.
|
||||
|
||||
Sums the ``additive`` magnitude of every boosted arm that surfaced the
|
||||
candidate. Flat by design: the bump does not depend on the candidate's rank
|
||||
within the arm, matching the post-rerank "additive boost" semantics.
|
||||
|
||||
Args:
|
||||
source_ranks: ``{"graph_rank": 3, "semantic_rank": 50, ...}`` from RRF.
|
||||
boosts: Map of strategy name -> priority level. Empty means no boost.
|
||||
|
||||
Returns:
|
||||
The additive boost (0.0 when no boosted arm surfaced this candidate).
|
||||
"""
|
||||
if not boosts:
|
||||
return 0.0
|
||||
return sum(BOOST_LEVELS[level].additive for strategy, level in boosts.items() if f"{strategy}_rank" in source_ranks)
|
||||
@@ -160,13 +160,29 @@ class CrossEncoderReranker:
|
||||
|
||||
import asyncio
|
||||
|
||||
from hindsight_api.config import ENV_MODEL_INIT_TIMEOUT, get_config
|
||||
|
||||
cross_encoder = self.cross_encoder
|
||||
# For local providers, run in thread pool to avoid blocking event loop
|
||||
if cross_encoder.provider_name == "local":
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, lambda: asyncio.run(cross_encoder.initialize()))
|
||||
init = loop.run_in_executor(None, lambda: asyncio.run(cross_encoder.initialize()))
|
||||
else:
|
||||
await cross_encoder.initialize()
|
||||
init = cross_encoder.initialize()
|
||||
|
||||
# Cap lazy init with the same wall-clock timeout used at startup so a
|
||||
# hung model download surfaces as a clear error on the request that
|
||||
# triggered it, rather than hanging the caller forever.
|
||||
init_timeout = get_config().model_init_timeout
|
||||
try:
|
||||
await asyncio.wait_for(init, timeout=init_timeout)
|
||||
except TimeoutError as e:
|
||||
raise RuntimeError(
|
||||
f"Cross-encoder initialization did not complete within {init_timeout:g}s. "
|
||||
f"The reranker model is likely blocked loading — e.g. an offline model "
|
||||
f"download. Increase {ENV_MODEL_INIT_TIMEOUT} if the first-time download "
|
||||
f"legitimately needs more time."
|
||||
) from e
|
||||
self._initialized = True
|
||||
|
||||
async def rerank(self, query: str, candidates: list[MergedCandidate]) -> list[ScoredResult]:
|
||||
|
||||
@@ -137,6 +137,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
"""
|
||||
result_dict: dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]] = {ft: ([], []) for ft in fact_types}
|
||||
|
||||
config = get_config()
|
||||
tokens = tokenize_query(query_text)
|
||||
|
||||
# Over-fetch for HNSW approximation; semantic results trimmed to limit in Python.
|
||||
@@ -148,8 +149,6 @@ async def retrieve_semantic_bm25_combined(
|
||||
)
|
||||
table = fq_table("memory_units")
|
||||
|
||||
config = get_config()
|
||||
|
||||
# Use the SQL dialect to build backend-specific query arms, avoiding
|
||||
# inline if/else branches for each database.
|
||||
# Use getattr for backward compat: raw asyncpg connections (used in some
|
||||
@@ -201,6 +200,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
embedding_param="$1",
|
||||
bank_id_param="$2",
|
||||
fetch_limit=hnsw_fetch,
|
||||
min_similarity=config.semantic_min_similarity,
|
||||
tags_clause=tags_clause,
|
||||
groups_clause=groups_clause,
|
||||
extra_where=created_range_clause,
|
||||
@@ -226,6 +226,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
arm_index=i,
|
||||
text_search_extension=text_ext,
|
||||
bm25_language=config.text_search_extension_native_language,
|
||||
bm25_min_score=config.bm25_min_score,
|
||||
extra_where=created_range_clause,
|
||||
)
|
||||
)
|
||||
@@ -273,6 +274,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
embedding_param="$1",
|
||||
bank_id_param="$2",
|
||||
fetch_limit=hnsw_fetch,
|
||||
min_similarity=config.semantic_min_similarity,
|
||||
tags_clause=fb_tags_clause,
|
||||
groups_clause=fb_groups_clause,
|
||||
extra_where=fb_created_clause,
|
||||
@@ -307,6 +309,66 @@ async def retrieve_semantic_bm25_combined(
|
||||
return result_dict
|
||||
|
||||
|
||||
# Temporal entry-point selection tuning.
|
||||
_TEMPORAL_POOL_SIZE = 60 # ANN candidates fetched per fact_type before coverage selection
|
||||
_TEMPORAL_ENTRY_POINTS = 10 # entry points kept per fact_type after coverage selection
|
||||
_TEMPORAL_COVERAGE_BUCKETS = 8 # time-buckets the window is divided into for coverage
|
||||
|
||||
|
||||
def _coalesce_date(row: Any) -> datetime | None:
|
||||
"""The unit's effective time — matches COALESCE(occurred_start, mentioned_at, occurred_end)."""
|
||||
return row["occurred_start"] or row["mentioned_at"] or row["occurred_end"]
|
||||
|
||||
|
||||
def _select_with_temporal_coverage(
|
||||
pool: list,
|
||||
start_date: datetime,
|
||||
end_date: datetime,
|
||||
limit: int,
|
||||
n_buckets: int,
|
||||
) -> list:
|
||||
"""Pick `limit` entry points from a similarity-ranked pool, spread across the window.
|
||||
|
||||
The window [start_date, end_date] is split into `n_buckets` equal time-buckets.
|
||||
Candidates are taken round-robin across the buckets that contain them — the
|
||||
best-similarity item from each populated bucket first, then the second-best from each,
|
||||
and so on — so every populated slice of the window is represented before any slice
|
||||
contributes a second item. Within a tier, higher-similarity items lead. When the
|
||||
in-window dates are degenerate (all in one bucket — e.g. a batch stamped with a single
|
||||
date) this collapses to plain similarity order.
|
||||
"""
|
||||
if len(pool) <= limit:
|
||||
return list(pool)
|
||||
|
||||
ranked = sorted(pool, key=lambda r: r["similarity"], reverse=True)
|
||||
span = (end_date - start_date).total_seconds()
|
||||
|
||||
def _bucket(row: Any) -> int:
|
||||
d = _coalesce_date(row)
|
||||
if d is None or span <= 0:
|
||||
return 0
|
||||
if d.tzinfo is None:
|
||||
d = d.replace(tzinfo=UTC)
|
||||
frac = (d - start_date).total_seconds() / span
|
||||
return max(0, min(int(frac * n_buckets), n_buckets - 1))
|
||||
|
||||
buckets: dict[int, list] = {}
|
||||
for row in ranked: # ranked is similarity-desc, so each bucket list inherits that order
|
||||
buckets.setdefault(_bucket(row), []).append(row)
|
||||
|
||||
selected: list = []
|
||||
tier = 0
|
||||
while len(selected) < limit and any(len(b) > tier for b in buckets.values()):
|
||||
# The tier-th best item from every bucket that still has one, strongest first.
|
||||
tier_rows = [b[tier] for b in buckets.values() if len(b) > tier]
|
||||
tier_rows.sort(key=lambda r: r["similarity"], reverse=True)
|
||||
for row in tier_rows:
|
||||
if len(selected) < limit:
|
||||
selected.append(row)
|
||||
tier += 1
|
||||
return selected
|
||||
|
||||
|
||||
async def retrieve_temporal_combined(
|
||||
conn,
|
||||
query_emb_str: str,
|
||||
@@ -350,9 +412,12 @@ async def retrieve_temporal_combined(
|
||||
end_date = end_date.replace(tzinfo=UTC)
|
||||
|
||||
# Build tags clause
|
||||
# Entry point query: fixed params are $1-$6, tags at $7
|
||||
tags_clause = build_tags_where_clause_simple(tags, 7, match=tags_match)
|
||||
tag_groups_param_start = 7 + (1 if tags else 0)
|
||||
# Entry-point query: fixed params are $1-$5 (emb, bank, start, end, threshold), tags at $6.
|
||||
# fact_type is inlined as a literal per UNION ALL arm (not a bind) — this avoids `unnest`,
|
||||
# which has no Oracle equivalent (the `<=>` operator and LIMIT are translated to Oracle by
|
||||
# the backend on execute, but `unnest` is not). Mirrors retrieve_semantic_bm25_combined.
|
||||
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
|
||||
tag_groups_param_start = 6 + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
# created_at time range filter (after tags/groups)
|
||||
@@ -368,69 +433,88 @@ async def retrieve_temporal_combined(
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
params: list = [query_emb_str, bank_id, fact_types, start_date, end_date, semantic_threshold]
|
||||
params: list = [query_emb_str, bank_id, start_date, end_date, semantic_threshold]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
# Two-phase entry point query:
|
||||
# Phase 1 (date_ranked): rank by date only — no embedding computation — for all units in
|
||||
# the temporal window. This lets the planner use date indexes for filtering.
|
||||
# Phase 2 (sim_ranked): join back to memory_units for only the top-50-per-type candidates
|
||||
# and compute embedding similarity for that small set (≤ 50 × len(fact_types) rows).
|
||||
# This avoids computing embedding distances for potentially thousands of date-range rows.
|
||||
entry_points = await conn.fetch(
|
||||
f"""
|
||||
WITH date_ranked AS MATERIALIZED (
|
||||
SELECT id, fact_type,
|
||||
ROW_NUMBER() OVER (
|
||||
PARTITION BY fact_type
|
||||
ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC NULLS LAST
|
||||
) AS rn
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = ANY($3)
|
||||
AND embedding IS NOT NULL
|
||||
AND (
|
||||
(occurred_start IS NOT NULL AND occurred_end IS NOT NULL
|
||||
AND occurred_start <= $5 AND occurred_end >= $4)
|
||||
OR
|
||||
(mentioned_at IS NOT NULL AND mentioned_at BETWEEN $4 AND $5)
|
||||
OR
|
||||
(occurred_start IS NOT NULL AND occurred_start BETWEEN $4 AND $5)
|
||||
OR
|
||||
(occurred_end IS NOT NULL AND occurred_end BETWEEN $4 AND $5)
|
||||
)
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
{created_range_clause}
|
||||
),
|
||||
sim_ranked AS (
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.proof_count, mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
|
||||
1 - (mu.embedding <=> $1::vector) AS similarity,
|
||||
ROW_NUMBER() OVER (PARTITION BY mu.fact_type ORDER BY mu.embedding <=> $1::vector) AS sim_rn
|
||||
FROM date_ranked dr
|
||||
JOIN {fq_table("memory_units")} mu ON mu.id = dr.id
|
||||
WHERE dr.rn <= 50
|
||||
AND (1 - (mu.embedding <=> $1::vector)) >= $6
|
||||
)
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, proof_count, document_id, chunk_id, tags, metadata, similarity
|
||||
FROM sim_ranked
|
||||
WHERE sim_rn <= 10
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
# Entry-point selection: similarity-gated, window-filtered, then narrowed for coverage.
|
||||
#
|
||||
# For each fact_type, ANN-rank the units whose time overlaps the window
|
||||
# (ORDER BY embedding <=> query) and keep a pool of the most relevant
|
||||
# (_TEMPORAL_POOL_SIZE). The planner serves this from the per-(bank, fact_type) vector
|
||||
# index when the window is broad — the dense-metadata case, where the window matches
|
||||
# most rows — and from the partial date indexes plus an exact sort when the window is
|
||||
# narrow. Either way the work is bounded; neither path is a scan-and-sort of the whole
|
||||
# match set.
|
||||
#
|
||||
# Selecting by *similarity* (not recency) is deliberate. The earlier form ranked the
|
||||
# entire match set by COALESCE(occurred_start, mentioned_at, occurred_end) and kept the
|
||||
# 50 most recent: that biased results toward the end of the window and, on banks with
|
||||
# dense/near-uniform dates (e.g. a retain batch stamped with one date), the date key was
|
||||
# degenerate so the "50 most recent" became a near-random sample that could drop the
|
||||
# single most relevant in-window memory — and it degraded to a full scan + disk-spilling
|
||||
# sort (30s+ on a 660k-row bank). The pool is then narrowed to _TEMPORAL_ENTRY_POINTS per
|
||||
# fact_type by _select_with_temporal_coverage so the entry points span the window's range
|
||||
# rather than clustering in one slice.
|
||||
if not fact_types:
|
||||
return {}
|
||||
|
||||
if not entry_points:
|
||||
# One similarity-ranked, window-filtered arm per fact_type, UNION ALL'd — each arm has its
|
||||
# own ORDER BY ... LIMIT so the per-(bank, fact_type) vector index can serve it. fact_type
|
||||
# is inlined as a literal (controlled internal enum, never user input), matching
|
||||
# retrieve_semantic_bm25_combined; this keeps the query free of `unnest`/LATERAL, which the
|
||||
# Oracle backend cannot translate.
|
||||
pool_cols = (
|
||||
"id, text, context, event_date, occurred_start, occurred_end, mentioned_at, "
|
||||
"fact_type, proof_count, document_id, chunk_id, tags, metadata"
|
||||
)
|
||||
table = fq_table("memory_units")
|
||||
arms = [
|
||||
f"""(
|
||||
SELECT {pool_cols}, 1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM {table}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = '{ft}'
|
||||
AND embedding IS NOT NULL
|
||||
AND (
|
||||
(occurred_start IS NOT NULL AND occurred_end IS NOT NULL
|
||||
AND occurred_start <= $4 AND occurred_end >= $3)
|
||||
OR
|
||||
(mentioned_at IS NOT NULL AND mentioned_at BETWEEN $3 AND $4)
|
||||
OR
|
||||
(occurred_start IS NOT NULL AND occurred_start BETWEEN $3 AND $4)
|
||||
OR
|
||||
(occurred_end IS NOT NULL AND occurred_end BETWEEN $3 AND $4)
|
||||
)
|
||||
AND (1 - (embedding <=> $1::vector)) >= $5
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
{created_range_clause}
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT {_TEMPORAL_POOL_SIZE}
|
||||
)"""
|
||||
for ft in fact_types
|
||||
]
|
||||
pool_rows = await conn.fetch("\nUNION ALL\n".join(arms), *params)
|
||||
|
||||
if not pool_rows:
|
||||
return {ft: [] for ft in fact_types}
|
||||
|
||||
# Group entry points by fact type
|
||||
entries_by_ft: dict[str, list] = {ft: [] for ft in fact_types}
|
||||
for ep in entry_points:
|
||||
ft = ep["fact_type"]
|
||||
if ft in entries_by_ft:
|
||||
entries_by_ft[ft].append(ep)
|
||||
# Group the ANN pool by fact type, then narrow each to coverage-spread entry points.
|
||||
pool_by_ft: dict[str, list] = {ft: [] for ft in fact_types}
|
||||
for row in pool_rows:
|
||||
ft = row["fact_type"]
|
||||
if ft in pool_by_ft:
|
||||
pool_by_ft[ft].append(row)
|
||||
|
||||
entries_by_ft: dict[str, list] = {
|
||||
ft: _select_with_temporal_coverage(
|
||||
rows, start_date, end_date, _TEMPORAL_ENTRY_POINTS, _TEMPORAL_COVERAGE_BUCKETS
|
||||
)
|
||||
for ft, rows in pool_by_ft.items()
|
||||
}
|
||||
|
||||
# Calculate shared temporal parameters
|
||||
total_days = (end_date - start_date).total_seconds() / 86400
|
||||
@@ -498,7 +582,13 @@ async def retrieve_temporal_combined(
|
||||
tag_groups, spreading_groups_param_start, table_alias="mu."
|
||||
)
|
||||
|
||||
while frontier and budget_remaining > 0 and iteration < max_iterations:
|
||||
# Multi-hop temporal spreading expands a batch of seed ids with
|
||||
# ``FROM unnest($2::uuid[])``, which has no Oracle equivalent. On backends
|
||||
# without unnest, skip the spread: the temporal entry points are still
|
||||
# returned above, and the semantic/keyword/graph retrievers cover the rest.
|
||||
supports_unnest = getattr(conn, "backend_type", "postgresql") != "oracle"
|
||||
|
||||
while frontier and budget_remaining > 0 and iteration < max_iterations and supports_unnest:
|
||||
iteration += 1
|
||||
batch_ids = frontier[:batch_size]
|
||||
frontier = frontier[batch_size:]
|
||||
|
||||
@@ -358,12 +358,15 @@ class SearchTracer:
|
||||
"""
|
||||
self.rrf_merged = []
|
||||
for rank, (doc_id, data, rrf_meta) in enumerate(merged_results, start=1):
|
||||
source_ranks = rrf_meta.get("source_ranks")
|
||||
if source_ranks is None:
|
||||
source_ranks = {key: value for key, value in rrf_meta.items() if key.endswith("_rank")}
|
||||
self.rrf_merged.append(
|
||||
RRFMergeResult(
|
||||
node_id=doc_id,
|
||||
text=data.get("text", ""),
|
||||
rrf_score=rrf_meta.get("rrf_score", 0.0),
|
||||
source_ranks=rrf_meta.get("source_ranks", {}),
|
||||
source_ranks=source_ranks,
|
||||
final_rrf_rank=rank,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -371,6 +371,7 @@ class SQLDialect(ABC):
|
||||
embedding_param: str,
|
||||
bank_id_param: str,
|
||||
fetch_limit: int,
|
||||
min_similarity: float,
|
||||
tags_clause: str = "",
|
||||
groups_clause: str = "",
|
||||
extra_where: str = "",
|
||||
@@ -387,6 +388,7 @@ class SQLDialect(ABC):
|
||||
embedding_param: Parameter placeholder for query embedding.
|
||||
bank_id_param: Parameter placeholder for bank_id.
|
||||
fetch_limit: Max rows to fetch (over-fetched for HNSW approximation).
|
||||
min_similarity: Minimum cosine similarity to include.
|
||||
tags_clause: Optional WHERE clause fragment for tag filtering.
|
||||
groups_clause: Optional WHERE clause fragment for tag group filtering.
|
||||
extra_where: Optional additional WHERE clause fragment (e.g. time range filter).
|
||||
@@ -408,6 +410,7 @@ class SQLDialect(ABC):
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
bm25_language: str = "english",
|
||||
bm25_min_score: float = 0.0,
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
"""Build a BM25/full-text search subquery arm.
|
||||
@@ -430,6 +433,11 @@ class SQLDialect(ABC):
|
||||
"pg_textsearch", "pgroonga"). Only relevant for PostgreSQL.
|
||||
bm25_language: PostgreSQL text search dictionary used by the native
|
||||
backend (e.g. "english", "french"). Ignored by other backends.
|
||||
bm25_min_score: Minimum BM25 relevance score a row must exceed to be
|
||||
returned. Gates out non-matching rows on backends whose
|
||||
operator (e.g. VectorChord) ranks every document instead
|
||||
of pre-filtering to query-term matches. Backends that
|
||||
already apply a boolean match gate ignore this.
|
||||
extra_where: Optional additional WHERE clause fragment (e.g. time range filter).
|
||||
"""
|
||||
...
|
||||
|
||||
@@ -234,6 +234,7 @@ class OracleDialect(SQLDialect):
|
||||
embedding_param: str,
|
||||
bank_id_param: str,
|
||||
fetch_limit: int,
|
||||
min_similarity: float,
|
||||
tags_clause: str = "",
|
||||
groups_clause: str = "",
|
||||
extra_where: str = "",
|
||||
@@ -249,7 +250,7 @@ class OracleDialect(SQLDialect):
|
||||
f" WHERE bank_id = {bank_id_param}"
|
||||
f" AND fact_type = '{fact_type}'"
|
||||
f" AND embedding IS NOT NULL"
|
||||
f" AND (1 - VECTOR_DISTANCE(embedding, {embedding_param}, COSINE)) >= 0.3"
|
||||
f" AND (1 - VECTOR_DISTANCE(embedding, {embedding_param}, COSINE)) >= {min_similarity}"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" {extra_where}"
|
||||
@@ -271,6 +272,7 @@ class OracleDialect(SQLDialect):
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
bm25_language: str = "english",
|
||||
bm25_min_score: float = 0.0,
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
# Oracle Text: CONTAINS() / SCORE() with the CTXSYS.CONTEXT index.
|
||||
@@ -285,7 +287,9 @@ class OracleDialect(SQLDialect):
|
||||
f" FROM {table}"
|
||||
f" WHERE bank_id = {bank_id_param}"
|
||||
f" AND fact_type = '{fact_type}'"
|
||||
f" AND CONTAINS(text, {text_param}, {label}) > 0"
|
||||
# CONTAINS already gates to genuine matches; the configurable floor
|
||||
# (default 0) keeps the threshold semantics uniform across backends.
|
||||
f" AND CONTAINS(text, {text_param}, {label}) > {bm25_min_score:g}"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" {extra_where}"
|
||||
|
||||
@@ -148,6 +148,7 @@ class PostgreSQLDialect(SQLDialect):
|
||||
embedding_param: str,
|
||||
bank_id_param: str,
|
||||
fetch_limit: int,
|
||||
min_similarity: float,
|
||||
tags_clause: str = "",
|
||||
groups_clause: str = "",
|
||||
extra_where: str = "",
|
||||
@@ -161,7 +162,7 @@ class PostgreSQLDialect(SQLDialect):
|
||||
f" WHERE bank_id = {bank_id_param}"
|
||||
f" AND fact_type = '{fact_type}'"
|
||||
f" AND embedding IS NOT NULL"
|
||||
f" AND (1 - (embedding <=> {embedding_param}::vector)) >= 0.3"
|
||||
f" AND (1 - (embedding <=> {embedding_param}::vector)) >= {min_similarity}"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" {extra_where}"
|
||||
@@ -183,25 +184,32 @@ class PostgreSQLDialect(SQLDialect):
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
bm25_language: str = "english",
|
||||
bm25_min_score: float = 0.0,
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
if text_search_extension == "vchord":
|
||||
# <&> returns a distance (lower = more relevant), negate for score
|
||||
# <&> returns the NEGATIVE BM25 score (lower = more relevant), negate
|
||||
# for a positive score where higher = more relevant.
|
||||
bm25_score_expr = f"-(search_vector <&> to_bm25query('idx_memory_units_text_search', tokenize({text_param}, 'llmlingua2')))"
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = ""
|
||||
# Unlike native tsvector (which has a boolean `@@` match gate), the
|
||||
# VectorChord operator ranks *every* document, so a bare ORDER BY ...
|
||||
# LIMIT pads the result with zero-score, non-matching rows. Gate on the
|
||||
# score so only genuine term matches survive into fusion/reranking.
|
||||
bm25_where_filter = f"AND -(search_vector <&> to_bm25query('idx_memory_units_text_search', tokenize({text_param}, 'llmlingua2'))) > {bm25_min_score:g}"
|
||||
elif text_search_extension == "pg_textsearch":
|
||||
bm25_score_expr = f"-({text_param} <@> to_bm25query({text_param}, 'idx_memory_units_text_search'))"
|
||||
bm25_order_by = f"text <@> to_bm25query({text_param}, 'idx_memory_units_text_search') ASC"
|
||||
bm25_where_filter = ""
|
||||
elif text_search_extension == "pgroonga":
|
||||
# &@~ accepts pgroonga's query syntax (raw query text). pgroonga_score
|
||||
# returns a non-negative relevance score (higher = better).
|
||||
# &@~ accepts pgroonga's query syntax. Escape the bind parameter so
|
||||
# literal memory text containing operators like ">" or "(" is not
|
||||
# parsed as a malformed query expression.
|
||||
bm25_score_expr = "pgroonga_score(tableoid, ctid)"
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = (
|
||||
f"AND (COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, '')) "
|
||||
f"&@~ {text_param}"
|
||||
f"&@~ pgroonga_query_escape({text_param})"
|
||||
)
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search: BM25 index over (id, text, context, text_signals)
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Document transfer: export/import documents between banks without re-running the LLM.
|
||||
|
||||
An export is a ZIP of already-extracted facts (text, entities by canonical name,
|
||||
causal relations, chunks) — never embeddings or DB ids. An import replays the
|
||||
deterministic half of the retain pipeline against the target bank: it re-embeds
|
||||
locally with the target bank's embedding model, re-resolves entities, and
|
||||
recreates temporal/semantic/causal links relative to the target bank's existing
|
||||
memories. No LLM fact-extraction is involved.
|
||||
|
||||
Consolidated observations (``fact_type='observation'``) are intentionally
|
||||
excluded from export — they are derived by consolidation and are regenerated in
|
||||
the target bank.
|
||||
"""
|
||||
|
||||
from .export import export_bank, export_documents
|
||||
from .importer import BankImportResult, ImportResult, import_bank, import_documents
|
||||
from .schema import (
|
||||
SCHEMA_VERSION,
|
||||
TransferCausalRelation,
|
||||
TransferChunk,
|
||||
TransferDocument,
|
||||
TransferFact,
|
||||
TransferManifest,
|
||||
TransferObservation,
|
||||
TransferObservationSource,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"SCHEMA_VERSION",
|
||||
"BankImportResult",
|
||||
"ImportResult",
|
||||
"TransferCausalRelation",
|
||||
"TransferChunk",
|
||||
"TransferDocument",
|
||||
"TransferFact",
|
||||
"TransferManifest",
|
||||
"TransferObservation",
|
||||
"TransferObservationSource",
|
||||
"export_bank",
|
||||
"export_documents",
|
||||
"import_bank",
|
||||
"import_documents",
|
||||
]
|
||||
@@ -0,0 +1,561 @@
|
||||
"""Export documents (with extracted facts, entities, causal links, chunks) to a ZIP archive.
|
||||
|
||||
Reads directly from the database via the backend connection. Embeddings and
|
||||
database ids are deliberately omitted — they are regenerated/re-resolved on
|
||||
import. Consolidated observations are excluded unless ``include_observations``
|
||||
is set, in which case they are written to ``observations.json``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import zipfile
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, date, datetime
|
||||
from decimal import Decimal
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..schema import fq_table
|
||||
from .schema import (
|
||||
SCHEMA_VERSION,
|
||||
TransferCausalRelation,
|
||||
TransferChunk,
|
||||
TransferDocument,
|
||||
TransferFact,
|
||||
TransferManifest,
|
||||
TransferObservation,
|
||||
TransferObservationSource,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Whole-bank export classification. Every bank-scoped table (admin.cli.BACKUP_TABLES)
|
||||
# must fall into exactly one bucket below; tests/test_document_transfer.py's
|
||||
# test_export_bank_covers_schema enforces this so a table added by a future
|
||||
# migration can't be silently dropped from a migration archive.
|
||||
|
||||
# NOT written to the archive — rebuilt on import by replaying the document/fact/
|
||||
# observation payload through the import pipeline:
|
||||
# * documents / chunks / memory_units carry their *text* in the logical document
|
||||
# payload (TransferDocument) and are re-embedded with the target model;
|
||||
# * entities / unit_entities / memory_links / entity_cooccurrences are derived
|
||||
# data — the pipeline re-resolves entities and rebuilds links/cooccurrence
|
||||
# stats against the target bank, so they are never exported.
|
||||
# Listed here only so the coverage guard can assert every table is classified.
|
||||
_REPLAYED_TABLES = frozenset(
|
||||
{
|
||||
"documents",
|
||||
"chunks",
|
||||
"memory_units",
|
||||
"entities",
|
||||
"unit_entities",
|
||||
"memory_links",
|
||||
"entity_cooccurrences",
|
||||
# observation_history FKs to a memory_units observation, but observations
|
||||
# are derived: they're regenerated with FRESH ids when consolidation is
|
||||
# replayed on import (see _EXPORTED_FACT_TYPES — observations are excluded).
|
||||
# There is no stable observation id to re-attach history to, so it is not
|
||||
# carried; the target rebuilds observation history as it re-consolidates.
|
||||
"observation_history",
|
||||
}
|
||||
)
|
||||
# Carried verbatim as JSON rows (bank config + synthesized state). Embedding-bearing
|
||||
# rows have their vector stripped (see _DERIVED_COLUMNS) and are re-embedded on import.
|
||||
_BANK_ROW_TABLES = ("banks", "mental_models", "directives", "webhooks")
|
||||
# Bank-scoped child-history carried verbatim. Unlike observations, mental models
|
||||
# keep their (id, bank_id) across export/import, so their refresh history can be
|
||||
# re-attached. The surrogate ``id`` is dropped on dump so the target reassigns it
|
||||
# (see _dump_history_rows); restored after its parent table (mental_models).
|
||||
_CARRIED_HISTORY_TABLES = ("mental_model_history",)
|
||||
# Operational history — only carried with include_history=True.
|
||||
_HISTORY_TABLES = ("audit_log", "llm_requests")
|
||||
# Intentionally never exported.
|
||||
_SKIP_TABLES = frozenset(
|
||||
{
|
||||
"async_operations", # in-flight ops; drain on the source before migrating
|
||||
"graph_maintenance_queue", # transient work queue; regenerated on import
|
||||
"file_storage", # raw uploads; documents.original_text is already carried
|
||||
# Curation archive of retired facts — local operational state, not part of
|
||||
# the live knowledge the export replays. Its rows mirror memory_units (stale
|
||||
# embedding) and snapshot source-bank entity ids that the import re-resolves
|
||||
# to fresh ids, so carrying them would only produce dangling associations.
|
||||
# Revert anything worth keeping on the source before migrating.
|
||||
"invalidated_memory_units",
|
||||
}
|
||||
)
|
||||
# Derived columns dropped from carried rows so the target regenerates them with
|
||||
# its own embedding model / text-search backend.
|
||||
_DERIVED_COLUMNS = ("embedding", "search_vector")
|
||||
|
||||
|
||||
@dataclass
|
||||
class _UnitLocation:
|
||||
"""Where a memory unit's fact lives in the assembled export (document + ordinal)."""
|
||||
|
||||
document_id: str
|
||||
ordinal: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class _LoadedFacts:
|
||||
"""Facts grouped by document plus an index from unit id to its location.
|
||||
|
||||
``facts_by_doc`` and ``unit_index`` share the same fixed ordering so that
|
||||
causal ``target_fact_index`` ordinals stay consistent across both.
|
||||
"""
|
||||
|
||||
facts_by_doc: dict[str, list[TransferFact]] = field(default_factory=dict)
|
||||
unit_index: dict[Any, _UnitLocation] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _LoadedExport:
|
||||
"""Assembled documents plus the unit-id → location index.
|
||||
|
||||
``unit_index`` is retained so observation source unit ids can be resolved to
|
||||
(document_id, fact_index) references when observations are exported.
|
||||
"""
|
||||
|
||||
documents: list[TransferDocument] = field(default_factory=list)
|
||||
unit_index: dict[Any, _UnitLocation] = field(default_factory=dict)
|
||||
|
||||
|
||||
# Causal link types that retain persists between facts. Only these travel in the
|
||||
# archive; temporal/semantic/entity links are regenerated against the target bank.
|
||||
_CAUSAL_LINK_TYPES = ("caused_by", "causes", "enables", "prevents")
|
||||
|
||||
# Facts of these types are exported; observations are derived and excluded.
|
||||
_EXPORTED_FACT_TYPES = ("world", "experience")
|
||||
|
||||
|
||||
def _as_jsonb(value: Any) -> Any:
|
||||
"""Coerce an asyncpg JSONB column (str or already-decoded) to a Python object."""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, str):
|
||||
return json.loads(value)
|
||||
return value
|
||||
|
||||
|
||||
def _chunk_index_from_chunk_id(chunk_id: str | None) -> int | None:
|
||||
"""Recover the chunk ordinal from a ``{bank_id}_{document_id}_{index}`` chunk_id.
|
||||
|
||||
The index is always the final underscore-delimited segment, so rsplit is
|
||||
correct even when bank/document ids themselves contain underscores.
|
||||
"""
|
||||
if not chunk_id:
|
||||
return None
|
||||
try:
|
||||
return int(chunk_id.rsplit("_", 1)[1])
|
||||
except (IndexError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
async def export_documents(
|
||||
backend: Any,
|
||||
bank_id: str,
|
||||
document_ids: list[str] | None = None,
|
||||
*,
|
||||
include_observations: bool = False,
|
||||
) -> bytes:
|
||||
"""Export documents from ``bank_id`` into an in-memory ZIP archive.
|
||||
|
||||
Args:
|
||||
backend: Database backend (provides ``acquire()``).
|
||||
bank_id: Source bank.
|
||||
document_ids: Specific document ids to export. ``None`` exports every
|
||||
document in the bank.
|
||||
include_observations: Also export consolidated observations (written to
|
||||
``observations.json``). Only valid for a whole-bank export.
|
||||
|
||||
Returns:
|
||||
The ZIP archive as bytes.
|
||||
|
||||
Raises:
|
||||
ValueError: if ``include_observations`` is combined with ``document_ids``.
|
||||
"""
|
||||
# Observations are bank-level and can be derived from facts spanning several
|
||||
# documents, so they're only coherent when the whole bank is exported. For a
|
||||
# document subset we'd have to silently drop every cross-document observation
|
||||
# — reject the combination instead so the caller isn't surprised.
|
||||
if include_observations and document_ids is not None:
|
||||
raise ValueError("include_observations is only supported when exporting the whole bank (omit document_id)")
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
loaded = await _load_documents(conn, bank_id, document_ids)
|
||||
documents = loaded.documents
|
||||
observations = await _load_observations(conn, bank_id, loaded.unit_index) if include_observations else []
|
||||
|
||||
archive = io.BytesIO()
|
||||
fact_total = 0
|
||||
with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as zf:
|
||||
for index, document in enumerate(documents):
|
||||
fact_total += len(document.facts)
|
||||
zf.writestr(
|
||||
f"documents/{index:06d}.json",
|
||||
document.model_dump_json(indent=2, exclude_none=False),
|
||||
)
|
||||
|
||||
if observations:
|
||||
payload = "[\n" + ",\n".join(o.model_dump_json(indent=2) for o in observations) + "\n]\n"
|
||||
zf.writestr("observations.json", payload)
|
||||
|
||||
manifest = TransferManifest(
|
||||
schema_version=SCHEMA_VERSION,
|
||||
source_bank_id=bank_id,
|
||||
exported_at=datetime.now(UTC),
|
||||
document_count=len(documents),
|
||||
fact_count=fact_total,
|
||||
observation_count=len(observations),
|
||||
)
|
||||
zf.writestr("manifest.json", manifest.model_dump_json(indent=2))
|
||||
|
||||
logger.info(
|
||||
"[transfer] Exported %d document(s), %d fact(s), %d observation(s) from bank %s",
|
||||
len(documents),
|
||||
fact_total,
|
||||
len(observations),
|
||||
bank_id,
|
||||
)
|
||||
return archive.getvalue()
|
||||
|
||||
|
||||
def _row_json_default(obj: Any) -> Any:
|
||||
"""JSON serializer for the value types asyncpg returns from bank rows."""
|
||||
if isinstance(obj, datetime):
|
||||
return obj.isoformat()
|
||||
if isinstance(obj, date):
|
||||
return obj.isoformat()
|
||||
if isinstance(obj, UUID):
|
||||
return str(obj)
|
||||
if isinstance(obj, Decimal):
|
||||
# str preserves precision; import casts back to numeric.
|
||||
return str(obj)
|
||||
if isinstance(obj, (bytes, bytearray, memoryview)):
|
||||
return base64.b64encode(bytes(obj)).decode("ascii")
|
||||
raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
|
||||
|
||||
|
||||
async def _dump_bank_rows(conn: Any, table: str, bank_id: str) -> list[dict]:
|
||||
"""Dump all rows of a bank-scoped table as JSON-ready dicts (derived columns stripped).
|
||||
|
||||
Embedding/search-vector columns are omitted so the target instance
|
||||
regenerates them with its own model/backend on import.
|
||||
"""
|
||||
rows = await conn.fetch(f"SELECT * FROM {fq_table(table)} WHERE bank_id = $1", bank_id)
|
||||
return [{k: v for k, v in dict(row).items() if k not in _DERIVED_COLUMNS} for row in rows]
|
||||
|
||||
|
||||
async def _dump_history_rows(conn: Any, table: str, bank_id: str) -> list[dict]:
|
||||
"""Dump a bank-scoped child-history table for carrying across instances.
|
||||
|
||||
Drops the surrogate ``id`` so the target reassigns it from its own IDENTITY
|
||||
sequence (carrying explicit ids would leave the sequence un-advanced and
|
||||
collide with later writes). Ordered oldest-first so the reassigned ids keep
|
||||
the same chronological tie-break order the read path relies on.
|
||||
"""
|
||||
rows = await conn.fetch(
|
||||
f"SELECT * FROM {fq_table(table)} WHERE bank_id = $1 ORDER BY changed_at, id",
|
||||
bank_id,
|
||||
)
|
||||
return [{k: v for k, v in dict(row).items() if k not in _DERIVED_COLUMNS and k != "id"} for row in rows]
|
||||
|
||||
|
||||
async def export_bank(conn: Any, bank_id: str, *, include_history: bool = False) -> bytes:
|
||||
"""Export an entire bank into a portable ZIP archive (no embeddings).
|
||||
|
||||
Produces a superset of the documents archive: the logical
|
||||
document/fact/observation export (replayed and re-embedded on import) plus
|
||||
the bank's config, mental models, directives and webhooks as JSON rows. With
|
||||
``include_history`` the operational tails (audit_log, llm_requests) are also
|
||||
carried. Intended for migrating a bank to a new instance configured with a
|
||||
different embedding model / vector / text-search backend — every vector is
|
||||
regenerated on the target, so nothing here is encoder-specific.
|
||||
|
||||
``conn`` is a live connection scoped to the bank's schema (the admin CLI sets
|
||||
``_current_schema`` and passes its raw connection; the engine acquires one
|
||||
after tenant auth).
|
||||
"""
|
||||
loaded = await _load_documents(conn, bank_id, None)
|
||||
documents = loaded.documents
|
||||
# Whole-bank export always carries observations (they're bank-level state).
|
||||
observations = await _load_observations(conn, bank_id, loaded.unit_index)
|
||||
|
||||
bank_rows = {table: await _dump_bank_rows(conn, table, bank_id) for table in _BANK_ROW_TABLES}
|
||||
for table in _CARRIED_HISTORY_TABLES:
|
||||
bank_rows[table] = await _dump_history_rows(conn, table, bank_id)
|
||||
history_rows: dict[str, list[dict]] = {}
|
||||
if include_history:
|
||||
history_rows = {table: await _dump_bank_rows(conn, table, bank_id) for table in _HISTORY_TABLES}
|
||||
|
||||
archive = io.BytesIO()
|
||||
fact_total = 0
|
||||
with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as zf:
|
||||
for index, document in enumerate(documents):
|
||||
fact_total += len(document.facts)
|
||||
zf.writestr(f"documents/{index:06d}.json", document.model_dump_json(indent=2, exclude_none=False))
|
||||
|
||||
if observations:
|
||||
payload = "[\n" + ",\n".join(o.model_dump_json(indent=2) for o in observations) + "\n]\n"
|
||||
zf.writestr("observations.json", payload)
|
||||
|
||||
for table, rows in bank_rows.items():
|
||||
zf.writestr(f"{table}.json", json.dumps(rows, indent=2, default=_row_json_default))
|
||||
for table, rows in history_rows.items():
|
||||
zf.writestr(f"history/{table}.json", json.dumps(rows, indent=2, default=_row_json_default))
|
||||
|
||||
manifest = TransferManifest(
|
||||
schema_version=SCHEMA_VERSION,
|
||||
source_bank_id=bank_id,
|
||||
exported_at=datetime.now(UTC),
|
||||
document_count=len(documents),
|
||||
fact_count=fact_total,
|
||||
observation_count=len(observations),
|
||||
archive_type="bank",
|
||||
mental_model_count=len(bank_rows.get("mental_models", [])),
|
||||
directive_count=len(bank_rows.get("directives", [])),
|
||||
webhook_count=len(bank_rows.get("webhooks", [])),
|
||||
includes_history=include_history,
|
||||
)
|
||||
zf.writestr("manifest.json", manifest.model_dump_json(indent=2))
|
||||
|
||||
logger.info(
|
||||
"[transfer] Exported bank %s: %d document(s), %d fact(s), %d observation(s), "
|
||||
"%d mental model(s), %d directive(s), %d webhook(s)%s",
|
||||
bank_id,
|
||||
len(documents),
|
||||
fact_total,
|
||||
len(observations),
|
||||
len(bank_rows.get("mental_models", [])),
|
||||
len(bank_rows.get("directives", [])),
|
||||
len(bank_rows.get("webhooks", [])),
|
||||
" (with history)" if include_history else "",
|
||||
)
|
||||
return archive.getvalue()
|
||||
|
||||
|
||||
async def _load_documents(
|
||||
conn: Any,
|
||||
bank_id: str,
|
||||
document_ids: list[str] | None,
|
||||
) -> _LoadedExport:
|
||||
"""Load and assemble TransferDocument payloads for the requested documents."""
|
||||
doc_filter = "AND id = ANY($2)" if document_ids else ""
|
||||
params: list[Any] = [bank_id]
|
||||
if document_ids:
|
||||
params.append(document_ids)
|
||||
doc_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, original_text, retain_params, tags, created_at
|
||||
FROM {fq_table("documents")}
|
||||
WHERE bank_id = $1 {doc_filter}
|
||||
ORDER BY created_at, id
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
if not doc_rows:
|
||||
return _LoadedExport()
|
||||
|
||||
selected_ids = [row["id"] for row in doc_rows]
|
||||
|
||||
chunks_by_doc = await _load_chunks(conn, bank_id, selected_ids)
|
||||
loaded = await _load_facts(conn, bank_id, selected_ids)
|
||||
await _attach_entities(conn, loaded)
|
||||
await _attach_causal_relations(conn, loaded)
|
||||
|
||||
documents: list[TransferDocument] = []
|
||||
for row in doc_rows:
|
||||
doc_id = row["id"]
|
||||
documents.append(
|
||||
TransferDocument(
|
||||
id=doc_id,
|
||||
original_text=row["original_text"],
|
||||
retain_params=_as_jsonb(row["retain_params"]),
|
||||
tags=list(row["tags"] or []),
|
||||
created_at=row["created_at"],
|
||||
chunks=chunks_by_doc.get(doc_id, []),
|
||||
facts=loaded.facts_by_doc.get(doc_id, []),
|
||||
)
|
||||
)
|
||||
return _LoadedExport(documents=documents, unit_index=loaded.unit_index)
|
||||
|
||||
|
||||
async def _load_observations(
|
||||
conn: Any,
|
||||
bank_id: str,
|
||||
unit_index: dict[Any, _UnitLocation],
|
||||
) -> list[TransferObservation]:
|
||||
"""Load observations whose source facts are all present in the exported set.
|
||||
|
||||
Each source unit id is rewritten to its (document_id, fact_index) reference
|
||||
via ``unit_index``. Only called for a whole-bank export, so every live source
|
||||
fact is present; an observation is skipped only if a source no longer exists
|
||||
(stale reference) — that keeps every exported observation resolvable on import.
|
||||
"""
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, tags, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, observation_scopes, proof_count, source_memory_ids
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY created_at, id
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
observations: list[TransferObservation] = []
|
||||
skipped = 0
|
||||
for row in rows:
|
||||
source_ids = list(row["source_memory_ids"] or [])
|
||||
locations = [unit_index.get(sid) for sid in source_ids]
|
||||
if not source_ids or any(loc is None for loc in locations):
|
||||
# An observation with sources outside the exported documents would be
|
||||
# incoherent on import — skip it rather than emit dangling refs.
|
||||
skipped += 1
|
||||
continue
|
||||
observations.append(
|
||||
TransferObservation(
|
||||
text=row["text"],
|
||||
tags=list(row["tags"] or []),
|
||||
event_date=row["event_date"],
|
||||
occurred_start=row["occurred_start"],
|
||||
occurred_end=row["occurred_end"],
|
||||
mentioned_at=row["mentioned_at"],
|
||||
observation_scopes=_as_jsonb(row["observation_scopes"]),
|
||||
proof_count=row["proof_count"] or len(source_ids),
|
||||
sources=[
|
||||
TransferObservationSource(document_id=loc.document_id, fact_index=loc.ordinal)
|
||||
for loc in locations
|
||||
if loc is not None
|
||||
],
|
||||
)
|
||||
)
|
||||
if skipped:
|
||||
logger.info("[transfer] Skipped %d observation(s) with sources outside the exported documents", skipped)
|
||||
return observations
|
||||
|
||||
|
||||
async def _load_chunks(conn: Any, bank_id: str, doc_ids: list[str]) -> dict[str, list[TransferChunk]]:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT document_id, chunk_index, chunk_text
|
||||
FROM {fq_table("chunks")}
|
||||
WHERE bank_id = $1 AND document_id = ANY($2)
|
||||
ORDER BY document_id, chunk_index
|
||||
""",
|
||||
bank_id,
|
||||
doc_ids,
|
||||
)
|
||||
chunks_by_doc: dict[str, list[TransferChunk]] = {}
|
||||
for row in rows:
|
||||
chunks_by_doc.setdefault(row["document_id"], []).append(
|
||||
TransferChunk(chunk_index=row["chunk_index"], chunk_text=row["chunk_text"])
|
||||
)
|
||||
return chunks_by_doc
|
||||
|
||||
|
||||
async def _load_facts(conn: Any, bank_id: str, doc_ids: list[str]) -> _LoadedFacts:
|
||||
"""Load non-observation facts grouped by document, with a unit-id location index.
|
||||
|
||||
The ordering is fixed (created_at, id) so that
|
||||
``causal_relations.target_fact_index`` ordinals stay consistent.
|
||||
"""
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, document_id, text, fact_type, context, event_date,
|
||||
occurred_start, occurred_end, mentioned_at, metadata,
|
||||
chunk_id, tags, observation_scopes
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND document_id = ANY($2)
|
||||
AND fact_type = ANY($3)
|
||||
ORDER BY document_id, created_at, id
|
||||
""",
|
||||
bank_id,
|
||||
doc_ids,
|
||||
list(_EXPORTED_FACT_TYPES),
|
||||
)
|
||||
|
||||
loaded = _LoadedFacts()
|
||||
for row in rows:
|
||||
doc_id = row["document_id"]
|
||||
bucket = loaded.facts_by_doc.setdefault(doc_id, [])
|
||||
ordinal = len(bucket)
|
||||
fact = TransferFact(
|
||||
text=row["text"],
|
||||
fact_type=row["fact_type"],
|
||||
context=row["context"],
|
||||
event_date=row["event_date"],
|
||||
occurred_start=row["occurred_start"],
|
||||
occurred_end=row["occurred_end"],
|
||||
mentioned_at=row["mentioned_at"],
|
||||
metadata=_as_jsonb(row["metadata"]) or {},
|
||||
tags=list(row["tags"] or []),
|
||||
observation_scopes=_as_jsonb(row["observation_scopes"]),
|
||||
chunk_index=_chunk_index_from_chunk_id(row["chunk_id"]),
|
||||
)
|
||||
bucket.append(fact)
|
||||
loaded.unit_index[row["id"]] = _UnitLocation(document_id=doc_id, ordinal=ordinal)
|
||||
return loaded
|
||||
|
||||
|
||||
async def _attach_entities(conn: Any, loaded: _LoadedFacts) -> None:
|
||||
"""Populate each fact's ``entities`` list with its entities' canonical names."""
|
||||
if not loaded.unit_index:
|
||||
return
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT ue.unit_id, e.canonical_name
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("entities")} e ON e.id = ue.entity_id
|
||||
WHERE ue.unit_id = ANY($1)
|
||||
ORDER BY e.canonical_name
|
||||
""",
|
||||
list(loaded.unit_index.keys()),
|
||||
)
|
||||
for row in rows:
|
||||
location = loaded.unit_index.get(row["unit_id"])
|
||||
if location is None:
|
||||
continue
|
||||
loaded.facts_by_doc[location.document_id][location.ordinal].entities.append(row["canonical_name"])
|
||||
|
||||
|
||||
async def _attach_causal_relations(conn: Any, loaded: _LoadedFacts) -> None:
|
||||
"""Reconstruct causal edges as fact ordinals within each document.
|
||||
|
||||
A memory_link (from_unit -> to_unit, link_type) means ``from_unit`` carries
|
||||
the relation pointing at ``to_unit``, so the edge is attached to the source
|
||||
fact with the target's ordinal. Edges spanning two documents are skipped
|
||||
(causal links are created within a single retain batch in practice).
|
||||
"""
|
||||
if not loaded.unit_index:
|
||||
return
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_unit_id, to_unit_id, link_type
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE link_type = ANY($1)
|
||||
AND from_unit_id = ANY($2)
|
||||
AND to_unit_id = ANY($2)
|
||||
""",
|
||||
list(_CAUSAL_LINK_TYPES),
|
||||
list(loaded.unit_index.keys()),
|
||||
)
|
||||
for row in rows:
|
||||
source = loaded.unit_index.get(row["from_unit_id"])
|
||||
target = loaded.unit_index.get(row["to_unit_id"])
|
||||
if source is None or target is None:
|
||||
continue
|
||||
if source.document_id != target.document_id:
|
||||
continue
|
||||
loaded.facts_by_doc[source.document_id][source.ordinal].causal_relations.append(
|
||||
TransferCausalRelation(
|
||||
relation_type=row["link_type"],
|
||||
target_fact_index=target.ordinal,
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,716 @@
|
||||
"""Import documents from a transfer archive by replaying the deterministic retain pipeline.
|
||||
|
||||
For each document the importer rebuilds the extracted facts, re-embeds them with
|
||||
the *target* bank's embedding model, then runs entity resolution (Phase 1) and
|
||||
the fact/link insert (Phase 2) — exactly the steps retain runs after LLM
|
||||
extraction. No LLM is called. Temporal/semantic/causal links and entity merges
|
||||
are therefore computed relative to the target bank's existing memories.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
import zipfile
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, date, datetime
|
||||
from typing import Any, Literal
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..retain import bank_utils, chunk_storage, embedding_processing, fact_storage, orchestrator
|
||||
from ..retain.types import (
|
||||
CausalRelation,
|
||||
ChunkMetadata,
|
||||
ExtractedFact,
|
||||
ProcessedFact,
|
||||
RetainContent,
|
||||
)
|
||||
from ..schema import fq_table
|
||||
from .schema import (
|
||||
SCHEMA_VERSION,
|
||||
TransferDocument,
|
||||
TransferFact,
|
||||
TransferManifest,
|
||||
TransferObservation,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
OnConflict = Literal["skip", "replace", "new-id"]
|
||||
_VALID_CONFLICT_MODES: tuple[OnConflict, ...] = ("skip", "replace", "new-id")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImportedDocument:
|
||||
"""A single document successfully imported, with the units it produced.
|
||||
|
||||
Carried back so the engine can fire the post-retain extension hook
|
||||
(usage tracking / metrics / notifications) once per imported document,
|
||||
mirroring how retain reports each completed document.
|
||||
"""
|
||||
|
||||
document_id: str
|
||||
unit_ids: list[str]
|
||||
content: str
|
||||
tags: list[str]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImportResult:
|
||||
"""Outcome of importing a transfer archive into a bank."""
|
||||
|
||||
documents_imported: int = 0
|
||||
documents_skipped: int = 0
|
||||
facts_imported: int = 0
|
||||
observations_imported: int = 0
|
||||
# Observations dropped because some source fact was not imported in this run.
|
||||
observations_skipped: int = 0
|
||||
skipped_document_ids: list[str] = field(default_factory=list)
|
||||
# Original id -> freshly generated id, for documents imported under "new-id".
|
||||
remapped_document_ids: dict[str, str] = field(default_factory=dict)
|
||||
# Per-document outcomes, for the engine's post-retain hook. Not serialized
|
||||
# into operation result_metadata (the worker handler writes counts only).
|
||||
imported_documents: list[ImportedDocument] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ObservationOutcome:
|
||||
"""Counts from the observation import pass."""
|
||||
|
||||
imported: int = 0
|
||||
skipped: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParsedArchive:
|
||||
"""A transfer archive after parsing/validation."""
|
||||
|
||||
manifest: TransferManifest
|
||||
documents: list[TransferDocument]
|
||||
observations: list[TransferObservation] = field(default_factory=list)
|
||||
|
||||
|
||||
def parse_archive(archive_bytes: bytes) -> ParsedArchive:
|
||||
"""Parse and validate a transfer ZIP archive produced by ``export_documents``."""
|
||||
with zipfile.ZipFile(io.BytesIO(archive_bytes), "r") as zf:
|
||||
names = set(zf.namelist())
|
||||
if "manifest.json" not in names:
|
||||
raise ValueError("Invalid transfer archive: manifest.json is missing")
|
||||
manifest = TransferManifest.model_validate_json(zf.read("manifest.json"))
|
||||
if manifest.schema_version != SCHEMA_VERSION:
|
||||
raise ValueError(
|
||||
f"Unsupported transfer archive schema version {manifest.schema_version} "
|
||||
f"(this build supports {SCHEMA_VERSION})"
|
||||
)
|
||||
doc_names = sorted(n for n in names if n.startswith("documents/") and n.endswith(".json"))
|
||||
documents = [TransferDocument.model_validate_json(zf.read(name)) for name in doc_names]
|
||||
observations: list[TransferObservation] = []
|
||||
if "observations.json" in names:
|
||||
observations = [TransferObservation.model_validate(o) for o in json.loads(zf.read("observations.json"))]
|
||||
return ParsedArchive(manifest=manifest, documents=documents, observations=observations)
|
||||
|
||||
|
||||
async def import_documents(
|
||||
*,
|
||||
backend: Any,
|
||||
embeddings_model: Any,
|
||||
entity_resolver: Any,
|
||||
config: Any,
|
||||
format_date_fn: Any,
|
||||
bank_id: str,
|
||||
archive_bytes: bytes,
|
||||
on_conflict: OnConflict = "skip",
|
||||
ops: Any = None,
|
||||
outbox_callback_factory: Any = None,
|
||||
) -> ImportResult:
|
||||
"""Import every document in ``archive_bytes`` into ``bank_id``.
|
||||
|
||||
Args:
|
||||
backend: Database backend (provides ``acquire()`` and ``ops``).
|
||||
embeddings_model: Target bank's embedding model (used to re-embed facts).
|
||||
entity_resolver: Shared entity resolver for the target bank.
|
||||
config: Resolved bank config for the target bank.
|
||||
format_date_fn: Date formatter used when augmenting fact text for embedding
|
||||
(must match retain so embeddings are consistent).
|
||||
bank_id: Target bank.
|
||||
archive_bytes: A ZIP archive produced by ``export_documents``.
|
||||
on_conflict: How to handle a document id that already exists in the target
|
||||
bank — ``skip`` (default), ``replace`` (delete old data and re-import),
|
||||
or ``new-id`` (import under a freshly generated id).
|
||||
ops: Backend ``DataAccessOps``. Defaults to ``backend.ops``.
|
||||
|
||||
Returns:
|
||||
An :class:`ImportResult` with per-document counts.
|
||||
"""
|
||||
if on_conflict not in _VALID_CONFLICT_MODES:
|
||||
raise ValueError(f"Invalid on_conflict '{on_conflict}'; expected one of {_VALID_CONFLICT_MODES}")
|
||||
if ops is None:
|
||||
ops = backend.ops
|
||||
|
||||
parsed = parse_archive(archive_bytes)
|
||||
result = ImportResult()
|
||||
|
||||
# (original document_id, fact ordinal) -> freshly inserted unit id. Used to
|
||||
# resolve observation source references after all facts exist.
|
||||
ref_map: dict[tuple[str, int], str] = {}
|
||||
|
||||
for document in parsed.documents:
|
||||
target_id = await _resolve_target_id(backend, bank_id, document.id, on_conflict)
|
||||
if target_id is None:
|
||||
result.documents_skipped += 1
|
||||
result.skipped_document_ids.append(document.id)
|
||||
continue
|
||||
if target_id != document.id:
|
||||
result.remapped_document_ids[document.id] = target_id
|
||||
|
||||
unit_ids = await _import_one_document(
|
||||
backend=backend,
|
||||
embeddings_model=embeddings_model,
|
||||
entity_resolver=entity_resolver,
|
||||
config=config,
|
||||
format_date_fn=format_date_fn,
|
||||
bank_id=bank_id,
|
||||
document=document,
|
||||
target_id=target_id,
|
||||
ops=ops,
|
||||
outbox_callback_factory=outbox_callback_factory,
|
||||
)
|
||||
result.documents_imported += 1
|
||||
result.facts_imported += len(unit_ids)
|
||||
result.imported_documents.append(
|
||||
ImportedDocument(
|
||||
document_id=target_id,
|
||||
unit_ids=unit_ids,
|
||||
content=document.original_text or "",
|
||||
tags=list(document.tags),
|
||||
)
|
||||
)
|
||||
for ordinal, unit_id in enumerate(unit_ids):
|
||||
ref_map[(document.id, ordinal)] = unit_id
|
||||
|
||||
if parsed.observations:
|
||||
outcome = await _import_observations(
|
||||
backend=backend,
|
||||
embeddings_model=embeddings_model,
|
||||
bank_id=bank_id,
|
||||
observations=parsed.observations,
|
||||
ref_map=ref_map,
|
||||
ops=ops,
|
||||
)
|
||||
result.observations_imported = outcome.imported
|
||||
result.observations_skipped = outcome.skipped
|
||||
|
||||
logger.info(
|
||||
"[transfer] Imported %d document(s), %d fact(s), %d observation(s) into bank %s "
|
||||
"(%d docs skipped, %d observations skipped)",
|
||||
result.documents_imported,
|
||||
result.facts_imported,
|
||||
result.observations_imported,
|
||||
bank_id,
|
||||
result.documents_skipped,
|
||||
result.observations_skipped,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
# Bank-level config/state tables restored verbatim from a whole-bank archive.
|
||||
# Order matters for foreign keys: banks (parent) is restored before any child.
|
||||
_BANK_CHILD_TABLES = ("mental_models", "directives", "webhooks")
|
||||
# Child-history carried verbatim; restored after its parent (mental_models) so the
|
||||
# foreign key resolves. Surrogate ids were dropped on export (the target reassigns
|
||||
# them), so these restore via fresh IDENTITY values.
|
||||
_CARRIED_HISTORY_TABLES = ("mental_model_history",)
|
||||
_HISTORY_TABLES = ("audit_log", "llm_requests")
|
||||
|
||||
|
||||
@dataclass
|
||||
class BankImportResult:
|
||||
"""Outcome of importing a whole-bank archive."""
|
||||
|
||||
bank_id: str
|
||||
documents_imported: int = 0
|
||||
facts_imported: int = 0
|
||||
observations_imported: int = 0
|
||||
mental_models_imported: int = 0
|
||||
mental_model_history_imported: int = 0
|
||||
directives_imported: int = 0
|
||||
webhooks_imported: int = 0
|
||||
history_rows_imported: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParsedBankArchive:
|
||||
"""The bank-level sections of a whole-bank archive (documents read separately)."""
|
||||
|
||||
manifest: TransferManifest
|
||||
# table name -> list of verbatim row dicts (banks, mental_models, directives, webhooks)
|
||||
bank_rows: dict[str, list[dict]] = field(default_factory=dict)
|
||||
# table name -> rows (audit_log, llm_requests), present only with --include-history
|
||||
history_rows: dict[str, list[dict]] = field(default_factory=dict)
|
||||
|
||||
|
||||
def parse_bank_archive(archive_bytes: bytes) -> ParsedBankArchive:
|
||||
"""Parse the bank-level sections of a whole-bank archive (``archive_type='bank'``)."""
|
||||
with zipfile.ZipFile(io.BytesIO(archive_bytes), "r") as zf:
|
||||
names = set(zf.namelist())
|
||||
if "manifest.json" not in names:
|
||||
raise ValueError("Invalid transfer archive: manifest.json is missing")
|
||||
manifest = TransferManifest.model_validate_json(zf.read("manifest.json"))
|
||||
if manifest.archive_type != "bank":
|
||||
raise ValueError(
|
||||
f"Not a whole-bank archive (archive_type={manifest.archive_type!r}); use import_documents instead"
|
||||
)
|
||||
bank_rows: dict[str, list[dict]] = {}
|
||||
for table in ("banks", *_BANK_CHILD_TABLES, *_CARRIED_HISTORY_TABLES):
|
||||
fname = f"{table}.json"
|
||||
bank_rows[table] = json.loads(zf.read(fname)) if fname in names else []
|
||||
history_rows: dict[str, list[dict]] = {}
|
||||
for table in _HISTORY_TABLES:
|
||||
fname = f"history/{table}.json"
|
||||
if fname in names:
|
||||
history_rows[table] = json.loads(zf.read(fname))
|
||||
return ParsedBankArchive(manifest=manifest, bank_rows=bank_rows, history_rows=history_rows)
|
||||
|
||||
|
||||
async def _restore_rows(conn: Any, table: str, rows: list[dict]) -> int:
|
||||
"""Insert verbatim rows into a bank-scoped table, coercing JSON-encoded values
|
||||
back to the column's type (timestamps, uuids, jsonb). ``ON CONFLICT DO NOTHING``
|
||||
keeps an import idempotent and safe to re-run against a partially-filled target."""
|
||||
if not rows:
|
||||
return 0
|
||||
from ..memory_engine import get_current_schema
|
||||
|
||||
schema = get_current_schema()
|
||||
col_types = {
|
||||
r["column_name"]: r["data_type"]
|
||||
for r in await conn.fetch(
|
||||
"SELECT column_name, data_type FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2",
|
||||
schema,
|
||||
table,
|
||||
)
|
||||
}
|
||||
inserted = 0
|
||||
for row in rows:
|
||||
cols = [c for c in row if c in col_types]
|
||||
placeholders: list[str] = []
|
||||
values: list[Any] = []
|
||||
for position, col in enumerate(cols, start=1):
|
||||
data_type = col_types[col]
|
||||
value = row[col]
|
||||
if data_type in ("jsonb", "json"):
|
||||
# asyncpg has no JSON codec on these raw connections; pass JSON
|
||||
# text and cast. Values may already be str (no codec on export) or
|
||||
# a Python object (codec on export) — normalize to text either way.
|
||||
values.append(value if isinstance(value, str) or value is None else json.dumps(value))
|
||||
placeholders.append(f"${position}::jsonb")
|
||||
continue
|
||||
if value is not None and isinstance(value, str):
|
||||
if data_type in ("timestamp with time zone", "timestamp without time zone"):
|
||||
value = datetime.fromisoformat(value)
|
||||
elif data_type == "date":
|
||||
value = date.fromisoformat(value)
|
||||
elif data_type == "uuid":
|
||||
value = uuid.UUID(value)
|
||||
placeholders.append(f"${position}")
|
||||
values.append(value)
|
||||
col_list = ", ".join(f'"{c}"' for c in cols)
|
||||
await conn.execute(
|
||||
f"INSERT INTO {fq_table(table)} ({col_list}) VALUES ({', '.join(placeholders)}) ON CONFLICT DO NOTHING",
|
||||
*values,
|
||||
)
|
||||
inserted += 1
|
||||
return inserted
|
||||
|
||||
|
||||
async def import_bank(
|
||||
*,
|
||||
backend: Any,
|
||||
embeddings_model: Any,
|
||||
entity_resolver: Any,
|
||||
config: Any,
|
||||
format_date_fn: Any,
|
||||
archive_bytes: bytes,
|
||||
target_bank_id: str | None = None,
|
||||
include_history: bool = False,
|
||||
ops: Any = None,
|
||||
) -> BankImportResult:
|
||||
"""Restore a whole bank from a ``export_bank`` archive into the target instance.
|
||||
|
||||
Re-embeds facts with the *target* instance's embedding model and rebuilds links,
|
||||
entities and search/vector indexes — the path for migrating a bank to an instance
|
||||
configured with a different embedding model / vector / text-search backend.
|
||||
|
||||
The **target bank must not already exist**: import restores a complete bank
|
||||
(config + facts + mental models + …) and is not a merge. If a bank with the
|
||||
target id is present, this raises — delete it first or pass ``target_bank_id``
|
||||
for a fresh id. A migration restores *exact* state, so unlike the document
|
||||
import it fires no retain webhooks and triggers no consolidation/graph
|
||||
maintenance: observations and mental models are restored as exported.
|
||||
"""
|
||||
if ops is None:
|
||||
ops = backend.ops
|
||||
parsed = parse_bank_archive(archive_bytes)
|
||||
source_bank_id = parsed.manifest.source_bank_id
|
||||
bank_id = target_bank_id or source_bank_id
|
||||
|
||||
# Remapping to a different id: rewrite the carried bank_id on every row so FKs
|
||||
# and PKs line up with the (also-remapped) documents/facts.
|
||||
if bank_id != source_bank_id:
|
||||
for rows in (*parsed.bank_rows.values(), *parsed.history_rows.values()):
|
||||
for row in rows:
|
||||
if "bank_id" in row:
|
||||
row["bank_id"] = bank_id
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
# Refuse to import into an existing bank — this restores a whole bank, it
|
||||
# does not merge. Merging would silently mix the archive's config/mental
|
||||
# models/webhooks with whatever is already there (and global-unique ids
|
||||
# like webhooks/directives would collide).
|
||||
if await conn.fetchval(f"SELECT 1 FROM {fq_table('banks')} WHERE bank_id = $1", bank_id):
|
||||
raise ValueError(
|
||||
f"Target bank '{bank_id}' already exists; import-bank restores into a fresh bank "
|
||||
f"(it is not a merge). Delete the bank first, or pass a different target bank id."
|
||||
)
|
||||
# Bank row first — children (documents, mental_models, …) FK to it.
|
||||
await _restore_rows(conn, "banks", parsed.bank_rows.get("banks", []))
|
||||
# Ensure the bank's per-bank vector indexes exist (no-op for global-index
|
||||
# extensions); idempotent and keeps the restored banks row (ON CONFLICT DO NOTHING).
|
||||
await bank_utils.get_or_create_bank_profile(backend, bank_id)
|
||||
|
||||
doc_result = await import_documents(
|
||||
backend=backend,
|
||||
embeddings_model=embeddings_model,
|
||||
entity_resolver=entity_resolver,
|
||||
config=config,
|
||||
format_date_fn=format_date_fn,
|
||||
bank_id=bank_id,
|
||||
archive_bytes=archive_bytes,
|
||||
ops=ops,
|
||||
outbox_callback_factory=None,
|
||||
)
|
||||
|
||||
result = BankImportResult(
|
||||
bank_id=bank_id,
|
||||
documents_imported=doc_result.documents_imported,
|
||||
facts_imported=doc_result.facts_imported,
|
||||
observations_imported=doc_result.observations_imported,
|
||||
)
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
result.mental_models_imported = await _restore_rows(
|
||||
conn, "mental_models", parsed.bank_rows.get("mental_models", [])
|
||||
)
|
||||
# Restored after mental_models so the (mental_model_id, bank_id) FK resolves.
|
||||
result.mental_model_history_imported = await _restore_rows(
|
||||
conn, "mental_model_history", parsed.bank_rows.get("mental_model_history", [])
|
||||
)
|
||||
result.directives_imported = await _restore_rows(conn, "directives", parsed.bank_rows.get("directives", []))
|
||||
result.webhooks_imported = await _restore_rows(conn, "webhooks", parsed.bank_rows.get("webhooks", []))
|
||||
if include_history:
|
||||
for table in _HISTORY_TABLES:
|
||||
result.history_rows_imported += await _restore_rows(conn, table, parsed.history_rows.get(table, []))
|
||||
|
||||
logger.info(
|
||||
"[transfer] Imported bank %s: %d doc(s), %d fact(s), %d observation(s), "
|
||||
"%d mental model(s), %d mm-history row(s), %d directive(s), %d webhook(s), %d history row(s)",
|
||||
bank_id,
|
||||
result.documents_imported,
|
||||
result.facts_imported,
|
||||
result.observations_imported,
|
||||
result.mental_models_imported,
|
||||
result.mental_model_history_imported,
|
||||
result.directives_imported,
|
||||
result.webhooks_imported,
|
||||
result.history_rows_imported,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
async def _resolve_target_id(backend: Any, bank_id: str, document_id: str, on_conflict: OnConflict) -> str | None:
|
||||
"""Decide the document id to write under, or ``None`` to skip.
|
||||
|
||||
Returns the original id when there is no conflict, a fresh id under
|
||||
``new-id``, the original id under ``replace`` (the insert path cascades the
|
||||
old data away), or ``None`` under ``skip`` when the document already exists.
|
||||
"""
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
exists = await conn.fetchval(
|
||||
f"SELECT 1 FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
if not exists:
|
||||
return document_id
|
||||
if on_conflict == "skip":
|
||||
return None
|
||||
if on_conflict == "new-id":
|
||||
return str(uuid.uuid4())
|
||||
return document_id # replace
|
||||
|
||||
|
||||
async def _import_one_document(
|
||||
*,
|
||||
backend: Any,
|
||||
embeddings_model: Any,
|
||||
entity_resolver: Any,
|
||||
config: Any,
|
||||
format_date_fn: Any,
|
||||
bank_id: str,
|
||||
document: TransferDocument,
|
||||
target_id: str,
|
||||
ops: Any,
|
||||
outbox_callback_factory: Any = None,
|
||||
) -> list[str]:
|
||||
"""Re-embed and insert a single document; returns the new unit ids in fact order."""
|
||||
log_buffer: list[str] = []
|
||||
|
||||
# Fire the same retain.completed webhook retain emits, transactionally inside
|
||||
# this document's insert. Factory returns None when no webhook manager exists.
|
||||
outbox_callback = (
|
||||
outbox_callback_factory([{"document_id": target_id, "tags": list(document.tags)}])
|
||||
if outbox_callback_factory
|
||||
else None
|
||||
)
|
||||
|
||||
extracted_facts = [_to_extracted_fact(fact) for fact in document.facts]
|
||||
|
||||
processed_facts: list[ProcessedFact] = []
|
||||
if extracted_facts:
|
||||
augmented = embedding_processing.augment_texts_with_dates(extracted_facts, format_date_fn)
|
||||
embeddings = await embedding_processing.generate_embeddings_batch(embeddings_model, augmented)
|
||||
processed_facts = [ProcessedFact.from_extracted_fact(ef, emb) for ef, emb in zip(extracted_facts, embeddings)]
|
||||
|
||||
contents = [RetainContent(content=document.original_text or "")]
|
||||
chunk_meta = [
|
||||
ChunkMetadata(chunk_text=chunk.chunk_text, fact_count=0, content_index=0, chunk_index=chunk.chunk_index)
|
||||
for chunk in document.chunks
|
||||
]
|
||||
|
||||
# Phase 1 (entity resolution + semantic ANN) on its own connection, outside
|
||||
# the write transaction — mirrors the retain pipeline.
|
||||
entity_resolver.discard_pending_stats()
|
||||
phase1 = await orchestrator._pre_resolve_phase1(
|
||||
backend,
|
||||
entity_resolver,
|
||||
bank_id,
|
||||
contents,
|
||||
processed_facts,
|
||||
config,
|
||||
log_buffer,
|
||||
skip_semantic_ann=False,
|
||||
)
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
# is_first_batch=True: cascade-delete any existing data for this id
|
||||
# (the "replace" path) and (re)insert the document row.
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn,
|
||||
bank_id,
|
||||
target_id,
|
||||
document.original_text or "",
|
||||
True,
|
||||
document.retain_params,
|
||||
document.tags,
|
||||
ops=ops,
|
||||
)
|
||||
|
||||
chunk_id_map: dict[int, str] = {}
|
||||
if chunk_meta:
|
||||
chunk_id_map = await chunk_storage.store_chunks_batch(conn, bank_id, target_id, chunk_meta, ops=ops)
|
||||
|
||||
for extracted, processed in zip(extracted_facts, processed_facts):
|
||||
processed.document_id = target_id
|
||||
if chunk_id_map and extracted.chunk_index is not None:
|
||||
chunk_id = chunk_id_map.get(extracted.chunk_index)
|
||||
if chunk_id:
|
||||
processed.chunk_id = chunk_id
|
||||
|
||||
result_unit_ids = await orchestrator._insert_facts_and_links(
|
||||
conn,
|
||||
entity_resolver,
|
||||
bank_id,
|
||||
contents,
|
||||
extracted_facts,
|
||||
processed_facts,
|
||||
config,
|
||||
log_buffer,
|
||||
resolved_entity_ids=phase1.entities.resolved_entity_ids,
|
||||
entity_to_unit=phase1.entities.entity_to_unit,
|
||||
unit_to_entity_ids=phase1.entities.unit_to_entity_ids,
|
||||
semantic_ann_links=phase1.semantic_ann_links,
|
||||
skip_semantic_links=False,
|
||||
outbox_callback=outbox_callback,
|
||||
ops=ops,
|
||||
)
|
||||
|
||||
try:
|
||||
await entity_resolver.flush_pending_stats()
|
||||
except Exception:
|
||||
logger.warning("[transfer] Entity stats flush failed for document %s", target_id, exc_info=True)
|
||||
|
||||
logger.debug("[transfer] Imported document %s:\n%s", target_id, "\n".join(log_buffer))
|
||||
# Single content item -> result_unit_ids[0] holds the new unit ids in fact order.
|
||||
return list(result_unit_ids[0]) if result_unit_ids else []
|
||||
|
||||
|
||||
async def _import_observations(
|
||||
*,
|
||||
backend: Any,
|
||||
embeddings_model: Any,
|
||||
bank_id: str,
|
||||
observations: list[TransferObservation],
|
||||
ref_map: dict[tuple[str, int], str],
|
||||
ops: Any,
|
||||
) -> _ObservationOutcome:
|
||||
"""Insert observations whose source facts were all imported in this run.
|
||||
|
||||
Observations carry no embedding, links, or entity rows — only the unit row
|
||||
plus ``source_memory_ids`` (remapped to the freshly inserted source units)
|
||||
and ``proof_count``. Their source facts are marked ``consolidated_at`` so the
|
||||
target bank's consolidator won't re-process them. Mirrors what consolidation
|
||||
writes, but driven from the archive instead of the LLM.
|
||||
|
||||
Inserted as-is: imported observations are NOT merged or deduplicated against
|
||||
observations that already exist in the target bank (unlike consolidation,
|
||||
which merges related observations). Importing into a bank that already has
|
||||
observations — or importing the same archive twice — can therefore produce
|
||||
overlapping observations over the same facts.
|
||||
"""
|
||||
outcome = _ObservationOutcome()
|
||||
|
||||
# Resolve each observation's sources to new unit ids; drop any whose sources
|
||||
# weren't all imported (e.g. a subset/skip import).
|
||||
resolved: list[tuple[TransferObservation, list[str]]] = []
|
||||
for obs in observations:
|
||||
source_ids = [ref_map.get((s.document_id, s.fact_index)) for s in obs.sources]
|
||||
if not source_ids or any(sid is None for sid in source_ids):
|
||||
outcome.skipped += 1
|
||||
continue
|
||||
resolved.append((obs, [sid for sid in source_ids if sid is not None]))
|
||||
|
||||
if not resolved:
|
||||
return outcome
|
||||
|
||||
# Observations embed the raw text (matching consolidation), not the
|
||||
# date-augmented text used for facts.
|
||||
embeddings = await embedding_processing.generate_embeddings_batch(
|
||||
embeddings_model, [obs.text for obs, _ in resolved]
|
||||
)
|
||||
processed = [
|
||||
ProcessedFact(
|
||||
fact_text=obs.text,
|
||||
fact_type="observation",
|
||||
embedding=embedding,
|
||||
occurred_start=obs.occurred_start,
|
||||
occurred_end=obs.occurred_end,
|
||||
mentioned_at=_observation_mentioned_at(obs),
|
||||
context="",
|
||||
metadata={},
|
||||
tags=list(obs.tags),
|
||||
observation_scopes=obs.observation_scopes,
|
||||
document_id=None,
|
||||
chunk_id=None,
|
||||
)
|
||||
for (obs, _sources), embedding in zip(resolved, embeddings)
|
||||
]
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
obs_unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, processed, ops=ops)
|
||||
|
||||
all_source_ids: set[uuid.UUID] = set()
|
||||
for (obs, sources), obs_unit_id in zip(resolved, obs_unit_ids):
|
||||
source_uuids = [uuid.UUID(s) for s in sources]
|
||||
all_source_ids.update(source_uuids)
|
||||
await _link_observation_sources(
|
||||
conn, ops, bank_id, uuid.UUID(obs_unit_id), source_uuids, obs.proof_count
|
||||
)
|
||||
|
||||
# Mark source facts consolidated so the target consolidator skips them.
|
||||
if all_source_ids:
|
||||
await conn.execute(
|
||||
f"UPDATE {fq_table('memory_units')} SET consolidated_at = now() "
|
||||
f"WHERE bank_id = $1 AND id = ANY($2)",
|
||||
bank_id,
|
||||
list(all_source_ids),
|
||||
)
|
||||
|
||||
outcome.imported = len(resolved)
|
||||
return outcome
|
||||
|
||||
|
||||
async def _link_observation_sources(
|
||||
conn: Any,
|
||||
ops: Any,
|
||||
bank_id: str,
|
||||
observation_id: uuid.UUID,
|
||||
source_ids: list[uuid.UUID],
|
||||
proof_count: int,
|
||||
) -> None:
|
||||
"""Attach source ids + proof_count to a freshly inserted observation row.
|
||||
|
||||
PG stores the sources in the ``source_memory_ids`` array column; Oracle uses
|
||||
the ``observation_sources`` junction table (same split as consolidation).
|
||||
"""
|
||||
if ops.uses_observation_sources_table:
|
||||
await conn.executemany(
|
||||
f"INSERT INTO {fq_table('observation_sources')} (observation_id, source_id) "
|
||||
f"VALUES ($1, $2) ON CONFLICT (observation_id, source_id) DO NOTHING",
|
||||
[(observation_id, sid) for sid in dict.fromkeys(source_ids)],
|
||||
)
|
||||
await conn.execute(
|
||||
f"UPDATE {fq_table('memory_units')} SET proof_count = $1 WHERE id = $2 AND bank_id = $3",
|
||||
proof_count,
|
||||
observation_id,
|
||||
bank_id,
|
||||
)
|
||||
else:
|
||||
await conn.execute(
|
||||
f"UPDATE {fq_table('memory_units')} SET source_memory_ids = $1, proof_count = $2 "
|
||||
f"WHERE id = $3 AND bank_id = $4",
|
||||
source_ids,
|
||||
proof_count,
|
||||
observation_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
def _observation_mentioned_at(obs: TransferObservation) -> datetime | None:
|
||||
"""event_date (NOT NULL) is derived from occurred_start or mentioned_at on
|
||||
insert; fall back so the column stays populated for observations too."""
|
||||
mentioned_at = obs.mentioned_at
|
||||
if obs.occurred_start is None and mentioned_at is None:
|
||||
mentioned_at = obs.event_date or datetime.now(UTC)
|
||||
return mentioned_at
|
||||
|
||||
|
||||
def _to_extracted_fact(fact: TransferFact) -> ExtractedFact:
|
||||
"""Rebuild the retain pipeline's ExtractedFact from a serialized transfer fact."""
|
||||
# event_date is NOT NULL in the schema and is derived from occurred_start or
|
||||
# mentioned_at on insert. When neither is present, fall back to the carried
|
||||
# event_date (or now) via mentioned_at so the column stays populated.
|
||||
mentioned_at = fact.mentioned_at
|
||||
if fact.occurred_start is None and mentioned_at is None:
|
||||
mentioned_at = fact.event_date or datetime.now(UTC)
|
||||
|
||||
return ExtractedFact(
|
||||
fact_text=fact.text,
|
||||
fact_type=fact.fact_type,
|
||||
entities=list(fact.entities),
|
||||
occurred_start=fact.occurred_start,
|
||||
occurred_end=fact.occurred_end,
|
||||
where=None,
|
||||
causal_relations=[
|
||||
CausalRelation(relation_type=rel.relation_type, target_fact_index=rel.target_fact_index)
|
||||
for rel in fact.causal_relations
|
||||
],
|
||||
content_index=0,
|
||||
chunk_index=fact.chunk_index,
|
||||
context=fact.context or "",
|
||||
mentioned_at=mentioned_at,
|
||||
metadata=dict(fact.metadata),
|
||||
tags=list(fact.tags),
|
||||
observation_scopes=fact.observation_scopes,
|
||||
)
|
||||
@@ -0,0 +1,138 @@
|
||||
"""Serialization schema for the document transfer archive (manifest + per-document payloads).
|
||||
|
||||
The archive is a ZIP:
|
||||
|
||||
manifest.json -- TransferManifest
|
||||
documents/000000.json -- TransferDocument (one file per document)
|
||||
documents/000001.json
|
||||
...
|
||||
|
||||
Documents are stored under a zero-padded index rather than their id so that
|
||||
arbitrary document ids (which may contain path-unsafe characters) never leak
|
||||
into archive entry names. The real id lives inside each payload.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
# Bump when the archive layout changes in a backward-incompatible way.
|
||||
SCHEMA_VERSION = 1
|
||||
|
||||
ObservationScopes = Literal["per_tag", "combined", "all_combinations"] | list[list[str]]
|
||||
|
||||
|
||||
class TransferCausalRelation(BaseModel):
|
||||
"""A causal edge from this fact to an earlier fact in the same document.
|
||||
|
||||
``target_fact_index`` is the ordinal of the target fact within the document's
|
||||
``facts`` list (not a database id), so it survives transfer to a new bank.
|
||||
"""
|
||||
|
||||
relation_type: str
|
||||
target_fact_index: int
|
||||
|
||||
|
||||
class TransferFact(BaseModel):
|
||||
"""One extracted fact (memory unit) without its embedding or database id.
|
||||
|
||||
Everything here is reused verbatim on import except the embedding, which is
|
||||
regenerated by the target bank's model, and the entity ids, which are
|
||||
re-resolved against the target bank by canonical name.
|
||||
"""
|
||||
|
||||
text: str
|
||||
fact_type: str
|
||||
context: str | None = None
|
||||
# event_date is a fallback used only when both occurred_start and
|
||||
# mentioned_at are absent, to satisfy the NOT NULL event_date column.
|
||||
event_date: datetime | None = None
|
||||
occurred_start: datetime | None = None
|
||||
occurred_end: datetime | None = None
|
||||
mentioned_at: datetime | None = None
|
||||
metadata: dict[str, str] = Field(default_factory=dict)
|
||||
tags: list[str] = Field(default_factory=list)
|
||||
observation_scopes: ObservationScopes | None = None
|
||||
# Ordinal of the source chunk within the document (parsed from chunk_id).
|
||||
chunk_index: int | None = None
|
||||
# Entity canonical names; re-resolved against the target bank on import.
|
||||
entities: list[str] = Field(default_factory=list)
|
||||
causal_relations: list[TransferCausalRelation] = Field(default_factory=list)
|
||||
|
||||
|
||||
class TransferChunk(BaseModel):
|
||||
"""A raw text chunk of the source document, reused verbatim."""
|
||||
|
||||
chunk_index: int
|
||||
chunk_text: str
|
||||
|
||||
|
||||
class TransferObservationSource(BaseModel):
|
||||
"""A reference to a source fact of an observation, by document + ordinal.
|
||||
|
||||
Observations span documents and reference their source facts by unit id;
|
||||
those ids don't survive transfer, so each source is carried as the
|
||||
(document_id, fact_index) of the fact within the exported document set.
|
||||
"""
|
||||
|
||||
document_id: str
|
||||
fact_index: int
|
||||
|
||||
|
||||
class TransferObservation(BaseModel):
|
||||
"""A consolidated observation (``fact_type='observation'``).
|
||||
|
||||
Observations are bank-level (not tied to one document), carry no embedding
|
||||
(re-generated on import) and no entity/link associations — retrieval reaches
|
||||
entities/links through their source facts. Only exported when explicitly
|
||||
requested, and only when every source resolves within the archive.
|
||||
"""
|
||||
|
||||
text: str
|
||||
tags: list[str] = Field(default_factory=list)
|
||||
event_date: datetime | None = None
|
||||
occurred_start: datetime | None = None
|
||||
occurred_end: datetime | None = None
|
||||
mentioned_at: datetime | None = None
|
||||
observation_scopes: ObservationScopes | None = None
|
||||
proof_count: int = 1
|
||||
sources: list[TransferObservationSource] = Field(default_factory=list)
|
||||
|
||||
|
||||
class TransferDocument(BaseModel):
|
||||
"""A single document plus its chunks and extracted facts."""
|
||||
|
||||
id: str
|
||||
original_text: str | None = None
|
||||
retain_params: dict | None = None
|
||||
tags: list[str] = Field(default_factory=list)
|
||||
created_at: datetime | None = None
|
||||
chunks: list[TransferChunk] = Field(default_factory=list)
|
||||
facts: list[TransferFact] = Field(default_factory=list)
|
||||
|
||||
|
||||
class TransferManifest(BaseModel):
|
||||
"""Top-level archive descriptor (``manifest.json``).
|
||||
|
||||
The bank-level fields default to a documents-only archive so older
|
||||
document-only archives (and the document import path) keep parsing
|
||||
unchanged; ``export_bank`` populates them for a whole-bank archive.
|
||||
"""
|
||||
|
||||
schema_version: int = SCHEMA_VERSION
|
||||
source_bank_id: str
|
||||
exported_at: datetime | None = None
|
||||
document_count: int = 0
|
||||
fact_count: int = 0
|
||||
observation_count: int = 0
|
||||
# "documents" = doc/fact/observation subset; "bank" = whole-bank export
|
||||
# (also carries bank config, mental models, directives, webhooks).
|
||||
archive_type: Literal["documents", "bank"] = "documents"
|
||||
mental_model_count: int = 0
|
||||
directive_count: int = 0
|
||||
webhook_count: int = 0
|
||||
# True when --include-history carried audit_log / llm_requests.
|
||||
includes_history: bool = False
|
||||
@@ -16,11 +16,24 @@ with the system (e.g., running migrations for tenant schemas).
|
||||
"""
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
from hindsight_api.extensions.builtin import ApiKeyTenantExtension, SupabaseTenantExtension
|
||||
from hindsight_api.extensions.builtin import (
|
||||
ApiKeyTenantExtension,
|
||||
MemoryDefenseRegexExtension,
|
||||
SupabaseTenantExtension,
|
||||
)
|
||||
from hindsight_api.extensions.context import DefaultExtensionContext, ExtensionContext
|
||||
from hindsight_api.extensions.http import HttpExtension
|
||||
from hindsight_api.extensions.loader import load_extension
|
||||
from hindsight_api.extensions.mcp import MCPExtension
|
||||
from hindsight_api.extensions.memory_defense import (
|
||||
DefenseAction,
|
||||
DefenseDecision,
|
||||
DefensePolicy,
|
||||
MemoryDefenseExtension,
|
||||
PolicyRule,
|
||||
apply_redaction,
|
||||
parse_policy,
|
||||
)
|
||||
from hindsight_api.extensions.operation_validator import (
|
||||
# Bank Management operations
|
||||
BankListContext,
|
||||
@@ -104,4 +117,13 @@ __all__ = [
|
||||
"Tenant",
|
||||
"TenantContext",
|
||||
"TenantExtension",
|
||||
# Memory Defense
|
||||
"DefenseAction",
|
||||
"DefenseDecision",
|
||||
"DefensePolicy",
|
||||
"MemoryDefenseExtension",
|
||||
"MemoryDefenseRegexExtension",
|
||||
"PolicyRule",
|
||||
"apply_redaction",
|
||||
"parse_policy",
|
||||
]
|
||||
|
||||
@@ -13,10 +13,12 @@ Example usage:
|
||||
HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.supabase_tenant:SupabaseTenantExtension
|
||||
"""
|
||||
|
||||
from hindsight_api.extensions.builtin.memory_defense_regex import MemoryDefenseRegexExtension
|
||||
from hindsight_api.extensions.builtin.supabase_tenant import SupabaseTenantExtension
|
||||
from hindsight_api.extensions.builtin.tenant import ApiKeyTenantExtension
|
||||
|
||||
__all__ = [
|
||||
"ApiKeyTenantExtension",
|
||||
"MemoryDefenseRegexExtension",
|
||||
"SupabaseTenantExtension",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Memory Defense (regex) — the default extension shipping with hindsight-api-slim.
|
||||
|
||||
Scrubs known secret/PII patterns from retained content via the
|
||||
``sensitive_data`` detector. Matching is pure regex (see ``apply_redaction``):
|
||||
no LLM call, no external dependency. A ``sensitive_data`` rule may either
|
||||
``redact`` matches in place or ``block`` the item entirely.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from hindsight_api.extensions.memory_defense import (
|
||||
DefenseAction,
|
||||
DefenseDecision,
|
||||
DefensePolicy,
|
||||
MemoryDefenseExtension,
|
||||
apply_redaction,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MemoryDefenseRegexExtension(MemoryDefenseExtension):
|
||||
"""Default Memory Defense — regex-based secret/PII redaction."""
|
||||
|
||||
async def screen(
|
||||
self,
|
||||
*,
|
||||
policy: DefensePolicy,
|
||||
bank_id: str,
|
||||
document_id: str | None,
|
||||
content: str,
|
||||
tags: list[str],
|
||||
) -> DefenseDecision:
|
||||
if not policy.enabled:
|
||||
return DefenseDecision(action=DefenseAction.ALLOW)
|
||||
|
||||
# The regex extension only runs the sensitive_data detector. If the
|
||||
# policy doesn't include a rule for it, there's nothing to do.
|
||||
rule = next((r for r in policy.rules if r.on == "sensitive_data"), None)
|
||||
if rule is None or rule.action is DefenseAction.ALLOW:
|
||||
return DefenseDecision(action=DefenseAction.ALLOW)
|
||||
|
||||
result = apply_redaction(content)
|
||||
if not result.matched_types:
|
||||
return DefenseDecision(action=DefenseAction.ALLOW)
|
||||
|
||||
return DefenseDecision(
|
||||
action=rule.action,
|
||||
detector="sensitive_data",
|
||||
message=f"Sensitive data pattern matched: {', '.join(result.matched_types)}",
|
||||
redacted_content=result.content if rule.action is DefenseAction.REDACT else None,
|
||||
matched_types=result.matched_types,
|
||||
)
|
||||
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.interface import MemoryEngineInterface
|
||||
from hindsight_api.webhooks.manager import WebhookManager
|
||||
|
||||
|
||||
class ExtensionContext(ABC):
|
||||
@@ -83,6 +84,8 @@ class DefaultExtensionContext(ExtensionContext):
|
||||
self,
|
||||
database_url: str,
|
||||
memory_engine: "MemoryEngineInterface | None" = None,
|
||||
webhook_manager: "WebhookManager | None" = None,
|
||||
current_schema: str | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize the context.
|
||||
@@ -90,9 +93,13 @@ class DefaultExtensionContext(ExtensionContext):
|
||||
Args:
|
||||
database_url: SQLAlchemy database URL for migrations.
|
||||
memory_engine: Optional MemoryEngine instance for memory operations.
|
||||
webhook_manager: Optional WebhookManager for firing webhooks.
|
||||
current_schema: Optional current schema name for tenant context.
|
||||
"""
|
||||
self._database_url = database_url
|
||||
self._memory_engine = memory_engine
|
||||
self.webhook_manager = webhook_manager
|
||||
self.current_schema = current_schema
|
||||
|
||||
async def run_migration(self, schema: str) -> None:
|
||||
"""Run migrations for a specific schema."""
|
||||
|
||||
@@ -0,0 +1,201 @@
|
||||
"""Memory Defense extension contract and shared policy types.
|
||||
|
||||
Lives in extensions/ (not engine/) because it defines the public contract
|
||||
between the retain orchestrator and any installed Memory Defense extension —
|
||||
the same shape as TenantExtension and OperationValidatorExtension.
|
||||
|
||||
api-slim ships the :class:`MemoryDefenseExtension` protocol and a regex default
|
||||
that scrubs known secret/PII patterns from retained content.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DefenseAction(str, Enum):
|
||||
ALLOW = "allow"
|
||||
REDACT = "redact"
|
||||
BLOCK = "block"
|
||||
|
||||
|
||||
_VALID_ACTIONS = {a.value for a in DefenseAction}
|
||||
|
||||
# Detector identifiers valid as ``policy.rules[*].on``. The OSS extension only
|
||||
# screens for sensitive data (secrets/PII), so that's the only accepted value.
|
||||
_VALID_DETECTORS = {"sensitive_data"}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PolicyRule:
|
||||
on: str
|
||||
action: DefenseAction
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DefensePolicy:
|
||||
enabled: bool = False
|
||||
rules: tuple[PolicyRule, ...] = ()
|
||||
|
||||
|
||||
@dataclass
|
||||
class DefenseDecision:
|
||||
action: DefenseAction
|
||||
detector: str | None = None
|
||||
message: str = ""
|
||||
redacted_content: str | None = None
|
||||
matched_types: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RedactionResult:
|
||||
content: str
|
||||
matched_types: list[str]
|
||||
|
||||
|
||||
def parse_policy(raw: dict | None) -> DefensePolicy:
|
||||
"""Parse a raw bank-config dict into a frozen DefensePolicy.
|
||||
|
||||
Raises ValueError for unknown detectors or actions; the HTTP layer
|
||||
converts those into a 422 response.
|
||||
"""
|
||||
if raw is None:
|
||||
return DefensePolicy()
|
||||
|
||||
rules: list[PolicyRule] = []
|
||||
for item in raw.get("rules", []) or []:
|
||||
on_raw = item.get("on")
|
||||
if on_raw not in _VALID_DETECTORS:
|
||||
raise ValueError(f"invalid on {on_raw!r}; must be one of {sorted(_VALID_DETECTORS)}")
|
||||
action_raw = item.get("action")
|
||||
if action_raw not in _VALID_ACTIONS:
|
||||
raise ValueError(f"invalid action {action_raw!r}; must be one of {sorted(_VALID_ACTIONS)}")
|
||||
rules.append(PolicyRule(on=on_raw, action=DefenseAction(action_raw)))
|
||||
|
||||
return DefensePolicy(
|
||||
enabled=bool(raw.get("enabled", False)),
|
||||
rules=tuple(rules),
|
||||
)
|
||||
|
||||
|
||||
# Secret/PII redaction patterns.
|
||||
#
|
||||
# Scope: high-confidence patterns with unambiguous prefixes (low false-positive
|
||||
# rate). Context-dependent matches (e.g. Cohere/Mistral keys that only stand
|
||||
# out near surrounding "cohere"/"mistral" tokens) are NOT covered by pure
|
||||
# regex — operators who need that should layer a context-aware secret
|
||||
# scanner (detect-secrets, trufflehog) on top.
|
||||
#
|
||||
# Order matters: more-specific patterns first so broader ones don't consume
|
||||
# substrings partially. Example: `sk-ant-...` and `sk-proj-...` must run
|
||||
# before the generic `sk-...` pattern.
|
||||
_REDACTION_PATTERNS: list[tuple[str, str]] = [
|
||||
# --- AI / LLM providers ---
|
||||
("anthropic_key", r"\bsk-ant-[A-Za-z0-9_-]{20,}\b"),
|
||||
("openai_project_key", r"\bsk-proj-[A-Za-z0-9_-]{48,}\b"),
|
||||
("openai_admin_key", r"\bsk-admin-[A-Za-z0-9_-]{40,}\b"),
|
||||
("openai_key", r"\bsk-[A-Za-z0-9_-]{20,}\b"),
|
||||
("google_api_key", r"\bAIza[0-9A-Za-z_-]{35}\b"),
|
||||
("google_oauth_token", r"\bya29\.[0-9A-Za-z_-]{20,}\b"),
|
||||
("xai_key", r"\bxai-[A-Za-z0-9]{40,}\b"),
|
||||
("groq_key", r"\bgsk_[A-Za-z0-9]{20,}\b"),
|
||||
("huggingface_token", r"\bhf_[A-Za-z0-9]{30,}\b"),
|
||||
("replicate_token", r"\br8_[A-Za-z0-9]{30,}\b"),
|
||||
("perplexity_key", r"\bpplx-[A-Za-z0-9]{40,}\b"),
|
||||
("databricks_token", r"\bdapi[A-Za-z0-9]{32}\b"),
|
||||
# --- Cloud providers ---
|
||||
("aws_access_key", r"\bAKIA[0-9A-Z]{16}\b"),
|
||||
("aws_session_token", r"\bASIA[0-9A-Z]{16}\b"),
|
||||
(
|
||||
"aws_secret_key",
|
||||
r"(?i)aws(.{0,20})?(secret|private)?[\s_-]?access[\s_-]?key[\s_-]?[:=][\s\"']*([A-Za-z0-9/+=]{40})",
|
||||
),
|
||||
("digitalocean_token", r"\bdop_v1_[a-f0-9]{64}\b"),
|
||||
# --- Source control & CI ---
|
||||
("github_fg_pat", r"\bgithub_pat_[A-Za-z0-9_]{60,}\b"),
|
||||
("github_token", r"\bghp_[A-Za-z0-9]{36}\b"),
|
||||
("github_app_token", r"\bghs_[A-Za-z0-9]{36}\b"),
|
||||
("github_user_token", r"\bghu_[A-Za-z0-9]{36}\b"),
|
||||
("github_refresh", r"\bghr_[A-Za-z0-9]{36}\b"),
|
||||
("github_oauth", r"\bgho_[A-Za-z0-9]{36}\b"),
|
||||
("gitlab_pat", r"\bglpat-[A-Za-z0-9_-]{20,}\b"),
|
||||
("npm_token", r"\bnpm_[A-Za-z0-9]{30,}\b"),
|
||||
("pypi_token", r"\bpypi-AgEIcHlwaS5vcmc[A-Za-z0-9_-]{20,}\b"),
|
||||
# --- Payment processors ---
|
||||
("stripe_secret", r"\bsk_(?:live|test)_[A-Za-z0-9]{20,}\b"),
|
||||
("stripe_restricted", r"\brk_(?:live|test)_[A-Za-z0-9]{20,}\b"),
|
||||
("square_token", r"\bsq0[a-z]{3}-[A-Za-z0-9_-]{22,}\b"),
|
||||
("braintree_token", r"\baccess_token\$production\$[a-z0-9]{16}\$[a-f0-9]{32}\b"),
|
||||
# --- Communication / email ---
|
||||
("slack_token", r"\bxox[abpr]-[0-9A-Za-z-]{10,}\b"),
|
||||
("slack_webhook", r"https://hooks\.slack\.com/services/T[A-Za-z0-9_]{8,}/B[A-Za-z0-9_]{8,}/[A-Za-z0-9_]{20,}"),
|
||||
("twilio_api_key", r"\bSK[0-9a-fA-F]{32}\b"),
|
||||
("twilio_account_sid", r"\bAC[0-9a-fA-F]{32}\b"),
|
||||
("sendgrid_key", r"\bSG\.[A-Za-z0-9_-]{22}\.[A-Za-z0-9_-]{43}\b"),
|
||||
("mailgun_key", r"\bkey-[A-Za-z0-9]{32}\b"),
|
||||
("discord_bot", r"\b[MNO][A-Za-z0-9]{23}\.[A-Za-z0-9_-]{6}\.[A-Za-z0-9_-]{27}\b"),
|
||||
("telegram_bot", r"\b[0-9]{8,10}:[A-Za-z0-9_-]{35}\b"),
|
||||
# --- Commerce ---
|
||||
("shopify_token", r"\bshpat_[a-fA-F0-9]{32}\b"),
|
||||
# --- Database connection strings (creds embedded in URL) ---
|
||||
("db_url_postgres", r"postgres(?:ql)?://[^\s:/@]+:[^\s/@]+@[^\s]+"),
|
||||
("db_url_mysql", r"mysql://[^\s:/@]+:[^\s/@]+@[^\s]+"),
|
||||
("db_url_mongodb", r"mongodb(?:\+srv)?://[^\s:/@]+:[^\s/@]+@[^\s]+"),
|
||||
# --- Private keys & generic credentials ---
|
||||
("private_key_pem", r"-----BEGIN (?:RSA |EC |DSA |OPENSSH |PGP )?PRIVATE KEY( BLOCK)?-----"),
|
||||
("jwt", r"\beyJ[A-Za-z0-9_-]{10,}\.eyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\b"),
|
||||
# --- PII (US-centric defaults; can be tuned per deployment) ---
|
||||
# NOTE: credit_card regex is intentionally narrowed to 13-19 digits with
|
||||
# exact separators to reduce false positives on long product IDs.
|
||||
("credit_card", r"\b(?:\d{4}[ -]?){3}\d{1,4}\b"),
|
||||
("ssn_us", r"\b\d{3}-\d{2}-\d{4}\b"),
|
||||
]
|
||||
_COMPILED_REDACTIONS: list[tuple[str, re.Pattern]] = [
|
||||
(label, re.compile(pattern)) for label, pattern in _REDACTION_PATTERNS
|
||||
]
|
||||
|
||||
|
||||
def apply_redaction(content: str) -> RedactionResult:
|
||||
"""Scrub known secret/PII patterns from content with [REDACTED:type] markers.
|
||||
|
||||
Returns the (possibly unchanged) content alongside the list of pattern
|
||||
labels that matched (empty when nothing matched).
|
||||
"""
|
||||
matched: list[str] = []
|
||||
for label, pattern in _COMPILED_REDACTIONS:
|
||||
new_content = pattern.sub(f"[REDACTED:{label}]", content)
|
||||
if new_content != content:
|
||||
matched.append(label)
|
||||
content = new_content
|
||||
return RedactionResult(content=content, matched_types=matched)
|
||||
|
||||
|
||||
class MemoryDefenseExtension(Extension, ABC):
|
||||
"""Abstract base for Memory Defense extensions.
|
||||
|
||||
Implementations decide whether to allow, redact, or block a given retain
|
||||
item by inspecting its content against a per-bank policy. The orchestrator
|
||||
applies the returned decision (redacts content / drops blocked items) and
|
||||
fires a webhook for non-allow decisions when one is configured.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def screen(
|
||||
self,
|
||||
*,
|
||||
policy: DefensePolicy,
|
||||
bank_id: str,
|
||||
document_id: str | None,
|
||||
content: str,
|
||||
tags: list[str],
|
||||
) -> DefenseDecision:
|
||||
"""Inspect content under the given policy and return a decision."""
|
||||
...
|
||||
@@ -40,6 +40,11 @@ class Tenant:
|
||||
"""
|
||||
|
||||
schema: str
|
||||
# Optional tenant identifier. When provided, background maintenance (e.g. the
|
||||
# consolidation reconcile sweep) can build a RequestContext carrying this id so
|
||||
# tenant-level config overrides are honored. Leave as None for single-tenant
|
||||
# setups or extensions that do not key config by tenant id.
|
||||
tenant_id: str | None = None
|
||||
|
||||
|
||||
class TenantExtension(Extension, ABC):
|
||||
|
||||
@@ -50,6 +50,8 @@ _ALL_TOOLS: frozenset[str] = frozenset(
|
||||
"delete_directive",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"update_memory",
|
||||
"invalidate_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"delete_document",
|
||||
@@ -228,6 +230,8 @@ def register_mcp_tools(
|
||||
"delete_directive",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"update_memory",
|
||||
"invalidate_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"delete_document",
|
||||
@@ -299,6 +303,12 @@ def register_mcp_tools(
|
||||
if "get_memory" in tools_to_register:
|
||||
_register_get_memory(mcp, memory, config)
|
||||
|
||||
if "update_memory" in tools_to_register:
|
||||
_register_update_memory(mcp, memory, config)
|
||||
|
||||
if "invalidate_memory" in tools_to_register:
|
||||
_register_invalidate_memory(mcp, memory, config)
|
||||
|
||||
# Document tools
|
||||
if "list_documents" in tools_to_register:
|
||||
_register_list_documents(mcp, memory, config)
|
||||
@@ -2293,6 +2303,206 @@ def _register_get_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the update_memory (edit) tool."""
|
||||
|
||||
_EDIT_DOC = """
|
||||
Edit a memory unit to correct what was extracted.
|
||||
|
||||
Pass any of text / context / occurred_start / occurred_end / fact_type /
|
||||
entities. For context and the dates, "" clears the field and omitting it
|
||||
leaves it unchanged; entities replaces the fact's entity set ([] detaches
|
||||
all). The memory is re-embedded and its derived observations, links, and
|
||||
graph are recomputed automatically.
|
||||
|
||||
Only raw world/experience facts can be edited; observations are derived.
|
||||
To retire or restore a fact, use invalidate_memory instead.
|
||||
"""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def update_memory(
|
||||
memory_id: str,
|
||||
text: str | None = None,
|
||||
context: str | None = None,
|
||||
occurred_start: str | None = None,
|
||||
occurred_end: str | None = None,
|
||||
fact_type: str | None = None,
|
||||
entities: list[str] | None = None,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
f"""{_EDIT_DOC}
|
||||
Args:
|
||||
memory_id: The ID of the memory unit to edit.
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
result = await memory.update_memory_unit(
|
||||
target_bank,
|
||||
memory_id,
|
||||
text=text,
|
||||
context=context,
|
||||
occurred_start=occurred_start,
|
||||
occurred_end=occurred_end,
|
||||
new_fact_type=fact_type,
|
||||
entities=entities,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if result is None:
|
||||
return json.dumps({"error": f"Memory '{memory_id}' not found"})
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except ValueError as e:
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating memory: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def update_memory(
|
||||
memory_id: str,
|
||||
text: str | None = None,
|
||||
context: str | None = None,
|
||||
occurred_start: str | None = None,
|
||||
occurred_end: str | None = None,
|
||||
fact_type: str | None = None,
|
||||
entities: list[str] | None = None,
|
||||
) -> dict:
|
||||
f"""{_EDIT_DOC}
|
||||
Args:
|
||||
memory_id: The ID of the memory unit to edit.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
result = await memory.update_memory_unit(
|
||||
target_bank,
|
||||
memory_id,
|
||||
text=text,
|
||||
context=context,
|
||||
occurred_start=occurred_start,
|
||||
occurred_end=occurred_end,
|
||||
new_fact_type=fact_type,
|
||||
entities=entities,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if result is None:
|
||||
return {"error": f"Memory '{memory_id}' not found"}
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except ValueError as e:
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating memory: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_invalidate_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the invalidate_memory (retire / restore) tool."""
|
||||
|
||||
_INVALIDATE_DOC = """
|
||||
Soft-retire a memory unit (or restore a previously retired one).
|
||||
|
||||
Invalidating moves the fact out of the active set: it's excluded from
|
||||
recall, consolidation, and the knowledge graph, its links are pruned, and
|
||||
its derived observations are recomputed without it — but it's kept for
|
||||
audit and is fully reversible. Pass restore=True to bring it back.
|
||||
|
||||
Only raw world/experience facts can be invalidated; observations are derived.
|
||||
"""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def invalidate_memory(
|
||||
memory_id: str,
|
||||
reason: str | None = None,
|
||||
restore: bool = False,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
f"""{_INVALIDATE_DOC}
|
||||
Args:
|
||||
memory_id: The ID of the memory unit to retire (or restore).
|
||||
reason: Optional free-text reason recorded when invalidating.
|
||||
restore: Set True to restore a previously invalidated fact.
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
result = await memory.update_memory_unit(
|
||||
target_bank,
|
||||
memory_id,
|
||||
state="valid" if restore else "invalidated",
|
||||
reason=reason,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if result is None:
|
||||
return json.dumps({"error": f"Memory '{memory_id}' not found"})
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except ValueError as e:
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error invalidating memory: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def invalidate_memory(
|
||||
memory_id: str,
|
||||
reason: str | None = None,
|
||||
restore: bool = False,
|
||||
) -> dict:
|
||||
f"""{_INVALIDATE_DOC}
|
||||
Args:
|
||||
memory_id: The ID of the memory unit to retire (or restore).
|
||||
reason: Optional free-text reason recorded when invalidating.
|
||||
restore: Set True to restore a previously invalidated fact.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
result = await memory.update_memory_unit(
|
||||
target_bank,
|
||||
memory_id,
|
||||
state="valid" if restore else "invalidated",
|
||||
reason=reason,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if result is None:
|
||||
return {"error": f"Memory '{memory_id}' not found"}
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except ValueError as e:
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error invalidating memory: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# DOCUMENT TOOLS
|
||||
# =========================================================================
|
||||
|
||||
@@ -184,6 +184,8 @@ class MetricsCollectorBase:
|
||||
input_tokens: int = 0,
|
||||
output_tokens: int = 0,
|
||||
success: bool = True,
|
||||
cached_input_tokens: int = 0,
|
||||
thoughts_tokens: int = 0,
|
||||
):
|
||||
"""
|
||||
Record metrics for an LLM call.
|
||||
@@ -193,9 +195,11 @@ class MetricsCollectorBase:
|
||||
model: Model name
|
||||
scope: Scope identifier (e.g., "memory", "reflect", "consolidation")
|
||||
duration: Call duration in seconds
|
||||
input_tokens: Number of input/prompt tokens
|
||||
output_tokens: Number of output/completion tokens
|
||||
input_tokens: Number of input/prompt tokens (total)
|
||||
output_tokens: Number of output/completion tokens visible in candidates
|
||||
success: Whether the call was successful
|
||||
cached_input_tokens: Subset of input_tokens billed at the cached rate
|
||||
thoughts_tokens: Reasoning tokens (billed as output, hidden from candidates)
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@@ -233,6 +237,8 @@ class NoOpMetricsCollector(MetricsCollectorBase):
|
||||
input_tokens: int = 0,
|
||||
output_tokens: int = 0,
|
||||
success: bool = True,
|
||||
cached_input_tokens: int = 0,
|
||||
thoughts_tokens: int = 0,
|
||||
):
|
||||
"""No-op LLM call recording."""
|
||||
pass
|
||||
@@ -287,6 +293,27 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
name="hindsight.llm.calls.total", description="Total number of LLM API calls", unit="calls"
|
||||
)
|
||||
|
||||
# Cached input tokens (subset of input_tokens billed at the cached rate).
|
||||
# Useful for tracking prompt-cache hit-rate independently of total
|
||||
# input volume. provider.scope.model labels matche llm_tokens_input.
|
||||
self.llm_tokens_cached_input = self.meter.create_counter(
|
||||
name="hindsight.llm.tokens.cached_input",
|
||||
description="Number of cached input tokens (billed at cached rate) for LLM calls",
|
||||
unit="tokens",
|
||||
)
|
||||
|
||||
# Thinking / reasoning tokens (Gemini 2.5+ family). Billed at the
|
||||
# output rate by the provider but invisible to candidates_token_count.
|
||||
# Surfacing them as a distinct counter is required for honest cost
|
||||
# attribution: a workload that "looks cheap" by output volume can be
|
||||
# silently expensive if the model is doing long reasoning chains.
|
||||
self.llm_tokens_thoughts = self.meter.create_counter(
|
||||
name="hindsight.llm.tokens.thoughts",
|
||||
description="Number of reasoning/thinking tokens emitted by the model "
|
||||
"(billed as output but not surfaced in candidates)",
|
||||
unit="tokens",
|
||||
)
|
||||
|
||||
# HTTP request metrics
|
||||
self.http_request_duration = self.meter.create_histogram(
|
||||
name="hindsight.http.duration", description="Duration of HTTP requests in seconds", unit="s"
|
||||
@@ -370,6 +397,8 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
input_tokens: int = 0,
|
||||
output_tokens: int = 0,
|
||||
success: bool = True,
|
||||
cached_input_tokens: int = 0,
|
||||
thoughts_tokens: int = 0,
|
||||
):
|
||||
"""
|
||||
Record metrics for an LLM call.
|
||||
@@ -379,9 +408,15 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
model: Model name
|
||||
scope: Scope identifier (e.g., "memory", "reflect", "consolidation")
|
||||
duration: Call duration in seconds
|
||||
input_tokens: Number of input/prompt tokens
|
||||
output_tokens: Number of output/completion tokens
|
||||
input_tokens: Number of input/prompt tokens (total, including cached portion)
|
||||
output_tokens: Number of output/completion tokens visible in candidates
|
||||
success: Whether the call was successful
|
||||
cached_input_tokens: Subset of input_tokens billed at the cached
|
||||
rate (Gemini context caching). Defaults to 0 when caching is
|
||||
disabled or the provider doesn't surface this field.
|
||||
thoughts_tokens: Reasoning/thinking tokens (Gemini 2.5+ family).
|
||||
Billed at the output rate but not counted in candidates.
|
||||
Defaults to 0 for providers that don't emit thoughts.
|
||||
"""
|
||||
# Base attributes for all metrics
|
||||
base_attributes = {
|
||||
@@ -413,6 +448,18 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
}
|
||||
self.llm_tokens_output.add(output_tokens, output_attributes)
|
||||
|
||||
if cached_input_tokens > 0:
|
||||
self.llm_tokens_cached_input.add(
|
||||
cached_input_tokens,
|
||||
{**base_attributes, "token_bucket": get_token_bucket(cached_input_tokens)},
|
||||
)
|
||||
|
||||
if thoughts_tokens > 0:
|
||||
self.llm_tokens_thoughts.add(
|
||||
thoughts_tokens,
|
||||
{**base_attributes, "token_bucket": get_token_bucket(thoughts_tokens)},
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def record_http_request(self, method: str, endpoint: str, status_code_getter: Callable[[], int]):
|
||||
"""
|
||||
|
||||
@@ -25,6 +25,7 @@ from pathlib import Path
|
||||
from alembic import command
|
||||
from alembic.config import Config
|
||||
from alembic.script.revision import ResolutionError
|
||||
from alembic.util.exc import CommandError
|
||||
from sqlalchemy import Connection, create_engine, text
|
||||
|
||||
from ._pg_search import normalize_pg_search_tokenizer, pg_search_bm25_columns
|
||||
@@ -131,7 +132,12 @@ def _run_migrations_internal(database_url: str, script_location: str, schema: st
|
||||
try:
|
||||
with _alembic_lock:
|
||||
command.upgrade(alembic_cfg, "heads")
|
||||
except ResolutionError as e:
|
||||
except (ResolutionError, CommandError) as e:
|
||||
# command.upgrade() wraps ResolutionError in CommandError via
|
||||
# ScriptDirectory._catch_revision_errors, so the wrapped form is what
|
||||
# actually reaches us; re-raise CommandErrors with any other cause.
|
||||
if isinstance(e, CommandError) and not isinstance(e.__cause__, ResolutionError):
|
||||
raise
|
||||
# This happens during rolling deployments when a newer version of the code
|
||||
# has already run migrations, and this older replica doesn't have the new
|
||||
# migration files. The database is already at a newer revision than we know.
|
||||
|
||||
@@ -4,8 +4,12 @@ SQLAlchemy models for the memory system.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING
|
||||
from uuid import UUID as PyUUID
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .cancellation import CancellationToken
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestContext:
|
||||
@@ -30,6 +34,21 @@ class RequestContext:
|
||||
# validators that want exponential backoff on repeated failures (e.g.
|
||||
# "defer for 2^retry_count minutes") without querying the DB themselves.
|
||||
retry_count: int = 0
|
||||
# Cooperative cancellation signal for long-running operations. The HTTP
|
||||
# layer sets this to a token that fires when the client disconnects; the
|
||||
# engine checks it at stage boundaries and aborts abandoned work so it stops
|
||||
# consuming CPU/DB resources (issue #2122). None means "never cancelled" —
|
||||
# every checkpoint is a no-op.
|
||||
cancellation: "CancellationToken | None" = None
|
||||
|
||||
def raise_if_cancelled(self) -> None:
|
||||
"""Abort the current operation if its cancellation token has fired.
|
||||
|
||||
A no-op when no token is attached, so engine code can call it at every
|
||||
stage boundary without caring whether the caller opted into cancellation.
|
||||
"""
|
||||
if self.cancellation is not None:
|
||||
self.cancellation.raise_if_cancelled()
|
||||
|
||||
|
||||
from pgvector.sqlalchemy import Vector
|
||||
|
||||
@@ -273,6 +273,8 @@ class LLMSpanRecorder:
|
||||
finish_reason: Optional[str] = None,
|
||||
error: Optional[Exception] = None,
|
||||
tool_calls: Optional[list[dict[str, Any]]] = None,
|
||||
cached_tokens: int = 0,
|
||||
**_extra: Any,
|
||||
) -> None:
|
||||
"""
|
||||
Record a completed LLM call as a span with GenAI semantic conventions.
|
||||
@@ -293,6 +295,8 @@ class LLMSpanRecorder:
|
||||
finish_reason: Reason the model stopped (stop, length, tool_calls, etc.)
|
||||
error: Exception if call failed
|
||||
tool_calls: List of tool calls made (for function calling)
|
||||
cached_tokens: Cached/cache-read prompt tokens, when reported by the provider.
|
||||
_extra: Tolerated forward-compatible kwargs from other recorders.
|
||||
"""
|
||||
try:
|
||||
# Map provider name to GenAI semantic convention
|
||||
@@ -326,6 +330,8 @@ class LLMSpanRecorder:
|
||||
span.set_attribute(GenAIAttributes.RESPONSE_MODEL, model)
|
||||
span.set_attribute(GenAIAttributes.USAGE_INPUT_TOKENS, input_tokens)
|
||||
span.set_attribute(GenAIAttributes.USAGE_OUTPUT_TOKENS, output_tokens)
|
||||
if cached_tokens:
|
||||
span.set_attribute("gen_ai.usage.cached_tokens", cached_tokens)
|
||||
|
||||
# Add custom attributes for Hindsight context
|
||||
span.set_attribute("hindsight.scope", scope)
|
||||
@@ -460,22 +466,61 @@ class NoOpLLMSpanRecorder:
|
||||
pass
|
||||
|
||||
|
||||
# Global span recorder instance
|
||||
class CompositeSpanRecorder:
|
||||
"""Fans out ``record_llm_call`` to every registered recorder.
|
||||
|
||||
This lets multiple GenAI consumers observe the same LLM calls — e.g. the
|
||||
OpenTelemetry span exporter and the per-bank DB tracer — through the single
|
||||
``record_llm_call`` chokepoint each provider already calls. A failure in one
|
||||
recorder never affects the others or the LLM call itself.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._recorders: list[Any] = []
|
||||
|
||||
def register(self, recorder: Any) -> None:
|
||||
if recorder not in self._recorders:
|
||||
self._recorders.append(recorder)
|
||||
|
||||
def unregister(self, recorder: Any) -> None:
|
||||
if recorder in self._recorders:
|
||||
self._recorders.remove(recorder)
|
||||
|
||||
def record_llm_call(self, **kwargs: Any) -> None:
|
||||
for recorder in self._recorders:
|
||||
try:
|
||||
recorder.record_llm_call(**kwargs)
|
||||
except Exception as e: # never let one recorder break others
|
||||
logger.debug(f"Span recorder {type(recorder).__name__} failed: {e}", exc_info=True)
|
||||
|
||||
|
||||
# Global composite recorder — always present; fans out to whatever is registered.
|
||||
_composite_recorder = CompositeSpanRecorder()
|
||||
# Backward-compat reference to the OTel recorder (if created).
|
||||
_span_recorder: Optional[LLMSpanRecorder] = None
|
||||
|
||||
|
||||
def get_span_recorder() -> LLMSpanRecorder | NoOpLLMSpanRecorder:
|
||||
"""Get the global span recorder (NoOp if tracing disabled)."""
|
||||
if _span_recorder is None:
|
||||
return NoOpLLMSpanRecorder()
|
||||
return _span_recorder
|
||||
def get_span_recorder() -> CompositeSpanRecorder:
|
||||
"""Get the global composite span recorder (fans out to all registered recorders)."""
|
||||
return _composite_recorder
|
||||
|
||||
|
||||
def register_span_recorder(recorder: Any) -> None:
|
||||
"""Register an additional GenAI recorder (e.g. the per-bank DB tracer)."""
|
||||
_composite_recorder.register(recorder)
|
||||
|
||||
|
||||
def unregister_span_recorder(recorder: Any) -> None:
|
||||
"""Remove a previously registered recorder."""
|
||||
_composite_recorder.unregister(recorder)
|
||||
|
||||
|
||||
def create_span_recorder() -> LLMSpanRecorder:
|
||||
"""Create and set the global span recorder."""
|
||||
"""Create and register the OpenTelemetry span recorder."""
|
||||
global _span_recorder
|
||||
tracer = get_tracer()
|
||||
if tracer is None:
|
||||
raise RuntimeError("Tracing not initialized. Call initialize_tracing() first.")
|
||||
_span_recorder = LLMSpanRecorder(tracer)
|
||||
register_span_recorder(_span_recorder)
|
||||
return _span_recorder
|
||||
|
||||
@@ -1,7 +1,14 @@
|
||||
"""Webhook system for Hindsight API event notifications."""
|
||||
|
||||
from .manager import WebhookManager
|
||||
from .models import ConsolidationEventData, RetainEventData, WebhookConfig, WebhookEvent, WebhookEventType
|
||||
from .models import (
|
||||
ConsolidationEventData,
|
||||
MemoryDefenseEventData,
|
||||
RetainEventData,
|
||||
WebhookConfig,
|
||||
WebhookEvent,
|
||||
WebhookEventType,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"WebhookManager",
|
||||
@@ -9,5 +16,6 @@ __all__ = [
|
||||
"WebhookEvent",
|
||||
"WebhookEventType",
|
||||
"ConsolidationEventData",
|
||||
"MemoryDefenseEventData",
|
||||
"RetainEventData",
|
||||
]
|
||||
|
||||
@@ -9,6 +9,7 @@ from pydantic import BaseModel, Field
|
||||
class WebhookEventType(StrEnum):
|
||||
CONSOLIDATION_COMPLETED = "consolidation.completed"
|
||||
RETAIN_COMPLETED = "retain.completed"
|
||||
MEMORY_DEFENSE_TRIGGERED = "memory_defense.triggered"
|
||||
|
||||
|
||||
class ConsolidationEventData(BaseModel):
|
||||
@@ -23,13 +24,23 @@ class RetainEventData(BaseModel):
|
||||
tags: list[str] | None = None
|
||||
|
||||
|
||||
class MemoryDefenseEventData(BaseModel):
|
||||
"""Payload for a memory_defense.triggered event (one item, one non-allow decision)."""
|
||||
|
||||
action: str # "redact" or "block"
|
||||
detector: str | None = None # e.g. "sensitive_data"
|
||||
document_id: str | None = None
|
||||
matched_types: list[str] | None = None # redaction pattern labels that fired
|
||||
message: str | None = None
|
||||
|
||||
|
||||
class WebhookEvent(BaseModel):
|
||||
event: WebhookEventType
|
||||
bank_id: str
|
||||
operation_id: str
|
||||
status: str # "completed" or "failed"
|
||||
status: str # "completed"/"failed" for retain/consolidation; the action ("redact"/"block") for memory_defense
|
||||
timestamp: datetime
|
||||
data: ConsolidationEventData | RetainEventData
|
||||
data: ConsolidationEventData | RetainEventData | MemoryDefenseEventData
|
||||
|
||||
|
||||
class WebhookHttpConfig(BaseModel):
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api-slim"
|
||||
version = "0.7.1"
|
||||
version = "0.8.1"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -81,6 +81,10 @@ local-ml = [
|
||||
# Local ML models for embeddings/reranking
|
||||
"sentence-transformers>=3.3.0",
|
||||
"transformers>=4.53.0", # Security fixes for ReDoS vulnerabilities
|
||||
# transformers (incl. latest 5.x) hard-requires tokenizers<=0.23.0 via a
|
||||
# runtime check; without this cap an in-place upgrade can pull tokenizers
|
||||
# 0.23.1 and break local embeddings/reranker startup. See issue #2055.
|
||||
"tokenizers>=0.22.0,<=0.23.0",
|
||||
"torch>=2.6.0", # CVE fix for remote code execution
|
||||
"einops>=0.8.2",
|
||||
"flashrank>=0.2.0",
|
||||
@@ -96,6 +100,14 @@ local-llm = [
|
||||
"llama-cpp-python[server]>=0.3.0",
|
||||
"huggingface-hub>=0.20.0",
|
||||
]
|
||||
local-onnx = [
|
||||
# In-process ONNX Runtime embeddings without an Ollama/TEI sidecar
|
||||
"onnxruntime>=1.17.0",
|
||||
"transformers>=4.53.0",
|
||||
"tokenizers>=0.22.0,<=0.23.0", # See issue #2055 (transformers caps tokenizers<=0.23.0)
|
||||
"huggingface-hub>=0.20.0",
|
||||
"numpy>=1.26.0",
|
||||
]
|
||||
embedded-db = [
|
||||
"pg0-embedded>=0.14.2",
|
||||
]
|
||||
@@ -103,7 +115,7 @@ oracle = [
|
||||
"oracledb>=2.5.0",
|
||||
]
|
||||
all = [
|
||||
"hindsight-api-slim[local-ml,embedded-db]",
|
||||
"hindsight-api-slim[local-ml,local-onnx,embedded-db]",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
@@ -175,12 +187,14 @@ dev = [
|
||||
[tool.ruff]
|
||||
line-length = 120
|
||||
target-version = "py311"
|
||||
exclude = [
|
||||
"tests/",
|
||||
"**/tests/",
|
||||
]
|
||||
|
||||
[tool.ruff.lint]
|
||||
# Tests are formatted (via `ruff format`) but excluded from lint rules, which
|
||||
# are too noisy for test code (unused imports/vars, import ordering).
|
||||
exclude = [
|
||||
"tests/**",
|
||||
"**/tests/**",
|
||||
]
|
||||
select = [
|
||||
"E", # pycodestyle errors
|
||||
"W", # pycodestyle warnings
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""
|
||||
Pytest configuration and shared fixtures.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from pathlib import Path
|
||||
@@ -20,6 +21,16 @@ from hindsight_api.pg0 import EmbeddedPostgres
|
||||
DEFAULT_PG0_INSTANCE_NAME = "hindsight-test"
|
||||
DEFAULT_PG0_PORT = int(os.environ.get("HINDSIGHT_TEST_PG_PORT", "5556"))
|
||||
|
||||
# Keep the background MaintenanceLoop from auto-starting during tests. In
|
||||
# production it sweeps retention and re-schedules consolidation, but its timers
|
||||
# would race shared-pg0 test data (e.g. delete llm_requests/audit_log rows a test
|
||||
# just inserted). Disabling the reconcile interval and llm-trace retention — with
|
||||
# audit retention already off by default — leaves no job enabled, so the loop
|
||||
# never starts. Tests that exercise it call MaintenanceLoop methods
|
||||
# (_run_reconcile / _purge_expired) directly.
|
||||
os.environ.setdefault("HINDSIGHT_API_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS", "0")
|
||||
os.environ.setdefault("HINDSIGHT_API_LLM_TRACE_RETENTION_DAYS", "-1")
|
||||
|
||||
|
||||
# Load environment variables from .env at the start of test session
|
||||
def pytest_configure(config):
|
||||
@@ -66,6 +77,7 @@ def pg0_db_url(db_url, tmp_path_factory, worker_id):
|
||||
if db_url and not _parse_pg0_url(db_url)[0]:
|
||||
# Plain postgresql:// URL - use it directly but still run migrations
|
||||
from hindsight_api.migrations import run_migrations
|
||||
|
||||
run_migrations(db_url)
|
||||
return db_url
|
||||
|
||||
@@ -117,6 +129,7 @@ def pg0_db_url(db_url, tmp_path_factory, worker_id):
|
||||
# Run migrations - uses PostgreSQL advisory lock internally,
|
||||
# so safe to call from multiple workers (only one will actually run migrations)
|
||||
from hindsight_api.migrations import run_migrations
|
||||
|
||||
run_migrations(url)
|
||||
|
||||
# Clean up stale test data from previous sessions. Per-bank vector indexes
|
||||
@@ -147,8 +160,7 @@ def _cleanup_stale_test_data(db_url: str) -> None:
|
||||
conn = await asyncpg.connect(db_url)
|
||||
try:
|
||||
idx_rows = await conn.fetch(
|
||||
"SELECT indexname FROM pg_indexes "
|
||||
"WHERE schemaname = 'public' AND indexname LIKE 'idx_mu_emb_%'"
|
||||
"SELECT indexname FROM pg_indexes WHERE schemaname = 'public' AND indexname LIKE 'idx_mu_emb_%'"
|
||||
)
|
||||
if idx_rows:
|
||||
for row in idx_rows:
|
||||
@@ -156,10 +168,20 @@ def _cleanup_stale_test_data(db_url: str) -> None:
|
||||
|
||||
# Truncate test data in dependency order
|
||||
for table in [
|
||||
"entity_cooccurrences", "unit_entities", "memory_links",
|
||||
"entities", "memory_units", "chunks", "documents",
|
||||
"mental_models", "directives", "async_operations",
|
||||
"audit_log", "webhooks", "file_storage", "banks",
|
||||
"entity_cooccurrences",
|
||||
"unit_entities",
|
||||
"memory_links",
|
||||
"entities",
|
||||
"memory_units",
|
||||
"chunks",
|
||||
"documents",
|
||||
"mental_models",
|
||||
"directives",
|
||||
"async_operations",
|
||||
"audit_log",
|
||||
"webhooks",
|
||||
"file_storage",
|
||||
"banks",
|
||||
]:
|
||||
try:
|
||||
await conn.execute(f"TRUNCATE {table} CASCADE")
|
||||
@@ -242,8 +264,7 @@ def oracle_db_url(_oracle_admin_dsn):
|
||||
# Create test user (idempotent — skip if already exists)
|
||||
try:
|
||||
cursor.execute(
|
||||
f'CREATE USER {test_user} IDENTIFIED BY "{test_pass}" '
|
||||
f"DEFAULT TABLESPACE USERS QUOTA UNLIMITED ON USERS"
|
||||
f'CREATE USER {test_user} IDENTIFIED BY "{test_pass}" DEFAULT TABLESPACE USERS QUOTA UNLIMITED ON USERS'
|
||||
)
|
||||
except oracledb.DatabaseError as e:
|
||||
if hasattr(e.args[0], "code") and e.args[0].code == 1920:
|
||||
@@ -410,13 +431,12 @@ def cross_encoder(tmp_path_factory, worker_id):
|
||||
|
||||
return ce
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def query_analyzer():
|
||||
return DateparserQueryAnalyzer()
|
||||
|
||||
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope="function")
|
||||
async def memory(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
"""
|
||||
@@ -512,3 +532,21 @@ async def memory_no_llm_verify(pg0_db_url, embeddings, cross_encoder, query_anal
|
||||
await mem.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client(memory):
|
||||
"""General-purpose HTTP test client over the `memory` fixture's app.
|
||||
|
||||
Use for any integration test that exercises the FastAPI surface without
|
||||
needing audit-logging side effects. See `audit_api_client` for the
|
||||
audit-enabled variant.
|
||||
"""
|
||||
import httpx
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
@@ -23,6 +23,7 @@ from urllib.parse import urlparse
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _log(step: int, total: int, msg: str) -> None:
|
||||
print(f" [{step}/{total}] {msg}")
|
||||
|
||||
@@ -64,8 +65,7 @@ def _bootstrap_test_user(admin_dsn: dict[str, str]) -> str:
|
||||
# Create user (skip if already exists - ORA-01920)
|
||||
try:
|
||||
cursor.execute(
|
||||
f'CREATE USER {test_user} IDENTIFIED BY "{test_pass}" '
|
||||
f"DEFAULT TABLESPACE USERS QUOTA UNLIMITED ON USERS"
|
||||
f'CREATE USER {test_user} IDENTIFIED BY "{test_pass}" DEFAULT TABLESPACE USERS QUOTA UNLIMITED ON USERS'
|
||||
)
|
||||
except oracledb.DatabaseError as e:
|
||||
if hasattr(e.args[0], "code") and e.args[0].code == 1920:
|
||||
@@ -100,6 +100,7 @@ def _bootstrap_test_user(admin_dsn: dict[str, str]) -> str:
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _run() -> None:
|
||||
total_steps = 8
|
||||
|
||||
@@ -290,6 +291,7 @@ def main() -> int:
|
||||
except Exception as exc:
|
||||
print(f"\nFAILED: {exc}", file=sys.stderr)
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
return 1
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user