Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ab10de7140 | ||
|
|
cfe640c357 | ||
|
|
742eb1315a | ||
|
|
e44da25794 |
@@ -1,7 +1,6 @@
|
||||
{
|
||||
"$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
|
||||
"name": "hindsight",
|
||||
"version": "0.7.2",
|
||||
"description": "Official Hindsight integrations for Claude Code",
|
||||
"owner": {
|
||||
"name": "vectorize-io"
|
||||
|
||||
@@ -73,11 +73,6 @@ 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.
|
||||
@@ -140,13 +135,6 @@ 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:
|
||||
@@ -154,12 +142,6 @@ 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:
|
||||
@@ -172,8 +154,7 @@ 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` 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.
|
||||
- **Release process** — check that the integration name is in the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh`. If missing, flag it.
|
||||
- **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
|
||||
@@ -192,18 +173,6 @@ 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:
|
||||
@@ -227,10 +196,7 @@ 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:
|
||||
|
||||
@@ -1,116 +0,0 @@
|
||||
---
|
||||
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.
|
||||
+3
-69
@@ -2,7 +2,7 @@
|
||||
# Copy this file to .env and fill in your values
|
||||
|
||||
# LLM Configuration (Required)
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, zai, atlas, volcano
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, zai, volcano
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
@@ -10,17 +10,6 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# Reasoning effort for providers/models that support it. Examples: low, medium, high, xhigh.
|
||||
# HINDSIGHT_API_LLM_REASONING_EFFORT=low
|
||||
|
||||
# Sampling temperature for internal LLM calls. Set a number in [0.0, 2.0], or `none`
|
||||
# to omit the temperature parameter entirely (required for models that reject explicit
|
||||
# temperatures, e.g. Azure gpt-5.5). The global override below applies to every operation;
|
||||
# per-operation overrides (defaults: verification=0.0, retain=0.1, reflect=0.9,
|
||||
# consolidation=0.0) take precedence.
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE=none
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_VERIFICATION=0.0
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_RETAIN=0.1
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_REFLECT=0.9
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_CONSOLIDATION=0.0
|
||||
|
||||
# Example: Anthropic Claude configuration
|
||||
# HINDSIGHT_API_LLM_PROVIDER=anthropic
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-anthropic-api-key
|
||||
@@ -36,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-M3 # or MiniMax-M2.7 for the previous generation
|
||||
# HINDSIGHT_API_LLM_MODEL=MiniMax-M2.7
|
||||
|
||||
# Example: DeepSeek configuration (https://api.deepseek.com)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=deepseek
|
||||
@@ -48,41 +37,16 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-zai-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=glm-4.5-flash # or glm-4.5-air for the paid tier
|
||||
|
||||
# Example: Atlas Cloud configuration (OpenAI-compatible, https://www.atlascloud.ai)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=atlas
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-atlascloud-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=deepseek-ai/deepseek-v4-pro # reasoning model; also Qwen / GLM / Kimi / MiniMax, etc.
|
||||
|
||||
# Example: LM Studio local configuration (Qwen 2.5 32B recommended)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=lmstudio
|
||||
# HINDSIGHT_API_LLM_API_KEY=lmstudio
|
||||
# HINDSIGHT_API_LLM_BASE_URL=http://localhost:1234/v1
|
||||
# HINDSIGHT_API_LLM_MODEL=qwen2.5-32b-instruct
|
||||
|
||||
# Multi-LLM strategies: configure extra LLMs by index alongside the primary above,
|
||||
# then pick a routing strategy. Unset = single primary LLM (default). Members are
|
||||
# numbered from 1; indices must be contiguous. Each operation can override with a
|
||||
# RETAIN_/REFLECT_/CONSOLIDATION_ prefix (e.g. HINDSIGHT_API_RETAIN_LLM_1_PROVIDER).
|
||||
# HINDSIGHT_API_LLM_1_PROVIDER=groq
|
||||
# HINDSIGHT_API_LLM_1_API_KEY=your-groq-api-key
|
||||
# HINDSIGHT_API_LLM_1_MODEL=openai/gpt-oss-120b
|
||||
# HINDSIGHT_API_LLM_2_PROVIDER=anthropic
|
||||
# HINDSIGHT_API_LLM_2_API_KEY=your-anthropic-api-key
|
||||
# Strategy JSON: {"mode": "failover"} or {"mode": "round-robin"}.
|
||||
# Round-robin accepts optional positive-int "weights" (one per member, primary first).
|
||||
# HINDSIGHT_API_LLM_STRATEGY={"mode": "failover"}
|
||||
|
||||
# API Configuration (Optional)
|
||||
HINDSIGHT_API_HOST=0.0.0.0
|
||||
HINDSIGHT_API_PORT=8888
|
||||
HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Optional retain chunking override for structured logs/transcripts.
|
||||
# Unset uses HINDSIGHT_API_RETAIN_CHUNK_SIZE as the structured-chunk limit.
|
||||
# HINDSIGHT_API_RETAIN_STRUCTURED_CHUNK_SIZE=
|
||||
|
||||
# Dry-run extraction preview endpoint (POST /memories/dry-run-extract). Enabled by default; it makes
|
||||
# a real LLM call but stores nothing. Set to false to remove the endpoint (returns 404).
|
||||
# HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT=true
|
||||
|
||||
# Base Path / Reverse Proxy Support (Optional)
|
||||
# Set these when deploying behind a reverse proxy with path-based routing
|
||||
@@ -95,7 +59,6 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_READ_DATABASE_URL= # Optional read-replica URL. When set, recall queries (semantic, BM25, graph, temporal) flow through a separate pool against this URL, offloading the primary. Typically points to a read-only endpoint (CNPG's <cluster>-ro service or Aurora reader endpoint).
|
||||
# HINDSIGHT_API_MIGRATION_DATABASE_URL= # Direct PostgreSQL URL for migrations (bypasses PgBouncer). Falls back to DATABASE_URL.
|
||||
# HINDSIGHT_API_DATABASE_SCHEMA=public # PostgreSQL schema name (default: public)
|
||||
# HINDSIGHT_API_MIGRATION_CONCURRENCY=1 # Tenant schemas to migrate concurrently (PG only, each in its own process; per-schema work stays sequential). Each worker has ~1-2s startup cost + uses ~3 DB connections, so it only pays off with many schemas (tens+) or slow migrations; keep concurrency*3 <= spare max_connections. Default: 1 (sequential).
|
||||
|
||||
# Vector Extension (Optional - uses pgvector by default)
|
||||
# Options: "pgvector" (default), "vchord", "pgvectorscale" (DiskANN)
|
||||
@@ -116,36 +79,11 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# korean_lindera/lindera(korean), ngram(min,max), edge_ngram(min,max)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER=
|
||||
|
||||
# File Parser (Optional - uses markitdown by default)
|
||||
# HINDSIGHT_API_FILE_PARSER=markitdown
|
||||
# Enable image OCR for MarkItDown using an OpenAI-compatible OCR/vision endpoint.
|
||||
# These OCR settings are independent from HINDSIGHT_API_LLM_* because MarkItDown
|
||||
# uses the OpenAI SDK directly and requires Chat Completions image input support.
|
||||
# When OCR is enabled, API_KEY, BASE_URL, and MODEL are required.
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED=false
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT=
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default), "onnx", "tei", "openai", "cohere", "google", "openrouter", "zeroentropy", "litellm", or "litellm-sdk"
|
||||
# Provider: "local" (default), "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:
|
||||
@@ -191,10 +129,6 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Custom service name and environment (optional, defaults: hindsight-api, development)
|
||||
# HINDSIGHT_API_OTEL_SERVICE_NAME=hindsight-production
|
||||
# HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT=production
|
||||
#
|
||||
# Expose async-operation queue + consolidation-backlog gauges on /metrics.
|
||||
# Runs periodic per-schema COUNT queries on a background task (disabled by default).
|
||||
# HINDSIGHT_API_METRICS_BACKLOG_ENABLED=true
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Control Plane (Optional)
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
version: 2
|
||||
updates:
|
||||
- package-ecosystem: "github-actions"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
@@ -22,8 +22,6 @@ 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,9 +23,7 @@ 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)."
|
||||
@@ -35,18 +33,6 @@ 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
|
||||
@@ -212,86 +198,3 @@ 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,11 +9,7 @@ jobs:
|
||||
publish:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
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).
|
||||
id-token: write # for PyPI trusted publishing
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -75,7 +71,7 @@ jobs:
|
||||
if: steps.type.outputs.type == 'plugin'
|
||||
run: |
|
||||
echo "Plugin integration ${{ steps.info.outputs.integration }} v${{ steps.info.outputs.version }} — no package to publish."
|
||||
echo "Users install via: claude plugin marketplace add vectorize-io/hindsight"
|
||||
echo "Users install via: claude plugin marketplace add vectorize-io/hindsight --sparse hindsight-integrations"
|
||||
|
||||
# ── TypeScript integrations (ai-sdk, chat, openclaw) ────────────────────
|
||||
|
||||
@@ -116,71 +112,6 @@ 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 }}
|
||||
@@ -190,12 +121,7 @@ jobs:
|
||||
EXIT_CODE=$?
|
||||
echo "$OUTPUT"
|
||||
if [ $EXIT_CODE -ne 0 ]; 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
|
||||
if echo "$OUTPUT" | grep -q "cannot publish over"; then
|
||||
echo "Package version already published, skipping..."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
@@ -266,7 +266,7 @@ jobs:
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- os: ubuntu-22.04
|
||||
- os: ubuntu-latest
|
||||
target: x86_64-unknown-linux-gnu
|
||||
artifact_name: hindsight
|
||||
asset_name: hindsight-linux-amd64
|
||||
@@ -278,7 +278,7 @@ jobs:
|
||||
target: aarch64-apple-darwin
|
||||
artifact_name: hindsight
|
||||
asset_name: hindsight-darwin-arm64
|
||||
- os: ubuntu-22.04-arm
|
||||
- os: ubuntu-24.04-arm
|
||||
target: aarch64-unknown-linux-gnu
|
||||
artifact_name: hindsight
|
||||
asset_name: hindsight-linux-arm64
|
||||
|
||||
+8
-982
File diff suppressed because it is too large
Load Diff
+1
-5
@@ -6,7 +6,6 @@ dist/
|
||||
wheels/
|
||||
*.egg-info
|
||||
.mcp.json
|
||||
.playwright-mcp/
|
||||
.osgrep
|
||||
# Virtual environments
|
||||
.venv
|
||||
@@ -16,8 +15,6 @@ node_modules/
|
||||
|
||||
# Environment variables and local config
|
||||
.env
|
||||
.env.bak*
|
||||
.env.*.bak
|
||||
docker-compose.yml
|
||||
docker-compose.override.yml
|
||||
|
||||
@@ -62,5 +59,4 @@ hindsight-integrations/_drafts/
|
||||
# Changelog is now tracked in hindsight-docs/src/pages/changelog.md
|
||||
# CHANGELOG.md
|
||||
|
||||
blog-post*
|
||||
.worktrees/
|
||||
blog-post*
|
||||
@@ -216,46 +216,10 @@ migration file dispatches through `run_for_dialect`, which calls either
|
||||
./scripts/hooks/lint.sh
|
||||
```
|
||||
|
||||
Dead-code detection runs in CI (the `check-unused-code` job) at two levels:
|
||||
- **Blocking:** unused imports (ruff `F401`) and variables (`F841`) — `lint.sh` auto-removes
|
||||
them and `verify-generated-files` fails on any leftover diff; and **knip** for orphaned
|
||||
control-plane files / unused (or unlisted) `package.json` dependencies.
|
||||
- **Advisory:** whole unused Python functions (vulture) and unused control-plane *exports*
|
||||
(the shadcn/ui surface is kept on purpose) — surfaced, not gated.
|
||||
|
||||
Run both locally with:
|
||||
```bash
|
||||
./scripts/hooks/check-unused.sh
|
||||
```
|
||||
|
||||
**After completing any implementation work, run `/code-review`** to verify your changes against project standards (missing tests, dead code, type safety, etc.). Fix any "must fix" issues before considering the task done.
|
||||
|
||||
**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
|
||||
@@ -327,10 +291,7 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
|
||||
```
|
||||
|
||||
2. **main.py** (`hindsight-api-slim/hindsight_api/main.py`):
|
||||
- No change is needed for ordinary environment-backed config fields. The CLI starts from `_get_raw_config()`,
|
||||
so new `HindsightConfig` fields are carried through automatically.
|
||||
- If the new field should be overridable by a CLI flag, add the argparse option in `_parse_cli_args()` and include
|
||||
that field in the `dataclasses.replace(config, ...)` call near the "CLI override" comment.
|
||||
- Add field to the manual `HindsightConfig()` constructor call (search for "CLI override")
|
||||
|
||||
3. **Use hierarchical config in MemoryEngine**:
|
||||
```python
|
||||
@@ -350,16 +311,6 @@ 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):
|
||||
@@ -376,7 +327,7 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
|
||||
|
||||
```bash
|
||||
cp .env.example .env
|
||||
# Edit .env with the LLM provider/model and credentials for your setup
|
||||
# Edit .env with LLM API key
|
||||
|
||||
# Python deps
|
||||
uv sync --directory hindsight-api-slim/
|
||||
@@ -385,10 +336,10 @@ uv sync --directory hindsight-api-slim/
|
||||
npm install
|
||||
```
|
||||
|
||||
Common LLM settings:
|
||||
Required env vars:
|
||||
- `HINDSIGHT_API_LLM_PROVIDER`: openai, anthropic, gemini, groq, minimax, ollama, lmstudio
|
||||
- `HINDSIGHT_API_LLM_API_KEY`: API key for providers that require one
|
||||
- `HINDSIGHT_API_LLM_MODEL`: Model name (defaults are provider-specific)
|
||||
- `HINDSIGHT_API_LLM_API_KEY`: Your API key
|
||||
- `HINDSIGHT_API_LLM_MODEL`: Model name (e.g., gpt-4o-mini, claude-sonnet-4-20250514)
|
||||
|
||||
Optional (uses local models by default):
|
||||
- `HINDSIGHT_API_EMBEDDINGS_PROVIDER`: local (default) or tei
|
||||
|
||||
+2
-25
@@ -9,36 +9,13 @@ Thanks for your interest in contributing to Hindsight!
|
||||
git clone [email protected]:vectorize-io/hindsight.git
|
||||
cd hindsight
|
||||
```
|
||||
|
||||
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:
|
||||
2. Set up your environment:
|
||||
```bash
|
||||
cp .env.example .env
|
||||
```
|
||||
Edit the .env to add LLM API key and config as required
|
||||
|
||||
2. Install dependencies:
|
||||
3. Install dependencies:
|
||||
```bash
|
||||
# Python dependencies
|
||||
uv sync --directory hindsight-api/
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
[](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml)
|
||||
[](https://join.slack.com/t/hindsight-space/shared_invite/zt-3nhbm4w29-LeSJ5Ixi6j8PdiYOCPlOgg)
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
[](https://gitcgr.com/vectorize-io/hindsight)
|
||||

|
||||

|
||||
<br/>
|
||||
@@ -61,16 +62,16 @@ If you need more control over how and when your agent stores and recalls memorie
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
|
||||
docker run -it --pull always --name hindsight --restart unless-stopped -p 8888:8888 -p 9999:9999 \
|
||||
docker run --rm -it --pull always -p 8888:8888 -p 9999:9999 \
|
||||
-e HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY \
|
||||
-v hindsight-data:/home/hindsight/.pg0 \
|
||||
-v $HOME/.hindsight-docker:/home/hindsight/.pg0 \
|
||||
ghcr.io/vectorize-io/hindsight:latest
|
||||
```
|
||||
|
||||
>API: http://localhost:8888
|
||||
>UI: http://localhost:9999
|
||||
|
||||
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio`, `minimax`, and `atlas` ([Atlas Cloud](https://www.atlascloud.ai/?utm_source=github&utm_medium=link&utm_campaign=hindsight)). The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio`, and `minimax`. The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
|
||||
|
||||
|
||||
@@ -142,8 +143,6 @@ 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
|
||||
@@ -250,7 +249,7 @@ Recall performs 4 retrieval strategies in parallel:
|
||||
- Graph: Entity/temporal/causal links
|
||||
- Temporal: Time range filtering
|
||||
|
||||

|
||||

|
||||
|
||||
The individual results from the retrievals are merged, then ordered by relevance using reciprocal rank fusion and a cross-encoder reranking model.
|
||||
|
||||
@@ -276,7 +275,7 @@ client = Hindsight(base_url="http://localhost:8888")
|
||||
client.reflect(bank_id="my-bank", query="What should I know about Alice?")
|
||||
```
|
||||
|
||||

|
||||

|
||||
|
||||
---
|
||||
|
||||
@@ -301,19 +300,6 @@ 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://hindsight.vectorize.io/developer/installation#supported-platforms) for details.
|
||||
|
||||
---
|
||||
|
||||
## Contributing
|
||||
|
||||
See [CONTRIBUTING.md](./CONTRIBUTING.md).
|
||||
|
||||
@@ -77,6 +77,7 @@
|
||||
"npm:@radix-ui/react-dropdown-menu@^2.1.16",
|
||||
"npm:@radix-ui/react-label@^2.1.8",
|
||||
"npm:@radix-ui/react-popover@^1.1.15",
|
||||
"npm:@radix-ui/react-radio-group@^1.3.8",
|
||||
"npm:@radix-ui/react-select@^2.2.6",
|
||||
"npm:@radix-ui/react-slider@^1.3.6",
|
||||
"npm:@radix-ui/react-slot@^1.2.4",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and vectorchord
|
||||
# docker compose -f docker/docker-compose/vchord/docker-compose.yaml down && sleep 2 && docker compose -f docker/docker-compose/vchord/docker-compose.yaml up -d
|
||||
# docker compose -f docker/docker-compose/docker-compose.yaml down && sleep 2 && docker compose -f docker/docker-compose/docker-compose.yaml up -d
|
||||
# Make sure to set the required environment variables before running:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see below in the hindsight service)
|
||||
|
||||
@@ -50,8 +50,6 @@ 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,56 +43,11 @@ 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 'chmod -R u+rwx "$TMP_DIR" 2>/dev/null || true; rm -rf "$TMP_DIR"' EXIT
|
||||
trap 'rm -rf "$TMP_DIR"' EXIT
|
||||
|
||||
assert_contains() {
|
||||
local output="$1"
|
||||
@@ -71,51 +71,3 @@ 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.8.4
|
||||
appVersion: "0.8.4"
|
||||
version: 0.7.1
|
||||
appVersion: "0.7.1"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -66,13 +66,13 @@ helm install hindsight ./helm/hindsight -n hindsight --create-namespace -f value
|
||||
|
||||
| Parameter | Description | Default |
|
||||
|-----------|-------------|---------|
|
||||
| `version` | Default image tag for all components | Chart `appVersion` |
|
||||
| `version` | Default image tag for all components | `0.1.0` |
|
||||
| `api.enabled` | Enable the API component | `true` |
|
||||
| `api.image.repository` | API image repository | `ghcr.io/vectorize-io/hindsight-api` |
|
||||
| `api.image.repository` | API image repository | `hindsight/api` |
|
||||
| `api.image.tag` | API image tag (defaults to `version`) | - |
|
||||
| `api.service.port` | API service port | `8888` |
|
||||
| `controlPlane.enabled` | Enable the control plane | `true` |
|
||||
| `controlPlane.image.repository` | Control plane image repository | `ghcr.io/vectorize-io/hindsight-control-plane` |
|
||||
| `controlPlane.image.repository` | Control plane image repository | `hindsight/control-plane` |
|
||||
| `controlPlane.image.tag` | Control plane image tag (defaults to `version`) | - |
|
||||
| `controlPlane.service.port` | Control plane service port | `3000` |
|
||||
| `postgresql.enabled` | Deploy PostgreSQL as subchart | `true` |
|
||||
|
||||
@@ -13,6 +13,9 @@
|
||||
# - Any other env vars you want to inject
|
||||
# existingSecret: "my-hindsight-secret"
|
||||
|
||||
# Global settings
|
||||
replicaCount: 1
|
||||
|
||||
# Image settings for api
|
||||
api:
|
||||
enabled: true
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.8.4",
|
||||
"version": "0.7.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.8.4"
|
||||
version = "0.7.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.8.4",
|
||||
"hindsight-api-slim==0.7.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.8.4"
|
||||
version = "0.7.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.8.4",
|
||||
"hindsight-api-slim[all]==0.7.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.8.4",
|
||||
"hindsight-api-slim[local-llm]==0.7.1",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -99,7 +99,7 @@ hindsight-api
|
||||
## Docker
|
||||
|
||||
```bash
|
||||
docker run -it --name hindsight --restart unless-stopped -p 8888:8888 \
|
||||
docker run --rm -it -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
|
||||
@@ -121,7 +121,7 @@ This runs a stdio-based MCP server that can be used directly with MCP-compatible
|
||||
- **Entity Graph** — Automatic entity extraction and relationship tracking
|
||||
- **Temporal Reasoning** — Native support for time-based queries
|
||||
- **Disposition Traits** — Configurable skepticism, literalism, and empathy influence opinion formation
|
||||
- **Three Memory Types** — World facts, experience facts (the bank's own actions), and observations
|
||||
- **Three Memory Types** — World facts, bank actions, and formed opinions with confidence scores
|
||||
|
||||
## Documentation
|
||||
|
||||
|
||||
@@ -4,13 +4,6 @@ Memory System for AI Agents.
|
||||
Temporal + Semantic Memory Architecture using PostgreSQL with pgvector.
|
||||
"""
|
||||
|
||||
# Cap native ML thread pools (OpenBLAS/OpenMP/MKL) before any import pulls in
|
||||
# numpy/torch/onnxruntime — they read these env vars only at load time. See
|
||||
# hindsight_api/_thread_limits.py for the rationale.
|
||||
from ._thread_limits import apply_default_thread_limits
|
||||
|
||||
apply_default_thread_limits()
|
||||
|
||||
from .config import HindsightConfig, get_config
|
||||
from .engine.cross_encoder import CrossEncoderModel, LocalSTCrossEncoder, RemoteTEICrossEncoder
|
||||
from .engine.embeddings import Embeddings, LocalSTEmbeddings, RemoteTEIEmbeddings
|
||||
@@ -53,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.8.4"
|
||||
__version__ = "0.7.1"
|
||||
|
||||
@@ -1,107 +0,0 @@
|
||||
"""Process-level caps for native ML thread pools.
|
||||
|
||||
OpenBLAS, OpenMP, and MKL each spawn a worker pool sized to the host CPU count
|
||||
the first time they are loaded (numpy pulls in OpenBLAS eagerly; torch and
|
||||
onnxruntime load their pools lazily on first inference). Hindsight already
|
||||
parallelizes at the request level via thread-pool executors (embeddings on the
|
||||
default executor, the reranker on its own pool), so these native intra-op pools
|
||||
oversubscribe the CPU: on a many-core host the process accumulates 100+ native
|
||||
threads, which inflates memory and, under contention, can degrade throughput.
|
||||
|
||||
We bound each pool to ``_MAX_NATIVE_THREADS`` (or the available CPU count, if
|
||||
smaller). "Available" is the CPU budget actually granted to the process, not
|
||||
``os.cpu_count()``: in a CPU-limited container ``os.cpu_count()`` still reports
|
||||
the host's cores, so sizing pools by it oversubscribes the container's real
|
||||
quota — the exact failure mode this guards against. We therefore take the
|
||||
smallest of the CPU-affinity set, the cgroup CPU quota, and ``os.cpu_count()``.
|
||||
|
||||
Every cap is applied with ``setdefault`` so an operator who has deliberately
|
||||
tuned one of these variables keeps their value. This must run *before* numpy,
|
||||
torch, or onnxruntime are imported — those libraries read the variables only at
|
||||
load time — which is why it is invoked at the very top of
|
||||
``hindsight_api/__init__.py``, ahead of the package's other imports.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
# Native threading env vars, each read by the respective library at load time.
|
||||
_NATIVE_THREAD_VARS = (
|
||||
"OMP_NUM_THREADS", # OpenMP — torch, onnxruntime, some BLAS builds
|
||||
"OPENBLAS_NUM_THREADS", # OpenBLAS — numpy's default BLAS
|
||||
"MKL_NUM_THREADS", # Intel MKL — numpy/torch when MKL-backed
|
||||
"NUMEXPR_NUM_THREADS", # numexpr expression engine
|
||||
)
|
||||
|
||||
# Upper bound on intra-op threads per native pool. Bounds runaway growth on
|
||||
# many-core hosts without serialising single-request inference.
|
||||
_MAX_NATIVE_THREADS = 16
|
||||
|
||||
|
||||
def _quota_to_cpus(quota: int, period: int) -> int | None:
|
||||
"""Whole CPUs from a CFS quota/period pair, or None if unlimited."""
|
||||
if quota > 0 and period > 0:
|
||||
# Floor (never round up) so we never exceed the granted budget.
|
||||
return max(1, quota // period)
|
||||
return None
|
||||
|
||||
|
||||
def _parse_cgroup_v2_cpu_max(text: str) -> int | None:
|
||||
"""Parse cgroup v2 ``cpu.max`` ("<quota> <period>", or "max <period>")."""
|
||||
parts = text.split()
|
||||
if len(parts) >= 2 and parts[0] != "max":
|
||||
try:
|
||||
return _quota_to_cpus(int(parts[0]), int(parts[1]))
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _cgroup_cpu_quota() -> int | None:
|
||||
"""Effective CPUs from the cgroup CPU quota, or None if unlimited/unknown."""
|
||||
try: # cgroup v2
|
||||
with open("/sys/fs/cgroup/cpu.max") as fh:
|
||||
return _parse_cgroup_v2_cpu_max(fh.read())
|
||||
except OSError:
|
||||
pass
|
||||
try: # cgroup v1
|
||||
with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as fh:
|
||||
quota = int(fh.read())
|
||||
with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as fh:
|
||||
period = int(fh.read())
|
||||
return _quota_to_cpus(quota, period)
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _available_cpu_count() -> int:
|
||||
"""CPUs actually available to this process.
|
||||
|
||||
The smallest of the CPU-affinity set (cpuset / ``--cpuset-cpus``), the
|
||||
cgroup CPU quota (``--cpus``), and ``os.cpu_count()`` — each captures a
|
||||
different way the budget can be constrained, and the last alone overcounts
|
||||
inside a limited container.
|
||||
"""
|
||||
candidates = [os.cpu_count() or 1]
|
||||
if hasattr(os, "sched_getaffinity"):
|
||||
try:
|
||||
candidates.append(len(os.sched_getaffinity(0)))
|
||||
except OSError:
|
||||
pass
|
||||
quota = _cgroup_cpu_quota()
|
||||
if quota is not None:
|
||||
candidates.append(quota)
|
||||
return max(1, min(candidates))
|
||||
|
||||
|
||||
def default_native_thread_count() -> int:
|
||||
"""Per-pool cap: ``_MAX_NATIVE_THREADS``, or available CPUs if fewer."""
|
||||
return min(_MAX_NATIVE_THREADS, _available_cpu_count())
|
||||
|
||||
|
||||
def apply_default_thread_limits() -> None:
|
||||
"""Cap native ML thread pools unless the operator has set the var already."""
|
||||
value = str(default_native_thread_count())
|
||||
for var in _NATIVE_THREAD_VARS:
|
||||
os.environ.setdefault(var, value)
|
||||
@@ -54,20 +54,23 @@ _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, 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.
|
||||
# - vchord exposes vchordrq.probes (no default; see VectorChord issue #392)
|
||||
# and vchordrq.epsilon (default 1.9). probes = 10 / 30 are starting
|
||||
# defaults pending a workload-specific sweep — vchordrq's recall curve
|
||||
# shape differs from HNSW's, so the pgvector numbers don't translate
|
||||
# directly. Revisit with a per-cluster benchmark once we have production
|
||||
# recall data; until then these are deliberately conservative on the
|
||||
# high-recall path. We leave epsilon at its default; tightening it is a
|
||||
# separate trade-off.
|
||||
# - pgvectorscale / pg_diskann / scann do not expose an equivalent per-statement
|
||||
# knob in the engine today, so the dispatcher returns no statements for them.
|
||||
_ANN_TUNING_LOW_LATENCY: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (("hnsw.ef_search", "60"),),
|
||||
"vchord": (("vchordrq.probes", "10"),),
|
||||
}
|
||||
_ANN_TUNING_HIGH_RECALL: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (("hnsw.ef_search", "200"),),
|
||||
"vchord": (("vchordrq.probes", "30"),),
|
||||
}
|
||||
|
||||
_EXTENSION_INSTALL_SQL = {
|
||||
|
||||
@@ -17,9 +17,7 @@ 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
|
||||
|
||||
@@ -49,43 +47,21 @@ BACKUP_TABLES = [
|
||||
"entities",
|
||||
"chunks",
|
||||
"memory_units",
|
||||
"invalidated_memory_units",
|
||||
"unit_entities",
|
||||
"entity_cooccurrences",
|
||||
"memory_links",
|
||||
"observation_history",
|
||||
"mental_models",
|
||||
"mental_model_history",
|
||||
"knowledge_pages",
|
||||
"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)
|
||||
@@ -257,10 +233,14 @@ async def _run_migration(
|
||||
schema: str | None = None,
|
||||
base_schema: str = DEFAULT_DATABASE_SCHEMA,
|
||||
embedding_dimension: int | None = None,
|
||||
ensure_extensions: bool = True,
|
||||
) -> list[str]:
|
||||
"""Resolve database URL and run migrations for one schema or all discovered schemas."""
|
||||
from ..migrations import run_migrations_for_schemas
|
||||
from ..migrations import (
|
||||
ensure_embedding_dimension,
|
||||
ensure_text_search_extension,
|
||||
ensure_vector_extension,
|
||||
run_migrations,
|
||||
)
|
||||
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
if is_pg0:
|
||||
@@ -281,21 +261,32 @@ async def _run_migration(
|
||||
# Preserve order while removing duplicates.
|
||||
schemas = list(dict.fromkeys(schemas))
|
||||
|
||||
# Migrate up to `migration_concurrency` schemas at once (each in its own
|
||||
# process); within a schema the work stays sequential. Run off the event
|
||||
# loop so the process pool's blocking joins don't stall it.
|
||||
await asyncio.to_thread(
|
||||
run_migrations_for_schemas,
|
||||
resolved_url,
|
||||
schemas,
|
||||
concurrency=config.migration_concurrency,
|
||||
migration_database_url=config.migration_database_url,
|
||||
embedding_dimension=embedding_dimension,
|
||||
vector_extension=config.vector_extension,
|
||||
text_search_extension=config.text_search_extension,
|
||||
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
|
||||
ensure_extensions=ensure_extensions,
|
||||
)
|
||||
for schema in schemas:
|
||||
run_migrations(resolved_url, schema=schema, migration_database_url=config.migration_database_url)
|
||||
|
||||
if embedding_dimension is not None:
|
||||
for schema in schemas:
|
||||
ensure_embedding_dimension(
|
||||
resolved_url,
|
||||
embedding_dimension,
|
||||
schema=schema,
|
||||
vector_extension=config.vector_extension,
|
||||
)
|
||||
|
||||
for schema in schemas:
|
||||
ensure_vector_extension(
|
||||
resolved_url,
|
||||
vector_extension=config.vector_extension,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
for schema in schemas:
|
||||
ensure_text_search_extension(
|
||||
resolved_url,
|
||||
text_search_extension=config.text_search_extension,
|
||||
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
return schemas
|
||||
|
||||
@@ -313,18 +304,6 @@ def run_db_migration(
|
||||
"--embedding-dimension",
|
||||
help="Expected embedding dimension to enforce after migrations. Omit to skip dimension sync.",
|
||||
),
|
||||
skip_extension_reconcile: bool = typer.Option(
|
||||
False,
|
||||
"--skip-extension-reconcile",
|
||||
help=(
|
||||
"Skip the post-migration vector / text-search index reconcile. This step only does "
|
||||
"work when the configured backend (HINDSIGHT_API_VECTOR_EXTENSION / "
|
||||
"HINDSIGHT_API_TEXT_SEARCH_EXTENSION) differs from a schema's existing indexes — a "
|
||||
"rare, operator-driven change. Skipping it makes a no-change re-migration over many "
|
||||
"tenant schemas much faster. Only use when you have NOT changed the backend; a "
|
||||
"backend change still needs a normal run to reshape the indexes."
|
||||
),
|
||||
),
|
||||
):
|
||||
"""Run database migrations to the latest version."""
|
||||
config = HindsightConfig.from_env()
|
||||
@@ -338,8 +317,6 @@ def run_db_migration(
|
||||
typer.echo(f"Running database migrations for schema: {schema}...")
|
||||
else:
|
||||
typer.echo("Running database migrations for base schema and all discovered tenant schemas...")
|
||||
if skip_extension_reconcile:
|
||||
typer.echo("Skipping post-migration extension reconcile (--skip-extension-reconcile).")
|
||||
|
||||
schemas = asyncio.run(
|
||||
_run_migration(
|
||||
@@ -347,130 +324,12 @@ def run_db_migration(
|
||||
schema=schema,
|
||||
base_schema=config.database_schema,
|
||||
embedding_dimension=embedding_dimension,
|
||||
ensure_extensions=not skip_extension_reconcile,
|
||||
)
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
-105
@@ -1,105 +0,0 @@
|
||||
"""Add a composite index on memory_links(bank_id, link_type) (PostgreSQL).
|
||||
|
||||
``bank_id`` was added to ``memory_links`` in ``c5d6e7f8a9b0`` precisely so that
|
||||
bank-scoped reads (e.g. the stats endpoint) could filter on the link table
|
||||
directly instead of joining ``memory_units`` — that JOIN took 18+ seconds on
|
||||
banks with millions of links. The column landed without an index, so every
|
||||
``bank_id = $1`` predicate still falls back to a sequential scan over the whole
|
||||
table.
|
||||
|
||||
This adds the missing btree. It is composite on ``(bank_id, link_type)`` rather
|
||||
than ``bank_id`` alone because the hot query is the stats endpoint's
|
||||
``SELECT link_type, COUNT(*) ... WHERE bank_id = $1 GROUP BY link_type``: a
|
||||
``(bank_id, link_type)`` index serves that filter, grouping and count as an
|
||||
index-only scan, never touching the heap, whereas a ``bank_id``-only index would
|
||||
still have to read every matching row to recover ``link_type``. ``link_type`` is
|
||||
low-cardinality (only ``temporal``/``semantic``/``caused_by`` are written —
|
||||
entity edges were dropped in ``e9b2c7d1f3a4``), so the trailing column adds
|
||||
little to the index size while removing the heap fetch.
|
||||
|
||||
The Oracle baseline (``o1a2b3c4d5e6``) already creates ``idx_ml_bank_id`` on
|
||||
``memory_links(bank_id)``; that single-column index already covers Oracle's
|
||||
bank-scoped filter, so the Oracle slot here is intentionally absent and only the
|
||||
PostgreSQL dialect gets the composite index.
|
||||
|
||||
``memory_links`` can hold tens of millions of rows, so the index is built
|
||||
CONCURRENTLY to avoid taking a write lock on the table. CONCURRENTLY cannot run
|
||||
inside a transaction block, so the statement runs in an ``autocommit_block()``;
|
||||
``IF NOT EXISTS`` keeps it idempotent across retries and re-migrated tenant
|
||||
schemas. A CONCURRENTLY build interrupted partway (lock conflict, disk
|
||||
pressure, signal) leaves the index behind as *invalid*; ``IF NOT EXISTS`` would
|
||||
then skip over it forever, so the upgrade first drops any invalid leftover of
|
||||
this name before (re)creating it.
|
||||
|
||||
Revision ID: 2071c7518f88
|
||||
Revises: a1d3f5b7c9e2
|
||||
Create Date: 2026-06-16
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "2071c7518f88"
|
||||
down_revision: str | Sequence[str] | None = "a1d3f5b7c9e2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_INDEX_NAME = "idx_memory_links_bank_id_link_type"
|
||||
|
||||
|
||||
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:
|
||||
bind = op.get_bind()
|
||||
# `or None` collapses an unset option and an explicit empty string into NULL
|
||||
# so the COALESCE below falls back to current_schema() in both cases.
|
||||
target_schema = context.config.get_main_option("target_schema") or None
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block; the
|
||||
# autocommit_block runs each statement outside Alembic's migration
|
||||
# transaction.
|
||||
with op.get_context().autocommit_block():
|
||||
# A CONCURRENTLY build that errored on a previous run leaves an INVALID
|
||||
# index of this name behind. `CREATE INDEX ... IF NOT EXISTS` would see
|
||||
# that relation and skip, so bank_id queries would keep seq-scanning.
|
||||
# Drop only the invalid leftover — never a healthy index — so the retry
|
||||
# actually rebuilds a usable one.
|
||||
leftover_invalid = bind.execute(
|
||||
text(
|
||||
"SELECT NOT i.indisvalid "
|
||||
"FROM pg_class c "
|
||||
"JOIN pg_index i ON c.oid = i.indexrelid "
|
||||
"JOIN pg_namespace n ON c.relnamespace = n.oid "
|
||||
"WHERE c.relname = :index_name "
|
||||
" AND n.nspname = COALESCE(:target_schema, current_schema())"
|
||||
),
|
||||
{"index_name": _INDEX_NAME, "target_schema": target_schema},
|
||||
).scalar()
|
||||
if leftover_invalid:
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}{_INDEX_NAME}")
|
||||
|
||||
# IF NOT EXISTS keeps the create idempotent across retries and schemas.
|
||||
op.execute(f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {_INDEX_NAME} ON {schema}memory_links(bank_id, link_type)")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}{_INDEX_NAME}")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-85
@@ -1,85 +0,0 @@
|
||||
"""Repair: widen the remaining live ``bank_id`` columns from VARCHAR(64) to TEXT on PostgreSQL.
|
||||
|
||||
Follow-up to ``c3e5a7b9d1f4`` (issue #2106), which widened the two *history*
|
||||
tables (``observation_history``, ``mental_model_history``) to ``TEXT`` after the
|
||||
narrow ``VARCHAR(64)`` declaration bricked startup. The same VARCHAR(64) / TEXT
|
||||
inconsistency still affects the live tables that store a user-supplied
|
||||
``bank_id``:
|
||||
|
||||
* ``directives`` -- created VARCHAR(64) in ``p1k2l3m4n5o6``
|
||||
* ``mental_models`` -- VARCHAR(64) (origin ``pinned_reflections`` in
|
||||
``n9i0j1k2l3m4``; recreated in ``h3c4d5e6f7g8``)
|
||||
|
||||
``mental_model_versions`` is intentionally *not* widened here: it is created in
|
||||
``j5e6f7g8h9i0`` but dropped (``DROP TABLE ... CASCADE``) in ``o0j1k2l3m4n5`` and
|
||||
never recreated on the upgrade path, so it does not exist at head. Issuing
|
||||
``ALTER TABLE mental_model_versions ...`` would raise ``UndefinedTable`` and --
|
||||
because migrations run inside the lifespan-startup transaction -- roll the whole
|
||||
migration back, bricking the API. (It is unrelated to the live
|
||||
``mental_model_history`` table widened by ``c3e5a7b9d1f4``.)
|
||||
|
||||
``banks.bank_id`` is ``TEXT`` (unbounded), so a deployment can create a bank
|
||||
whose id exceeds 64 chars -- the 78-char hierarchical org-unit shape reported in
|
||||
issue #2106 -- and the bank insert succeeds. The next write that propagates that
|
||||
id (``create_directive``, ``create_mental_model`` / consolidation, or
|
||||
mental-model versioning) then aborts with::
|
||||
|
||||
psycopg2.errors.StringDataRightTruncation: value too long for type
|
||||
character varying(64)
|
||||
|
||||
i.e. a 500 on core write endpoints, instead of the startup brick that
|
||||
``c3e5a7b9d1f4`` already repaired.
|
||||
|
||||
``ALTER COLUMN ... TYPE TEXT`` is a no-op on a column that is already ``TEXT``,
|
||||
so every upgrade path converges on ``TEXT``. These 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 -- the same mechanism as
|
||||
``c3e5a7b9d1f4``.
|
||||
|
||||
PostgreSQL only: these tables are created by PostgreSQL-only migrations
|
||||
(``run_for_dialect(pg=...)``); on Oracle they are absent or already
|
||||
``VARCHAR2(256)`` (consistent, never truncates), so the Oracle slot is
|
||||
intentionally absent -- mirroring ``c3e5a7b9d1f4``.
|
||||
|
||||
Revision ID: a1d3f5b7c9e2
|
||||
Revises: c3e5a7b9d1f4
|
||||
Create Date: 2026-06-13
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a1d3f5b7c9e2"
|
||||
down_revision: str | Sequence[str] | None = "c3e5a7b9d1f4"
|
||||
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}directives ALTER COLUMN bank_id TYPE TEXT")
|
||||
op.execute(f"ALTER TABLE {schema}mental_models 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 types are owned by
|
||||
# the migrations that created the tables.
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-52
@@ -1,52 +0,0 @@
|
||||
"""Add managed flag to knowledge_pages.
|
||||
|
||||
The knowledge base is managed by clients (CRUD over folders/pages). ``managed``
|
||||
lets a client tag a node as system-owned vs. hand-authored; it carries no
|
||||
server-side behaviour.
|
||||
|
||||
Revision ID: a5b6c7d8e9f0
|
||||
Revises: a9b8c7d6e5f4
|
||||
Create Date: 2026-06-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a5b6c7d8e9f0"
|
||||
down_revision: str | Sequence[str] | None = "a9b8c7d6e5f4"
|
||||
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()
|
||||
op.execute(f"ALTER TABLE {schema}knowledge_pages ADD COLUMN IF NOT EXISTS managed BOOLEAN NOT NULL DEFAULT false")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}knowledge_pages DROP COLUMN IF EXISTS managed")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
op.execute("ALTER TABLE knowledge_pages ADD (managed NUMBER(1) DEFAULT 0 NOT NULL)")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("ALTER TABLE knowledge_pages DROP COLUMN managed")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
-253
@@ -1,253 +0,0 @@
|
||||
"""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)
|
||||
-110
@@ -1,110 +0,0 @@
|
||||
"""Add knowledge_pages table (knowledge-base hierarchy).
|
||||
|
||||
The knowledge base organizes synthesized mental models into a navigable tree of
|
||||
**folders** and **pages**. A page references the mental model that holds its
|
||||
content (``mental_model_id``); a folder is a pure container (``mental_model_id``
|
||||
NULL). Hierarchy is a single self-referential ``parent_id`` so folders can nest
|
||||
arbitrarily. Content stays in ``mental_models`` — this table is metadata + tree
|
||||
structure only.
|
||||
|
||||
Revision ID: a9b8c7d6e5f4
|
||||
Revises: b57a7c9e0d13
|
||||
Create Date: 2026-06-25
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a9b8c7d6e5f4"
|
||||
down_revision: str | Sequence[str] | None = "b57a7c9e0d13"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# parent_id self-FK cascades so deleting a folder row removes its whole
|
||||
# subtree of rows in one shot. The mental_model FK is composite (matches the
|
||||
# mental_models (id, bank_id) PK) and cascades too, so deleting a page's
|
||||
# mental model removes the page row — folders skip the FK because a NULL
|
||||
# column in a composite FK is not enforced (MATCH SIMPLE).
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}knowledge_pages (
|
||||
id VARCHAR(64) NOT NULL,
|
||||
bank_id TEXT NOT NULL,
|
||||
parent_id VARCHAR(64),
|
||||
kind VARCHAR(16) NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
mental_model_id VARCHAR(64),
|
||||
sort_order INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
CONSTRAINT pk_knowledge_pages PRIMARY KEY (id),
|
||||
CONSTRAINT ck_knowledge_pages_kind CHECK (kind IN ('folder', 'page')),
|
||||
CONSTRAINT fk_kp_bank FOREIGN KEY (bank_id)
|
||||
REFERENCES {schema}banks(bank_id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_parent FOREIGN KEY (parent_id)
|
||||
REFERENCES {schema}knowledge_pages(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_mm FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES {schema}mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_kp_bank_parent ON {schema}knowledge_pages (bank_id, parent_id, sort_order)"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_kp_bank_parent")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}knowledge_pages")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
op.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS knowledge_pages (
|
||||
id VARCHAR2(64) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
parent_id VARCHAR2(64),
|
||||
kind VARCHAR2(16) NOT NULL,
|
||||
name CLOB NOT NULL,
|
||||
mental_model_id VARCHAR2(64),
|
||||
sort_order NUMBER DEFAULT 0 NOT NULL,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_knowledge_pages PRIMARY KEY (id),
|
||||
CONSTRAINT ck_knowledge_pages_kind CHECK (kind IN ('folder', 'page')),
|
||||
CONSTRAINT fk_kp_bank FOREIGN KEY (bank_id)
|
||||
REFERENCES banks(bank_id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_parent FOREIGN KEY (parent_id)
|
||||
REFERENCES knowledge_pages(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_mm FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute("CREATE INDEX idx_kp_bank_parent ON knowledge_pages (bank_id, parent_id, sort_order)")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("DROP TABLE knowledge_pages 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
@@ -1,156 +0,0 @@
|
||||
"""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)
|
||||
-61
@@ -1,61 +0,0 @@
|
||||
"""Add bank_stats_cache table for distributed get_bank_stats caching
|
||||
|
||||
Revision ID: b57a7c9e0d13
|
||||
Revises: c3f7a1b9d2e4
|
||||
Create Date: 2026-07-01
|
||||
|
||||
get_bank_stats aggregates over memory_links / unit_entities — a multi-second scan
|
||||
on banks with millions of rows. The result was cached per-process (in-memory), so
|
||||
every API worker recomputed it once per TTL and the first caller after expiry
|
||||
stalled. This table backs a shared, cross-process TTL cache: one worker's compute
|
||||
is written here and served to all the others.
|
||||
|
||||
PostgreSQL only. Oracle keeps the in-process cache (the runtime picks the backing
|
||||
store by dialect), so the Oracle upgrade slot is intentionally absent.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b57a7c9e0d13"
|
||||
down_revision: str | Sequence[str] | None = "c3f7a1b9d2e4"
|
||||
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()
|
||||
# One row per bank: payload is the full get_bank_stats result, computed_at
|
||||
# drives logical TTL expiry. Rows are overwritten in place (ON CONFLICT), so
|
||||
# the table never grows beyond the number of banks and needs no purge job.
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}bank_stats_cache (
|
||||
bank_id TEXT PRIMARY KEY,
|
||||
payload JSONB NOT NULL,
|
||||
computed_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}bank_stats_cache")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade) # oracle slot intentionally absent → no-op
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-71
@@ -1,71 +0,0 @@
|
||||
"""Unique page name per folder in knowledge_pages.
|
||||
|
||||
The folder curator can fire concurrently (folder-create trigger + the
|
||||
post-consolidation sweep), and an in-process lock can't serialize runs that
|
||||
execute in different threads/loops. A partial unique index on
|
||||
(bank_id, parent, lower(name)) for pages makes duplicate-named pages in the same
|
||||
folder impossible at the DB level — the second concurrent insert fails and the
|
||||
curator treats it as "already exists".
|
||||
|
||||
PostgreSQL only: the Oracle ``name`` column is a CLOB and cannot back a
|
||||
functional unique index; Oracle relies on the in-process serialization instead.
|
||||
|
||||
Revision ID: c3d4e5f6a7b8
|
||||
Revises: a5b6c7d8e9f0
|
||||
Create Date: 2026-06-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c3d4e5f6a7b8"
|
||||
down_revision: str | Sequence[str] | None = "a5b6c7d8e9f0"
|
||||
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()
|
||||
# First drop any pre-existing duplicate pages (created by the racy curator
|
||||
# before this guard existed), keeping the earliest row of each duplicate set,
|
||||
# so the unique index can be built. Their backing mental models are left in
|
||||
# place (harmless orphans).
|
||||
op.execute(
|
||||
f"""
|
||||
DELETE FROM {schema}knowledge_pages a
|
||||
USING {schema}knowledge_pages b
|
||||
WHERE a.kind = 'page' AND b.kind = 'page'
|
||||
AND a.bank_id = b.bank_id
|
||||
AND COALESCE(a.parent_id, '') = COALESCE(b.parent_id, '')
|
||||
AND lower(a.name) = lower(b.name)
|
||||
AND a.ctid > b.ctid
|
||||
"""
|
||||
)
|
||||
# COALESCE(parent_id, '') so root-level pages (NULL parent) are also unique by
|
||||
# name — NULLs would otherwise compare distinct and allow duplicates.
|
||||
op.execute(
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS uq_kp_folder_pagename "
|
||||
f"ON {schema}knowledge_pages (bank_id, COALESCE(parent_id, ''), lower(name)) "
|
||||
"WHERE kind = 'page'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}uq_kp_folder_pagename")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade) # oracle slot intentionally absent (CLOB name)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-75
@@ -1,75 +0,0 @@
|
||||
"""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)
|
||||
-144
@@ -1,144 +0,0 @@
|
||||
"""Backfill search_vector for native-backend observations.
|
||||
|
||||
Observations created or updated by the consolidator landed with a NULL
|
||||
``search_vector`` under the ``native`` text-search backend: the
|
||||
single-row INSERT/UPDATE paths in ``consolidator.py`` never populated the
|
||||
tsvector (only the batch raw-fact path in ``ops_postgresql.insert_facts_batch``
|
||||
did). Those observations were therefore invisible to the BM25 retrieval arm
|
||||
until they were re-written by a later consolidation pass. The writer is fixed
|
||||
in the same change set (all four consolidator sites now call
|
||||
``to_tsvector($lang, COALESCE(text, ''))``); this migration repairs the
|
||||
historical residue so existing observations become BM25-searchable without a
|
||||
re-ingest.
|
||||
|
||||
Scope mirrors the writer fix exactly:
|
||||
* Only the ``native`` backend is touched. The gate is the column *type*:
|
||||
under ``native`` ``search_vector`` is a regular (non-generated) tsvector
|
||||
column; under ``vchord`` it is a ``bm25vector`` and under
|
||||
``pg_textsearch`` / ``pgroonga`` / ``pg_search`` it is a dummy ``text``
|
||||
column. ``_is_regular_tsvector`` is true only for ``native``, so every
|
||||
other backend is a no-op.
|
||||
* The tsvector is built from the observation's own ``text`` only — matching
|
||||
the consolidator INSERT/UPDATE paths (entity / source / temporal signals
|
||||
are intentionally excluded; the other retrieval arms cover those).
|
||||
* Only ``fact_type = 'observation'`` rows with a NULL ``search_vector`` are
|
||||
rewritten. Raw facts already carry a populated tsvector, and the
|
||||
``IS NULL`` predicate makes the migration idempotent and re-runnable.
|
||||
|
||||
The configured ``HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE`` is used
|
||||
so backfilled rows are lexically identical to newly-created observations. The
|
||||
value is validated as a PG identifier (mirroring
|
||||
``HindsightConfig.validate``) before being embedded as a SQL literal.
|
||||
|
||||
This is a single UPDATE per schema: it locks the targeted observation rows for
|
||||
its duration. It is one-time and only touches unpopulated rows, so subsequent
|
||||
online writes (which now carry the tsvector via the writer fix) are unaffected.
|
||||
|
||||
Oracle slot is intentionally absent: the consolidator INSERT/UPDATE paths that
|
||||
this repairs are PostgreSQL-specific (``ops_postgresql``), and the native
|
||||
tsvector ``search_vector`` column only exists on PostgreSQL. There is no Oracle
|
||||
residue to repair.
|
||||
|
||||
Revision ID: c3f7a1b9d2e4
|
||||
Revises: f4d1c2b3a5e6
|
||||
Create Date: 2026-06-29
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import Connection, text
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
from hindsight_api.config import (
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
)
|
||||
|
||||
revision: str = "c3f7a1b9d2e4"
|
||||
down_revision: str | Sequence[str] | None = "f4d1c2b3a5e6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
# Matches HindsightConfig.validate(): a tsvector regconfig name embedded as a
|
||||
# SQL literal must be a bare PG identifier.
|
||||
_PG_IDENTIFIER = re.compile(r"[a-zA-Z_][a-zA-Z0-9_]*")
|
||||
|
||||
|
||||
def _schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _schema_name() -> str:
|
||||
return (context.config.get_main_option("target_schema") or "public").strip('"')
|
||||
|
||||
|
||||
def _native_language() -> str:
|
||||
"""Configured native tsvector language, validated as a PG identifier."""
|
||||
lang = os.getenv(
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
)
|
||||
if not _PG_IDENTIFIER.fullmatch(lang):
|
||||
return DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE
|
||||
return lang
|
||||
|
||||
|
||||
def _is_regular_tsvector(conn: Connection, schema: str, table: str) -> bool:
|
||||
"""True iff ``schema.table.search_vector`` is a non-generated tsvector column.
|
||||
|
||||
This is the ``native`` backend signature. ``vchord`` (bm25vector) and
|
||||
``pg_textsearch`` / ``pgroonga`` / ``pg_search`` (dummy text column) all
|
||||
fail this check, so the backfill is a no-op for them.
|
||||
"""
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT is_generated, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table
|
||||
AND column_name = 'search_vector'
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
is_generated, udt_name = row[0], row[1]
|
||||
return udt_name == "tsvector" and is_generated != "ALWAYS"
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
schema_name = _schema_name()
|
||||
if not _is_regular_tsvector(conn, schema_name, "memory_units"):
|
||||
# Non-native backend (or column absent) — nothing to backfill.
|
||||
return
|
||||
schema_prefix = _schema_prefix()
|
||||
lang = _native_language()
|
||||
op.execute(
|
||||
f"""
|
||||
UPDATE {schema_prefix}memory_units
|
||||
SET search_vector = to_tsvector('{lang}'::regconfig, COALESCE(text, ''))
|
||||
WHERE fact_type = 'observation' AND search_vector IS NULL
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# No-op: backfilled rows are indistinguishable from observations that were
|
||||
# populated by the post-fix writer, and reverting either to NULL would
|
||||
# re-break BM25 retrieval. The column simply stays populated.
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-158
@@ -1,158 +0,0 @@
|
||||
"""Make maintenance routines resilient to schemas that vanish mid-scan.
|
||||
|
||||
``public.banks_needing_consolidation()`` and
|
||||
``public.schemas_with_expired_rows(...)`` snapshot the set of schemas owning a
|
||||
target table from ``pg_class`` and then run a dynamic query against each schema
|
||||
in turn. That is a time-of-check/time-of-use race: a schema (or its tables) can
|
||||
be dropped — a tenant being deleted, or a tenant migration that recreates
|
||||
tables — between the snapshot and the per-schema query, which then aborts the
|
||||
whole routine with::
|
||||
|
||||
relation "<schema>.memory_units" does not exist
|
||||
relation "<schema>.audit_log" does not exist
|
||||
|
||||
In the test suite this surfaces as cross-worker contamination: the multi-tenant
|
||||
maintenance test creates and drops ~100 ``mt<hash>_NNN`` schemas while
|
||||
``test_maintenance_routines`` (on another xdist worker, same DB) calls the
|
||||
routines. In production the background maintenance loop hits the same race when
|
||||
a tenant is removed or mid-migration.
|
||||
|
||||
Wrap each per-schema query in its own ``BEGIN ... EXCEPTION`` block so a schema
|
||||
that disappears (``undefined_table`` / ``invalid_schema_name`` /
|
||||
``undefined_column``) is skipped instead of aborting the scan. The routines stay
|
||||
``CREATE OR REPLACE`` and PostgreSQL-only, and are (re)installed only on the run
|
||||
that targets the shared ``public`` schema — same gating as the original
|
||||
install (``e5f6a7b8c9d0``) and its repair (``b2d4f6a8c1e3``).
|
||||
|
||||
Revision ID: c7e9f1a3b5d2
|
||||
Revises: e1f2a3b4c5d6
|
||||
Create Date: 2026-06-19
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c7e9f1a3b5d2"
|
||||
down_revision: str | Sequence[str] | None = "e1f2a3b4c5d6"
|
||||
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``, so they are installed exactly
|
||||
once — on the base run (no ``target_schema``) or the run that explicitly
|
||||
targets ``public``. Mirrors ``b2d4f6a8c1e3``.
|
||||
"""
|
||||
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
|
||||
|
||||
# Same body as b2d4f6a8c1e3, but each per-schema query runs in its own
|
||||
# subtransaction so a schema dropped mid-scan is skipped, not fatal.
|
||||
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
|
||||
BEGIN
|
||||
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);
|
||||
EXCEPTION
|
||||
-- Schema or its tables vanished between the pg_class
|
||||
-- snapshot and this query (tenant dropped or migrating).
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
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
|
||||
BEGIN
|
||||
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;
|
||||
EXCEPTION
|
||||
-- Schema or its table vanished mid-scan; skip it.
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
IF has_expired THEN
|
||||
RETURN NEXT sch;
|
||||
END IF;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# No-op: e5f6a7b8c9d0 owns these functions' lifecycle and drops them on its
|
||||
# own downgrade. This migration only re-installs them (the resilient body is
|
||||
# a strict superset of the previous behaviour), 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)
|
||||
-108
@@ -1,108 +0,0 @@
|
||||
"""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 — except ``embedding``,
|
||||
which it never keeps: the archive is cold storage, never a recall surface, and
|
||||
revert recomputes the embedding from the unit's text/dates/entities. Keeping no
|
||||
archive vector also means a later embedding-model switch (which re-dimensions
|
||||
``memory_units``) can't trip a dimension mismatch on the move (#2209). 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.
|
||||
# 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)"
|
||||
)
|
||||
# ...then drop the inherited embedding: the archive never stores one (revert
|
||||
# recomputes it), so it isn't created here only to be dropped again later by
|
||||
# d4f6a8c2e1b3. That migration still runs as a no-op (DROP ... IF EXISTS) on
|
||||
# fresh DBs and does the real drop on DBs created before this column was removed.
|
||||
op.execute(f"ALTER TABLE {schema}invalidated_memory_units DROP COLUMN IF EXISTS embedding")
|
||||
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
@@ -1,96 +0,0 @@
|
||||
"""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)
|
||||
-93
@@ -1,93 +0,0 @@
|
||||
"""Drop the embedding column from the curation archive (invalidated_memory_units).
|
||||
|
||||
The archive is cold storage, never a recall surface, so it has no business
|
||||
keeping an embedding. Earlier curation code copied the live row's embedding into
|
||||
``invalidated_memory_units`` on invalidate; the engine now leaves it out on
|
||||
invalidate and recomputes it on revert, so the column is dead weight.
|
||||
|
||||
Dropping it makes "the archive holds no embedding" a schema-enforced invariant
|
||||
rather than a convention the move queries have to honour, and removes a latent
|
||||
failure mode (#2209): after an embedding-model switch the live tables are
|
||||
re-dimensioned but the archive was not, so a stale old-dimension embedding in
|
||||
the archive tripped a vector-dimension mismatch on the INSERT … SELECT
|
||||
round-trip. With no column at all, there is nothing to mismatch.
|
||||
|
||||
The creation sites no longer add the column (the PG ``LIKE`` clone in
|
||||
c9a1b2d3e4f5 drops it; the Oracle baseline omits it), so on a fresh database
|
||||
this migration is a no-op (DROP ... IF EXISTS / Oracle ORA-00904 swallow). It
|
||||
does the real work on databases created before the column was removed there.
|
||||
|
||||
DROP COLUMN is a metadata-only operation on both PostgreSQL and Oracle 23ai (no
|
||||
table rewrite), so it is cheap even across many tenant schemas. The downgrade
|
||||
re-adds an unconstrained vector column (any dimension) — empty, since the
|
||||
embeddings are intentionally discarded.
|
||||
|
||||
Revision ID: d4f6a8c2e1b3
|
||||
Revises: a1d3f5b7c9e2
|
||||
Create Date: 2026-06-15
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "d4f6a8c2e1b3"
|
||||
down_revision: str | Sequence[str] | None = "a1d3f5b7c9e2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}invalidated_memory_units DROP COLUMN IF EXISTS embedding")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# Unconstrained `vector` (no dimension) so the re-added column accepts any
|
||||
# model's embeddings; it comes back empty regardless.
|
||||
op.execute(f"ALTER TABLE {schema}invalidated_memory_units ADD COLUMN IF NOT EXISTS embedding vector")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
# Oracle has no `DROP COLUMN IF EXISTS`; swallow ORA-00904 (column does not
|
||||
# exist) so the migration is idempotent and safe on a fresh schema whose
|
||||
# baseline already omits the column.
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units DROP COLUMN embedding';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -904 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
# Swallow ORA-01430 (column already exists) for idempotency.
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units ADD (embedding VECTOR)';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -1430 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
-39
@@ -1,39 +0,0 @@
|
||||
"""Merge two divergent migration heads.
|
||||
|
||||
``d4f6a8c2e1b3`` (drop the curation-archive embedding column) and
|
||||
``2071c7518f88`` (add the memory_links(bank_id, link_type) index) were authored
|
||||
in parallel off the same parent (``a1d3f5b7c9e2``) and merged independently,
|
||||
leaving the DAG with two heads. This is a no-op merge that re-unifies them so
|
||||
``alembic upgrade head`` is unambiguous again (enforced by
|
||||
``tests/test_alembic_dag.py::test_single_head``).
|
||||
|
||||
Revision ID: e1f2a3b4c5d6
|
||||
Revises: d4f6a8c2e1b3, 2071c7518f88
|
||||
Create Date: 2026-06-16
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e1f2a3b4c5d6"
|
||||
down_revision: str | Sequence[str] | None = ("d4f6a8c2e1b3", "2071c7518f88")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
# Pure DAG merge — both parents already applied their schema changes.
|
||||
pass
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-153
@@ -1,153 +0,0 @@
|
||||
"""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)
|
||||
-110
@@ -1,110 +0,0 @@
|
||||
"""Add server-side routine for cron-scheduled mental model refresh.
|
||||
|
||||
Installs ``public.mental_models_with_cron()`` — a discovery routine that returns
|
||||
every mental model carrying a non-empty ``trigger->>'refresh_cron'`` across all
|
||||
tenant schemas in one round-trip (the same per-schema scan as the other
|
||||
maintenance routines from ``e5f6a7b8c9d0``). The maintenance loop evaluates each
|
||||
candidate's cron expression in Python (``croniter``) against ``last_refreshed_at``
|
||||
to decide whether a scheduled refresh is due — cron arithmetic isn't expressible
|
||||
in plain SQL — and only the cron *candidate set* is discovered here.
|
||||
|
||||
Models that already have a ``refresh_mental_model`` operation pending/processing
|
||||
are excluded so a slow refresh isn't double-queued (mirrors the in-flight guard
|
||||
in ``banks_needing_consolidation``). Each per-schema query runs in its own
|
||||
``BEGIN ... EXCEPTION`` subtransaction so a schema dropped mid-scan (tenant
|
||||
deletion / migration) is skipped, not fatal — same resilience as
|
||||
``c7e9f1a3b5d2``.
|
||||
|
||||
Read-only (STABLE) discovery routine — the caller performs the refresh enqueue —
|
||||
so installing it never mutates data. PostgreSQL only: the worker poller and the
|
||||
maintenance loop are PG-only (Oracle slot intentionally absent, mirroring
|
||||
``e5f6a7b8c9d0``). The routine lives in ``public`` and is CREATE OR REPLACE, so
|
||||
it is installed exactly once (base / ``public`` run) to avoid the
|
||||
``tuple concurrently updated`` race on concurrent per-tenant runs.
|
||||
|
||||
Revision ID: f4d1c2b3a5e6
|
||||
Revises: c7e9f1a3b5d2
|
||||
Create Date: 2026-06-23
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "f4d1c2b3a5e6"
|
||||
down_revision: str | Sequence[str] | None = "c7e9f1a3b5d2"
|
||||
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.*`` routine.
|
||||
|
||||
The routine physically lives in ``public``, so it is installed exactly once —
|
||||
on the base run (no ``target_schema``) or the run that explicitly targets
|
||||
``public``. Mirrors ``c7e9f1a3b5d2``.
|
||||
"""
|
||||
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
|
||||
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.mental_models_with_cron()
|
||||
RETURNS TABLE(schema_name text, bank_id text, mental_model_id text,
|
||||
refresh_cron text, last_refreshed_at timestamptz)
|
||||
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 = 'mental_models' AND c.relkind = 'r'
|
||||
LOOP
|
||||
BEGIN
|
||||
RETURN QUERY EXECUTE format($q$
|
||||
SELECT %1$L::text, mm.bank_id::text, mm.id::text,
|
||||
mm.trigger->>'refresh_cron', mm.last_refreshed_at
|
||||
FROM %1$I.mental_models mm
|
||||
WHERE COALESCE(mm.trigger->>'refresh_cron', '') <> ''
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM %1$I.async_operations o
|
||||
WHERE o.bank_id = mm.bank_id
|
||||
AND o.operation_type = 'refresh_mental_model'
|
||||
AND o.status IN ('pending', 'processing')
|
||||
AND o.task_payload->>'mental_model_id' = mm.id::text
|
||||
)
|
||||
$q$, sch);
|
||||
EXCEPTION
|
||||
-- Schema or its tables vanished between the pg_class
|
||||
-- snapshot and this query (tenant dropped or migrating).
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
if not _should_install_public_routines(context.config.get_main_option("target_schema")):
|
||||
return
|
||||
op.execute("DROP FUNCTION IF EXISTS public.mental_models_with_cron()")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
@@ -122,7 +122,6 @@ _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),
|
||||
@@ -139,50 +138,6 @@ _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.
|
||||
# No `embedding` column: the archive is cold storage and revert recomputes the
|
||||
# embedding, so there is no archive vector to fall out of sync with the live
|
||||
# model's dimension on a model switch (#2209).
|
||||
"""
|
||||
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,
|
||||
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,
|
||||
|
||||
+2
@@ -16,7 +16,9 @@ retention parameters, retrieval settings, etc.) in Python field name format.
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import context, op
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
|
||||
@@ -1,97 +0,0 @@
|
||||
"""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
@@ -9,7 +9,7 @@ from fastmcp import FastMCP
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api import __version__ as HINDSIGHT_VERSION
|
||||
from hindsight_api.config import DEFAULT_MCP_RECALL_DESCRIPTION, DEFAULT_MCP_RETAIN_DESCRIPTION, _get_raw_config
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.memory_engine import _current_schema
|
||||
from hindsight_api.extensions import MCPExtension, load_extension
|
||||
from hindsight_api.extensions.tenant import AuthenticationError
|
||||
@@ -78,19 +78,6 @@ def get_current_mcp_authenticated() -> bool:
|
||||
return _current_mcp_authenticated.get()
|
||||
|
||||
|
||||
def _build_mcp_tool_descriptions(extra_instructions: str | None) -> tuple[str | None, str | None]:
|
||||
"""Return custom retain/recall descriptions when server-level MCP instructions are set."""
|
||||
if not isinstance(extra_instructions, str):
|
||||
return None, None
|
||||
|
||||
extra_instructions = extra_instructions.strip()
|
||||
if not extra_instructions:
|
||||
return None, None
|
||||
|
||||
suffix = f"\n\nAdditional instructions: {extra_instructions}"
|
||||
return DEFAULT_MCP_RETAIN_DESCRIPTION + suffix, DEFAULT_MCP_RECALL_DESCRIPTION + suffix
|
||||
|
||||
|
||||
def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"""
|
||||
Create and configure the Hindsight MCP server.
|
||||
@@ -126,8 +113,6 @@ 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",
|
||||
@@ -148,10 +133,6 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
allowed = frozenset(global_config.mcp_enabled_tools)
|
||||
base_tools = (base_tools if base_tools is not None else _ALL_TOOLS) & allowed
|
||||
|
||||
retain_description, recall_description = _build_mcp_tool_descriptions(
|
||||
getattr(global_config, "mcp_instructions", None)
|
||||
)
|
||||
|
||||
# Configure and register tools using shared module
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=get_current_bank_id,
|
||||
@@ -161,8 +142,6 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
mcp_authenticated_resolver=get_current_mcp_authenticated, # Propagate MCP pre-auth flag
|
||||
include_bank_id_param=multi_bank,
|
||||
tools=base_tools,
|
||||
retain_description=retain_description,
|
||||
recall_description=recall_description,
|
||||
)
|
||||
|
||||
register_mcp_tools(mcp, memory, config)
|
||||
|
||||
@@ -1,263 +0,0 @@
|
||||
"""Open Knowledge Format (OKF) projection for knowledge pages.
|
||||
|
||||
Knowledge pages are a *read-only* OKF view over the existing mental models: each
|
||||
mental model is projected into an OKF document — a markdown body with YAML
|
||||
frontmatter (``type`` required; ``title``/``description``/``tags``/``timestamp``
|
||||
optional) — and pages are linked into a constellation graph via shared tags.
|
||||
|
||||
See the Open Knowledge Format spec:
|
||||
https://github.com/GoogleCloudPlatform/knowledge-catalog/tree/main/okf
|
||||
|
||||
This module is intentionally pure: every function transforms the mental-model
|
||||
dicts returned by ``MemoryEngine.list_mental_models`` / ``get_mental_model`` and
|
||||
never touches the database. That keeps the OKF contract unit-testable without a
|
||||
DB or LLM and lets the HTTP layer stay a thin wrapper.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
# OKF requires exactly one frontmatter field — ``type``. We default to this when
|
||||
# a page does not declare one via a ``type:<x>`` tag.
|
||||
DEFAULT_PAGE_TYPE = "knowledge-page"
|
||||
|
||||
# A page declares its OKF ``type`` through a tag of the form ``type:runbook``.
|
||||
# This keeps the projection schema-free (no new mental_models column): the type
|
||||
# is lifted from the existing tags array.
|
||||
TYPE_TAG_PREFIX = "type:"
|
||||
|
||||
INDEX_FILENAME = "index.md"
|
||||
|
||||
# Deterministic, colour-blind-friendly palette. Type → colour is stable across
|
||||
# requests so the constellation keeps the same colours between reloads.
|
||||
_PALETTE = (
|
||||
"#0074d9", # blue
|
||||
"#2ecc40", # green
|
||||
"#b10dc9", # purple
|
||||
"#ff851b", # orange
|
||||
"#39cccc", # teal
|
||||
"#f012be", # magenta
|
||||
"#3d9970", # olive
|
||||
"#ff4136", # red
|
||||
)
|
||||
|
||||
_EDGE_COLOR = "#9aa5b1"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PageType:
|
||||
"""A page's OKF ``type`` and the tags that remain after the type tag is split off."""
|
||||
|
||||
type: str
|
||||
display_tags: list[str]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class KnowledgeGraph:
|
||||
"""Cytoscape-style node/edge graph of knowledge pages linked by shared tags."""
|
||||
|
||||
nodes: list[dict[str, Any]] = field(default_factory=list)
|
||||
edges: list[dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
|
||||
def _color_for(key: str) -> str:
|
||||
"""Stable colour for a string key (FNV-ish hash into the fixed palette)."""
|
||||
h = 0
|
||||
for ch in key:
|
||||
h = (h * 31 + ord(ch)) & 0xFFFFFFFF
|
||||
return _PALETTE[h % len(_PALETTE)]
|
||||
|
||||
|
||||
def _scalar(value: Any) -> str:
|
||||
"""Emit a YAML-safe double-quoted scalar.
|
||||
|
||||
We always double-quote so arbitrary page names / source queries can't be
|
||||
misread as YAML special forms (``true``, ``2026-01-01``, ``- x``, etc.).
|
||||
"""
|
||||
text = str(value)
|
||||
escaped = text.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n").replace("\r", "")
|
||||
return f'"{escaped}"'
|
||||
|
||||
|
||||
def page_type(tags: list[str] | None) -> PageType:
|
||||
"""Split an OKF ``type`` out of the tag list.
|
||||
|
||||
The first ``type:<x>`` tag wins; all ``type:`` tags are removed from the
|
||||
returned ``display_tags`` so they don't pollute the constellation's
|
||||
shared-tag edges. Falls back to :data:`DEFAULT_PAGE_TYPE`.
|
||||
"""
|
||||
resolved = DEFAULT_PAGE_TYPE
|
||||
display: list[str] = []
|
||||
for tag in tags or []:
|
||||
if tag.startswith(TYPE_TAG_PREFIX):
|
||||
suffix = tag[len(TYPE_TAG_PREFIX) :].strip()
|
||||
if suffix and resolved == DEFAULT_PAGE_TYPE:
|
||||
resolved = suffix
|
||||
continue
|
||||
display.append(tag)
|
||||
return PageType(type=resolved, display_tags=display)
|
||||
|
||||
|
||||
def _timestamp(mm: dict[str, Any]) -> str | None:
|
||||
return mm.get("last_refreshed_at") or mm.get("created_at")
|
||||
|
||||
|
||||
def frontmatter(mm: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Build the ordered OKF frontmatter mapping for a mental model.
|
||||
|
||||
``None``/empty values are dropped by :func:`render_frontmatter`.
|
||||
"""
|
||||
pt = page_type(mm.get("tags"))
|
||||
return {
|
||||
"id": mm.get("id"),
|
||||
"type": pt.type,
|
||||
"title": mm.get("name"),
|
||||
"description": mm.get("source_query"),
|
||||
"tags": pt.display_tags,
|
||||
"timestamp": _timestamp(mm),
|
||||
}
|
||||
|
||||
|
||||
def render_frontmatter(fm: dict[str, Any]) -> str:
|
||||
"""Render a frontmatter mapping into a ``---`` fenced YAML block."""
|
||||
lines = ["---"]
|
||||
for key, value in fm.items():
|
||||
if value is None:
|
||||
continue
|
||||
if isinstance(value, list):
|
||||
if not value:
|
||||
continue
|
||||
lines.append(f"{key}:")
|
||||
lines.extend(f" - {_scalar(item)}" for item in value)
|
||||
else:
|
||||
lines.append(f"{key}: {_scalar(value)}")
|
||||
lines.append("---")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def render_document(mm: dict[str, Any]) -> str:
|
||||
"""Render a full OKF document: frontmatter block + markdown body."""
|
||||
body = (mm.get("content") or "").strip()
|
||||
return f"{render_frontmatter(frontmatter(mm))}\n\n{body}\n" if body else f"{render_frontmatter(frontmatter(mm))}\n"
|
||||
|
||||
|
||||
def page_filename(page_id: str) -> str:
|
||||
"""OKF bundle filename for a page id."""
|
||||
return f"{page_id}.md"
|
||||
|
||||
|
||||
def log_filename(page_id: str) -> str:
|
||||
"""OKF reserved per-page history filename."""
|
||||
return f"{page_id}.log.md"
|
||||
|
||||
|
||||
def render_index(nodes: list[dict[str, Any]]) -> str:
|
||||
"""Render the reserved ``index.md`` — nested OKF navigation over the tree.
|
||||
|
||||
``nodes`` is the flat folder/page list (each with ``id``, ``kind``, ``name``,
|
||||
``parent_id``); folders nest their children, pages link to their ``.md``.
|
||||
"""
|
||||
fm = render_frontmatter({"type": "index", "title": "Knowledge base"})
|
||||
lines = [fm, "", "# Knowledge base", ""]
|
||||
|
||||
children: dict[Any, list[dict[str, Any]]] = {}
|
||||
for node in nodes:
|
||||
children.setdefault(node.get("parent_id"), []).append(node)
|
||||
|
||||
def walk(parent: Any, depth: int) -> None:
|
||||
ordered = sorted(children.get(parent, []), key=lambda n: (n.get("sort_order", 0), n.get("name") or ""))
|
||||
for node in ordered:
|
||||
indent = " " * depth
|
||||
if node.get("kind") == "folder":
|
||||
lines.append(f"{indent}- **{node['name']}/**")
|
||||
walk(node["id"], depth + 1)
|
||||
else:
|
||||
description = node.get("source_query") or node.get("description")
|
||||
link = f"{indent}- [{node['name']}](./{page_filename(node['id'])})"
|
||||
lines.append(f"{link} — {description}" if description else link)
|
||||
|
||||
walk(None, 0)
|
||||
if len(lines) == 4:
|
||||
lines.append("_No knowledge pages yet._")
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def render_log(mm: dict[str, Any], history: list[dict[str, Any]]) -> str:
|
||||
"""Render the reserved per-page ``log.md`` from refresh history.
|
||||
|
||||
Each history entry is ``{previous_content, previous_reflect_response,
|
||||
changed_at}`` (newest first), capturing the content *before* a refresh.
|
||||
"""
|
||||
name = mm.get("name") or mm.get("id")
|
||||
fm = render_frontmatter({"type": "log", "title": f"{name} — history"})
|
||||
lines = [fm, "", f"# {name} — history", ""]
|
||||
if not history:
|
||||
lines.append("_No refresh history._")
|
||||
return "\n".join(lines) + "\n"
|
||||
for entry in history:
|
||||
changed_at = entry.get("changed_at") or "unknown"
|
||||
previous = (entry.get("previous_content") or "").strip()
|
||||
lines.append(f"## {changed_at}")
|
||||
lines.append("")
|
||||
lines.append(previous if previous else "_(empty)_")
|
||||
lines.append("")
|
||||
return "\n".join(lines).rstrip() + "\n"
|
||||
|
||||
|
||||
def knowledge_graph(
|
||||
pages: list[dict[str, Any]],
|
||||
cluster_for: "Callable[[dict[str, Any]], str] | None" = None,
|
||||
) -> KnowledgeGraph:
|
||||
"""Derive the constellation graph: pages as nodes, shared tags as edges.
|
||||
|
||||
Two pages are linked when they share at least one (non-``type:``) tag; the
|
||||
edge weight is the number of shared tags. Each node's cluster (``type`` field
|
||||
+ colour) comes from ``cluster_for(page)`` — the knowledge base groups by
|
||||
parent folder; the default groups by OKF ``type``.
|
||||
"""
|
||||
nodes: list[dict[str, Any]] = []
|
||||
tag_sets: list[tuple[str, frozenset[str]]] = []
|
||||
for mm in pages:
|
||||
page_id = mm["id"]
|
||||
pt = page_type(mm.get("tags"))
|
||||
cluster = cluster_for(mm) if cluster_for else pt.type
|
||||
tag_sets.append((page_id, frozenset(pt.display_tags)))
|
||||
nodes.append(
|
||||
{
|
||||
"data": {
|
||||
"id": page_id,
|
||||
"label": mm.get("name") or page_id,
|
||||
"type": cluster,
|
||||
"tagCount": len(pt.display_tags),
|
||||
"color": _color_for(cluster),
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
edges: list[dict[str, Any]] = []
|
||||
for i in range(len(tag_sets)):
|
||||
source_id, source_tags = tag_sets[i]
|
||||
if not source_tags:
|
||||
continue
|
||||
for j in range(i + 1, len(tag_sets)):
|
||||
target_id, target_tags = tag_sets[j]
|
||||
shared = source_tags & target_tags
|
||||
if not shared:
|
||||
continue
|
||||
edges.append(
|
||||
{
|
||||
"data": {
|
||||
"id": f"{source_id}--{target_id}",
|
||||
"source": source_id,
|
||||
"target": target_id,
|
||||
"sharedTags": sorted(shared),
|
||||
"weight": len(shared),
|
||||
"color": _EDGE_COLOR,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
return KnowledgeGraph(nodes=nodes, edges=edges)
|
||||
@@ -1,85 +0,0 @@
|
||||
"""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()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -8,7 +8,6 @@ Config values are resolved on every request to ensure consistency across
|
||||
multiple API servers.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import asdict, replace
|
||||
@@ -19,8 +18,6 @@ from hindsight_api.config import (
|
||||
HindsightConfig,
|
||||
_get_raw_config,
|
||||
normalize_config_dict,
|
||||
validate_retain_chunking_config,
|
||||
validate_retain_completion_token_budget,
|
||||
)
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
@@ -32,35 +29,6 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _validate_retain_strategy_chunking(base_config: HindsightConfig, strategies: Any) -> None:
|
||||
"""Validate retain strategy chunking with the same semantics as apply_strategy()."""
|
||||
if not isinstance(strategies, dict):
|
||||
return
|
||||
configurable = HindsightConfig.get_configurable_fields()
|
||||
for strategy_name, overrides in strategies.items():
|
||||
if not isinstance(overrides, dict):
|
||||
raise ValueError(f"Invalid retain strategy {strategy_name!r}: must be an object")
|
||||
filtered = {k: v for k, v in overrides.items() if k in configurable}
|
||||
if not filtered:
|
||||
continue
|
||||
try:
|
||||
resolved = replace(base_config, **filtered)
|
||||
validate_retain_chunking_config(
|
||||
resolved.retain_chunk_size,
|
||||
resolved.retain_structured_chunk_size,
|
||||
)
|
||||
validate_retain_completion_token_budget(
|
||||
llm_provider=resolved.llm_provider,
|
||||
retain_max_completion_tokens=resolved.retain_max_completion_tokens,
|
||||
retain_chunk_size=resolved.retain_chunk_size,
|
||||
retain_llm_model=resolved.retain_llm_model,
|
||||
llm_model=resolved.llm_model,
|
||||
retain_llm_provider=resolved.retain_llm_provider,
|
||||
)
|
||||
except ValueError as e:
|
||||
raise ValueError(f"Invalid retain strategy {strategy_name!r}: {e}") from e
|
||||
|
||||
|
||||
class ConfigResolver:
|
||||
"""Resolves hierarchical configuration with tenant/bank overrides."""
|
||||
|
||||
@@ -78,26 +46,6 @@ class ConfigResolver:
|
||||
self._configurable_fields = HindsightConfig.get_configurable_fields()
|
||||
self._credential_fields = HindsightConfig.get_credential_fields()
|
||||
|
||||
async def _resolve_parent_config_dict(self, bank_id: str, context: RequestContext | None = None) -> dict[str, Any]:
|
||||
"""Resolve global + tenant config before bank-level overrides."""
|
||||
config_dict = asdict(self._global_config)
|
||||
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
tenant_overrides = await self.tenant_extension.get_tenant_config(context)
|
||||
if tenant_overrides:
|
||||
# Normalize keys and filter to configurable fields only
|
||||
normalized_tenant = normalize_config_dict(tenant_overrides)
|
||||
configurable_tenant = {k: v for k, v in normalized_tenant.items() if k in self._configurable_fields}
|
||||
config_dict.update(configurable_tenant)
|
||||
logger.debug(
|
||||
f"Applied tenant config overrides for bank {bank_id}: {list(configurable_tenant.keys())}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load tenant config for bank {bank_id}: {e}")
|
||||
|
||||
return config_dict
|
||||
|
||||
async def resolve_full_config(self, bank_id: str, context: RequestContext | None = None) -> HindsightConfig:
|
||||
"""
|
||||
Resolve full HindsightConfig for a bank with hierarchical overrides applied.
|
||||
@@ -117,7 +65,23 @@ class ConfigResolver:
|
||||
Returns:
|
||||
Complete HindsightConfig with hierarchical overrides applied
|
||||
"""
|
||||
config_dict = await self._resolve_parent_config_dict(bank_id, context)
|
||||
# Start with global config (all fields)
|
||||
config_dict = asdict(self._global_config)
|
||||
|
||||
# Load tenant config overrides (if tenant extension available)
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
tenant_overrides = await self.tenant_extension.get_tenant_config(context)
|
||||
if tenant_overrides:
|
||||
# Normalize keys and filter to configurable fields only
|
||||
normalized_tenant = normalize_config_dict(tenant_overrides)
|
||||
configurable_tenant = {k: v for k, v in normalized_tenant.items() if k in self._configurable_fields}
|
||||
config_dict.update(configurable_tenant)
|
||||
logger.debug(
|
||||
f"Applied tenant config overrides for bank {bank_id}: {list(configurable_tenant.keys())}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load tenant config for bank {bank_id}: {e}")
|
||||
|
||||
# Load bank config overrides
|
||||
bank_overrides = await self._load_bank_config(bank_id)
|
||||
@@ -128,25 +92,6 @@ class ConfigResolver:
|
||||
# Return full config object (dataclass doesn't have __init__ that accepts kwargs, so we update the object)
|
||||
# Create a new config instance by copying the global config and updating fields
|
||||
resolved_config = HindsightConfig(**config_dict)
|
||||
# Multi-LLM chains are static credential fields (never tenant/bank-overridable),
|
||||
# but asdict() above flattened their member dataclasses into plain dicts. Restore
|
||||
# the original typed objects from the global config so the resolved object stays
|
||||
# well-typed for any consumer that reads them.
|
||||
resolved_config = replace(
|
||||
resolved_config,
|
||||
llm_members=self._global_config.llm_members,
|
||||
llm_strategy=self._global_config.llm_strategy,
|
||||
retain_llm_members=self._global_config.retain_llm_members,
|
||||
retain_llm_strategy=self._global_config.retain_llm_strategy,
|
||||
reflect_llm_members=self._global_config.reflect_llm_members,
|
||||
reflect_llm_strategy=self._global_config.reflect_llm_strategy,
|
||||
consolidation_llm_members=self._global_config.consolidation_llm_members,
|
||||
consolidation_llm_strategy=self._global_config.consolidation_llm_strategy,
|
||||
)
|
||||
validate_retain_chunking_config(
|
||||
resolved_config.retain_chunk_size,
|
||||
resolved_config.retain_structured_chunk_size,
|
||||
)
|
||||
return resolved_config
|
||||
|
||||
async def get_bank_config(self, bank_id: str, context: RequestContext | None = None) -> dict[str, Any]:
|
||||
@@ -177,83 +122,26 @@ class ConfigResolver:
|
||||
resolved_config = await self.resolve_full_config(bank_id, context)
|
||||
config_dict = asdict(resolved_config)
|
||||
|
||||
# SECURITY: drop static/infrastructure + credential fields, then permission-filter.
|
||||
filtered = self._strip_static_and_credential_fields(config_dict)
|
||||
return await self._apply_permission_filter(filtered, bank_id, context)
|
||||
# SECURITY: Filter to only configurable fields (exclude static/infrastructure)
|
||||
filtered = {k: v for k, v in config_dict.items() if k in self._configurable_fields}
|
||||
|
||||
def _strip_static_and_credential_fields(self, config_dict: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Keep only configurable, non-credential fields.
|
||||
# SECURITY: Remove ALL credential fields (API keys, base URLs, etc.)
|
||||
filtered = {k: v for k, v in filtered.items() if k not in self._credential_fields}
|
||||
|
||||
SECURITY: excludes static/infrastructure fields and ALL credential fields
|
||||
(API keys, base URLs, etc.) so a resolved config is safe to return over the API.
|
||||
"""
|
||||
return {
|
||||
k: v for k, v in config_dict.items() if k in self._configurable_fields and k not in self._credential_fields
|
||||
}
|
||||
|
||||
async def _apply_permission_filter(
|
||||
self, filtered: dict[str, Any], bank_id: str, context: RequestContext | None
|
||||
) -> dict[str, Any]:
|
||||
"""Further restrict already-stripped config to the tenant/bank permission allow-list.
|
||||
|
||||
On extension error, leaves ``filtered`` unchanged (parity with the historical
|
||||
single-bank path: a permissions lookup failure must not leak or drop fields).
|
||||
"""
|
||||
if not (self.tenant_extension and context):
|
||||
return filtered
|
||||
try:
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
filtered = {k: v for k, v in filtered.items() if k in allowed_fields}
|
||||
logger.debug(
|
||||
f"Applied permission filter for bank {bank_id}: allowed={len(allowed_fields)} fields, "
|
||||
f"returned={len(filtered)} fields"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load permissions for bank {bank_id}: {e}")
|
||||
return filtered
|
||||
|
||||
async def get_bank_configs(
|
||||
self, bank_ids: list[str], context: RequestContext | None = None
|
||||
) -> dict[str, dict[str, Any]]:
|
||||
"""Batch variant of :meth:`get_bank_config` for many banks.
|
||||
|
||||
Equivalent to calling ``get_bank_config`` per bank, but resolves the
|
||||
global + tenant base once and loads every bank's ``banks.config`` JSONB
|
||||
in a single query, instead of one config round-trip per bank. Used by
|
||||
``list_banks`` to overlay disposition + mission without an N+1.
|
||||
|
||||
Returns a mapping of bank_id -> filtered configurable-field dict. A bank
|
||||
with no config row still appears, mapped to the global+tenant base.
|
||||
"""
|
||||
if not bank_ids:
|
||||
return {}
|
||||
|
||||
# Global + tenant base, resolved once (tenant override is per-request, not per-bank).
|
||||
base_dict = asdict(self._global_config)
|
||||
# PERMISSIONS: Further filter based on tenant/bank permissions
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
tenant_overrides = await self.tenant_extension.get_tenant_config(context)
|
||||
if tenant_overrides:
|
||||
normalized_tenant = normalize_config_dict(tenant_overrides)
|
||||
base_dict.update({k: v for k, v in normalized_tenant.items() if k in self._configurable_fields})
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
filtered = {k: v for k, v in filtered.items() if k in allowed_fields}
|
||||
logger.debug(
|
||||
f"Applied permission filter for bank {bank_id}: allowed={len(allowed_fields)} fields, "
|
||||
f"returned={len(filtered)} fields"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load tenant config for bulk resolve: {e}")
|
||||
logger.warning(f"Failed to load permissions for bank {bank_id}: {e}")
|
||||
|
||||
# All bank overrides in one query, then merge + strip per bank.
|
||||
bank_overrides = await self._load_bank_configs(bank_ids)
|
||||
stripped = {
|
||||
bank_id: self._strip_static_and_credential_fields({**base_dict, **bank_overrides.get(bank_id, {})})
|
||||
for bank_id in bank_ids
|
||||
}
|
||||
|
||||
# Permission filter is per-bank; resolve concurrently when an extension is present.
|
||||
if not (self.tenant_extension and context):
|
||||
return stripped
|
||||
permission_filtered = await asyncio.gather(
|
||||
*(self._apply_permission_filter(stripped[bank_id], bank_id, context) for bank_id in bank_ids)
|
||||
)
|
||||
return dict(zip(bank_ids, permission_filtered, strict=True))
|
||||
return filtered
|
||||
|
||||
async def _load_bank_config(self, bank_id: str) -> dict[str, Any]:
|
||||
"""
|
||||
@@ -292,45 +180,6 @@ class ConfigResolver:
|
||||
|
||||
return {}
|
||||
|
||||
async def _load_bank_configs(self, bank_ids: list[str]) -> dict[str, dict[str, Any]]:
|
||||
"""Bulk variant of :meth:`_load_bank_config`: load many banks' overrides in one query.
|
||||
|
||||
Returns a mapping of bank_id -> normalized active overrides. Banks with no row
|
||||
(or an empty/all-tombstone config) are simply absent from the mapping.
|
||||
"""
|
||||
result: dict[str, dict[str, Any]] = {}
|
||||
if not bank_ids:
|
||||
return result
|
||||
try:
|
||||
async with self._backend.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT bank_id, config FROM {fq_table("banks")} WHERE bank_id = ANY($1)
|
||||
""",
|
||||
bank_ids,
|
||||
)
|
||||
for row in rows:
|
||||
config_data = row["config"]
|
||||
if not config_data:
|
||||
continue
|
||||
# Handle case where JSONB is returned as JSON string
|
||||
if isinstance(config_data, str):
|
||||
config_data = json.loads(config_data)
|
||||
|
||||
# Normalize keys (handle both env var format and Python field format)
|
||||
normalized = normalize_config_dict(config_data)
|
||||
|
||||
# Only active overrides for configurable fields. JSON null is a tombstone
|
||||
# for "Server Default" in the bank-config UI and must not override defaults.
|
||||
overrides = {
|
||||
k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None
|
||||
}
|
||||
if overrides:
|
||||
result[row["bank_id"]] = overrides
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to bulk-load bank configs: {e}")
|
||||
return result
|
||||
|
||||
async def update_bank_config(
|
||||
self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None
|
||||
) -> None:
|
||||
@@ -417,46 +266,12 @@ class ConfigResolver:
|
||||
# Validate recall budget fields
|
||||
_validate_recall_budget_updates(normalized_updates)
|
||||
|
||||
# Validate disposition trait fields (1-5 integer scale)
|
||||
_validate_disposition_updates(normalized_updates)
|
||||
|
||||
chunking_fields_updated = (
|
||||
"retain_chunk_size" in normalized_updates
|
||||
or "retain_structured_chunk_size" in normalized_updates
|
||||
or "retain_strategies" in normalized_updates
|
||||
)
|
||||
if chunking_fields_updated:
|
||||
config_dict = await self._resolve_parent_config_dict(bank_id, context)
|
||||
active_bank_overrides = await self._load_bank_config(bank_id)
|
||||
for key, value in normalized_updates.items():
|
||||
if key not in self._configurable_fields:
|
||||
continue
|
||||
if value is None:
|
||||
active_bank_overrides.pop(key, None)
|
||||
else:
|
||||
active_bank_overrides[key] = value
|
||||
config_dict.update(active_bank_overrides)
|
||||
base_config = HindsightConfig(**config_dict)
|
||||
validate_retain_chunking_config(
|
||||
base_config.retain_chunk_size,
|
||||
base_config.retain_structured_chunk_size,
|
||||
)
|
||||
_validate_retain_strategy_chunking(base_config, base_config.retain_strategies)
|
||||
|
||||
# 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
|
||||
|
||||
# Merge with existing config (JSONB || operator)
|
||||
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 = COALESCE(config, '{{}}'::jsonb) || $1::jsonb,
|
||||
SET config = config || $1::jsonb,
|
||||
updated_at = now()
|
||||
WHERE bank_id = $2
|
||||
""",
|
||||
@@ -534,31 +349,6 @@ def _validate_recall_budget_updates(updates: dict[str, Any]) -> None:
|
||||
)
|
||||
|
||||
|
||||
_DISPOSITION_KEYS = (
|
||||
"disposition_skepticism",
|
||||
"disposition_literalism",
|
||||
"disposition_empathy",
|
||||
)
|
||||
|
||||
|
||||
def _validate_disposition_updates(updates: dict[str, Any]) -> None:
|
||||
"""Validate disposition trait config updates. Raises ValueError on invalid input.
|
||||
|
||||
Each trait is an integer on a 1-5 scale (or None to clear the per-bank
|
||||
override). The read overlay injects the stored value verbatim into a strict
|
||||
``DispositionTraits(int, ge=1, le=5)``; an out-of-contract value (a float, a
|
||||
0-1 scale, or an int outside 1-5) accepted here would later 500 the whole
|
||||
bank list when any bank profile is serialized (issue #2348).
|
||||
"""
|
||||
for key in _DISPOSITION_KEYS:
|
||||
if key in updates:
|
||||
value = updates[key]
|
||||
if value is None:
|
||||
continue
|
||||
if not isinstance(value, int) or isinstance(value, bool) or not (1 <= value <= 5):
|
||||
raise ValueError(f"{key} must be an integer between 1 and 5, got {value!r}")
|
||||
|
||||
|
||||
def apply_strategy(config: HindsightConfig, strategy_name: str) -> HindsightConfig:
|
||||
"""
|
||||
Apply a named retain strategy's overrides on top of a resolved config.
|
||||
@@ -566,8 +356,7 @@ def apply_strategy(config: HindsightConfig, strategy_name: str) -> HindsightConf
|
||||
A strategy is a named set of hierarchical field overrides stored in
|
||||
config.retain_strategies. Any field in _HIERARCHICAL_FIELDS can be
|
||||
overridden, including retain_extraction_mode, retain_chunk_size,
|
||||
retain_structured_chunk_size, entity_labels,
|
||||
entities_allow_free_form, etc.
|
||||
entity_labels, entities_allow_free_form, etc.
|
||||
|
||||
Unknown strategy names log a warning and return config unchanged.
|
||||
Unknown or non-hierarchical fields in the strategy are silently ignored.
|
||||
@@ -589,17 +378,4 @@ def apply_strategy(config: HindsightConfig, strategy_name: str) -> HindsightConf
|
||||
return config
|
||||
|
||||
logger.debug(f"Applying retain strategy '{strategy_name}': {list(filtered.keys())}")
|
||||
resolved = replace(config, **filtered)
|
||||
validate_retain_chunking_config(
|
||||
resolved.retain_chunk_size,
|
||||
resolved.retain_structured_chunk_size,
|
||||
)
|
||||
validate_retain_completion_token_budget(
|
||||
llm_provider=resolved.llm_provider,
|
||||
retain_max_completion_tokens=resolved.retain_max_completion_tokens,
|
||||
retain_chunk_size=resolved.retain_chunk_size,
|
||||
retain_llm_model=resolved.retain_llm_model,
|
||||
llm_model=resolved.llm_model,
|
||||
retain_llm_provider=resolved.retain_llm_provider,
|
||||
)
|
||||
return resolved
|
||||
return replace(config, **filtered)
|
||||
|
||||
@@ -16,59 +16,11 @@ 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."""
|
||||
@@ -107,11 +59,11 @@ def _safe_json(data: Any) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
class AuditLogger:
|
||||
"""Fire-and-forget audit log writer.
|
||||
_SWEEP_INTERVAL_SECONDS = 3600 # Run retention sweep every hour
|
||||
|
||||
Retention of old rows is handled by the background :class:`MaintenanceLoop`.
|
||||
"""
|
||||
|
||||
class AuditLogger:
|
||||
"""Fire-and-forget audit log writer with optional retention sweep."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -119,11 +71,14 @@ 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."""
|
||||
@@ -173,6 +128,48 @@ 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(
|
||||
|
||||
@@ -1,34 +0,0 @@
|
||||
"""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
|
||||
@@ -1,254 +0,0 @@
|
||||
"""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 json
|
||||
import logging
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .db.base import DatabaseBackend
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
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]]],
|
||||
*,
|
||||
force_refresh: bool = False,
|
||||
) -> 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. When ``force_refresh`` is set the cached value is
|
||||
ignored: the loader runs and its result replaces the cached entry.
|
||||
"""
|
||||
if not self.enabled:
|
||||
return await loader()
|
||||
|
||||
key = (schema, bank_id)
|
||||
|
||||
if force_refresh:
|
||||
value = await loader()
|
||||
async with self._lock:
|
||||
self._store_unlocked(key, value)
|
||||
# Supersede any loader that was in flight for this key.
|
||||
self._in_flight.pop(key, None)
|
||||
return value
|
||||
|
||||
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:
|
||||
# Invalidation may have detached this loader and allowed a new
|
||||
# one to claim the key. Never remove that newer loader's slot.
|
||||
if self._in_flight.get(key) is in_flight:
|
||||
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:
|
||||
# Only the loader that still owns the key may populate the cache.
|
||||
# An invalidated loader can finish for its original callers, but its
|
||||
# pre-invalidation result must not overwrite a newer load.
|
||||
if self._in_flight.get(key) is in_flight:
|
||||
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:
|
||||
key = (schema, bank_id)
|
||||
self._entries.pop(key, None)
|
||||
# Detach rather than cancel: existing callers may finish with the
|
||||
# snapshot they requested, while post-invalidation callers reload.
|
||||
self._in_flight.pop(key, None)
|
||||
|
||||
async def clear(self) -> None:
|
||||
async with self._lock:
|
||||
self._entries.clear()
|
||||
self._in_flight.clear()
|
||||
|
||||
|
||||
class DistributedBankStatsCache:
|
||||
"""Table-backed (cross-process) TTL cache for `get_bank_stats`.
|
||||
|
||||
Same ``get_or_load`` / ``invalidate`` / ``clear`` contract as
|
||||
:class:`BankStatsCache`, but the store is the per-schema ``bank_stats_cache``
|
||||
table instead of a per-process dict — so one worker's computation is shared
|
||||
with every other worker, and no caller recomputes while a fresh row exists.
|
||||
|
||||
On a hit, a call is a single primary-key ``SELECT`` (sub-millisecond); only a
|
||||
miss runs the (expensive) ``loader`` and writes the row back. Concurrent
|
||||
misses are *not* coalesced across processes (that would need a lock): they
|
||||
each compute and ``UPSERT``, last write wins — all results are correct, at the
|
||||
cost of a brief redundant compute at expiry.
|
||||
|
||||
Every DB touch is best-effort: if the cache table is unreachable or missing
|
||||
(e.g. a schema mid-migration), the call degrades to computing without caching
|
||||
rather than failing ``get_bank_stats``. PostgreSQL only — the engine keeps the
|
||||
in-process :class:`BankStatsCache` for Oracle.
|
||||
"""
|
||||
|
||||
def __init__(self, *, backend: "DatabaseBackend", ttl_seconds: float) -> None:
|
||||
self._backend = backend
|
||||
self._ttl = float(ttl_seconds)
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self._ttl > 0
|
||||
|
||||
@staticmethod
|
||||
def _qualified(schema: str) -> str:
|
||||
return f'"{schema}".bank_stats_cache' if schema else "bank_stats_cache"
|
||||
|
||||
async def get_or_load(
|
||||
self,
|
||||
schema: str,
|
||||
bank_id: str,
|
||||
loader: Callable[[], Awaitable[dict[str, Any]]],
|
||||
*,
|
||||
force_refresh: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
if not self.enabled:
|
||||
return await loader()
|
||||
|
||||
table = self._qualified(schema)
|
||||
|
||||
# 1. Fresh row? Single PK lookup; ``payload::text`` sidesteps any
|
||||
# jsonb->object codec so we always decode the same way. Skipped when
|
||||
# the caller forces a refresh — then we recompute and overwrite below.
|
||||
if not force_refresh:
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT payload::text AS payload FROM {table} "
|
||||
f"WHERE bank_id = $1 AND computed_at > now() - make_interval(secs => $2::double precision)",
|
||||
bank_id,
|
||||
self._ttl,
|
||||
)
|
||||
if row is not None:
|
||||
return json.loads(row["payload"])
|
||||
except Exception as exc: # noqa: BLE001 — cache read must never break the endpoint
|
||||
logger.debug("bank_stats_cache read failed for %s.%s (%s); computing uncached", schema, bank_id, exc)
|
||||
return await loader()
|
||||
|
||||
# 2. Miss — compute, then write the row back (best-effort).
|
||||
value = await loader()
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
await conn.execute(
|
||||
f"INSERT INTO {table} (bank_id, payload, computed_at) VALUES ($1, $2::jsonb, now()) "
|
||||
f"ON CONFLICT (bank_id) DO UPDATE SET payload = EXCLUDED.payload, computed_at = now()",
|
||||
bank_id,
|
||||
json.dumps(value),
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 — a failed write just means no caching this round
|
||||
logger.warning("bank_stats_cache write failed for %s.%s (%s)", schema, bank_id, exc)
|
||||
return value
|
||||
|
||||
async def invalidate(self, schema: str, bank_id: str) -> None:
|
||||
"""Drop the cached row so the next read recomputes."""
|
||||
if not self.enabled:
|
||||
return
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
await conn.execute(f"DELETE FROM {self._qualified(schema)} WHERE bank_id = $1", bank_id)
|
||||
except Exception as exc: # noqa: BLE001 — invalidation must never break the write path
|
||||
logger.debug("bank_stats_cache invalidate failed for %s.%s (%s)", schema, bank_id, exc)
|
||||
|
||||
async def clear(self) -> None:
|
||||
"""Drop all cached rows in the current schema (best-effort)."""
|
||||
if not self.enabled:
|
||||
return
|
||||
from .memory_engine import get_current_schema
|
||||
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
await conn.execute(f"DELETE FROM {self._qualified(get_current_schema())}")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("bank_stats_cache clear failed (%s)", exc)
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -37,33 +37,6 @@ _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
|
||||
|
||||
@@ -92,7 +65,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`. Every entry must include a `reason`.
|
||||
Return a JSON object with three arrays: `creates`, `updates`, `deletes`.
|
||||
|
||||
### Example 1 — Merging recurring claims into an existing observation
|
||||
|
||||
@@ -106,7 +79,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"], "reason": "Both new facts restate the same sovereignty decision already captured by obs 1111 — merged as evidence rather than creating siblings."}}],
|
||||
"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"]}}],
|
||||
"deletes": []}}
|
||||
|
||||
### Example 2 — State change updates one observation; unrelated fact creates a new one
|
||||
@@ -120,8 +93,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"], "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."}}],
|
||||
{{"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"]}}],
|
||||
"deletes": []}}
|
||||
|
||||
### Observation text rules
|
||||
@@ -137,7 +110,6 @@ 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."""
|
||||
|
||||
@@ -173,55 +145,3 @@ 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)
|
||||
|
||||
@@ -27,21 +27,35 @@ from ..config import (
|
||||
DEFAULT_RERANKER_LITELLM_MODEL,
|
||||
DEFAULT_RERANKER_LITELLM_SDK_MODEL,
|
||||
DEFAULT_RERANKER_LOCAL_BATCH_SIZE,
|
||||
DEFAULT_RERANKER_LOCAL_FORCE_CPU,
|
||||
DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_LOCAL_MODEL,
|
||||
DEFAULT_RERANKER_LOCAL_TRUST_REMOTE_CODE,
|
||||
DEFAULT_RERANKER_PROVIDER,
|
||||
DEFAULT_RERANKER_SILICONFLOW_BASE_URL,
|
||||
DEFAULT_RERANKER_SILICONFLOW_MODEL,
|
||||
DEFAULT_RERANKER_TEI_BATCH_SIZE,
|
||||
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL,
|
||||
DEFAULT_ZEROENTROPY_BASE_URL,
|
||||
ENV_RERANKER_ALIBABA_API_KEY,
|
||||
ENV_RERANKER_COHERE_API_KEY,
|
||||
ENV_RERANKER_COHERE_MODEL,
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA,
|
||||
ENV_RERANKER_FLASHRANK_MODEL,
|
||||
ENV_RERANKER_GOOGLE_PROJECT_ID,
|
||||
ENV_RERANKER_LITELLM_SDK_API_KEY,
|
||||
ENV_RERANKER_LOCAL_FORCE_CPU,
|
||||
ENV_RERANKER_LOCAL_MAX_CONCURRENT,
|
||||
ENV_RERANKER_LOCAL_MODEL,
|
||||
ENV_RERANKER_LOCAL_TRUST_REMOTE_CODE,
|
||||
ENV_RERANKER_PROVIDER,
|
||||
ENV_RERANKER_SILICONFLOW_API_KEY,
|
||||
ENV_RERANKER_TEI_BATCH_SIZE,
|
||||
ENV_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
ENV_RERANKER_TEI_MAX_CONCURRENT,
|
||||
ENV_RERANKER_TEI_URL,
|
||||
ENV_RERANKER_ZEROENTROPY_API_KEY,
|
||||
)
|
||||
@@ -212,7 +226,7 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
device = "cpu"
|
||||
logger.info("Reranker: forcing CPU mode (HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU=1)")
|
||||
else:
|
||||
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
|
||||
# Check for GPU (CUDA) or Apple Silicon (MPS)
|
||||
# Wrap in try-except to gracefully handle any device detection issues
|
||||
# (e.g., in CI environments or when PyTorch is built without GPU support)
|
||||
device = "cpu" # Default to CPU
|
||||
@@ -220,13 +234,10 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
has_gpu = torch.cuda.is_available() or (
|
||||
hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
)
|
||||
# Intel Arc XPU support — torch.xpu is available when the XPU build is loaded
|
||||
if not has_gpu and hasattr(torch, "xpu"):
|
||||
has_gpu = torch.xpu.is_available()
|
||||
if has_gpu:
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS/XPU
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
|
||||
logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}")
|
||||
|
||||
# Patch transformers 5.x compatibility for models using XLM-RoBERTa
|
||||
# (e.g., jina-reranker-v2-base-multilingual). transformers 5.x removed
|
||||
@@ -293,6 +304,7 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
- bucket_batching: sort pairs by token length to reduce padding waste (36-54% speedup)
|
||||
- batch_size: explicit batch size for predict() calls (MPS optimal: 32)
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
if self.bucket_batching and len(pairs) > 1:
|
||||
@@ -1187,7 +1199,7 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str | None = None,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_RERANKER_LITELLM_SDK_MODEL,
|
||||
api_base: str | None = None,
|
||||
timeout: float = 60.0,
|
||||
@@ -1197,8 +1209,7 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
Initialize LiteLLM SDK cross-encoder client.
|
||||
|
||||
Args:
|
||||
api_key: API key for the reranking provider (optional — omit for
|
||||
providers that use ambient credentials, e.g. AWS Bedrock with IAM)
|
||||
api_key: API key for the reranking provider
|
||||
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)
|
||||
@@ -1273,9 +1284,8 @@ 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
|
||||
|
||||
@@ -1668,7 +1678,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
return CohereCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_openrouter_model,
|
||||
base_url=config.reranker_openrouter_base_url,
|
||||
base_url="https://openrouter.ai/api/v1/rerank",
|
||||
timeout=config.reranker_openrouter_timeout,
|
||||
)
|
||||
elif provider == "flashrank":
|
||||
@@ -1687,8 +1697,13 @@ 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=config.reranker_litellm_sdk_api_key or None,
|
||||
api_key=api_key,
|
||||
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,
|
||||
|
||||
@@ -19,6 +19,7 @@ and mirrors Django's ``DatabaseOperations`` architecture.
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from .base import DatabaseConnection
|
||||
from .result import ResultRow
|
||||
@@ -71,30 +72,6 @@ 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,
|
||||
|
||||
@@ -8,6 +8,8 @@ columns can't appear in GROUP BY).
|
||||
import json
|
||||
import uuid as uuid_mod
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from .base import DatabaseConnection
|
||||
from .ops import DataAccessOps, TagListingParts
|
||||
@@ -45,37 +47,6 @@ 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,
|
||||
@@ -256,21 +227,13 @@ class OracleOps(DataAccessOps):
|
||||
# Oracle doesn't support ON CONFLICT; rely on the PK and the
|
||||
# IGNORE_ROW_ON_DUPKEY_INDEX hint to skip duplicates server-side.
|
||||
# The hint name must match the PK constraint exactly.
|
||||
#
|
||||
# Sort to enforce a global lock-acquisition order on the
|
||||
# (bank_id, unit_id) PK. Without this, two concurrent
|
||||
# transactions inserting overlapping unit_id sets in different
|
||||
# orders can deadlock on the unique-check row locks. Sorting
|
||||
# gives every concurrent caller the same lock order, so
|
||||
# conflicting inserts queue cleanly instead of cycling.
|
||||
sorted_unit_ids = sorted(unit_ids)
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT /*+ IGNORE_ROW_ON_DUPKEY_INDEX({table}, pk_graph_maintenance_queue) */
|
||||
INTO {table} (bank_id, unit_id)
|
||||
VALUES ($1, $2)
|
||||
""",
|
||||
[(bank_id, uid) for uid in sorted_unit_ids],
|
||||
[(bank_id, uid) for uid in unit_ids],
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
|
||||
@@ -4,6 +4,11 @@ Uses unnest(), LATERAL, DISTINCT ON, and native array operations for
|
||||
efficient batch operations.
|
||||
"""
|
||||
|
||||
import json
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from .base import DatabaseConnection
|
||||
from .ops import DataAccessOps, TagListingParts
|
||||
from .result import ResultRow
|
||||
@@ -44,30 +49,6 @@ 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,
|
||||
@@ -348,15 +329,6 @@ class PostgreSQLOps(DataAccessOps):
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
# Sort to enforce a global lock-acquisition order on the
|
||||
# (bank_id, unit_id) unique-key. Without this, two concurrent
|
||||
# transactions inserting overlapping unit_id sets in different
|
||||
# orders can deadlock on the ON CONFLICT row locks — Postgres
|
||||
# acquires a short-lived lock per row being checked, and cycle
|
||||
# detection then aborts one transaction. Sorting gives every
|
||||
# concurrent caller the same lock order, so conflicting inserts
|
||||
# queue cleanly instead of cycling.
|
||||
sorted_unit_ids = sorted(unit_ids)
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (bank_id, unit_id)
|
||||
@@ -364,7 +336,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
ON CONFLICT (bank_id, unit_id) DO NOTHING
|
||||
""",
|
||||
bank_id,
|
||||
sorted_unit_ids,
|
||||
unit_ids,
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
@@ -624,6 +596,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
per_entity_limit: int,
|
||||
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
|
||||
# v0.5.6 array ops: unnest, &&, COUNT(DISTINCT) on source_memory_ids.
|
||||
from ..schema import fq_table
|
||||
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
|
||||
@@ -106,15 +106,6 @@ SCHEMAS_WITH_PENDING_WORK = OptionalRoutine(
|
||||
deployment.
|
||||
* Should be cheap and idempotent — called every poll cycle (~30s).
|
||||
|
||||
The poller trusts the result wholesale: any schema the routine does
|
||||
not return is treated as having no work this cycle. It does NOT
|
||||
second-guess omissions with a per-schema scan — that would re-run the
|
||||
exact queries this routine exists to avoid. Consequently the routine
|
||||
is *only* appropriate for multi-tenant deployments. Single-schema
|
||||
(default/public only) installs should NOT create it: the per-schema
|
||||
fallback below is a single cheap EXISTS check that covers ``public``
|
||||
correctly and cannot starve.
|
||||
|
||||
Fallback when the routine is absent: per-schema ``EXISTS`` queries
|
||||
from Python (~4ms per schema). The server-side path is a single-
|
||||
round-trip optimisation worth ~200ms in deployments with thousands
|
||||
|
||||
@@ -14,7 +14,6 @@ Supports multi-tenant schema isolation via ALTER SESSION SET CURRENT_SCHEMA.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
@@ -156,23 +155,6 @@ _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:
|
||||
@@ -703,11 +685,6 @@ 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)
|
||||
|
||||
@@ -885,7 +862,7 @@ class OracleConnection(DatabaseConnection):
|
||||
|
||||
return query, params
|
||||
|
||||
async def _read_returning_values(self, returning_cols: list[str], params: dict[str, Any]) -> dict[str, Any] | None:
|
||||
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):
|
||||
@@ -895,14 +872,6 @@ 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()
|
||||
@@ -1090,7 +1059,7 @@ class OracleConnection(DatabaseConnection):
|
||||
raise
|
||||
|
||||
if ret_cols is not None:
|
||||
row_dict = await self._read_returning_values(ret_cols, params)
|
||||
row_dict = 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 []]
|
||||
@@ -1128,7 +1097,7 @@ class OracleConnection(DatabaseConnection):
|
||||
raise
|
||||
|
||||
if ret_cols is not None:
|
||||
row_dict = await self._read_returning_values(ret_cols, params)
|
||||
row_dict = 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 []]
|
||||
@@ -1161,7 +1130,7 @@ class OracleConnection(DatabaseConnection):
|
||||
await cursor.execute(query, params)
|
||||
|
||||
if ret_cols is not None:
|
||||
row_dict = await self._read_returning_values(ret_cols, params)
|
||||
row_dict = self._read_returning_values(ret_cols, params)
|
||||
if row_dict is None:
|
||||
return None
|
||||
vals = list(row_dict.values())
|
||||
|
||||
@@ -26,8 +26,11 @@ from ..config import (
|
||||
DEFAULT_EMBEDDINGS_GEMINI_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LITELLM_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
DEFAULT_EMBEDDINGS_PROVIDER,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
@@ -37,6 +40,9 @@ from ..config import (
|
||||
DEFAULT_ZEROENTROPY_BASE_URL,
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY,
|
||||
ENV_EMBEDDINGS_GEMINI_API_KEY,
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL,
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
ENV_EMBEDDINGS_OPENAI_API_KEY,
|
||||
ENV_EMBEDDINGS_OPENAI_BASE_URL,
|
||||
ENV_EMBEDDINGS_OPENAI_MODEL,
|
||||
@@ -47,7 +53,6 @@ from ..config import (
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
ENV_LLM_API_KEY,
|
||||
)
|
||||
from .bank_attribution import apply_bank_attribution
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -190,7 +195,7 @@ class LocalSTEmbeddings(Embeddings):
|
||||
device = "cpu"
|
||||
logger.info("Embeddings: forcing CPU mode")
|
||||
else:
|
||||
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
|
||||
# Check for GPU (CUDA) or Apple Silicon (MPS)
|
||||
# Wrap in try-except to gracefully handle any device detection issues
|
||||
# (e.g., in CI environments or when PyTorch is built without GPU support)
|
||||
device = "cpu" # Default to CPU
|
||||
@@ -198,13 +203,10 @@ class LocalSTEmbeddings(Embeddings):
|
||||
has_gpu = torch.cuda.is_available() or (
|
||||
hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
)
|
||||
# Intel Arc XPU support — torch.xpu is available when the XPU build is loaded
|
||||
if not has_gpu and hasattr(torch, "xpu"):
|
||||
has_gpu = torch.xpu.is_available()
|
||||
if has_gpu:
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS/XPU
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
|
||||
logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}")
|
||||
|
||||
# Suppress verbose transformers warnings during model loading
|
||||
# This suppresses the "UNEXPECTED" warnings from BertModel which are harmless
|
||||
@@ -250,172 +252,6 @@ 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.
|
||||
@@ -699,7 +535,6 @@ class OpenAIEmbeddings(Embeddings):
|
||||
}
|
||||
if self.dimensions is not None:
|
||||
request["dimensions"] = self.dimensions
|
||||
apply_bank_attribution(request)
|
||||
|
||||
response = self._client.embeddings.create(**request)
|
||||
|
||||
@@ -712,8 +547,7 @@ class OpenAIEmbeddings(Embeddings):
|
||||
|
||||
class CodexOAuthEmbeddings(OpenAIEmbeddings):
|
||||
"""
|
||||
OpenAI embeddings using the Codex/ChatGPT OAuth token from the Codex
|
||||
``auth.json`` (``$CODEX_HOME/auth.json``, or ``~/.codex/auth.json`` when unset).
|
||||
OpenAI embeddings using the Codex/ChatGPT OAuth token from ``~/.codex/auth.json``.
|
||||
|
||||
Codex OAuth is an LLM-provider auth path in Hindsight, but the same bearer token
|
||||
can also authenticate against the standard OpenAI embeddings endpoint. This keeps
|
||||
@@ -1343,21 +1177,6 @@ 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.
|
||||
@@ -1367,10 +1186,6 @@ 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__(
|
||||
@@ -1525,13 +1340,9 @@ 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), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
|
||||
embed_kwargs = {"model": self.model, "contents": batch}
|
||||
if self._embed_config is not None:
|
||||
@@ -1539,13 +1350,7 @@ class GeminiEmbeddings(Embeddings):
|
||||
|
||||
result = self._client.models.embed_content(**embed_kwargs)
|
||||
|
||||
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])
|
||||
all_embeddings.extend([emb.values for emb in result.embeddings])
|
||||
|
||||
# L2-normalize when output_dimensionality is set — Gemini only returns
|
||||
# normalized vectors at full 3072 dims; truncated dims need re-normalization
|
||||
@@ -1586,20 +1391,6 @@ 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)
|
||||
@@ -1638,20 +1429,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "requesty":
|
||||
api_key = config.embeddings_requesty_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_EMBEDDINGS_REQUESTY_API_KEY, HINDSIGHT_API_REQUESTY_API_KEY, "
|
||||
f"or {ENV_LLM_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'requesty'"
|
||||
)
|
||||
return OpenAIEmbeddings(
|
||||
api_key=api_key,
|
||||
model=config.embeddings_requesty_model,
|
||||
base_url="https://router.requesty.ai/v1",
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "zeroentropy":
|
||||
api_key = config.embeddings_zeroentropy_api_key
|
||||
if not api_key:
|
||||
@@ -1715,6 +1492,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown embeddings provider: {provider}. "
|
||||
f"Supported: 'local', 'onnx', 'tei', 'openai', 'openai-codex', 'openrouter', 'requesty', 'cohere', 'google', "
|
||||
f"Supported: 'local', 'tei', 'openai', 'openai-codex', 'openrouter', 'cohere', 'google', "
|
||||
f"'zeroentropy', 'litellm', 'litellm-sdk'"
|
||||
)
|
||||
|
||||
@@ -782,6 +782,239 @@ class EntityResolver:
|
||||
|
||||
return entity_ids
|
||||
|
||||
async def resolve_entity(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_text: str,
|
||||
context: str,
|
||||
nearby_entities: list[dict],
|
||||
unit_event_date,
|
||||
) -> str:
|
||||
"""
|
||||
Resolve an entity to a canonical entity ID.
|
||||
|
||||
Args:
|
||||
bank_id: bank ID (entities are scoped to agents)
|
||||
entity_text: Entity text ("Alice", "Google", etc.)
|
||||
context: Context where entity appears
|
||||
nearby_entities: Other entities in the same unit
|
||||
unit_event_date: When this unit was created
|
||||
|
||||
Returns:
|
||||
Entity ID (creates new entity if needed)
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
# Find candidate entities with similar name
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, canonical_name, metadata, last_seen
|
||||
FROM {fq_table("entities")}
|
||||
WHERE bank_id = $1
|
||||
AND (
|
||||
canonical_name ILIKE $2
|
||||
OR canonical_name ILIKE $3
|
||||
OR $2 ILIKE canonical_name || '%%'
|
||||
)
|
||||
ORDER BY mention_count DESC
|
||||
""",
|
||||
bank_id,
|
||||
entity_text,
|
||||
f"%{entity_text}%",
|
||||
)
|
||||
|
||||
if not candidates:
|
||||
# New entity - create it
|
||||
return await self._create_entity(conn, bank_id, entity_text, unit_event_date)
|
||||
|
||||
# Score candidates based on:
|
||||
# 1. Name similarity
|
||||
# 2. Context overlap (TODO: could use embeddings)
|
||||
# 3. Co-occurring entities
|
||||
# 4. Temporal proximity
|
||||
|
||||
best_candidate = None
|
||||
best_score = 0.0
|
||||
best_name_similarity = 0.0
|
||||
|
||||
nearby_entity_set = {e["text"].lower() for e in nearby_entities if e["text"] != entity_text}
|
||||
|
||||
for row in candidates:
|
||||
candidate_id = row["id"]
|
||||
canonical_name = row["canonical_name"]
|
||||
metadata = row["metadata"]
|
||||
last_seen = row["last_seen"]
|
||||
score = 0.0
|
||||
|
||||
# 1. Name similarity (0-1)
|
||||
name_similarity = SequenceMatcher(None, entity_text.lower(), canonical_name.lower()).ratio()
|
||||
score += name_similarity * 0.5
|
||||
|
||||
# 2. Co-occurring entities (0-0.5)
|
||||
# Get entities that co-occurred with this candidate before
|
||||
# Use the materialized co-occurrence cache for fast lookup
|
||||
co_entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.canonical_name, ec.cooccurrence_count
|
||||
FROM {fq_table("entity_cooccurrences")} ec
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
CASE
|
||||
WHEN ec.entity_id_1 = $1 THEN ec.entity_id_2
|
||||
WHEN ec.entity_id_2 = $1 THEN ec.entity_id_1
|
||||
END = e.id
|
||||
)
|
||||
WHERE ec.entity_id_1 = $1 OR ec.entity_id_2 = $1
|
||||
""",
|
||||
candidate_id,
|
||||
)
|
||||
co_entities = {r["canonical_name"].lower() for r in co_entity_rows}
|
||||
|
||||
# Check overlap with nearby entities
|
||||
overlap = len(nearby_entity_set & co_entities)
|
||||
if nearby_entity_set:
|
||||
co_entity_score = overlap / len(nearby_entity_set)
|
||||
score += co_entity_score * 0.3
|
||||
|
||||
# 3. Temporal proximity (0-0.2)
|
||||
if last_seen:
|
||||
# Normalize both to UTC-aware to avoid naive/aware mismatch
|
||||
# (Oracle returns naive datetimes from fromisoformat)
|
||||
_evt = unit_event_date if unit_event_date.tzinfo else unit_event_date.replace(tzinfo=UTC)
|
||||
_seen = last_seen if last_seen.tzinfo else last_seen.replace(tzinfo=UTC)
|
||||
days_diff = abs((_evt - _seen).total_seconds() / 86400)
|
||||
if days_diff < 7: # Within a week
|
||||
temporal_score = max(0, 1.0 - (days_diff / 7))
|
||||
score += temporal_score * 0.2
|
||||
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_candidate = candidate_id
|
||||
best_name_similarity = name_similarity
|
||||
|
||||
# Threshold for considering it the same entity
|
||||
threshold = 0.6
|
||||
|
||||
if best_score > threshold:
|
||||
# Update entity
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("entities")}
|
||||
SET mention_count = mention_count + 1,
|
||||
last_seen = $1
|
||||
WHERE id = $2
|
||||
""",
|
||||
unit_event_date,
|
||||
best_candidate,
|
||||
)
|
||||
return best_candidate
|
||||
else:
|
||||
# Not confident - create new entity
|
||||
return await self._create_entity(conn, bank_id, entity_text, unit_event_date)
|
||||
|
||||
async def _create_entity(
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entity_text: str,
|
||||
event_date,
|
||||
) -> str:
|
||||
"""
|
||||
Create a new entity or get existing one if it already exists.
|
||||
|
||||
Uses INSERT ... ON CONFLICT to handle race conditions where
|
||||
two concurrent transactions try to create the same entity.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
bank_id: bank ID
|
||||
entity_text: Entity text
|
||||
event_date: When first seen
|
||||
|
||||
Returns:
|
||||
Entity ID
|
||||
"""
|
||||
entity_id = await conn.fetchval(
|
||||
f"""
|
||||
INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
VALUES ($1, $2, COALESCE($3, now()), COALESCE($4, now()), 1)
|
||||
ON CONFLICT (bank_id, LOWER(canonical_name))
|
||||
DO UPDATE SET
|
||||
mention_count = {fq_table("entities")}.mention_count + 1,
|
||||
last_seen = EXCLUDED.last_seen
|
||||
RETURNING id
|
||||
""",
|
||||
bank_id,
|
||||
entity_text,
|
||||
event_date,
|
||||
event_date,
|
||||
)
|
||||
return entity_id
|
||||
|
||||
async def link_unit_to_entity(self, unit_id: str, entity_id: str):
|
||||
"""
|
||||
Link a memory unit to an entity.
|
||||
Also updates co-occurrence cache with other entities in the same unit.
|
||||
|
||||
Args:
|
||||
unit_id: Memory unit ID
|
||||
entity_id: Entity ID
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
# Insert unit-entity link
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
unit_id,
|
||||
entity_id,
|
||||
)
|
||||
|
||||
# Update co-occurrence cache: find other entities in this unit
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT entity_id
|
||||
FROM {fq_table("unit_entities")}
|
||||
WHERE unit_id = $1 AND entity_id != $2
|
||||
""",
|
||||
unit_id,
|
||||
entity_id,
|
||||
)
|
||||
|
||||
other_entities = [row["entity_id"] for row in rows]
|
||||
|
||||
# Update co-occurrences for each pair
|
||||
for other_entity_id in other_entities:
|
||||
await self._update_cooccurrence(conn, entity_id, other_entity_id)
|
||||
|
||||
async def _update_cooccurrence(self, conn, entity_id_1: str, entity_id_2: str):
|
||||
"""
|
||||
Update the co-occurrence cache for two entities.
|
||||
|
||||
Uses CHECK constraint ordering (entity_id_1 < entity_id_2) to avoid duplicates.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
entity_id_1: First entity ID
|
||||
entity_id_2: Second entity ID
|
||||
"""
|
||||
# Ensure consistent ordering (smaller UUID first)
|
||||
if entity_id_1 > entity_id_2:
|
||||
entity_id_1, entity_id_2 = entity_id_2, entity_id_1
|
||||
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("entity_cooccurrences")} (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
VALUES ($1, $2, 1, NOW())
|
||||
ON CONFLICT (entity_id_1, entity_id_2)
|
||||
DO UPDATE SET
|
||||
cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + 1,
|
||||
last_cooccurred = NOW()
|
||||
""",
|
||||
entity_id_1,
|
||||
entity_id_2,
|
||||
)
|
||||
|
||||
async def link_units_to_entities_batch(
|
||||
self,
|
||||
unit_entity_pairs: list[tuple[str, str]] | list[tuple[str, str, datetime | None]],
|
||||
|
||||
@@ -10,7 +10,7 @@ from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.memory_engine import BankLlmHealthInfo, Budget
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.engine.response_models import RecallResult, ReflectResult
|
||||
from hindsight_api.engine.search.tags import TagsMatch
|
||||
from hindsight_api.models import RequestContext
|
||||
@@ -449,7 +449,6 @@ class MemoryEngineInterface(ABC):
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
force_refresh: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Get statistics about memory nodes and links for a bank.
|
||||
@@ -457,46 +456,10 @@ class MemoryEngineInterface(ABC):
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
force_refresh: Bypass the cached value and recompute (also refreshes
|
||||
the cache for subsequent callers).
|
||||
|
||||
Returns:
|
||||
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.
|
||||
Dict with node_counts, link_counts, link_counts_by_fact_type,
|
||||
link_breakdown, and operations stats.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@@ -6,10 +6,9 @@ enabling support for multiple LLM backends (OpenAI, Anthropic, Gemini, Codex, et
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from .response_models import LLMToolCallResult
|
||||
from .response_models import LLMToolCallResult, TokenUsage
|
||||
|
||||
|
||||
class LLMInterface(ABC):
|
||||
@@ -70,7 +69,6 @@ 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.
|
||||
@@ -85,13 +83,8 @@ 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: 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.
|
||||
strict_schema: Use strict JSON schema enforcement (OpenAI only).
|
||||
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.
|
||||
@@ -115,7 +108,6 @@ 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.
|
||||
@@ -145,46 +137,6 @@ 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]],
|
||||
@@ -253,11 +205,3 @@ class OutputTooLongError(Exception):
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ProviderRateLimitResetError(Exception):
|
||||
"""Raised when an upstream provider says quota will reopen at a known time."""
|
||||
|
||||
def __init__(self, retry_at: datetime, message: str = "") -> None:
|
||||
self.retry_at = retry_at
|
||||
super().__init__(message)
|
||||
|
||||
@@ -1,585 +0,0 @@
|
||||
"""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)
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMResponseUsage:
|
||||
"""Provider-reported token usage for the in-flight LLM call.
|
||||
|
||||
Stashed by provider implementations as soon as a response is received —
|
||||
*before* local JSON parsing / schema validation, which may still fail. The
|
||||
wrapper reads it to attach real token counts to an error trace when the
|
||||
provider call itself succeeded but the structured output couldn't be parsed
|
||||
or validated (providers charge for those tokens regardless). See #2387.
|
||||
"""
|
||||
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cached_tokens: int = 0
|
||||
|
||||
|
||||
# Per-call provider usage, set by providers right after a response is received.
|
||||
_response_usage_ctx: ContextVar[LLMResponseUsage | None] = ContextVar("hindsight_llm_response_usage_ctx", default=None)
|
||||
|
||||
|
||||
def set_response_usage(usage: LLMResponseUsage | None) -> Token:
|
||||
"""Bind provider-reported usage for the current call. Returns a reset token."""
|
||||
return _response_usage_ctx.set(usage)
|
||||
|
||||
|
||||
def stash_response_usage(usage: LLMResponseUsage | None) -> None:
|
||||
"""Record provider-reported usage so an error trace can attach it later.
|
||||
|
||||
Called by provider implementations once a response (with usage) is in hand,
|
||||
before parsing/validation that may raise. Overwrites any prior value from an
|
||||
earlier retry attempt so the last attempt's usage wins.
|
||||
"""
|
||||
_response_usage_ctx.set(usage)
|
||||
|
||||
|
||||
def reset_response_usage(token: Token) -> None:
|
||||
"""Unwind a binding made by :func:`set_response_usage`."""
|
||||
_response_usage_ctx.reset(token)
|
||||
|
||||
|
||||
def current_response_usage() -> LLMResponseUsage | None:
|
||||
"""Return the active call's provider-reported usage, or None."""
|
||||
return _response_usage_ctx.get()
|
||||
|
||||
|
||||
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}")
|
||||
@@ -10,10 +10,15 @@ import re
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import AsyncExitStack
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from openai import APIConnectionError, APIStatusError, AsyncOpenAI, LengthFinishReasonError
|
||||
|
||||
# Vertex AI imports (conditional - for LLMProvider to pass credentials to GeminiLLM)
|
||||
try:
|
||||
import google.auth
|
||||
from google.oauth2 import service_account
|
||||
|
||||
VERTEXAI_AVAILABLE = True
|
||||
@@ -22,14 +27,16 @@ except ImportError:
|
||||
|
||||
from ..config import (
|
||||
DEFAULT_LLM_MAX_CONCURRENT,
|
||||
DEFAULT_LLM_TIMEOUT,
|
||||
ENV_CONSOLIDATION_LLM_MAX_CONCURRENT,
|
||||
ENV_LLM_GROQ_SERVICE_TIER,
|
||||
ENV_LLM_MAX_CONCURRENT,
|
||||
ENV_LLM_TIMEOUT,
|
||||
ENV_REFLECT_LLM_MAX_CONCURRENT,
|
||||
ENV_RETAIN_LLM_MAX_CONCURRENT,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .response_models import LLMToolCallResult
|
||||
from ..metrics import get_metrics_collector
|
||||
from .response_models import TokenUsage
|
||||
|
||||
# Seed applied to every Groq request for deterministic behavior.
|
||||
DEFAULT_LLM_SEED = 4242
|
||||
@@ -107,32 +114,6 @@ 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.
|
||||
@@ -225,7 +206,6 @@ _PROVIDERS_WITHOUT_API_KEY = frozenset(
|
||||
"litellm",
|
||||
"litellmrouter",
|
||||
"bedrock",
|
||||
"nous",
|
||||
}
|
||||
)
|
||||
|
||||
@@ -243,17 +223,13 @@ 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,
|
||||
gemini_service_tier: str | None = None,
|
||||
timeout: float | None = None,
|
||||
) -> Any: # Returns LLMInterface
|
||||
"""
|
||||
Factory function to create the appropriate LLM provider implementation.
|
||||
@@ -266,31 +242,18 @@ 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).
|
||||
bedrock_service_tier: Bedrock service tier (for Bedrock provider) - None (default), "flex", "priority", or "reserved".
|
||||
gemini_service_tier: Gemini service tier (for Gemini provider) - None (default) or "flex" (50% cheaper).
|
||||
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 to provider SDK clients (used by operators
|
||||
routing through proxies / request-tracing middleware). Wired into the Anthropic
|
||||
provider (SDK ``default_headers``) and the LiteLLM-backed providers — ``litellm``,
|
||||
``litellmrouter`` and ``bedrock`` — as the LiteLLM ``extra_headers`` completion
|
||||
kwarg; other providers may opt in as needed.
|
||||
extra_body: Extra body params merged into OpenAI-compatible API calls.
|
||||
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.
|
||||
vertexai_project_id: Vertex AI project ID (for VertexAI provider).
|
||||
vertexai_region: Vertex AI region (for VertexAI provider).
|
||||
vertexai_credentials: Vertex AI credentials object (for VertexAI provider).
|
||||
timeout: Per-request LLM timeout in seconds (resolved by the caller from the
|
||||
per-operation/global config). Threaded into the providers that honour a
|
||||
configurable request timeout (LiteLLM, LiteLLM Router, OpenAI-compatible,
|
||||
Nous). ``None`` lets each provider fall back to its own default
|
||||
(``HINDSIGHT_API_LLM_TIMEOUT`` / ``DEFAULT_LLM_TIMEOUT`` for those four;
|
||||
Anthropic and Gemini keep their provider-specific defaults).
|
||||
|
||||
Returns:
|
||||
LLMInterface implementation for the specified provider.
|
||||
"""
|
||||
from .llm_interface import LLMInterface
|
||||
from .providers import (
|
||||
AnthropicLLM,
|
||||
ClaudeCodeLLM,
|
||||
@@ -306,12 +269,6 @@ def create_llm_provider(
|
||||
)
|
||||
|
||||
provider_lower = provider.lower()
|
||||
if provider_lower == "gemini":
|
||||
from ..config import parse_gemini_service_tier
|
||||
|
||||
gemini_service_tier = parse_gemini_service_tier(gemini_service_tier)
|
||||
else:
|
||||
gemini_service_tier = None
|
||||
|
||||
if provider_lower == "openai-codex":
|
||||
return CodexLLM(
|
||||
@@ -360,9 +317,6 @@ def create_llm_provider(
|
||||
vertexai_region=vertexai_region,
|
||||
vertexai_credentials=vertexai_credentials,
|
||||
gemini_safety_settings=gemini_safety_settings,
|
||||
gemini_service_tier=gemini_service_tier,
|
||||
prompt_cache_enabled=prompt_cache_enabled,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower == "anthropic":
|
||||
@@ -373,7 +327,6 @@ def create_llm_provider(
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
default_headers=default_headers,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower == "litellm":
|
||||
@@ -383,9 +336,6 @@ def create_llm_provider(
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
elif provider_lower == "litellmrouter":
|
||||
@@ -403,9 +353,6 @@ def create_llm_provider(
|
||||
model=model,
|
||||
config=litellmrouter_config,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
elif provider_lower == "bedrock":
|
||||
@@ -417,10 +364,6 @@ def create_llm_provider(
|
||||
base_url=base_url,
|
||||
model=bedrock_model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
bedrock_service_tier=bedrock_service_tier,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
elif provider_lower == "llamacpp":
|
||||
@@ -454,22 +397,6 @@ 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,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
elif provider_lower in (
|
||||
"openai",
|
||||
"groq",
|
||||
@@ -480,10 +407,8 @@ def create_llm_provider(
|
||||
"deepseek",
|
||||
"volcano",
|
||||
"openrouter",
|
||||
"requesty",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
):
|
||||
return OpenAICompatibleLLM(
|
||||
provider=provider,
|
||||
@@ -494,7 +419,6 @@ def create_llm_provider(
|
||||
groq_service_tier=groq_service_tier,
|
||||
openai_service_tier=openai_service_tier,
|
||||
extra_body=extra_body,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
else:
|
||||
@@ -517,20 +441,10 @@ 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,
|
||||
gemini_service_tier: str | None = None,
|
||||
vertexai_project_id: str | None = None,
|
||||
vertexai_region: str | None = None,
|
||||
vertexai_service_account_key: str | None = None,
|
||||
timeout: float | None = None,
|
||||
max_retries: int | None = None,
|
||||
initial_backoff: float | None = None,
|
||||
max_backoff: float | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize LLM provider.
|
||||
@@ -543,74 +457,42 @@ 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_service_tier: Gemini service tier (None or "flex") - from config.
|
||||
gemini_safety_settings: Safety settings for Gemini/VertexAI providers.
|
||||
extra_body: Extra request-body params merged into the provider's native call
|
||||
(OpenAI-compatible, Fireworks, Anthropic, Gemini/VertexAI, LiteLLM).
|
||||
extra_body: Extra body params merged into OpenAI-compatible API calls.
|
||||
default_headers: Custom headers passed as ``default_headers`` to provider SDK clients.
|
||||
Used by operators routing through proxies / request-tracing middleware.
|
||||
Used by operators routing through proxies / request-tracing middleware. Falls
|
||||
back to ``HindsightConfig.llm_default_headers`` (env: ``HINDSIGHT_API_LLM_DEFAULT_HEADERS``)
|
||||
when ``None``.
|
||||
litellmrouter_config: Provider-specific config for ``provider="litellmrouter"``.
|
||||
JSON object passed verbatim to ``litellm.Router(**config)`` — see
|
||||
https://docs.litellm.ai/docs/routing. Ignored unless ``provider == "litellmrouter"``.
|
||||
vertexai_project_id: Vertex AI project ID for ``provider="vertexai"`` (required for
|
||||
that provider).
|
||||
vertexai_region: Vertex AI region for ``provider="vertexai"`` (defaults to
|
||||
``"us-central1"`` when ``None``).
|
||||
vertexai_service_account_key: Path to a Vertex AI service-account key file for
|
||||
``provider="vertexai"`` (uses ADC when ``None``).
|
||||
timeout: Per-request LLM timeout in seconds. Resolved by the caller from the
|
||||
per-operation/global config (``retain_llm_timeout`` falling back to
|
||||
``llm_timeout``, etc.). ``None`` lets each provider apply its own default.
|
||||
max_retries: Default retry-attempt budget for ``call`` / ``call_with_tools``
|
||||
when the per-call argument is omitted. Resolved by the caller from the
|
||||
per-operation/global config (``reflect_llm_max_retries`` falling back to
|
||||
``llm_max_retries``, etc.). ``None`` keeps each method's own fallback.
|
||||
initial_backoff: Default initial retry backoff (seconds), same resolution as
|
||||
``max_retries``. ``None`` keeps each method's own fallback.
|
||||
max_backoff: Default maximum retry backoff (seconds), same resolution as
|
||||
``max_retries``. ``None`` keeps each method's own fallback.
|
||||
|
||||
This constructor uses every argument as passed and does not read global
|
||||
``HindsightConfig``: resolving the server-level default for a ``None`` argument is the
|
||||
caller's responsibility (see ``MemoryEngine``'s per-op builds, ``_member_to_llm``, and
|
||||
``LLMProvider.from_env``). Keeping it config-free makes a provider's effective settings a
|
||||
pure function of its arguments — which is what lets each member of a multi-LLM chain be
|
||||
configured independently.
|
||||
When None and the provider is ``litellmrouter``, falls back to
|
||||
``HindsightConfig.llm_litellmrouter_config``.
|
||||
"""
|
||||
self.provider = provider.lower()
|
||||
self.api_key = api_key
|
||||
self.base_url = base_url
|
||||
self.model = model
|
||||
self.reasoning_effort = reasoning_effort
|
||||
# Per-request timeout (seconds). Used verbatim — the caller resolves the
|
||||
# per-operation/global fallback. ``None`` defers to the provider default.
|
||||
self.timeout = timeout
|
||||
# Default retry policy for call()/call_with_tools(). The caller resolves the
|
||||
# per-operation/global fallback; ``None`` keeps each method's own fallback so
|
||||
# providers built without a resolved config (from_env, tests) are unchanged.
|
||||
self.max_retries = max_retries
|
||||
self.initial_backoff = initial_backoff
|
||||
self.max_backoff = max_backoff
|
||||
self.litellmrouter_config = litellmrouter_config
|
||||
# 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
|
||||
self.gemini_service_tier = gemini_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).
|
||||
# Used verbatim — callers resolve the global fallback (see _member_to_llm /
|
||||
# the per-op builds in MemoryEngine, and LLMProvider.from_env).
|
||||
# Same pattern as ``gemini_safety_settings``: explicit override wins; otherwise read
|
||||
# the static server-level default from ``HindsightConfig`` via ``_get_raw_config()``.
|
||||
self.default_headers = default_headers
|
||||
if self.default_headers is None:
|
||||
from ..config import _get_raw_config
|
||||
|
||||
try:
|
||||
self.default_headers = _get_raw_config().llm_default_headers
|
||||
except Exception:
|
||||
pass # Config may not be initialized in test environments
|
||||
|
||||
# Validate provider
|
||||
valid_providers = [
|
||||
@@ -634,12 +516,9 @@ class LLMProvider:
|
||||
"bedrock",
|
||||
"volcano",
|
||||
"openrouter",
|
||||
"requesty",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
"fireworks",
|
||||
"nous",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
raise ValueError(f"Invalid LLM provider: {self.provider}. Must be one of: {', '.join(valid_providers)}")
|
||||
@@ -660,31 +539,30 @@ class LLMProvider:
|
||||
self.base_url = "https://api.deepseek.com"
|
||||
elif self.provider == "openrouter":
|
||||
self.base_url = "https://openrouter.ai/api/v1"
|
||||
elif self.provider == "requesty":
|
||||
self.base_url = "https://router.requesty.ai/v1"
|
||||
elif self.provider == "zai":
|
||||
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 == "atlas":
|
||||
self.base_url = "https://api.atlascloud.ai/v1"
|
||||
elif self.provider == "nous":
|
||||
self.base_url = "https://inference-api.nousresearch.com/v1"
|
||||
|
||||
# Prepare Vertex AI config (if applicable). Values are used as passed; the
|
||||
# caller resolves the global-config fallback (MemoryEngine builds /
|
||||
# _member_to_llm / from_env). The region keeps a constant default here.
|
||||
# Prepare Vertex AI config (if applicable)
|
||||
vertexai_project_id = None
|
||||
vertexai_region = None
|
||||
vertexai_credentials = None
|
||||
|
||||
if self.provider == "vertexai":
|
||||
from ..config import get_config
|
||||
|
||||
config = get_config()
|
||||
|
||||
vertexai_project_id = config.llm_vertexai_project_id
|
||||
if not vertexai_project_id:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID is required for Vertex AI provider. "
|
||||
"Set it to your GCP project ID."
|
||||
)
|
||||
|
||||
vertexai_region = vertexai_region or "us-central1"
|
||||
service_account_key = vertexai_service_account_key
|
||||
vertexai_region = config.llm_vertexai_region or "us-central1"
|
||||
service_account_key = config.llm_vertexai_service_account_key
|
||||
|
||||
# Load explicit service account credentials if provided
|
||||
if service_account_key:
|
||||
@@ -708,20 +586,30 @@ class LLMProvider:
|
||||
f"model={self.model}, auth={'service_account' if service_account_key else 'ADC'}"
|
||||
)
|
||||
|
||||
# Normalize the Gemini service tier (pure: maps/validates the passed value,
|
||||
# no global config read). Non-Gemini providers never carry a tier. The
|
||||
# server-level default is resolved by the caller, like the other fields.
|
||||
if self.provider == "gemini":
|
||||
from ..config import parse_gemini_service_tier
|
||||
# For Gemini/VertexAI providers: read safety settings from global config if not explicitly provided
|
||||
# Use _get_raw_config() to bypass StaticConfigProxy (which blocks configurable fields),
|
||||
# since LLMProvider initialization legitimately needs the server-level default.
|
||||
if self.provider in ("gemini", "vertexai") and self.gemini_safety_settings is None:
|
||||
from ..config import _get_raw_config
|
||||
|
||||
self.gemini_service_tier = parse_gemini_service_tier(self.gemini_service_tier)
|
||||
else:
|
||||
self.gemini_service_tier = None
|
||||
try:
|
||||
raw_config = _get_raw_config()
|
||||
self.gemini_safety_settings = raw_config.llm_gemini_safety_settings
|
||||
except Exception:
|
||||
pass # Config may not be initialized in test environments
|
||||
|
||||
# gemini_safety_settings / prompt_cache_enabled / litellmrouter_config are
|
||||
# used as passed — the caller resolves the global-config fallback. Providers
|
||||
# that don't support prompt caching ignore the flag.
|
||||
# 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
|
||||
# ad-hoc constructions (e.g. ``LLMProvider.from_env()``) keep working.
|
||||
router_config: dict[str, Any] | None = self.litellmrouter_config
|
||||
if self.provider == "litellmrouter" and router_config is None:
|
||||
from ..config import _get_raw_config
|
||||
|
||||
try:
|
||||
router_config = _get_raw_config().llm_litellmrouter_config
|
||||
except Exception:
|
||||
router_config = None
|
||||
|
||||
# Create provider implementation using factory
|
||||
self._provider_impl = create_llm_provider(
|
||||
@@ -732,17 +620,13 @@ 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,
|
||||
gemini_service_tier=self.gemini_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,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
# Backward compatibility: Keep mock provider properties
|
||||
@@ -799,13 +683,12 @@ class LLMProvider:
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "memory",
|
||||
max_retries: int | None = None,
|
||||
initial_backoff: float | None = None,
|
||||
max_backoff: float | None = None,
|
||||
max_retries: int = 10,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 60.0,
|
||||
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.
|
||||
@@ -816,17 +699,11 @@ class LLMProvider:
|
||||
max_completion_tokens: Maximum tokens in response.
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts. ``None`` uses the provider's configured
|
||||
default (per-operation/global ``llm_max_retries``), else 10.
|
||||
initial_backoff: Initial backoff time in seconds. ``None`` uses the provider's
|
||||
configured default (``llm_initial_backoff``), else 1.0.
|
||||
max_backoff: Maximum backoff time in seconds. ``None`` uses the provider's
|
||||
configured default (``llm_max_backoff``), else 60.0.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
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.
|
||||
strict_schema: Use strict JSON schema enforcement (OpenAI only). Guarantees all required fields.
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
|
||||
Returns:
|
||||
@@ -846,112 +723,33 @@ class LLMProvider:
|
||||
structured = "+structured" if response_format is not None else ""
|
||||
set_stage(f"llm.{self.provider}.{scope}{structured}")
|
||||
|
||||
# Resolve the retry policy: explicit per-call arg wins, else the provider's
|
||||
# configured per-operation/global default, else this method's own fallback.
|
||||
max_retries = (
|
||||
max_retries if max_retries is not None else (self.max_retries if self.max_retries is not None else 10)
|
||||
)
|
||||
initial_backoff = (
|
||||
initial_backoff
|
||||
if initial_backoff is not None
|
||||
else (self.initial_backoff if self.initial_backoff is not None else 1.0)
|
||||
)
|
||||
max_backoff = (
|
||||
max_backoff if max_backoff is not None else (self.max_backoff if self.max_backoff is not None else 60.0)
|
||||
)
|
||||
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
|
||||
|
||||
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 (
|
||||
current_response_usage,
|
||||
reset_request_context,
|
||||
reset_response_usage,
|
||||
set_request_context,
|
||||
set_response_usage,
|
||||
)
|
||||
|
||||
call_start = time.monotonic()
|
||||
request_token = set_request_context(
|
||||
_request_params(
|
||||
# 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,
|
||||
response_format=response_format,
|
||||
max_retries=max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
skip_validation=skip_validation,
|
||||
strict_schema=strict_schema,
|
||||
return_usage=return_usage,
|
||||
)
|
||||
)
|
||||
# Cleared per call; the provider stashes real usage once a response is in
|
||||
# hand so the error path below can attach it if parsing/validation fails.
|
||||
usage_token = set_response_usage(None)
|
||||
try:
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
|
||||
# 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:
|
||||
# The provider call may have succeeded (and incurred token
|
||||
# cost) before local parsing/validation raised; attach the
|
||||
# provider-reported usage to the error trace when available.
|
||||
usage = current_response_usage()
|
||||
get_span_recorder().record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=None,
|
||||
input_tokens=usage.input_tokens if usage else 0,
|
||||
output_tokens=usage.output_tokens if usage else 0,
|
||||
cached_tokens=usage.cached_tokens if usage else 0,
|
||||
duration=time.monotonic() - call_start,
|
||||
error=e,
|
||||
)
|
||||
raise
|
||||
# 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
|
||||
|
||||
# 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)
|
||||
reset_response_usage(usage_token)
|
||||
if isinstance(self._provider_impl, MockLLM):
|
||||
# Sync the mock calls from provider implementation to wrapper
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
|
||||
return result
|
||||
|
||||
@@ -962,11 +760,10 @@ class LLMProvider:
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "tools",
|
||||
max_retries: int | None = None,
|
||||
initial_backoff: float | None = None,
|
||||
max_backoff: float | None = None,
|
||||
max_retries: int = 5,
|
||||
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.
|
||||
@@ -977,12 +774,9 @@ class LLMProvider:
|
||||
max_completion_tokens: Maximum tokens in response.
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts. ``None`` uses the provider's configured
|
||||
default (per-operation/global ``llm_max_retries``), else 5.
|
||||
initial_backoff: Initial backoff time in seconds. ``None`` uses the provider's
|
||||
configured default (``llm_initial_backoff``), else 1.0.
|
||||
max_backoff: Maximum backoff time in seconds. ``None`` uses the provider's
|
||||
configured default (``llm_max_backoff``), else 30.0.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or {"type": "function", "function": {"name": "..."}}
|
||||
|
||||
Returns:
|
||||
@@ -992,95 +786,31 @@ class LLMProvider:
|
||||
|
||||
set_stage(f"llm.{self.provider}.{scope}+tools")
|
||||
|
||||
# Resolve the retry policy: explicit per-call arg wins, else the provider's
|
||||
# configured per-operation/global default, else this method's own fallback.
|
||||
max_retries = (
|
||||
max_retries if max_retries is not None else (self.max_retries if self.max_retries is not None else 5)
|
||||
)
|
||||
initial_backoff = (
|
||||
initial_backoff
|
||||
if initial_backoff is not None
|
||||
else (self.initial_backoff if self.initial_backoff is not None else 1.0)
|
||||
)
|
||||
max_backoff = (
|
||||
max_backoff if max_backoff is not None else (self.max_backoff if self.max_backoff is not None else 30.0)
|
||||
)
|
||||
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 (
|
||||
current_response_usage,
|
||||
reset_request_context,
|
||||
reset_response_usage,
|
||||
set_request_context,
|
||||
set_response_usage,
|
||||
)
|
||||
|
||||
call_start = time.monotonic()
|
||||
request_token = set_request_context(
|
||||
_request_params(
|
||||
# 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,
|
||||
)
|
||||
)
|
||||
# Cleared per call; the provider stashes real usage once a response is in
|
||||
# hand so the error path below can attach it if parsing/validation fails.
|
||||
usage_token = set_response_usage(None)
|
||||
try:
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
|
||||
# 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:
|
||||
# The provider call may have succeeded (and incurred token
|
||||
# cost) before local parsing/validation raised; attach the
|
||||
# provider-reported usage to the error trace when available.
|
||||
usage = current_response_usage()
|
||||
get_span_recorder().record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=None,
|
||||
input_tokens=usage.input_tokens if usage else 0,
|
||||
output_tokens=usage.output_tokens if usage else 0,
|
||||
cached_tokens=usage.cached_tokens if usage else 0,
|
||||
duration=time.monotonic() - call_start,
|
||||
error=e,
|
||||
)
|
||||
raise
|
||||
# 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
|
||||
|
||||
# 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)
|
||||
reset_response_usage(usage_token)
|
||||
if isinstance(self._provider_impl, MockLLM):
|
||||
# Sync the mock calls from provider implementation to wrapper
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
|
||||
return result
|
||||
|
||||
@@ -1124,9 +854,7 @@ class LLMProvider:
|
||||
|
||||
def _load_codex_auth(self) -> tuple[str, str]:
|
||||
"""
|
||||
Load OAuth credentials from the Codex ``auth.json``.
|
||||
|
||||
Honors ``CODEX_HOME`` (falling back to ``~/.codex``).
|
||||
Load OAuth credentials from ~/.codex/auth.json.
|
||||
|
||||
Returns:
|
||||
Tuple of (access_token, account_id).
|
||||
@@ -1135,9 +863,7 @@ class LLMProvider:
|
||||
FileNotFoundError: If auth file doesn't exist.
|
||||
ValueError: If auth file is invalid.
|
||||
"""
|
||||
from .providers.codex_auth import default_codex_auth_file
|
||||
|
||||
auth_file = default_codex_auth_file()
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(
|
||||
@@ -1188,14 +914,7 @@ 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,
|
||||
*,
|
||||
bank_id: str | None = None,
|
||||
operation: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> "ConfiguredLLMProvider":
|
||||
def with_config(self, config: Any) -> "ConfiguredLLMProvider":
|
||||
"""
|
||||
Return a configured wrapper for a specific bank operation.
|
||||
|
||||
@@ -1205,31 +924,12 @@ 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.
|
||||
"""
|
||||
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)
|
||||
return ConfiguredLLMProvider(self, config.llm_gemini_safety_settings)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (e.g. stop llamacpp subprocess)."""
|
||||
@@ -1239,38 +939,17 @@ class LLMProvider:
|
||||
@classmethod
|
||||
def from_env(cls) -> "LLMProvider":
|
||||
"""Create provider from environment variables using config.py constants."""
|
||||
# Read every field straight from the environment. The constructor no longer
|
||||
# resolves global-config fallbacks, so this factory must supply them — and it
|
||||
# does so without building the full HindsightConfig, keeping from_env() a
|
||||
# lightweight env-only loader (see test_llm_provider_from_env_keeps_lightweight_loader).
|
||||
from ..config import (
|
||||
DEFAULT_LLM_GROQ_SERVICE_TIER,
|
||||
DEFAULT_LLM_OPENAI_SERVICE_TIER,
|
||||
DEFAULT_LLM_PROMPT_CACHE_ENABLED,
|
||||
DEFAULT_LLM_PROVIDER,
|
||||
DEFAULT_LLM_REASONING_EFFORT,
|
||||
DEFAULT_LLM_TIMEOUT,
|
||||
ENV_LLM_API_KEY,
|
||||
ENV_LLM_BASE_URL,
|
||||
ENV_LLM_BEDROCK_SERVICE_TIER,
|
||||
ENV_LLM_DEFAULT_HEADERS,
|
||||
ENV_LLM_EXTRA_BODY,
|
||||
ENV_LLM_GEMINI_SAFETY_SETTINGS,
|
||||
ENV_LLM_GEMINI_SERVICE_TIER,
|
||||
ENV_LLM_GROQ_SERVICE_TIER,
|
||||
ENV_LLM_LITELLMROUTER_CONFIG,
|
||||
ENV_LLM_MODEL,
|
||||
ENV_LLM_OPENAI_SERVICE_TIER,
|
||||
ENV_LLM_PROMPT_CACHE_ENABLED,
|
||||
ENV_LLM_PROVIDER,
|
||||
ENV_LLM_REASONING_EFFORT,
|
||||
ENV_LLM_TIMEOUT,
|
||||
ENV_LLM_VERTEXAI_PROJECT_ID,
|
||||
ENV_LLM_VERTEXAI_REGION,
|
||||
ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY,
|
||||
_get_default_model_for_provider,
|
||||
_parse_llm_router_config,
|
||||
parse_gemini_service_tier,
|
||||
)
|
||||
|
||||
provider = os.getenv(ENV_LLM_PROVIDER, DEFAULT_LLM_PROVIDER)
|
||||
@@ -1287,14 +966,6 @@ class LLMProvider:
|
||||
model = os.getenv(ENV_LLM_MODEL) or _get_default_model_for_provider(provider)
|
||||
extra_body = json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null"))
|
||||
default_headers = json.loads(os.getenv(ENV_LLM_DEFAULT_HEADERS, "null"))
|
||||
prompt_cache_enabled = os.getenv(
|
||||
ENV_LLM_PROMPT_CACHE_ENABLED, str(DEFAULT_LLM_PROMPT_CACHE_ENABLED)
|
||||
).lower() in (
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
"on",
|
||||
)
|
||||
|
||||
return cls(
|
||||
provider=provider,
|
||||
@@ -1304,21 +975,6 @@ class LLMProvider:
|
||||
reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
|
||||
openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
|
||||
bedrock_service_tier=os.getenv(ENV_LLM_BEDROCK_SERVICE_TIER) or None,
|
||||
gemini_service_tier=(
|
||||
parse_gemini_service_tier(os.getenv(ENV_LLM_GEMINI_SERVICE_TIER))
|
||||
if provider.lower() == "gemini"
|
||||
else None
|
||||
),
|
||||
gemini_safety_settings=json.loads(os.getenv(ENV_LLM_GEMINI_SAFETY_SETTINGS, "null")),
|
||||
prompt_cache_enabled=prompt_cache_enabled,
|
||||
litellmrouter_config=_parse_llm_router_config(ENV_LLM_LITELLMROUTER_CONFIG),
|
||||
vertexai_project_id=os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID) or None,
|
||||
vertexai_region=os.getenv(ENV_LLM_VERTEXAI_REGION) or None,
|
||||
vertexai_service_account_key=os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY) or None,
|
||||
timeout=float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT))),
|
||||
)
|
||||
|
||||
|
||||
@@ -1337,16 +993,10 @@ class ConfiguredLLMProvider:
|
||||
any changes.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: "LLMProvider",
|
||||
gemini_safety_settings: list | None,
|
||||
trace_ctx: Any | None = None,
|
||||
) -> None:
|
||||
def __init__(self, provider: "LLMProvider", gemini_safety_settings: list | 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 ──────────────────────────────────────────────────
|
||||
|
||||
@@ -1359,12 +1009,10 @@ 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,
|
||||
@@ -1375,38 +1023,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_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
|
||||
|
||||
@@ -1,347 +0,0 @@
|
||||
"""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.
|
||||
- **Scheduled mental model refresh** (configurable check cadence, default 60s):
|
||||
refresh mental models whose ``trigger.refresh_cron`` schedule is due, but only
|
||||
when the model is stale (new memories in its scope since its last refresh), so
|
||||
a scheduled tick never burns an LLM call to regenerate identical content. The
|
||||
per-model schedule lives in the cron expression; this loop only decides when to
|
||||
*check*.
|
||||
|
||||
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 collections.abc import Coroutine
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..config import HindsightConfig, get_config
|
||||
from ..models import RequestContext
|
||||
from .db_utils import acquire_with_retry
|
||||
from .schema import _is_oracle, fq_table
|
||||
|
||||
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
|
||||
mm_refresh_on = cfg.mental_model_refresh_tick_seconds > 0
|
||||
return reconcile_on or audit_on or llm_on or mm_refresh_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_timed("retention", self._run_retention(cfg))
|
||||
interval = cfg.consolidation_reconcile_interval_seconds
|
||||
if interval > 0 and self._is_due("reconcile", interval):
|
||||
await self._run_timed("consolidation reconcile", self._run_reconcile())
|
||||
mm_interval = cfg.mental_model_refresh_tick_seconds
|
||||
if mm_interval > 0 and self._is_due("mm_refresh", mm_interval):
|
||||
await self._run_timed("scheduled mental model refresh", self._run_scheduled_mm_refresh())
|
||||
|
||||
async def _run_timed(self, name: str, coro: Coroutine[Any, Any, None]) -> None:
|
||||
"""Run a maintenance job and emit one timing line for it.
|
||||
|
||||
Each job keeps its own summary log (counts of work done); this adds a
|
||||
single, uniform line per run so the cost of every sweep is observable.
|
||||
"""
|
||||
start = time.monotonic()
|
||||
try:
|
||||
await coro
|
||||
finally:
|
||||
logger.info(f"Maintenance: {name} took {time.monotonic() - start:.3f}s")
|
||||
|
||||
# ── 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 "")
|
||||
)
|
||||
|
||||
# ── scheduled mental model refresh ───────────────────────────────────────
|
||||
|
||||
async def _run_scheduled_mm_refresh(self) -> None:
|
||||
"""Refresh mental models whose ``trigger.refresh_cron`` is due.
|
||||
|
||||
Discovery (the set of cron-scheduled models, minus any with an in-flight
|
||||
refresh) is one cross-tenant round-trip via
|
||||
``public.mental_models_with_cron()``. Cron *due-ness* is evaluated here in
|
||||
Python — a scheduled fire has elapsed when the most recent cron boundary at
|
||||
or before now is later than ``last_refreshed_at`` — because cron arithmetic
|
||||
isn't expressible in plain SQL. Each due model is refreshed only when it is
|
||||
actually stale, so a schedule that fires while nothing changed costs a
|
||||
cheap staleness query, not an LLM call.
|
||||
"""
|
||||
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, mental_model_id, refresh_cron, last_refreshed_at "
|
||||
"FROM public.mental_models_with_cron()"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Scheduled mental model refresh discovery failed: {e}")
|
||||
return
|
||||
if not rows:
|
||||
return
|
||||
|
||||
from croniter import croniter
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
due = []
|
||||
for row in rows:
|
||||
cron = row["refresh_cron"]
|
||||
last = row["last_refreshed_at"]
|
||||
try:
|
||||
prev_fire = croniter(cron, now).get_prev(datetime)
|
||||
except (ValueError, KeyError) as e:
|
||||
logger.warning(
|
||||
f"Scheduled mental model refresh: skipping invalid cron {cron!r} for "
|
||||
f"{row['schema_name']}/{row['mental_model_id']}: {e}"
|
||||
)
|
||||
continue
|
||||
if last is None or prev_fire > last:
|
||||
due.append(row)
|
||||
if not due:
|
||||
return
|
||||
|
||||
# Only enqueue into schemas the worker actually polls (tenant discovery),
|
||||
# otherwise the op would never be claimed. The tenant_id (when provided)
|
||||
# lets config resolution honor tenant-level overrides.
|
||||
try:
|
||||
tenants = await engine._tenant_extension.list_tenants()
|
||||
except Exception as e:
|
||||
logger.warning(f"Scheduled mental model refresh 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
|
||||
skipped_fresh = 0
|
||||
for row in due:
|
||||
schema = row["schema_name"]
|
||||
bank_id = row["bank_id"]
|
||||
mm_id = row["mental_model_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)
|
||||
# Skip if nothing in the model's scope changed since its last
|
||||
# refresh — a scheduled refresh must not regenerate identical
|
||||
# content. compute_mental_model_is_stale needs the model's tags +
|
||||
# trigger, which the discovery routine doesn't return, so re-read
|
||||
# the row under the bank's schema context.
|
||||
async with acquire_with_retry(engine._backend, max_retries=1) as conn:
|
||||
mm_row = await conn.fetchrow(
|
||||
f"SELECT id, tags, trigger, last_refreshed_at FROM {fq_table('mental_models')} "
|
||||
"WHERE bank_id = $1 AND id = $2",
|
||||
bank_id,
|
||||
mm_id,
|
||||
)
|
||||
if mm_row is None:
|
||||
continue
|
||||
is_stale = await engine.compute_mental_model_is_stale(conn, bank_id, mm_row)
|
||||
if not is_stale:
|
||||
skipped_fresh += 1
|
||||
continue
|
||||
await engine.submit_async_refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm_id, request_context=context
|
||||
)
|
||||
submitted += 1
|
||||
except Exception as e:
|
||||
logger.warning(f"Scheduled mental model refresh failed for {mm_id} in {schema}: {e}")
|
||||
finally:
|
||||
_current_schema.reset(token)
|
||||
|
||||
if submitted or skipped_unknown or skipped_fresh:
|
||||
logger.info(
|
||||
f"Scheduled mental model refresh: scheduled {submitted} model(s)"
|
||||
+ (f", {skipped_fresh} up-to-date" if skipped_fresh else "")
|
||||
+ (f", skipped {skipped_unknown} in unrecognized schema(s)" if skipped_unknown else "")
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,204 +0,0 @@
|
||||
"""Multi-LLM routing: failover and (weighted) round-robin across N providers.
|
||||
|
||||
``MultiLLMProvider`` wraps an ordered list of :class:`LLMProvider` members and a
|
||||
:class:`~hindsight_api.config.LLMStrategyConfig`, exposing the same public surface
|
||||
as a single ``LLMProvider`` so it drops into every existing call path (including
|
||||
``with_config()`` / ``ConfiguredLLMProvider``).
|
||||
|
||||
Member 0 is the **primary** (the operation's unindexed/base LLM); members 1..N are
|
||||
the indexed extras (``HINDSIGHT_API_<OP>LLM_<n>_*``). Each member keeps its own
|
||||
internal retry budget, so we only advance to the next member after a member has
|
||||
exhausted its retries and raised.
|
||||
|
||||
Strategies:
|
||||
- ``failover``: try members in declared order ``[0..N]``.
|
||||
- ``round-robin``: rotate the starting member per request (optionally weighted),
|
||||
then fall through the remaining members on error.
|
||||
|
||||
Batch retain and any direct ``_provider_impl`` access operate on the **primary
|
||||
member only** (via attribute passthrough) — failover/round-robin apply to the
|
||||
interactive ``call`` / ``call_with_tools`` paths.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..config import LLM_STRATEGY_FAILOVER, LLMStrategyConfig
|
||||
from .llm_wrapper import LLMProvider, OutputTooLongError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .llm_wrapper import ConfiguredLLMProvider, LLMToolCallResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _should_failover(exc: BaseException) -> bool:
|
||||
"""Whether ``exc`` from one member should trigger a try on the next member.
|
||||
|
||||
Generic ``Exception`` instances (network errors, provider 5xx, timeouts after
|
||||
a member's own retries) fail over. ``OutputTooLongError`` is propagated — a
|
||||
different provider won't fit an over-length output either. ``CancelledError``,
|
||||
``KeyboardInterrupt`` and ``SystemExit`` are ``BaseException`` (not
|
||||
``Exception``) and therefore propagate unchanged.
|
||||
"""
|
||||
if isinstance(exc, OutputTooLongError):
|
||||
return False
|
||||
return isinstance(exc, Exception)
|
||||
|
||||
|
||||
class _WeightedRoundRobin:
|
||||
"""Smooth weighted round-robin scheduler (nginx SWRR).
|
||||
|
||||
Produces a starting member index per request such that, over time, member
|
||||
``i`` is chosen in proportion to ``weights[i]`` while keeping selections
|
||||
interleaved rather than bursty. Uniform weights degrade to plain round-robin.
|
||||
The tiny selection critical section is mutex-guarded so concurrent callers
|
||||
don't corrupt the running totals (they may still interleave, which only
|
||||
affects distribution, never correctness).
|
||||
"""
|
||||
|
||||
def __init__(self, weights: list[int]) -> None:
|
||||
self._weights = list(weights)
|
||||
self._current = [0] * len(weights)
|
||||
self._total = sum(weights)
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def next(self) -> int:
|
||||
with self._lock:
|
||||
best = 0
|
||||
for i, w in enumerate(self._weights):
|
||||
self._current[i] += w
|
||||
if self._current[i] > self._current[best]:
|
||||
best = i
|
||||
self._current[best] -= self._total
|
||||
return best
|
||||
|
||||
|
||||
class MultiLLMProvider:
|
||||
"""Route LLM calls across multiple members per a failover / round-robin strategy."""
|
||||
|
||||
def __init__(self, members: list[LLMProvider], strategy: LLMStrategyConfig) -> None:
|
||||
if not members:
|
||||
raise ValueError("MultiLLMProvider requires at least one member")
|
||||
self._members = members
|
||||
self._strategy = strategy
|
||||
|
||||
weights = strategy.weights or [1] * len(members)
|
||||
if len(weights) != len(members):
|
||||
raise ValueError(
|
||||
f"LLM strategy 'weights' has {len(weights)} entries but the chain has "
|
||||
f"{len(members)} members (primary + indexed); they must match."
|
||||
)
|
||||
self._scheduler = _WeightedRoundRobin(weights)
|
||||
|
||||
# ── routing ────────────────────────────────────────────────────────────────
|
||||
|
||||
def _member_order(self) -> list[int]:
|
||||
"""Indices to try, in order, for one request."""
|
||||
n = len(self._members)
|
||||
if self._strategy.mode == LLM_STRATEGY_FAILOVER:
|
||||
return list(range(n))
|
||||
start = self._scheduler.next()
|
||||
return [(start + i) % n for i in range(n)]
|
||||
|
||||
async def _dispatch(self, method_name: str, **kwargs: Any) -> Any:
|
||||
last_exc: BaseException | None = None
|
||||
order = self._member_order()
|
||||
for position, idx in enumerate(order):
|
||||
member = self._members[idx]
|
||||
try:
|
||||
return await getattr(member, method_name)(**kwargs)
|
||||
except BaseException as e: # noqa: BLE001 - re-raised unless it should fail over
|
||||
if not _should_failover(e):
|
||||
raise
|
||||
last_exc = e
|
||||
remaining = len(order) - position - 1
|
||||
logger.warning(
|
||||
"LLM member %d (%s/%s) failed on %s: %s%s",
|
||||
idx,
|
||||
member.provider,
|
||||
member.model,
|
||||
method_name,
|
||||
e,
|
||||
f"; trying next member ({remaining} left)" if remaining else "; no members left",
|
||||
)
|
||||
# All members failed; surface the last error (loop ran at least once).
|
||||
assert last_exc is not None
|
||||
raise last_exc
|
||||
|
||||
async def call(self, messages: list[dict[str, Any]], **kwargs: Any) -> Any:
|
||||
return await self._dispatch("call", messages=messages, **kwargs)
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]],
|
||||
**kwargs: Any,
|
||||
) -> "LLMToolCallResult":
|
||||
return await self._dispatch("call_with_tools", messages=messages, tools=tools, **kwargs)
|
||||
|
||||
# ── lifecycle ────────────────────────────────────────────────────────────────
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""Strictly verify the primary; soft-verify the rest (warn, don't fail).
|
||||
|
||||
A failover member being unreachable at startup must not block the server —
|
||||
it may come back before it's needed. The primary is the steady-state path,
|
||||
so its failure is still surfaced (the caller already wraps this in a
|
||||
warn-only try/except at startup).
|
||||
"""
|
||||
await self._members[0].verify_connection()
|
||||
for member in self._members[1:]:
|
||||
try:
|
||||
await member.verify_connection()
|
||||
except Exception as e: # noqa: BLE001 - soft verification
|
||||
logger.warning(
|
||||
"Failover LLM member %s/%s failed connection verification: %s. "
|
||||
"It will be tried at request time if the primary fails.",
|
||||
member.provider,
|
||||
member.model,
|
||||
e,
|
||||
)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
for member in self._members:
|
||||
await member.cleanup()
|
||||
|
||||
def with_config(
|
||||
self,
|
||||
config: Any,
|
||||
*,
|
||||
bank_id: str | None = None,
|
||||
operation: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> "ConfiguredLLMProvider":
|
||||
"""Mirror ``LLMProvider.with_config`` so the strategy runs inside the
|
||||
per-operation configured wrapper (gemini-safety + trace contextvars wrap
|
||||
every member call)."""
|
||||
from .llm_trace import LLMTraceContext
|
||||
from .llm_wrapper import ConfiguredLLMProvider
|
||||
|
||||
trace_ctx = None
|
||||
if bank_id is not None or operation is not None or metadata:
|
||||
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)
|
||||
|
||||
# ── attribute passthrough ────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def members(self) -> list[LLMProvider]:
|
||||
return self._members
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
# Anything not defined here (provider, model, api_key, base_url,
|
||||
# _provider_impl, mock helpers, batch helpers, ...) delegates to the
|
||||
# primary member so existing call sites keep working unchanged.
|
||||
return getattr(object.__getattribute__(self, "_members")[0], name)
|
||||
@@ -5,10 +5,8 @@ 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, field
|
||||
from typing import Any, Mapping
|
||||
|
||||
MAX_EXTRACTION_ERROR_SAMPLES = 5
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -50,79 +48,6 @@ 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."""
|
||||
|
||||
@@ -3,138 +3,43 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from hindsight_api.config import DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
|
||||
|
||||
from .base import FileParser
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from markitdown import StreamInfo
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Extensions whose markitdown converters decode the raw bytes as text. markitdown
|
||||
# samples only the first chunk for charset detection, so a UTF-8 file with a long
|
||||
# ASCII-only prefix is mis-detected as ASCII; the JSON/ipynb converter then crashes
|
||||
# decoding the first multibyte byte. Passing an explicit UTF-8 hint when the bytes
|
||||
# are valid UTF-8 sidesteps the faulty detection without affecting other encodings.
|
||||
_TEXT_EXTENSIONS = {
|
||||
".json",
|
||||
".jsonl",
|
||||
".ipynb",
|
||||
".txt",
|
||||
".text",
|
||||
".md",
|
||||
".markdown",
|
||||
".csv",
|
||||
".html",
|
||||
".htm",
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MarkitdownOcrOptions:
|
||||
"""OpenAI-compatible OCR options passed through to MarkItDown."""
|
||||
|
||||
# Keep this typed as object so the OpenAI SDK import stays lazy for non-OCR users.
|
||||
llm_client: object
|
||||
llm_model: str
|
||||
llm_prompt: str
|
||||
|
||||
|
||||
class MarkitdownParser(FileParser):
|
||||
"""
|
||||
Markitdown file parser.
|
||||
|
||||
Uses Microsoft's markitdown library to convert various file formats
|
||||
to markdown including PDF, Office docs, images with optional OCR,
|
||||
audio, HTML.
|
||||
to markdown including PDF, Office docs, images (via OCR), audio, HTML.
|
||||
|
||||
Supported formats:
|
||||
- PDF (.pdf)
|
||||
- Word (.docx, .doc)
|
||||
- PowerPoint (.pptx, .ppt)
|
||||
- Excel (.xlsx, .xls)
|
||||
- Images (.jpg, .jpeg, .png) - optional OCR
|
||||
- Images (.jpg, .jpeg, .png) - with OCR
|
||||
- HTML (.html, .htm)
|
||||
- Text (.txt, .md)
|
||||
- Audio (.mp3, .wav) - with transcription
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
ocr_enabled: bool = False,
|
||||
ocr_api_key: str | None = None,
|
||||
ocr_base_url: str | None = None,
|
||||
ocr_model: str | None = None,
|
||||
ocr_prompt: str | None = None,
|
||||
):
|
||||
def __init__(self):
|
||||
"""Initialize markitdown parser."""
|
||||
# Lazy import to avoid requiring markitdown for all users
|
||||
try:
|
||||
from markitdown import MarkItDown
|
||||
|
||||
self._markitdown = MarkItDown()
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"markitdown package is required for file parsing. Install with: pip install markitdown"
|
||||
) from e
|
||||
|
||||
self._ocr_enabled = ocr_enabled
|
||||
if ocr_enabled:
|
||||
ocr_options = self._build_ocr_options(
|
||||
api_key=ocr_api_key,
|
||||
base_url=ocr_base_url,
|
||||
model=ocr_model,
|
||||
prompt=ocr_prompt,
|
||||
)
|
||||
self._markitdown = MarkItDown(
|
||||
llm_client=ocr_options.llm_client,
|
||||
llm_model=ocr_options.llm_model,
|
||||
llm_prompt=ocr_options.llm_prompt,
|
||||
)
|
||||
else:
|
||||
self._markitdown = MarkItDown()
|
||||
|
||||
def _build_ocr_options(
|
||||
self,
|
||||
*,
|
||||
api_key: str | None,
|
||||
base_url: str | None,
|
||||
model: str | None,
|
||||
prompt: str | None,
|
||||
) -> MarkitdownOcrOptions:
|
||||
"""Build MarkItDown options for OpenAI-compatible image OCR."""
|
||||
if not model or not model.strip():
|
||||
raise ValueError(
|
||||
"Markitdown OCR is enabled but no model is configured. "
|
||||
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL to an OpenAI-compatible OCR/vision model "
|
||||
"with image-input support."
|
||||
)
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"Markitdown OCR is enabled but no API key is configured. "
|
||||
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY."
|
||||
)
|
||||
if not base_url or not base_url.strip():
|
||||
raise ValueError(
|
||||
"Markitdown OCR is enabled but no base URL is configured. "
|
||||
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL to an OpenAI-compatible OCR/vision endpoint."
|
||||
)
|
||||
|
||||
try:
|
||||
from openai import OpenAI
|
||||
except ImportError as e:
|
||||
raise RuntimeError("openai package is required when Markitdown OCR is enabled.") from e
|
||||
|
||||
return MarkitdownOcrOptions(
|
||||
llm_client=OpenAI(api_key=api_key, base_url=base_url.strip()),
|
||||
llm_model=model.strip(),
|
||||
llm_prompt=prompt or DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
|
||||
)
|
||||
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
"""Parse file to markdown using markitdown."""
|
||||
# markitdown is synchronous, so we run it in executor to avoid blocking
|
||||
@@ -143,22 +48,14 @@ class MarkitdownParser(FileParser):
|
||||
|
||||
def _convert_sync(self, file_data: bytes, filename: str) -> str:
|
||||
"""Synchronous parsing (runs in thread pool)."""
|
||||
if self._is_image_file(filename) and not self._ocr_enabled:
|
||||
raise RuntimeError(
|
||||
"Image OCR is not enabled for the markitdown parser. "
|
||||
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED=true and configure an OpenAI-compatible "
|
||||
"OCR/vision endpoint with image-input support, or choose an OCR-capable parser."
|
||||
)
|
||||
|
||||
# Write to temp file (markitdown requires file path)
|
||||
with tempfile.NamedTemporaryFile(suffix=Path(filename).suffix, delete=False) as tmp:
|
||||
tmp.write(file_data)
|
||||
tmp_path = tmp.name
|
||||
|
||||
try:
|
||||
# Parse using markitdown, passing an explicit charset hint for text
|
||||
# files to avoid markitdown's sample-based (and crash-prone) detection.
|
||||
result = self._markitdown.convert(tmp_path, stream_info=self._utf8_stream_info(file_data, filename))
|
||||
# Parse using markitdown
|
||||
result = self._markitdown.convert(tmp_path)
|
||||
|
||||
if not result or not result.text_content:
|
||||
raise RuntimeError(f"No content extracted from '{filename}'")
|
||||
@@ -176,28 +73,6 @@ class MarkitdownParser(FileParser):
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _utf8_stream_info(file_data: bytes, filename: str) -> "StreamInfo | None":
|
||||
"""Return a UTF-8 charset hint for text files that decode cleanly as UTF-8.
|
||||
|
||||
Returns None for binary files or non-UTF-8 text so markitdown falls back
|
||||
to its own detection.
|
||||
"""
|
||||
if Path(filename).suffix.lower() not in _TEXT_EXTENSIONS:
|
||||
return None
|
||||
try:
|
||||
file_data.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return None
|
||||
from markitdown import StreamInfo
|
||||
|
||||
return StreamInfo(charset="utf-8")
|
||||
|
||||
@staticmethod
|
||||
def _is_image_file(filename: str) -> bool:
|
||||
"""Return whether the file type needs OCR to extract useful text."""
|
||||
return Path(filename).suffix.lower() in {".jpg", ".jpeg", ".png"}
|
||||
|
||||
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||
"""Check if markitdown supports this file type."""
|
||||
# Supported extensions (from markitdown docs)
|
||||
@@ -210,7 +85,7 @@ class MarkitdownParser(FileParser):
|
||||
".ppt",
|
||||
".xlsx",
|
||||
".xls",
|
||||
# Images (optional OCR)
|
||||
# Images (with OCR)
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".png",
|
||||
|
||||
@@ -14,26 +14,13 @@ import logging
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _usage_from_anthropic_response(response: Any) -> LLMResponseUsage:
|
||||
"""Extract input/output/cached token counts from an Anthropic usage block."""
|
||||
usage = getattr(response, "usage", None)
|
||||
if not usage:
|
||||
return LLMResponseUsage()
|
||||
return LLMResponseUsage(
|
||||
input_tokens=usage.input_tokens or 0,
|
||||
output_tokens=usage.output_tokens or 0,
|
||||
cached_tokens=getattr(usage, "cache_read_input_tokens", 0) or 0,
|
||||
)
|
||||
|
||||
|
||||
class AnthropicLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using Anthropic's Claude models.
|
||||
@@ -51,7 +38,6 @@ 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,
|
||||
):
|
||||
"""
|
||||
@@ -68,10 +54,6 @@ 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)
|
||||
@@ -79,9 +61,6 @@ 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
|
||||
@@ -149,9 +128,7 @@ class AnthropicLLM(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: Route structured output through a forced tool_use tool for
|
||||
native constrained decoding (issue #1002). When False, falls back to
|
||||
schema-in-prompt + JSON parse.
|
||||
strict_schema: Use strict JSON schema enforcement (not supported by Anthropic).
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
|
||||
Returns:
|
||||
@@ -182,21 +159,14 @@ class AnthropicLLM(LLMInterface):
|
||||
else:
|
||||
anthropic_messages.append({"role": role, "content": content})
|
||||
|
||||
# Structured output: prefer Anthropic-native constrained decoding via a single
|
||||
# forced tool_use tool (strict_schema) over text-injecting the schema and
|
||||
# parsing the reply. Native constrained decoding guarantees schema-valid JSON,
|
||||
# eliminating the invalid-JSON retry storm (issue #1002). When strict_schema is
|
||||
# off we keep the text-inject + json.loads fallback for backward compatibility.
|
||||
schema = None
|
||||
use_forced_tool = False
|
||||
_tool_name = "structured_response"
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
if strict_schema:
|
||||
use_forced_tool = True
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
if system_prompt:
|
||||
system_prompt += schema_msg
|
||||
else:
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
system_prompt = (system_prompt + schema_msg) if system_prompt else schema_msg
|
||||
system_prompt = schema_msg
|
||||
|
||||
# Prepare parameters
|
||||
call_params: dict[str, Any] = {
|
||||
@@ -208,77 +178,44 @@ class AnthropicLLM(LLMInterface):
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if use_forced_tool:
|
||||
# Single tool whose input_schema IS the response schema; force the model to
|
||||
# emit it via tool_choice so the SDK does constrained decoding for us.
|
||||
call_params["tools"] = [
|
||||
{"name": _tool_name, "description": "Return the structured response.", "input_schema": schema}
|
||||
]
|
||||
call_params["tool_choice"] = {"type": "tool", "name": _tool_name}
|
||||
|
||||
if self._extra_body:
|
||||
call_params["extra_body"] = self._extra_body
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await self._client.messages.create(**call_params)
|
||||
# Stash usage before parse/validate, which may raise locally
|
||||
# even though the provider charged for these tokens (#2387).
|
||||
stash_response_usage(_usage_from_anthropic_response(response))
|
||||
|
||||
if use_forced_tool:
|
||||
# Forced tool_use → the validated args are already a dict; no parsing,
|
||||
# no markdown-strip, no JSON-decode retry possible.
|
||||
tool_input = None
|
||||
for block in response.content:
|
||||
if block.type == "tool_use" and block.name == _tool_name:
|
||||
tool_input = block.input or {}
|
||||
break
|
||||
if tool_input is None:
|
||||
# Model ignored the forced tool (rare, e.g. a gateway that drops
|
||||
# tool_choice). Fall back to text parse so we don't hard-fail; the
|
||||
# existing retry loop still covers genuine errors.
|
||||
content = "".join(b.text for b in response.content if b.type == "text")
|
||||
tool_input = json.loads(content)
|
||||
content = json.dumps(tool_input)
|
||||
result = tool_input if skip_validation else response_format.model_validate(tool_input)
|
||||
else:
|
||||
# Anthropic response content is a list of blocks
|
||||
content = ""
|
||||
for block in response.content:
|
||||
if block.type == "text":
|
||||
content += block.text
|
||||
# Anthropic response content is a list of blocks
|
||||
content = ""
|
||||
for block in response.content:
|
||||
if block.type == "text":
|
||||
content += block.text
|
||||
|
||||
if response_format is not None:
|
||||
# Models may wrap JSON in markdown code blocks
|
||||
clean_content = content
|
||||
if "```json" in content:
|
||||
clean_content = content.split("```json")[1].split("```")[0].strip()
|
||||
elif "```" in content:
|
||||
clean_content = content.split("```")[1].split("```")[0].strip()
|
||||
if response_format is not None:
|
||||
# Models may wrap JSON in markdown code blocks
|
||||
clean_content = content
|
||||
if "```json" in content:
|
||||
clean_content = content.split("```json")[1].split("```")[0].strip()
|
||||
elif "```" in content:
|
||||
clean_content = content.split("```")[1].split("```")[0].strip()
|
||||
|
||||
try:
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError:
|
||||
# Fallback to parsing raw content if markdown stripping failed
|
||||
json_data = json.loads(content)
|
||||
try:
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError:
|
||||
# Fallback to parsing raw content if markdown stripping failed
|
||||
json_data = json.loads(content)
|
||||
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
else:
|
||||
result = response_format.model_validate(json_data)
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
else:
|
||||
result = content
|
||||
result = response_format.model_validate(json_data)
|
||||
else:
|
||||
result = content
|
||||
|
||||
# Record metrics and log slow calls
|
||||
duration = time.time() - start_time
|
||||
response_usage = _usage_from_anthropic_response(response)
|
||||
input_tokens = response_usage.input_tokens
|
||||
output_tokens = response_usage.output_tokens
|
||||
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 = response_usage.cached_tokens
|
||||
|
||||
# Record LLM metrics
|
||||
metrics = get_metrics_collector()
|
||||
@@ -308,7 +245,6 @@ class AnthropicLLM(LLMInterface):
|
||||
duration=duration,
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
@@ -324,7 +260,6 @@ 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
|
||||
@@ -459,14 +394,10 @@ 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:
|
||||
response = await self._client.messages.create(**call_params)
|
||||
stash_response_usage(_usage_from_anthropic_response(response))
|
||||
|
||||
# Extract content and tool calls
|
||||
content_parts = []
|
||||
|
||||
@@ -15,8 +15,7 @@ from typing import Any
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
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
|
||||
|
||||
@@ -119,14 +118,12 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
Raises:
|
||||
RuntimeError: If the connection test fails.
|
||||
"""
|
||||
from ...config import get_config
|
||||
|
||||
try:
|
||||
test_messages = [{"role": "user", "content": "test"}]
|
||||
await self.call(
|
||||
messages=test_messages,
|
||||
max_completion_tokens=10,
|
||||
temperature=get_config().llm_temperature_verification,
|
||||
temperature=0.0,
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
@@ -229,16 +226,6 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
if isinstance(block, TextBlock):
|
||||
full_text += block.text
|
||||
|
||||
# The Claude Agent SDK doesn't report exact counts; stash the same
|
||||
# char/4 estimate the success path traces so a later parse/validate
|
||||
# failure records consistent (estimated) tokens, not zero (#2387).
|
||||
stash_response_usage(
|
||||
LLMResponseUsage(
|
||||
input_tokens=sum(len(m.get("content", "")) for m in messages) // 4,
|
||||
output_tokens=len(full_text) // 4,
|
||||
)
|
||||
)
|
||||
|
||||
# Handle structured output
|
||||
if response_format is not None:
|
||||
# Models may wrap JSON in markdown
|
||||
|
||||
@@ -60,22 +60,6 @@ _CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
|
||||
)
|
||||
|
||||
|
||||
def default_codex_auth_file() -> Path:
|
||||
"""Return the path to Codex's ``auth.json``.
|
||||
|
||||
Honors the ``CODEX_HOME`` environment variable — the same variable the
|
||||
canonical ``@openai/codex`` CLI uses to relocate its config/credentials
|
||||
directory — and falls back to ``~/.codex`` when it is unset or empty.
|
||||
|
||||
Resolved lazily on each call (rather than cached at import time) so that
|
||||
the environment is read at the point of use.
|
||||
"""
|
||||
codex_home = os.environ.get("CODEX_HOME")
|
||||
if codex_home:
|
||||
return Path(codex_home) / "auth.json"
|
||||
return Path.home() / ".codex" / "auth.json"
|
||||
|
||||
|
||||
class CodexRefreshExpiredError(RuntimeError):
|
||||
"""Raised when the Codex refresh_token itself is no longer valid.
|
||||
|
||||
@@ -102,7 +86,7 @@ class CodexAuthManager:
|
||||
The OAuth refresh token. May be ``None`` when the auth file omits it;
|
||||
the provider still works as a one-shot loader in that case.
|
||||
auth_file:
|
||||
Path to the Codex ``auth.json``. Used for re-reading the refresh token
|
||||
Path to ``~/.codex/auth.json``. Used for re-reading the refresh token
|
||||
on demand and for atomic persistence of rotated credentials.
|
||||
"""
|
||||
|
||||
@@ -131,8 +115,7 @@ class CodexAuthManager:
|
||||
Parameters
|
||||
----------
|
||||
auth_file:
|
||||
Defaults to ``$CODEX_HOME/auth.json`` (or ``~/.codex/auth.json``
|
||||
when ``CODEX_HOME`` is unset).
|
||||
Defaults to ``~/.codex/auth.json``.
|
||||
|
||||
Raises
|
||||
------
|
||||
@@ -143,7 +126,7 @@ class CodexAuthManager:
|
||||
``auth_mode``.
|
||||
"""
|
||||
if auth_file is None:
|
||||
auth_file = default_codex_auth_file()
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(f"Codex auth file not found: {auth_file}. Run 'codex auth login' to authenticate.")
|
||||
|
||||
@@ -2,9 +2,8 @@
|
||||
OpenAI Codex LLM provider using ChatGPT Plus/Pro OAuth authentication.
|
||||
|
||||
This provider enables using ChatGPT Plus/Pro subscriptions for API calls
|
||||
without separate OpenAI Platform API credits. It uses OAuth tokens from the
|
||||
Codex ``auth.json`` (``$CODEX_HOME/auth.json``, or ``~/.codex/auth.json`` when
|
||||
``CODEX_HOME`` is unset) and communicates with the ChatGPT backend API.
|
||||
without separate OpenAI Platform API credits. It uses OAuth tokens from
|
||||
~/.codex/auth.json and communicates with the ChatGPT backend API.
|
||||
|
||||
Tokens are refreshed automatically: the provider decodes the access_token
|
||||
JWT's ``exp`` claim and proactively refreshes via
|
||||
@@ -25,8 +24,7 @@ from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
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
|
||||
|
||||
@@ -37,7 +35,6 @@ from .codex_auth import (
|
||||
_CODEX_TOKEN_REFRESH_SKEW_SECONDS,
|
||||
CodexAuthManager,
|
||||
CodexRefreshExpiredError,
|
||||
default_codex_auth_file,
|
||||
)
|
||||
|
||||
# Re-export for backward compatibility (tests import from this module).
|
||||
@@ -58,15 +55,14 @@ class CodexLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using OpenAI Codex OAuth authentication.
|
||||
|
||||
Authenticates using ChatGPT Plus/Pro credentials stored in the Codex
|
||||
``auth.json`` (honoring ``CODEX_HOME``, default ``~/.codex``) and makes API
|
||||
calls to chatgpt.com/backend-api/codex/responses.
|
||||
Authenticates using ChatGPT Plus/Pro credentials stored in ~/.codex/auth.json
|
||||
and makes API calls to chatgpt.com/backend-api/codex/responses.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str, # Will be ignored, reads from the Codex auth.json (CODEX_HOME or ~/.codex)
|
||||
api_key: str, # Will be ignored, reads from ~/.codex/auth.json
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
@@ -85,14 +81,12 @@ class CodexLLM(LLMInterface):
|
||||
refresh_token = self._load_codex_refresh_token()
|
||||
logger.info(f"Loaded Codex OAuth credentials for account: {account_id}")
|
||||
except Exception as e:
|
||||
auth_file = default_codex_auth_file()
|
||||
raise RuntimeError(
|
||||
f"Failed to load Codex OAuth credentials from {auth_file}: {e}\n\n"
|
||||
f"Failed to load Codex OAuth credentials from ~/.codex/auth.json: {e}\n\n"
|
||||
"To set up Codex authentication:\n"
|
||||
"1. Install Codex CLI: npm install -g @openai/codex\n"
|
||||
"2. Login: codex auth login\n"
|
||||
f"3. Verify: ls {auth_file}\n\n"
|
||||
"(Set CODEX_HOME to use a credentials directory other than ~/.codex.)\n\n"
|
||||
"3. Verify: ls ~/.codex/auth.json\n\n"
|
||||
"Or use a different provider (openai, anthropic, gemini) with API keys."
|
||||
) from e
|
||||
|
||||
@@ -100,7 +94,7 @@ class CodexLLM(LLMInterface):
|
||||
access_token=access_token,
|
||||
account_id=account_id,
|
||||
refresh_token=refresh_token,
|
||||
auth_file=default_codex_auth_file(),
|
||||
auth_file=Path.home() / ".codex" / "auth.json",
|
||||
)
|
||||
|
||||
# Use ChatGPT backend API endpoint. Codex auth is tied to
|
||||
@@ -162,7 +156,7 @@ class CodexLLM(LLMInterface):
|
||||
|
||||
def _load_codex_auth(self) -> tuple[str, str]:
|
||||
"""
|
||||
Load OAuth credentials from the Codex ``auth.json`` (CODEX_HOME or ~/.codex).
|
||||
Load OAuth credentials from ~/.codex/auth.json.
|
||||
|
||||
Returns:
|
||||
Tuple of (access_token, account_id).
|
||||
@@ -171,7 +165,7 @@ class CodexLLM(LLMInterface):
|
||||
FileNotFoundError: If auth file doesn't exist.
|
||||
ValueError: If auth file is invalid.
|
||||
"""
|
||||
auth_file = default_codex_auth_file()
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(
|
||||
@@ -203,7 +197,9 @@ class CodexLLM(LLMInterface):
|
||||
pre- and post-``__init__`` because it does not depend on
|
||||
``_auth_manager`` being constructed yet.
|
||||
"""
|
||||
auth_file = self._auth_manager._auth_file if hasattr(self, "_auth_manager") else default_codex_auth_file()
|
||||
auth_file = (
|
||||
self._auth_manager._auth_file if hasattr(self, "_auth_manager") else Path.home() / ".codex" / "auth.json"
|
||||
)
|
||||
return CodexAuthManager.load_refresh_token_from_file(auth_file)
|
||||
|
||||
@staticmethod
|
||||
@@ -401,6 +397,7 @@ class CodexLLM(LLMInterface):
|
||||
}
|
||||
|
||||
url = f"{self.base_url}/codex/responses"
|
||||
last_exception = None
|
||||
|
||||
# Manual attempt tracking instead of ``for attempt in range(...)`` so
|
||||
# that the reactive-refresh path can retry once without consuming a
|
||||
@@ -415,16 +412,6 @@ class CodexLLM(LLMInterface):
|
||||
# Parse SSE stream
|
||||
content = await self._parse_sse_stream(response)
|
||||
|
||||
# Codex SSE carries no usage block; stash the same char/4 estimate
|
||||
# the success path traces so a later parse/validate failure records
|
||||
# consistent (estimated) token counts rather than zero (#2387).
|
||||
stash_response_usage(
|
||||
LLMResponseUsage(
|
||||
input_tokens=sum(len(m.get("content", "")) for m in messages) // 4,
|
||||
output_tokens=len(content) // 4,
|
||||
)
|
||||
)
|
||||
|
||||
# Handle structured output
|
||||
if response_format is not None:
|
||||
# Models may wrap JSON in markdown
|
||||
@@ -441,6 +428,7 @@ class CodexLLM(LLMInterface):
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
last_exception = e
|
||||
attempt += 1
|
||||
continue
|
||||
raise
|
||||
@@ -502,6 +490,7 @@ class CodexLLM(LLMInterface):
|
||||
return result
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
last_exception = e
|
||||
status_code = e.response.status_code
|
||||
|
||||
# Auth error: try one OAuth refresh + retry before giving up.
|
||||
@@ -560,6 +549,7 @@ class CodexLLM(LLMInterface):
|
||||
raise
|
||||
|
||||
except httpx.RequestError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
logger.warning(f"Codex connection error (attempt {attempt + 1}/{max_retries + 1}): {e}")
|
||||
@@ -574,6 +564,10 @@ class CodexLLM(LLMInterface):
|
||||
logger.error(f"Unexpected Codex error: {type(e).__name__}: {e}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("Codex call failed after all retries")
|
||||
|
||||
async def _parse_sse_stream(self, response: httpx.Response) -> str:
|
||||
"""
|
||||
Parse Server-Sent Events (SSE) stream from Codex API.
|
||||
|
||||
@@ -1,316 +0,0 @@
|
||||
"""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,9 +8,9 @@ This provider supports both:
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from contextvars import ContextVar
|
||||
from typing import Any
|
||||
@@ -19,8 +19,7 @@ from google import genai
|
||||
from google.genai import errors as genai_errors
|
||||
from google.genai import types as genai_types
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.llm_wrapper import parse_llm_json
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
@@ -36,6 +35,7 @@ _safety_settings_ctx: ContextVar[list | None] = ContextVar("gemini_safety_settin
|
||||
|
||||
# Vertex AI imports (optional)
|
||||
try:
|
||||
import google.auth
|
||||
from google.oauth2 import service_account
|
||||
|
||||
VERTEXAI_AVAILABLE = True
|
||||
@@ -43,26 +43,6 @@ 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
|
||||
|
||||
|
||||
def _usage_from_gemini_response(response: Any) -> LLMResponseUsage:
|
||||
"""Extract prompt/candidate/cached token counts from a Gemini usage_metadata block."""
|
||||
usage = getattr(response, "usage_metadata", None)
|
||||
if not usage:
|
||||
return LLMResponseUsage()
|
||||
return LLMResponseUsage(
|
||||
input_tokens=usage.prompt_token_count or 0,
|
||||
output_tokens=usage.candidates_token_count or 0,
|
||||
cached_tokens=getattr(usage, "cached_content_token_count", 0) or 0,
|
||||
)
|
||||
|
||||
|
||||
class GeminiLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider for Google Gemini and Vertex AI.
|
||||
@@ -89,23 +69,6 @@ class GeminiLLM(LLMInterface):
|
||||
|
||||
# Safety settings: None means use Gemini's defaults
|
||||
self._safety_settings: list | None = kwargs.get("gemini_safety_settings")
|
||||
self._service_tier: str | None = kwargs.get("gemini_service_tier")
|
||||
|
||||
# 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)
|
||||
@@ -120,16 +83,6 @@ class GeminiLLM(LLMInterface):
|
||||
self._client = genai.Client(api_key=self.api_key)
|
||||
logger.info(f"Gemini API: model={self.model}")
|
||||
|
||||
def _apply_service_tier(self, config_kwargs: dict[str, Any]) -> None:
|
||||
if not self._service_tier:
|
||||
return
|
||||
|
||||
http_options = dict(config_kwargs.get("http_options") or {})
|
||||
extra_body = dict(http_options.get("extra_body") or {})
|
||||
extra_body.setdefault("service_tier", self._service_tier)
|
||||
http_options["extra_body"] = extra_body
|
||||
config_kwargs["http_options"] = http_options
|
||||
|
||||
def _init_vertexai(self, **kwargs: Any) -> None:
|
||||
"""Initialize Vertex AI client with project, region, and credentials."""
|
||||
# Extract Vertex AI config from kwargs
|
||||
@@ -215,7 +168,6 @@ 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.
|
||||
@@ -230,17 +182,8 @@ 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: Ignored — Gemini always grammar-enforces structured output via its
|
||||
native response_schema, so it is strict regardless of this flag.
|
||||
strict_schema: Use strict JSON schema enforcement (not supported by Gemini).
|
||||
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.
|
||||
@@ -248,14 +191,9 @@ class GeminiLLM(LLMInterface):
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# 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.
|
||||
# Convert OpenAI-style messages to Gemini format
|
||||
system_instruction = None
|
||||
gemini_contents = []
|
||||
using_cache = cached_prefix is not None
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
@@ -271,60 +209,41 @@ class GeminiLLM(LLMInterface):
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
|
||||
|
||||
def _system_instruction_with_schema() -> str:
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
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"
|
||||
f"{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
)
|
||||
return (system_instruction + schema_msg) if system_instruction else schema_msg
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
if system_instruction:
|
||||
system_instruction += schema_msg
|
||||
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
|
||||
]
|
||||
|
||||
# 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)
|
||||
self._apply_service_tier(config_kwargs)
|
||||
if use_cache:
|
||||
config_kwargs["cached_content"] = cached_prefix
|
||||
elif (
|
||||
use_schema_prompt_fallback
|
||||
and response_format is not None
|
||||
and hasattr(response_format, "model_json_schema")
|
||||
):
|
||||
config_kwargs["system_instruction"] = _system_instruction_with_schema()
|
||||
elif system_instruction:
|
||||
config_kwargs["system_instruction"] = system_instruction
|
||||
if response_format is not None and not use_schema_prompt_fallback:
|
||||
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
|
||||
use_schema_prompt_fallback = False
|
||||
generation_config = _build_generation_config(cache_active)
|
||||
generation_config = genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
|
||||
|
||||
last_exception = None
|
||||
|
||||
@@ -340,9 +259,6 @@ class GeminiLLM(LLMInterface):
|
||||
),
|
||||
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
|
||||
)
|
||||
# Stash usage before parse/validate, which may raise locally
|
||||
# even though the provider charged for these tokens (#2387).
|
||||
stash_response_usage(_usage_from_gemini_response(response))
|
||||
|
||||
content = response.text
|
||||
|
||||
@@ -372,24 +288,13 @@ class GeminiLLM(LLMInterface):
|
||||
else:
|
||||
result = content
|
||||
|
||||
# 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.
|
||||
# Extract token usage
|
||||
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
|
||||
@@ -402,8 +307,6 @@ 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
|
||||
@@ -427,7 +330,6 @@ class GeminiLLM(LLMInterface):
|
||||
duration=duration,
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
@@ -443,27 +345,12 @@ class GeminiLLM(LLMInterface):
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=input_tokens + output_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
last_exception = e
|
||||
if (
|
||||
attempt < max_retries
|
||||
and response_format is not None
|
||||
and hasattr(response_format, "model_json_schema")
|
||||
and not cache_active
|
||||
and not use_schema_prompt_fallback
|
||||
):
|
||||
logger.warning("Gemini returned invalid JSON, retrying with prompt-side schema guidance...")
|
||||
cache_active = False
|
||||
use_schema_prompt_fallback = True
|
||||
generation_config = _build_generation_config(cache_active)
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
if attempt < max_retries:
|
||||
logger.warning("Gemini returned invalid JSON, retrying...")
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
@@ -479,20 +366,6 @@ 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
|
||||
@@ -526,7 +399,6 @@ 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.
|
||||
@@ -541,39 +413,27 @@ 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. 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``.
|
||||
# Convert tools to Gemini format
|
||||
gemini_tools = []
|
||||
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"),
|
||||
)
|
||||
]
|
||||
)
|
||||
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
|
||||
@@ -586,10 +446,6 @@ 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":
|
||||
@@ -637,64 +493,49 @@ 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
|
||||
]
|
||||
|
||||
# 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)
|
||||
self._apply_service_tier(config_kwargs)
|
||||
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)
|
||||
config = genai_types.GenerateContentConfig(**config_kwargs)
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
@@ -709,7 +550,6 @@ class GeminiLLM(LLMInterface):
|
||||
),
|
||||
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
|
||||
)
|
||||
stash_response_usage(_usage_from_gemini_response(response))
|
||||
|
||||
# Extract content and tool calls
|
||||
content = None
|
||||
@@ -738,18 +578,12 @@ class GeminiLLM(LLMInterface):
|
||||
|
||||
finish_reason = "tool_calls" if tool_calls else "stop"
|
||||
|
||||
# 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.
|
||||
# Extract token usage
|
||||
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
|
||||
@@ -762,8 +596,6 @@ 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
|
||||
@@ -788,7 +620,6 @@ class GeminiLLM(LLMInterface):
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
tool_calls=tool_calls_dict,
|
||||
cached_tokens=cached_input_tokens,
|
||||
)
|
||||
|
||||
return LLMToolCallResult(
|
||||
@@ -797,8 +628,6 @@ class GeminiLLM(LLMInterface):
|
||||
finish_reason=finish_reason,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cached_tokens=cached_input_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
|
||||
except genai_errors.APIError as e:
|
||||
@@ -807,18 +636,6 @@ 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:
|
||||
@@ -835,330 +652,6 @@ 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
|
||||
|
||||
@@ -15,15 +15,10 @@ is handled automatically by LiteLLM.
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from litellm.exceptions import Timeout as LiteLLMTimeout
|
||||
|
||||
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
@@ -31,22 +26,6 @@ from hindsight_api.worker.stage import set_stage
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _usage_from_litellm_response(response: Any) -> LLMResponseUsage:
|
||||
"""Extract prompt/completion/cached token counts from a LiteLLM (OpenAI-shaped) usage block."""
|
||||
usage = getattr(response, "usage", None)
|
||||
if not usage:
|
||||
return LLMResponseUsage()
|
||||
cached_tokens = 0
|
||||
details = getattr(usage, "prompt_tokens_details", None)
|
||||
if details:
|
||||
cached_tokens = getattr(details, "cached_tokens", 0) or 0
|
||||
return LLMResponseUsage(
|
||||
input_tokens=getattr(usage, "prompt_tokens", 0) or 0,
|
||||
output_tokens=getattr(usage, "completion_tokens", 0) or 0,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
|
||||
|
||||
class LiteLLMLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using the LiteLLM SDK for universal model support.
|
||||
@@ -68,31 +47,12 @@ class LiteLLMLLM(LLMInterface):
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
bedrock_service_tier: str | None = None,
|
||||
default_headers: dict[str, Any] | None = None,
|
||||
timeout: float = 300.0,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
# ``None`` falls back to HINDSIGHT_API_LLM_TIMEOUT, then DEFAULT_LLM_TIMEOUT — never None,
|
||||
# so the hard ``asyncio.wait_for`` backstop in ``call`` is always bounded.
|
||||
self.timeout = timeout if timeout is not None else float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT)))
|
||||
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 {}
|
||||
# Operator-configured default headers forwarded to litellm.acompletion as
|
||||
# ``extra_headers`` (used by deployments routing through proxies / request-
|
||||
# tracing middleware). Mirrors the Anthropic provider's default_headers
|
||||
# wiring. Sourced from llm_default_headers (env: HINDSIGHT_API_LLM_DEFAULT_HEADERS).
|
||||
# Copied so a caller-owned dict can't be mutated through us, and a fresh
|
||||
# copy is handed to each call below to avoid cross-request contamination.
|
||||
self._default_headers: dict[str, Any] = dict(default_headers or {})
|
||||
self.bedrock_service_tier = bedrock_service_tier
|
||||
|
||||
try:
|
||||
import litellm
|
||||
@@ -108,14 +68,12 @@ class LiteLLMLLM(LLMInterface):
|
||||
raise RuntimeError("LiteLLM SDK not installed. Run: uv add litellm or pip install litellm") from e
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
from ...config import get_config
|
||||
|
||||
try:
|
||||
test_messages = [{"role": "user", "content": "test"}]
|
||||
await self.call(
|
||||
messages=test_messages,
|
||||
max_completion_tokens=50,
|
||||
temperature=get_config().llm_temperature_verification,
|
||||
temperature=0.0,
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
@@ -149,22 +107,6 @@ 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)
|
||||
|
||||
# Forward operator-configured default headers as ``extra_headers`` so they
|
||||
# reach the provider behind LiteLLM (proxies / request-tracing middleware).
|
||||
# ``setdefault`` keeps any explicit per-call ``extra_headers`` authoritative;
|
||||
# a per-call copy prevents LiteLLM/downstream from mutating the stored dict.
|
||||
if self._default_headers:
|
||||
kwargs.setdefault("extra_headers", dict(self._default_headers))
|
||||
|
||||
# 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) ────────────
|
||||
@@ -249,14 +191,7 @@ class LiteLLMLLM(LLMInterface):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self._stage_label}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
timeout=self.timeout,
|
||||
)
|
||||
# Stash usage before the length check and parse/validate below,
|
||||
# which may raise locally even though the provider charged for
|
||||
# these tokens (#2387).
|
||||
stash_response_usage(_usage_from_litellm_response(response))
|
||||
response = await self._acompletion(**call_kwargs)
|
||||
|
||||
content = response.choices[0].message.content or ""
|
||||
finish_reason = response.choices[0].finish_reason
|
||||
@@ -287,9 +222,8 @@ class LiteLLMLLM(LLMInterface):
|
||||
result = content
|
||||
|
||||
# Extract usage
|
||||
response_usage = _usage_from_litellm_response(response)
|
||||
input_tokens = response_usage.input_tokens
|
||||
output_tokens = response_usage.output_tokens
|
||||
input_tokens = getattr(response.usage, "prompt_tokens", 0) or 0
|
||||
output_tokens = getattr(response.usage, "completion_tokens", 0) or 0
|
||||
total_tokens = input_tokens + output_tokens
|
||||
|
||||
# Record metrics
|
||||
@@ -352,25 +286,6 @@ class LiteLLMLLM(LLMInterface):
|
||||
logger.error(f"LiteLLM returned invalid JSON after {max_retries + 1} attempts")
|
||||
raise
|
||||
|
||||
except (TimeoutError, asyncio.TimeoutError, LiteLLMTimeout) as e:
|
||||
# litellm/httpx don't always honor their own ``timeout=`` (e.g. a connection held
|
||||
# open with no token progress), so ``wait_for`` is the hard cap that cancels a hung
|
||||
# call regardless — otherwise one straggler pins a worker slot and stalls its gather.
|
||||
last_exception = e
|
||||
exc_name = type(e).__name__
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"LiteLLM call exceeded timeout={self.timeout}s ({exc_name}, scope={scope}), retrying..."
|
||||
)
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
logger.error(
|
||||
f"LiteLLM call timed out after {self.timeout}s on {attempt + 1} attempts "
|
||||
f"({exc_name}, scope={scope})"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
error_str = str(e).lower()
|
||||
# Fast fail on auth errors
|
||||
@@ -421,18 +336,7 @@ class LiteLLMLLM(LLMInterface):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self._stage_label}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
timeout=self.timeout,
|
||||
)
|
||||
# Stash usage before the tool-call argument parse below, which
|
||||
# can raise json.JSONDecodeError locally even though the provider
|
||||
# already billed for these tokens; without this the error trace
|
||||
# records 0/0 tokens (#2387). Mirrors call() and the anthropic/
|
||||
# gemini call_with_tools paths so the litellm tool path (and the
|
||||
# LiteLLMRouterLLM subclass that inherits this method) completes
|
||||
# the #2396 usage-on-error coverage.
|
||||
stash_response_usage(_usage_from_litellm_response(response))
|
||||
response = await self._acompletion(**call_kwargs)
|
||||
|
||||
message = response.choices[0].message
|
||||
content = message.content
|
||||
@@ -502,23 +406,6 @@ class LiteLLMLLM(LLMInterface):
|
||||
output_tokens=output_tokens,
|
||||
)
|
||||
|
||||
except (TimeoutError, asyncio.TimeoutError, LiteLLMTimeout) as e:
|
||||
# See ``call`` — hard cap so a hung completion cannot block
|
||||
# forever and pin a worker slot / concurrency permit.
|
||||
last_exception = e
|
||||
exc_name = type(e).__name__
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"LiteLLM tool call exceeded timeout={self.timeout}s ({exc_name}, scope={scope}), retrying..."
|
||||
)
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
logger.error(
|
||||
f"LiteLLM tool call timed out after {self.timeout}s on {attempt + 1} attempts "
|
||||
f"({exc_name}, scope={scope})"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
error_str = str(e).lower()
|
||||
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
|
||||
|
||||
@@ -67,7 +67,7 @@ class LiteLLMRouterLLM(LiteLLMLLM):
|
||||
model: str,
|
||||
config: dict[str, Any],
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float | None = None,
|
||||
timeout: float = 300.0,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(
|
||||
@@ -146,28 +146,16 @@ class LiteLLMRouterLLM(LiteLLMLLM):
|
||||
kwargs["max_completion_tokens"] = self._cap_max_completion_tokens(max_completion_tokens)
|
||||
if temperature is not None:
|
||||
kwargs["temperature"] = temperature
|
||||
|
||||
# Forward operator-configured default headers as ``extra_headers`` so they
|
||||
# reach the provider behind the Router (proxies / request-tracing middleware).
|
||||
# This override deliberately omits api_key/base_url/extra_body (those live in
|
||||
# the per-deployment Router config), but headers are a cross-cutting operator
|
||||
# concern, so we inject them here too — mirroring the base provider.
|
||||
# ``setdefault`` keeps any explicit per-call ``extra_headers`` authoritative;
|
||||
# a per-call copy prevents LiteLLM/downstream from mutating the stored dict.
|
||||
if self._default_headers:
|
||||
kwargs.setdefault("extra_headers", dict(self._default_headers))
|
||||
return kwargs
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
from hindsight_api.engine.llm_interface import OutputTooLongError
|
||||
|
||||
from ...config import get_config
|
||||
|
||||
try:
|
||||
await self.call(
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
max_completion_tokens=50,
|
||||
temperature=get_config().llm_temperature_verification,
|
||||
temperature=0.0,
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
@@ -101,7 +101,7 @@ class MockLLM(LLMInterface):
|
||||
messages: List of message dicts with 'role' and 'content'.
|
||||
response_format: Optional Pydantic model for structured output.
|
||||
max_completion_tokens: Not used in mock.
|
||||
temperature: Recorded on the call record for test assertions.
|
||||
temperature: Not used in mock.
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Not used in mock.
|
||||
initial_backoff: Not used in mock.
|
||||
@@ -123,9 +123,6 @@ class MockLLM(LLMInterface):
|
||||
if response_format and hasattr(response_format, "__name__")
|
||||
else str(response_format),
|
||||
"scope": scope,
|
||||
# Record the temperature so tests can assert per-operation temperature
|
||||
# wiring (None means the parameter was omitted from the call).
|
||||
"temperature": temperature,
|
||||
}
|
||||
self._mock_calls.append(call_record)
|
||||
logger.debug(f"Mock LLM call recorded: scope={scope}, model={self.model}")
|
||||
@@ -165,12 +162,6 @@ 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."
|
||||
@@ -211,7 +202,7 @@ class MockLLM(LLMInterface):
|
||||
messages: List of message dicts. Can include tool results with role='tool'.
|
||||
tools: List of tool definitions in OpenAI format.
|
||||
max_completion_tokens: Not used in mock.
|
||||
temperature: Recorded on the call record for test assertions.
|
||||
temperature: Not used in mock.
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Not used in mock.
|
||||
initial_backoff: Not used in mock.
|
||||
@@ -228,9 +219,6 @@ class MockLLM(LLMInterface):
|
||||
"messages": messages,
|
||||
"tools": [t.get("function", {}).get("name") for t in tools],
|
||||
"scope": scope,
|
||||
# Record the temperature so tests can assert per-operation temperature
|
||||
# wiring (None means the parameter was omitted from the call).
|
||||
"temperature": temperature,
|
||||
}
|
||||
self._mock_calls.append(call_record)
|
||||
|
||||
|
||||
@@ -1,463 +0,0 @@
|
||||
"""
|
||||
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()
|
||||
@@ -1,167 +0,0 @@
|
||||
"""
|
||||
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-M3 / MiniMax-M2.7 models with 1M context window
|
||||
- MiniMax: 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
|
||||
|
||||
@@ -26,8 +26,6 @@ import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from email.utils import parsedate_to_datetime
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qs, urlparse, urlunparse
|
||||
|
||||
@@ -35,9 +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, ProviderRateLimitResetError
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
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
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
@@ -48,16 +44,6 @@ 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."""
|
||||
@@ -86,49 +72,6 @@ def _strip_code_fences(content: str) -> str:
|
||||
return content
|
||||
|
||||
|
||||
# Reasoning/thinking tags emitted by extended-thinking models. Some providers
|
||||
# (e.g. MiniMax-M3) leak the chain-of-thought wrapped in these tags into the
|
||||
# response body instead of a separate reasoning_content field. Each entry is
|
||||
# (open_tag, close_tag); the open tag also matches when the close tag is missing
|
||||
# (truncated output) so a dangling block is removed to end-of-string.
|
||||
_REASONING_TAG_PAIRS: tuple[tuple[str, str], ...] = (
|
||||
("<think>", "</think>"),
|
||||
("<thinking>", "</thinking>"),
|
||||
("<thought>", "</thought>"),
|
||||
("<reasoning>", "</reasoning>"),
|
||||
("|startthink|", "|endthink|"),
|
||||
)
|
||||
|
||||
|
||||
def _strip_reasoning_tags(text: str) -> str:
|
||||
"""Strip extended-thinking/reasoning blocks from an LLM response.
|
||||
|
||||
Removes the full set of tag styles emitted by reasoning models:
|
||||
``<think>``, ``<thinking>``, ``<thought>``, ``<reasoning>`` and the
|
||||
``|startthink|...|endthink|`` markers. Both the structured (JSON) path and
|
||||
the free-form path must call this — otherwise a non-structured response
|
||||
(e.g. a mental-model markdown blob from MiniMax-M3) leaks the raw
|
||||
``<think>...</think>`` verbatim into stored memories.
|
||||
|
||||
Handles two cases:
|
||||
1. Closed blocks: ``<think>...</think>`` removed wherever they appear.
|
||||
2. Unclosed blocks: a dangling ``<think>`` with no closing tag (model output
|
||||
truncated mid-thought) is removed from the open tag to end-of-string.
|
||||
|
||||
Returns the input unchanged (modulo surrounding whitespace) when no tags are
|
||||
present.
|
||||
"""
|
||||
if not text:
|
||||
return text
|
||||
for open_tag, close_tag in _REASONING_TAG_PAIRS:
|
||||
open_re = re.escape(open_tag)
|
||||
close_re = re.escape(close_tag)
|
||||
# Closed blocks first, then any remaining unclosed (truncated) block.
|
||||
text = re.sub(rf"{open_re}.*?{close_re}", "", text, flags=re.DOTALL)
|
||||
text = re.sub(rf"{open_re}.*", "", text, flags=re.DOTALL)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def _response_get(response: Any, key: str, default: Any = None) -> Any:
|
||||
if isinstance(response, dict):
|
||||
return response.get(key, default)
|
||||
@@ -233,21 +176,6 @@ def _content_or_error(response: Any, *, provider: str, model: str, scope: str) -
|
||||
return content, choice
|
||||
|
||||
|
||||
def _usage_from_openai_response(response: Any) -> LLMResponseUsage:
|
||||
"""Extract prompt/completion/cached token counts from an OpenAI-shaped usage block."""
|
||||
usage = getattr(response, "usage", None)
|
||||
input_tokens = (usage.prompt_tokens or 0) if usage else 0
|
||||
output_tokens = (usage.completion_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
|
||||
return LLMResponseUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
|
||||
|
||||
def _ensure_json_word_in_user_message(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Some OpenAI-compatible gateways require 'json' in a user message for json_object mode."""
|
||||
|
||||
@@ -295,122 +223,6 @@ def _summarize_status_error(e: APIStatusError, body_max: int = 400) -> str:
|
||||
return f"HTTP {e.status_code}: {body_str or '<no body>'}"
|
||||
|
||||
|
||||
_RATE_LIMIT_RESET_AT_RE = re.compile(
|
||||
r"\breset at\s+"
|
||||
r"(?P<reset_at>\d{4}-\d{2}-\d{2}[ T]\d{2}:\d{2}:\d{2}(?:\s*(?:Z|[+-]\d{2}:?\d{2}))?)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_RATE_LIMIT_WINDOW_RE = re.compile(
|
||||
r"\b(?:for|in)\s+(?P<amount>\d+)\s*(?P<unit>second|minute|hour|day)s?\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def _status_error_body_text(e: APIStatusError) -> str:
|
||||
body: Any = getattr(e, "body", None)
|
||||
if body is None:
|
||||
try:
|
||||
body = e.response.text
|
||||
except Exception:
|
||||
body = None
|
||||
if isinstance(body, (dict, list)):
|
||||
try:
|
||||
return json.dumps(body, default=str, ensure_ascii=False)
|
||||
except Exception:
|
||||
return str(body)
|
||||
return str(body or "").strip()
|
||||
|
||||
|
||||
def _parse_retry_after_header(value: str | None, now: datetime) -> datetime | None:
|
||||
if not value:
|
||||
return None
|
||||
raw = value.strip()
|
||||
try:
|
||||
seconds = float(raw)
|
||||
except ValueError:
|
||||
seconds = -1.0
|
||||
if seconds >= 0:
|
||||
return now + timedelta(seconds=seconds)
|
||||
|
||||
try:
|
||||
parsed = parsedate_to_datetime(raw)
|
||||
except (TypeError, ValueError, IndexError, OverflowError):
|
||||
return None
|
||||
if parsed.tzinfo is None:
|
||||
parsed = parsed.replace(tzinfo=UTC)
|
||||
return parsed.astimezone(UTC)
|
||||
|
||||
|
||||
def _parse_reset_at_datetime(value: str) -> datetime | None:
|
||||
raw = value.strip().replace(" ", "T")
|
||||
if raw.endswith("Z"):
|
||||
raw = f"{raw[:-1]}+00:00"
|
||||
elif re.search(r"[+-]\d{4}$", raw):
|
||||
raw = f"{raw[:-2]}:{raw[-2:]}"
|
||||
try:
|
||||
parsed = datetime.fromisoformat(raw)
|
||||
except ValueError:
|
||||
return None
|
||||
if parsed.tzinfo is None:
|
||||
# Some providers (z.ai included) return a wall-clock reset timestamp
|
||||
# without a zone. Interpret it in the host's local zone so logs, status
|
||||
# pages, and the queued next_retry_at describe the same operator-facing
|
||||
# clock instead of silently shifting by UTC offset.
|
||||
parsed = parsed.astimezone()
|
||||
return parsed.astimezone(UTC)
|
||||
|
||||
|
||||
def _rate_limit_retry_at(e: APIStatusError) -> datetime | None:
|
||||
now = datetime.now(UTC)
|
||||
response = getattr(e, "response", None)
|
||||
headers = getattr(response, "headers", None)
|
||||
if headers is not None:
|
||||
retry_at = _parse_retry_after_header(headers.get("retry-after") or headers.get("Retry-After"), now)
|
||||
if retry_at is not None and retry_at > now:
|
||||
return retry_at
|
||||
|
||||
body_text = _status_error_body_text(e)
|
||||
reset_match = _RATE_LIMIT_RESET_AT_RE.search(body_text)
|
||||
if reset_match:
|
||||
retry_at = _parse_reset_at_datetime(reset_match.group("reset_at"))
|
||||
if retry_at is not None and retry_at > now:
|
||||
return retry_at
|
||||
|
||||
window_match = _RATE_LIMIT_WINDOW_RE.search(body_text)
|
||||
if not window_match:
|
||||
return None
|
||||
amount = int(window_match.group("amount"))
|
||||
unit = window_match.group("unit").lower()
|
||||
if unit == "second":
|
||||
seconds = amount
|
||||
elif unit == "minute":
|
||||
seconds = amount * 60
|
||||
elif unit == "hour":
|
||||
seconds = amount * 3600
|
||||
else:
|
||||
seconds = amount * 86400
|
||||
return now + timedelta(seconds=seconds)
|
||||
|
||||
|
||||
def _raise_provider_quota_defer(
|
||||
e: APIStatusError, *, provider: str, model: str, scope: str, max_backoff: float
|
||||
) -> None:
|
||||
if e.status_code != 429:
|
||||
return
|
||||
retry_at = _rate_limit_retry_at(e)
|
||||
if retry_at is None:
|
||||
return
|
||||
if (retry_at - datetime.now(UTC)).total_seconds() <= max_backoff:
|
||||
return
|
||||
summary = _summarize_status_error(e)
|
||||
raise ProviderRateLimitResetError(
|
||||
retry_at=retry_at,
|
||||
message=(
|
||||
f"Provider quota exhausted ({provider}/{model}, scope={scope}); retry at {retry_at.isoformat()}: {summary}"
|
||||
),
|
||||
) from e
|
||||
|
||||
|
||||
class OpenAICompatibleLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider for OpenAI-compatible APIs.
|
||||
@@ -420,7 +232,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-M3 / MiniMax-M2.7 models via OpenAI-compatible API (https://api.minimax.io/v1)
|
||||
- MiniMax: 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
|
||||
"""
|
||||
@@ -446,7 +258,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
base_url: Base URL for the API (uses defaults for groq/ollama/lmstudio if empty).
|
||||
model: Model name.
|
||||
reasoning_effort: Reasoning effort level for supported models ("low", "medium", "high").
|
||||
timeout: Request timeout in seconds (uses env var or 120s default).
|
||||
timeout: Request timeout in seconds (uses env var or 300s default).
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto").
|
||||
extra_body: Extra body params merged into every API call.
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
@@ -465,10 +277,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"deepseek",
|
||||
"volcano",
|
||||
"openrouter",
|
||||
"requesty",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
"fireworks",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
@@ -490,14 +300,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
self.base_url = "https://api.deepseek.com"
|
||||
elif self.provider == "openrouter":
|
||||
self.base_url = "https://openrouter.ai/api/v1"
|
||||
elif self.provider == "requesty":
|
||||
self.base_url = "https://router.requesty.ai/v1"
|
||||
elif self.provider == "zai":
|
||||
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 == "atlas":
|
||||
self.base_url = "https://api.atlascloud.ai/v1"
|
||||
elif self.provider == "fireworks":
|
||||
# OpenAI-compatible inference host (online path). The batch API
|
||||
# lives on a separate control-plane host — see FireworksLLM.
|
||||
@@ -516,10 +322,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"minimax",
|
||||
"deepseek",
|
||||
"openrouter",
|
||||
"requesty",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
"ollama-cloud",
|
||||
)
|
||||
and not self.api_key
|
||||
@@ -556,21 +360,6 @@ 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.
|
||||
@@ -671,9 +460,7 @@ 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 (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.
|
||||
strict_schema: Use strict JSON schema enforcement (OpenAI only).
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
|
||||
Returns:
|
||||
@@ -781,8 +568,6 @@ 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):
|
||||
@@ -794,9 +579,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
try:
|
||||
if response_format is not None:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
# Stash usage before parse/validate, which may raise locally
|
||||
# even though the provider charged for these tokens (#2387).
|
||||
stash_response_usage(_usage_from_openai_response(response))
|
||||
|
||||
content, first_choice = _content_or_error(
|
||||
response,
|
||||
@@ -805,10 +587,15 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
scope=scope,
|
||||
)
|
||||
|
||||
# Strip reasoning model thinking tags (closed and unclosed).
|
||||
# Strip reasoning model thinking tags
|
||||
# Supports: <think>, <thinking>, <thought>, <reasoning>, |startthink|/|endthink|
|
||||
original_len = len(content)
|
||||
content = _strip_reasoning_tags(content)
|
||||
content = re.sub(r"<think>.*?</think>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"<thinking>.*?</thinking>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"<thought>.*?</thought>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"<reasoning>.*?</reasoning>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"\|startthink\|.*?\|endthink\|", "", content, flags=re.DOTALL)
|
||||
content = content.strip()
|
||||
if len(content) < original_len:
|
||||
logger.debug(f"Stripped {original_len - len(content)} chars of reasoning tokens")
|
||||
|
||||
@@ -850,7 +637,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
result = response_format.model_validate(json_data)
|
||||
else:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
stash_response_usage(_usage_from_openai_response(response))
|
||||
result, first_choice = _content_or_error(
|
||||
response,
|
||||
provider=self.provider,
|
||||
@@ -858,33 +644,12 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
scope=scope,
|
||||
)
|
||||
|
||||
# Free-form (non-structured) output also leaks reasoning tags:
|
||||
# reasoning models like MiniMax-M3 wrap their chain-of-thought
|
||||
# in <think>...</think> in the response body. Without this strip
|
||||
# a mental-model markdown blob is stored verbatim with the raw
|
||||
# thinking tags. Mirrors the structured-output path above.
|
||||
result = _strip_reasoning_tags(result)
|
||||
|
||||
# Record token usage metrics
|
||||
duration = time.time() - start_time
|
||||
usage = response.usage
|
||||
response_usage = _usage_from_openai_response(response)
|
||||
input_tokens = response_usage.input_tokens
|
||||
output_tokens = response_usage.output_tokens
|
||||
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 = response_usage.cached_tokens
|
||||
thoughts_tokens = 0
|
||||
if usage and getattr(usage, "completion_tokens_details", None):
|
||||
thoughts_tokens = getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0
|
||||
# OpenAI-compatible providers fold reasoning tokens into
|
||||
# ``completion_tokens`` (and thus ``total_tokens``), but the
|
||||
# TokenUsage contract — and the Gemini provider — treat
|
||||
# ``output_tokens``/``total_tokens`` as visible-only, surfacing
|
||||
# reasoning separately in ``thoughts_tokens``. Subtract so the
|
||||
# two fields don't double-count reasoning (cost over-attribution).
|
||||
if thoughts_tokens:
|
||||
output_tokens = max(0, output_tokens - thoughts_tokens)
|
||||
total_tokens = max(0, total_tokens - thoughts_tokens)
|
||||
|
||||
# Record LLM metrics
|
||||
metrics = get_metrics_collector()
|
||||
@@ -914,12 +679,14 @@ 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}, "
|
||||
@@ -932,8 +699,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=total_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
@@ -964,10 +729,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
logger.error(f"Auth error (HTTP {e.status_code}), not retrying: {str(e)}")
|
||||
raise
|
||||
|
||||
_raise_provider_quota_defer(
|
||||
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
|
||||
)
|
||||
|
||||
# Handle tool_use_failed error - model outputted in tool call format
|
||||
if e.status_code == 400 and response_format is not None:
|
||||
try:
|
||||
@@ -1021,6 +782,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
f"scope={scope}): {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except ProviderResponseError as e:
|
||||
last_exception = e
|
||||
if e.retryable and attempt < max_retries:
|
||||
@@ -1105,16 +867,6 @@ 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
|
||||
@@ -1154,8 +906,6 @@ 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):
|
||||
@@ -1184,17 +934,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
usage = response.usage
|
||||
input_tokens = usage.prompt_tokens or 0 if usage else 0
|
||||
output_tokens = usage.completion_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
|
||||
thoughts_tokens = 0
|
||||
if usage and getattr(usage, "completion_tokens_details", None):
|
||||
thoughts_tokens = getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0
|
||||
# See ``call()``: OpenAI-compatible ``completion_tokens`` includes
|
||||
# reasoning, so make ``output_tokens`` visible-only to avoid
|
||||
# double-counting it against ``thoughts_tokens``.
|
||||
if thoughts_tokens:
|
||||
output_tokens = max(0, output_tokens - thoughts_tokens)
|
||||
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
@@ -1237,8 +976,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
finish_reason=finish_reason,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
|
||||
except APIConnectionError as e:
|
||||
@@ -1266,10 +1003,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
f"not retrying: {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
_raise_provider_quota_defer(
|
||||
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
|
||||
)
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
@@ -1283,6 +1016,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
f"({self.provider}/{self.model}, scope={scope}): {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
|
||||
@@ -6,17 +6,12 @@ structured information like temporal constraints.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from hindsight_api.engine.temporal_periods import (
|
||||
NO_TEMPORAL_CONSTRAINT,
|
||||
extract_period,
|
||||
is_embedded_cjk_dateparser_match,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -128,12 +123,9 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
|
||||
# Check for period expressions first (these need special handling)
|
||||
query_lower = query.lower()
|
||||
period_result = extract_period(query_lower, reference_date)
|
||||
if period_result is NO_TEMPORAL_CONSTRAINT:
|
||||
return QueryAnalysis(temporal_constraint=None)
|
||||
if isinstance(period_result, tuple):
|
||||
start_date, end_date = period_result
|
||||
return QueryAnalysis(temporal_constraint=TemporalConstraint(start_date=start_date, end_date=end_date))
|
||||
period_result = self._extract_period(query_lower, reference_date)
|
||||
if period_result is not None:
|
||||
return QueryAnalysis(temporal_constraint=period_result)
|
||||
|
||||
# Lazy load dateparser (only imports on first call, then cached)
|
||||
self.load()
|
||||
@@ -166,12 +158,7 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
|
||||
# Filter out false positives (common words parsed as dates)
|
||||
false_positives = {"do", "may", "march", "will", "can", "sat", "sun", "mon", "tue", "wed", "thu", "fri"}
|
||||
valid_results = [
|
||||
(text, date)
|
||||
for text, date in results
|
||||
if (text.lower() not in false_positives or len(text) > 3)
|
||||
and not is_embedded_cjk_dateparser_match(query, text)
|
||||
]
|
||||
valid_results = [(text, date) for text, date in results if text.lower() not in false_positives or len(text) > 3]
|
||||
|
||||
if not valid_results:
|
||||
return QueryAnalysis(temporal_constraint=None)
|
||||
@@ -185,6 +172,127 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
|
||||
return QueryAnalysis(temporal_constraint=TemporalConstraint(start_date=start_date, end_date=end_date))
|
||||
|
||||
def _extract_period(self, query: str, reference_date: datetime) -> TemporalConstraint | None:
|
||||
"""
|
||||
Extract period-based temporal expressions (week, month, year, weekend).
|
||||
|
||||
These need special handling as they represent date ranges, not single dates.
|
||||
Supports multiple languages.
|
||||
"""
|
||||
|
||||
def constraint(start: datetime, end: datetime) -> TemporalConstraint:
|
||||
return TemporalConstraint(
|
||||
start_date=start.replace(hour=0, minute=0, second=0, microsecond=0),
|
||||
end_date=end.replace(hour=23, minute=59, second=59, microsecond=999999),
|
||||
)
|
||||
|
||||
# Yesterday patterns (English, Spanish, Italian, French, German)
|
||||
if re.search(r"\b(yesterday|ayer|ieri|hier|gestern)\b", query, re.IGNORECASE):
|
||||
d = reference_date - timedelta(days=1)
|
||||
return constraint(d, d)
|
||||
|
||||
# Today patterns
|
||||
if re.search(r"\b(today|hoy|oggi|aujourd\'?hui|heute)\b", query, re.IGNORECASE):
|
||||
return constraint(reference_date, reference_date)
|
||||
|
||||
# "a couple of days ago" / "a few days ago" patterns
|
||||
# These are imprecise so we create a range
|
||||
if re.search(r"\b(a\s+)?couple\s+(of\s+)?days?\s+ago\b", query, re.IGNORECASE):
|
||||
# "a couple of days" = approximately 2 days, give range of 1-3 days
|
||||
return constraint(reference_date - timedelta(days=3), reference_date - timedelta(days=1))
|
||||
|
||||
if re.search(r"\b(a\s+)?few\s+days?\s+ago\b", query, re.IGNORECASE):
|
||||
# "a few days" = approximately 3-4 days, give range of 2-5 days
|
||||
return constraint(reference_date - timedelta(days=5), reference_date - timedelta(days=2))
|
||||
|
||||
# "a couple of weeks ago" / "a few weeks ago" patterns
|
||||
if re.search(r"\b(a\s+)?couple\s+(of\s+)?weeks?\s+ago\b", query, re.IGNORECASE):
|
||||
# "a couple of weeks" = approximately 2 weeks, give range of 1-3 weeks
|
||||
return constraint(reference_date - timedelta(weeks=3), reference_date - timedelta(weeks=1))
|
||||
|
||||
if re.search(r"\b(a\s+)?few\s+weeks?\s+ago\b", query, re.IGNORECASE):
|
||||
# "a few weeks" = approximately 3-4 weeks, give range of 2-5 weeks
|
||||
return constraint(reference_date - timedelta(weeks=5), reference_date - timedelta(weeks=2))
|
||||
|
||||
# "a couple of months ago" / "a few months ago" patterns
|
||||
if re.search(r"\b(a\s+)?couple\s+(of\s+)?months?\s+ago\b", query, re.IGNORECASE):
|
||||
# "a couple of months" = approximately 2 months, give range of 1-3 months
|
||||
return constraint(reference_date - timedelta(days=90), reference_date - timedelta(days=30))
|
||||
|
||||
if re.search(r"\b(a\s+)?few\s+months?\s+ago\b", query, re.IGNORECASE):
|
||||
# "a few months" = approximately 3-4 months, give range of 2-5 months
|
||||
return constraint(reference_date - timedelta(days=150), reference_date - timedelta(days=60))
|
||||
|
||||
# Last week patterns (English, Spanish, Italian, French, German)
|
||||
if re.search(
|
||||
r"\b(last\s+week|la\s+semana\s+pasada|la\s+settimana\s+scorsa|la\s+semaine\s+derni[eè]re|letzte\s+woche)\b",
|
||||
query,
|
||||
re.IGNORECASE,
|
||||
):
|
||||
start = reference_date - timedelta(days=reference_date.weekday() + 7)
|
||||
return constraint(start, start + timedelta(days=6))
|
||||
|
||||
# Last month patterns
|
||||
if re.search(
|
||||
r"\b(last\s+month|el\s+mes\s+pasado|il\s+mese\s+scorso|le\s+mois\s+dernier|letzten?\s+monat)\b",
|
||||
query,
|
||||
re.IGNORECASE,
|
||||
):
|
||||
first = reference_date.replace(day=1)
|
||||
end = first - timedelta(days=1)
|
||||
start = end.replace(day=1)
|
||||
return constraint(start, end)
|
||||
|
||||
# Last year patterns
|
||||
if re.search(
|
||||
r"\b(last\s+year|el\s+a[ñn]o\s+pasado|l\'anno\s+scorso|l\'ann[ée]e\s+derni[eè]re|letztes?\s+jahr)\b",
|
||||
query,
|
||||
re.IGNORECASE,
|
||||
):
|
||||
year = reference_date.year - 1
|
||||
return constraint(datetime(year, 1, 1), datetime(year, 12, 31))
|
||||
|
||||
# Last weekend patterns
|
||||
if re.search(
|
||||
r"\b(last\s+weekend|el\s+fin\s+de\s+semana\s+pasado|lo\s+scorso\s+fine\s+settimana|le\s+week-?end\s+dernier|letztes?\s+wochenende)\b",
|
||||
query,
|
||||
re.IGNORECASE,
|
||||
):
|
||||
days_since_sat = (reference_date.weekday() + 2) % 7
|
||||
if days_since_sat == 0:
|
||||
days_since_sat = 7
|
||||
sat = reference_date - timedelta(days=days_since_sat)
|
||||
return constraint(sat, sat + timedelta(days=1))
|
||||
|
||||
# Month + Year patterns (e.g., "June 2024", "junio 2024", "giugno 2024")
|
||||
month_patterns = {
|
||||
"january|enero|gennaio|janvier|januar": 1,
|
||||
"february|febrero|febbraio|f[ée]vrier|februar": 2,
|
||||
"march|marzo|mars|m[äa]rz": 3,
|
||||
"april|abril|aprile|avril": 4,
|
||||
"may|mayo|maggio|mai": 5,
|
||||
"june|junio|giugno|juin|juni": 6,
|
||||
"july|julio|luglio|juillet|juli": 7,
|
||||
"august|agosto|ao[uû]t": 8,
|
||||
"september|septiembre|settembre|septembre": 9,
|
||||
"october|octubre|ottobre|octobre|oktober": 10,
|
||||
"november|noviembre|novembre": 11,
|
||||
"december|diciembre|dicembre|d[ée]cembre|dezember": 12,
|
||||
}
|
||||
|
||||
for pattern, month_num in month_patterns.items():
|
||||
match = re.search(rf"\b({pattern})\s+(\d{{4}})\b", query, re.IGNORECASE)
|
||||
if match:
|
||||
year = int(match.group(2))
|
||||
start = datetime(year, month_num, 1)
|
||||
if month_num == 12:
|
||||
end = datetime(year, 12, 31)
|
||||
else:
|
||||
end = datetime(year, month_num + 1, 1) - timedelta(days=1)
|
||||
return constraint(start, end)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class TransformerQueryAnalyzer(QueryAnalyzer):
|
||||
"""
|
||||
|
||||
@@ -14,8 +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, StructuredOutputResult, TokenUsageSummary, ToolCall
|
||||
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, TokenUsageSummary, ToolCall
|
||||
from .prompts import (
|
||||
_extract_directive_rules,
|
||||
build_final_prompt,
|
||||
@@ -90,87 +89,12 @@ _LEAKED_JSON_SUFFIX = re.compile(
|
||||
r'\s*```(?:json)?\s*\{[^}]*(?:"(?:observation_ids|memory_ids|mental_model_ids)"|\})\s*```\s*$',
|
||||
re.DOTALL | re.IGNORECASE,
|
||||
)
|
||||
_LEAKED_JSON_OBJECT = re.compile(
|
||||
r'\s*\{[^{]*"(?:observation_ids|memory_ids|mental_model_ids|answer)"[^}]*\}\s*$', re.DOTALL
|
||||
)
|
||||
_TRAILING_IDS_PATTERN = re.compile(
|
||||
r"\s*(?:observation_ids|memory_ids|mental_model_ids)\s*[=:]\s*\[.*?\]\s*$", re.DOTALL | re.IGNORECASE
|
||||
)
|
||||
_JSON_CODE_FENCE_PATTERN = re.compile(r"^\s*```(?:json)?\s*(\{.*\})\s*```\s*$", re.DOTALL | re.IGNORECASE)
|
||||
|
||||
_DONE_ARGUMENT_KEYS = frozenset(
|
||||
{
|
||||
"answer",
|
||||
"directive_compliance",
|
||||
"memory_ids",
|
||||
"mental_model_ids",
|
||||
"observation_ids",
|
||||
"model_ids",
|
||||
}
|
||||
)
|
||||
_DONE_ARGUMENT_MARKER_KEYS = _DONE_ARGUMENT_KEYS - {"answer"}
|
||||
_LEAKED_JSON_ID_KEYS = frozenset({"memory_ids", "mental_model_ids", "observation_ids", "model_ids"})
|
||||
|
||||
|
||||
def _unwrap_leaked_done_arguments(text: str) -> str | None:
|
||||
"""Return the answer when a done tool call was rendered as JSON text.
|
||||
|
||||
Some providers leak the done tool's argument object instead of surfacing it
|
||||
as a native tool call, e.g. {"answer": "...", "memory_ids": [...]}. Only
|
||||
unwrap objects that match the done argument shape so normal JSON answers
|
||||
stay intact.
|
||||
"""
|
||||
candidate = text.strip()
|
||||
if not candidate:
|
||||
return None
|
||||
|
||||
fenced = _JSON_CODE_FENCE_PATTERN.match(candidate)
|
||||
if fenced:
|
||||
candidate = fenced.group(1).strip()
|
||||
|
||||
try:
|
||||
payload = json.loads(candidate)
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
answer = payload.get("answer")
|
||||
if not isinstance(answer, str) or not answer.strip():
|
||||
return None
|
||||
|
||||
keys = set(payload)
|
||||
if not keys.intersection(_DONE_ARGUMENT_MARKER_KEYS):
|
||||
return None
|
||||
if not keys.issubset(_DONE_ARGUMENT_KEYS):
|
||||
return None
|
||||
|
||||
for key in ("memory_ids", "mental_model_ids", "observation_ids", "model_ids"):
|
||||
value = payload.get(key)
|
||||
if value is not None and not isinstance(value, list):
|
||||
return None
|
||||
|
||||
return answer.strip()
|
||||
|
||||
|
||||
def _strip_trailing_id_json_object(text: str) -> str:
|
||||
stripped = text.rstrip()
|
||||
if not stripped.endswith("}"):
|
||||
return text.strip()
|
||||
|
||||
start = stripped.rfind("{")
|
||||
if start < 0:
|
||||
return text.strip()
|
||||
|
||||
try:
|
||||
payload = json.loads(stripped[start:])
|
||||
except json.JSONDecodeError:
|
||||
return text.strip()
|
||||
|
||||
if not isinstance(payload, dict) or not payload:
|
||||
return text.strip()
|
||||
keys = set(payload)
|
||||
if not keys.issubset(_LEAKED_JSON_ID_KEYS):
|
||||
return text.strip()
|
||||
|
||||
return stripped[:start].strip()
|
||||
|
||||
|
||||
def _clean_answer_text(text: str) -> str:
|
||||
@@ -179,10 +103,6 @@ def _clean_answer_text(text: str) -> str:
|
||||
Some LLMs output the done() call as text instead of a proper tool call.
|
||||
This strips out patterns like: done({"answer": "...", ...})
|
||||
"""
|
||||
unwrapped = _unwrap_leaked_done_arguments(text)
|
||||
if unwrapped is not None:
|
||||
return unwrapped
|
||||
|
||||
# Remove done() call pattern from the end of the text
|
||||
cleaned = _DONE_CALL_PATTERN.sub("", text).strip()
|
||||
return cleaned if cleaned else text
|
||||
@@ -201,17 +121,13 @@ def _clean_done_answer(text: str) -> str:
|
||||
if not text:
|
||||
return text
|
||||
|
||||
unwrapped = _unwrap_leaked_done_arguments(text)
|
||||
if unwrapped is not None:
|
||||
return unwrapped
|
||||
|
||||
cleaned = text
|
||||
|
||||
# Remove leaked JSON in code blocks at the end
|
||||
cleaned = _LEAKED_JSON_SUFFIX.sub("", cleaned).strip()
|
||||
|
||||
# Remove leaked raw JSON objects at the end
|
||||
cleaned = _strip_trailing_id_json_object(cleaned)
|
||||
cleaned = _LEAKED_JSON_OBJECT.sub("", cleaned).strip()
|
||||
|
||||
# Remove trailing ID patterns
|
||||
cleaned = _TRAILING_IDS_PATTERN.sub("", cleaned).strip()
|
||||
@@ -224,7 +140,7 @@ async def _generate_structured_output(
|
||||
response_schema: dict,
|
||||
llm_config: "LLMProvider",
|
||||
reflect_id: str,
|
||||
) -> StructuredOutputResult:
|
||||
) -> tuple[dict[str, Any] | None, int, int]:
|
||||
"""Generate structured output from an answer using the provided JSON schema.
|
||||
|
||||
Args:
|
||||
@@ -234,8 +150,8 @@ async def _generate_structured_output(
|
||||
reflect_id: Reflect ID for logging
|
||||
|
||||
Returns:
|
||||
A StructuredOutputResult carrying the structured output (None if
|
||||
generation fails) and the call's token usage.
|
||||
Tuple of (structured_output, input_tokens, output_tokens).
|
||||
structured_output is None if generation fails.
|
||||
"""
|
||||
try:
|
||||
from typing import Any as TypingAny
|
||||
@@ -269,7 +185,7 @@ async def _generate_structured_output(
|
||||
|
||||
if not fields:
|
||||
logger.warning(f"[REFLECT {reflect_id}] No fields found in response_schema, skipping structured output")
|
||||
return StructuredOutputResult()
|
||||
return None, 0, 0
|
||||
|
||||
DynamicModel = create_model("StructuredResponse", **fields)
|
||||
|
||||
@@ -322,9 +238,6 @@ OUTPUT:"""
|
||||
],
|
||||
response_format=DynamicModel,
|
||||
scope="reflect_structured",
|
||||
max_retries=1,
|
||||
initial_backoff=0.25,
|
||||
max_backoff=1.0,
|
||||
skip_validation=True, # We'll handle the dict ourselves
|
||||
return_usage=True,
|
||||
)
|
||||
@@ -345,17 +258,11 @@ OUTPUT:"""
|
||||
logger.warning(f"[REFLECT {reflect_id}] Required field '{field_name}' is empty in structured output")
|
||||
|
||||
logger.info(f"[REFLECT {reflect_id}] Generated structured output with {len(structured_output)} fields")
|
||||
return StructuredOutputResult(
|
||||
structured_output=structured_output,
|
||||
input_tokens=usage.input_tokens,
|
||||
output_tokens=usage.output_tokens,
|
||||
cached_tokens=usage.cached_tokens,
|
||||
thoughts_tokens=usage.thoughts_tokens,
|
||||
)
|
||||
return structured_output, usage.input_tokens, usage.output_tokens
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"[REFLECT {reflect_id}] Failed to generate structured output: {e}")
|
||||
return StructuredOutputResult()
|
||||
return None, 0, 0
|
||||
|
||||
|
||||
def _count_messages_tokens(messages: list[dict[str, Any]]) -> int:
|
||||
@@ -395,24 +302,6 @@ 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,
|
||||
@@ -433,7 +322,6 @@ 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.
|
||||
@@ -470,16 +358,12 @@ 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).
|
||||
# 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
|
||||
# Get tools for this agent (with directive compliance field if directives exist)
|
||||
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")
|
||||
@@ -498,28 +382,6 @@ 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] = []
|
||||
@@ -527,14 +389,9 @@ async def run_reflect_agent(
|
||||
llm_trace: list[dict[str, Any]] = []
|
||||
context_history: list[dict[str, Any]] = [] # For final prompt fallback
|
||||
|
||||
# Token usage tracking - accumulate across all LLM calls.
|
||||
# cached_tokens and thoughts_tokens are surfaced for cost attribution
|
||||
# and prompt-cache tuning. Both are subsets of (or parallel to) the
|
||||
# input/output counts and are NOT double-counted in total_tokens.
|
||||
# Token usage tracking - accumulate across all LLM calls
|
||||
total_input_tokens = 0
|
||||
total_output_tokens = 0
|
||||
total_cached_tokens = 0
|
||||
total_thoughts_tokens = 0
|
||||
|
||||
# Track available IDs for validation (prevents hallucinated citations)
|
||||
available_memory_ids: set[str] = set()
|
||||
@@ -557,8 +414,6 @@ async def run_reflect_agent(
|
||||
input_tokens=total_input_tokens,
|
||||
output_tokens=total_output_tokens,
|
||||
total_tokens=total_input_tokens + total_output_tokens,
|
||||
cached_tokens=total_cached_tokens,
|
||||
thoughts_tokens=total_thoughts_tokens,
|
||||
)
|
||||
|
||||
def _log_completion(answer: str, iterations: int, forced: bool = False):
|
||||
@@ -587,19 +442,7 @@ 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:
|
||||
@@ -612,9 +455,7 @@ async def run_reflect_agent(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(
|
||||
bank_profile.get("mission"), llm_output_language, directives
|
||||
),
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
@@ -625,8 +466,6 @@ async def run_reflect_agent(
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
@@ -640,12 +479,11 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
@@ -677,9 +515,7 @@ async def run_reflect_agent(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(
|
||||
bank_profile.get("mission"), llm_output_language, directives
|
||||
),
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
@@ -690,8 +526,6 @@ async def run_reflect_agent(
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
@@ -704,12 +538,11 @@ async def run_reflect_agent(
|
||||
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
@@ -737,37 +570,22 @@ async def run_reflect_agent(
|
||||
if include_recall:
|
||||
forced_sequence.append("recall")
|
||||
|
||||
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]}}
|
||||
if iteration < len(forced_sequence):
|
||||
iter_tool_choice: str | dict = {"type": "function", "function": {"name": forced_sequence[iteration]}}
|
||||
else:
|
||||
iter_tool_choice = "auto"
|
||||
|
||||
try:
|
||||
ct_kwargs: dict[str, Any] = dict(
|
||||
result = await llm_config.call_with_tools(
|
||||
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
|
||||
total_output_tokens += result.output_tokens
|
||||
total_cached_tokens += getattr(result, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(result, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": f"agent_{iteration + 1}",
|
||||
@@ -803,9 +621,7 @@ async def run_reflect_agent(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(
|
||||
bank_profile.get("mission"), llm_output_language, directives
|
||||
),
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
@@ -816,8 +632,6 @@ async def run_reflect_agent(
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
@@ -831,12 +645,11 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
@@ -893,8 +706,6 @@ async def run_reflect_agent(
|
||||
)
|
||||
total_input_tokens += rewrite_usage.input_tokens
|
||||
total_output_tokens += rewrite_usage.output_tokens
|
||||
total_cached_tokens += getattr(rewrite_usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(rewrite_usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final_rewrite",
|
||||
@@ -908,12 +719,11 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
|
||||
_log_completion(answer, iteration + 1)
|
||||
return ReflectAgentResult(
|
||||
@@ -935,9 +745,7 @@ async def run_reflect_agent(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(
|
||||
bank_profile.get("mission"), llm_output_language, directives
|
||||
),
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
@@ -948,8 +756,6 @@ async def run_reflect_agent(
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
@@ -963,12 +769,11 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
@@ -1045,9 +850,7 @@ async def run_reflect_agent(
|
||||
hallucinated_tools = []
|
||||
for tc in other_tools:
|
||||
norm = _normalize_tool_name(tc.name)
|
||||
# "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":
|
||||
if enabled_tools is not None and norm not in enabled_tools and norm not in ("done", "expand"):
|
||||
hallucinated_tools.append(tc)
|
||||
else:
|
||||
allowed_tools.append(tc)
|
||||
@@ -1121,25 +924,6 @@ 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"
|
||||
@@ -1263,15 +1047,14 @@ async def _process_done_tool(
|
||||
structured_output = None
|
||||
final_usage = usage
|
||||
if response_schema and llm_config and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
structured_output = struct.structured_output
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
# Add structured output tokens to usage
|
||||
final_usage = TokenUsageSummary(
|
||||
input_tokens=usage.input_tokens + struct.input_tokens,
|
||||
output_tokens=usage.output_tokens + struct.output_tokens,
|
||||
total_tokens=usage.total_tokens + struct.input_tokens + struct.output_tokens,
|
||||
cached_tokens=usage.cached_tokens + struct.cached_tokens,
|
||||
thoughts_tokens=usage.thoughts_tokens + struct.thoughts_tokens,
|
||||
input_tokens=usage.input_tokens + struct_in,
|
||||
output_tokens=usage.output_tokens + struct_out,
|
||||
total_tokens=usage.total_tokens + struct_in + struct_out,
|
||||
)
|
||||
|
||||
log_completion(answer, iterations)
|
||||
@@ -1376,10 +1159,8 @@ 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.
|
||||
# "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":
|
||||
# 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"):
|
||||
return {"error": f"Tool '{tool_name}' is not available. Use only the tools provided to you."}
|
||||
|
||||
if tool_name == "search_mental_models":
|
||||
|
||||
@@ -26,13 +26,10 @@ or stay the same per refresh, never get worse.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Annotated, Any, Literal, Union
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
|
||||
|
||||
from hindsight_api.engine.llm_wrapper import parse_llm_json
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from .structured_doc import (
|
||||
Block,
|
||||
@@ -147,27 +144,6 @@ Operation = Annotated[
|
||||
Field(discriminator="op"),
|
||||
]
|
||||
|
||||
_OPERATION_ADAPTER: TypeAdapter[Operation] = TypeAdapter(Operation)
|
||||
|
||||
|
||||
def _validate_operations_list(raw_ops: Any) -> tuple[list[Operation], list[dict[str, Any]]]:
|
||||
"""Validate each operation independently; drop invalid ops instead of failing the batch."""
|
||||
if not isinstance(raw_ops, list):
|
||||
raise TypeError(f"operations must be a list, got {type(raw_ops)!r}")
|
||||
valid: list[Operation] = []
|
||||
skipped: list[dict[str, Any]] = []
|
||||
for i, item in enumerate(raw_ops):
|
||||
try:
|
||||
valid.append(_OPERATION_ADAPTER.validate_python(item))
|
||||
except ValidationError as exc:
|
||||
skipped.append({"index": i, "op": item, "error": exc.errors(include_url=False)})
|
||||
logger.warning(
|
||||
"[STRUCTURED_DELTA] skipping invalid operation at index %s: %s",
|
||||
i,
|
||||
exc.errors(include_url=False),
|
||||
)
|
||||
return valid, skipped
|
||||
|
||||
|
||||
class DeltaOperationList(BaseModel):
|
||||
"""Container for the operations produced by an LLM delta call."""
|
||||
@@ -176,104 +152,6 @@ class DeltaOperationList(BaseModel):
|
||||
operations: list[Operation] = Field(default_factory=list)
|
||||
|
||||
|
||||
class DeltaAllOpsInvalidError(ValueError):
|
||||
"""Raised when the model emitted operations but none survived validation.
|
||||
|
||||
Distinct from an empty ``operations`` array (a legitimate no-op): here every
|
||||
op was malformed, so returning zero valid ops would make the caller apply
|
||||
nothing and silently drop this refresh's new facts. Raising instead lets the
|
||||
caller fall back to a full rewrite, which still integrates the new facts.
|
||||
"""
|
||||
|
||||
|
||||
def _finalize_operations(valid: list[Operation], skipped: list[dict[str, Any]]) -> DeltaOperationList:
|
||||
"""Build the result, but refuse a wholesale validation failure as a silent no-op."""
|
||||
if skipped and not valid:
|
||||
raise DeltaAllOpsInvalidError(f"all {len(skipped)} delta operation(s) failed validation")
|
||||
return DeltaOperationList(operations=valid)
|
||||
|
||||
|
||||
def _extract_balanced_json_object(text: str) -> str | None:
|
||||
"""Return the first top-level ``{...}`` slice, ignoring trailing junk."""
|
||||
start = text.find("{")
|
||||
if start < 0:
|
||||
return None
|
||||
depth = 0
|
||||
in_string = False
|
||||
escape = False
|
||||
for i in range(start, len(text)):
|
||||
ch = text[i]
|
||||
if in_string:
|
||||
if escape:
|
||||
escape = False
|
||||
elif ch == "\\":
|
||||
escape = True
|
||||
elif ch == '"':
|
||||
in_string = False
|
||||
continue
|
||||
if ch == '"':
|
||||
in_string = True
|
||||
elif ch == "{":
|
||||
depth += 1
|
||||
elif ch == "}":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
return text[start : i + 1]
|
||||
return None
|
||||
|
||||
|
||||
def parse_delta_operation_list(raw: Any) -> DeltaOperationList:
|
||||
"""Parse structured-delta LLM output into a validated operation list."""
|
||||
if isinstance(raw, DeltaOperationList):
|
||||
return raw
|
||||
if isinstance(raw, dict):
|
||||
ops_raw = raw.get("operations", [])
|
||||
valid, skipped = _validate_operations_list(ops_raw)
|
||||
if skipped:
|
||||
logger.info(
|
||||
"[STRUCTURED_DELTA] parsed %s op(s), skipped %s invalid op(s) from dict payload",
|
||||
len(valid),
|
||||
len(skipped),
|
||||
)
|
||||
return _finalize_operations(valid, skipped)
|
||||
|
||||
text = (raw or "").strip()
|
||||
if not text:
|
||||
return DeltaOperationList()
|
||||
|
||||
candidates: list[str] = [text]
|
||||
extracted = _extract_balanced_json_object(text)
|
||||
if extracted and extracted != text:
|
||||
candidates.append(extracted)
|
||||
|
||||
last_error: Exception | None = None
|
||||
for candidate in candidates:
|
||||
try:
|
||||
payload = parse_llm_json(candidate)
|
||||
except json.JSONDecodeError as exc:
|
||||
last_error = exc
|
||||
continue
|
||||
if not isinstance(payload, dict) or "operations" not in payload:
|
||||
last_error = ValueError("delta payload must be an object with an operations array")
|
||||
continue
|
||||
try:
|
||||
valid, skipped = _validate_operations_list(payload["operations"])
|
||||
except TypeError as exc:
|
||||
last_error = exc
|
||||
continue
|
||||
if skipped:
|
||||
logger.info(
|
||||
"[STRUCTURED_DELTA] parsed %s op(s), skipped %s invalid op(s)",
|
||||
len(valid),
|
||||
len(skipped),
|
||||
)
|
||||
return _finalize_operations(valid, skipped)
|
||||
|
||||
if last_error is not None:
|
||||
raise last_error
|
||||
return DeltaOperationList()
|
||||
|
||||
|
||||
# Application ---------------------------------------------------------------
|
||||
|
||||
|
||||
|
||||
@@ -78,32 +78,9 @@ class DirectiveInfo(BaseModel):
|
||||
class TokenUsageSummary(BaseModel):
|
||||
"""Total token usage across all LLM calls."""
|
||||
|
||||
input_tokens: int = Field(default=0, description="Total input tokens used (includes any cached prefix tokens)")
|
||||
output_tokens: int = Field(default=0, description="Total visible output tokens used (excludes reasoning/thoughts)")
|
||||
total_tokens: int = Field(default=0, description="Total tokens (input + output, excludes thoughts)")
|
||||
cached_tokens: int = Field(
|
||||
default=0,
|
||||
description="Cached/cache-read prompt tokens summed across calls. Subset of input_tokens.",
|
||||
)
|
||||
thoughts_tokens: int = Field(
|
||||
default=0,
|
||||
description=(
|
||||
"Reasoning/thinking tokens summed across calls. Billed at the output rate by some providers "
|
||||
"but not part of visible output."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class StructuredOutputResult(BaseModel):
|
||||
"""Result of structured-output generation, including token usage for the call."""
|
||||
|
||||
structured_output: dict[str, Any] | None = Field(
|
||||
default=None, description="Generated structured output, or None if generation failed"
|
||||
)
|
||||
input_tokens: int = Field(default=0, description="Input tokens used")
|
||||
output_tokens: int = Field(default=0, description="Visible output tokens used")
|
||||
cached_tokens: int = Field(default=0, description="Cached prefix tokens. Subset of input_tokens.")
|
||||
thoughts_tokens: int = Field(default=0, description="Reasoning/thinking tokens, when reported by the provider")
|
||||
input_tokens: int = Field(default=0, description="Total input tokens used")
|
||||
output_tokens: int = Field(default=0, description="Total output tokens used")
|
||||
total_tokens: int = Field(default=0, description="Total tokens (input + output)")
|
||||
|
||||
|
||||
class ReflectAgentResult(BaseModel):
|
||||
|
||||
@@ -604,44 +604,16 @@ Just provide the direct answer with proper markdown formatting.
|
||||
CRITICAL: This is a NON-CONVERSATIONAL system. NEVER ask follow-up questions, offer to search again, suggest alternatives, or end with anything like "Would you like me to..." or "Let me know if...". The user cannot reply. Your answer must be complete and self-contained."""
|
||||
|
||||
|
||||
# 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:
|
||||
def build_final_system_prompt(mission: str | None = None, llm_output_language: str | None = None) -> str:
|
||||
"""Build the final synthesis system prompt, using mission as role when set.
|
||||
|
||||
``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).
|
||||
When ``llm_output_language`` is set, the response is forced into that
|
||||
language regardless of the query/source language.
|
||||
"""
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt, output_language_directive
|
||||
|
||||
role_section = escape_for_prompt(mission.strip()) if mission else _DEFAULT_FINAL_ROLE
|
||||
|
||||
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)
|
||||
return _FINAL_SYSTEM_PROMPT_BASE.format(role_section=role_section) + output_language_directive(llm_output_language)
|
||||
|
||||
|
||||
# Backward-compatible constant for non-identity missions
|
||||
@@ -734,65 +706,7 @@ Examples
|
||||
``{"operations": [{"op": "replace_block", "section_id": "overview",
|
||||
"index": 0, "block": {"type": "paragraph", "text": "Updated summary."}}]}``
|
||||
- Remove an obsolete block →
|
||||
``{"operations": [{"op": "remove_block", "section_id": "status", "index": 2}]}``
|
||||
|
||||
JSON STRING RULES (critical)
|
||||
- Every ``text`` and ``items`` string must be valid JSON: escape ``"`` as ``\\"``,
|
||||
backslashes as ``\\\\``, and newlines as ``\\n``. Do not use raw backticks inside
|
||||
strings unless needed; prefer plain quotes for file paths.
|
||||
- ``replace_block``, ``insert_block``, and ``remove_block`` MUST include ``index`` (0-based block position in that section). Use ``replace_section_blocks`` only when replacing every block in a section.
|
||||
|
||||
- Do not append extra ``]`` or ``}`` after the closing ``}`` of the root object."""
|
||||
|
||||
_STRUCTURED_DELTA_DEFAULT_MAX_INPUT_TOKENS = 24_000
|
||||
|
||||
|
||||
def _truncate_cl100k(text: str, max_tokens: int) -> str:
|
||||
"""Truncate text to at most max_tokens using cl100k_base."""
|
||||
if max_tokens <= 0:
|
||||
return ""
|
||||
from .tokenization import count_cl100k_tokens
|
||||
|
||||
if count_cl100k_tokens(text) <= max_tokens:
|
||||
return text
|
||||
enc = __import__("tiktoken").get_encoding("cl100k_base")
|
||||
return enc.decode(enc.encode(text)[:max_tokens])
|
||||
|
||||
|
||||
def _fit_structured_delta_prompt_parts(
|
||||
*,
|
||||
source_query: str,
|
||||
current_document_json: str,
|
||||
candidate_markdown: str,
|
||||
facts_block: str,
|
||||
budget_hint: str,
|
||||
task_footer: str,
|
||||
max_input_tokens: int,
|
||||
) -> tuple[str, str, str, bool]:
|
||||
"""Shrink large prompt sections to fit within max_input_tokens (cl100k estimate)."""
|
||||
from .tokenization import count_cl100k_tokens
|
||||
|
||||
fixed = (
|
||||
f"## Topic\n{source_query}\n\n"
|
||||
f"## CURRENT DOCUMENT (apply ops to this; reference section ids as listed)\n"
|
||||
f"```json\n\n```\n\n"
|
||||
f"## NEW INFORMATION SYNTHESIS (context for how new facts relate to the topic)\n"
|
||||
f"```markdown\n\n```\n\n"
|
||||
f"## SUPPORTING FACTS (new since last refresh — integrate these)\n"
|
||||
f"{budget_hint}\n\n"
|
||||
f"{task_footer}"
|
||||
)
|
||||
facts_header = "## SUPPORTING FACTS (new since last refresh — integrate these)\n"
|
||||
facts_prefix_tokens = count_cl100k_tokens(facts_header)
|
||||
reserved_facts = min(4096, max(512, max_input_tokens // 8))
|
||||
doc_budget = max(1024, (max_input_tokens - count_cl100k_tokens(fixed) - reserved_facts) * 55 // 100)
|
||||
cand_budget = max(512, (max_input_tokens - count_cl100k_tokens(fixed) - reserved_facts) * 30 // 100)
|
||||
facts_budget = max(256, reserved_facts - facts_prefix_tokens)
|
||||
doc_json = _truncate_cl100k(current_document_json, doc_budget)
|
||||
candidate = _truncate_cl100k(candidate_markdown, cand_budget)
|
||||
facts_body = _truncate_cl100k(facts_block, facts_budget)
|
||||
truncated = doc_json != current_document_json or candidate != candidate_markdown or facts_body != facts_block
|
||||
return doc_json, candidate, facts_body, truncated
|
||||
``{"operations": [{"op": "remove_block", "section_id": "status", "index": 2}]}``"""
|
||||
|
||||
|
||||
def build_structured_delta_prompt(
|
||||
@@ -802,7 +716,6 @@ def build_structured_delta_prompt(
|
||||
supporting_facts: list[dict[str, Any]],
|
||||
source_query: str,
|
||||
max_output_tokens: int | None = None,
|
||||
max_input_tokens: int | None = None,
|
||||
) -> str:
|
||||
"""Build the user prompt for a structured-delta mental model refresh.
|
||||
|
||||
@@ -833,39 +746,19 @@ def build_structured_delta_prompt(
|
||||
"block-level ops) so the response always parses as valid JSON."
|
||||
)
|
||||
|
||||
task_footer = (
|
||||
return (
|
||||
f"## Topic\n{source_query}\n\n"
|
||||
f"## CURRENT DOCUMENT (apply ops to this; reference section ids as listed)\n"
|
||||
f"```json\n{current_document_json}\n```\n\n"
|
||||
f"## NEW INFORMATION SYNTHESIS (context for how new facts relate to the topic)\n"
|
||||
f"```markdown\n{candidate_markdown}\n```\n\n"
|
||||
f"## SUPPORTING FACTS (new since last refresh — integrate these)\n{facts_block}"
|
||||
f"{budget_hint}\n\n"
|
||||
"## Task\n"
|
||||
"Output a JSON object matching the operations schema. Integrate the new "
|
||||
"supporting facts into CURRENT DOCUMENT. Add, update, or remove content "
|
||||
"as needed. Preserve unchanged sections and blocks by not mentioning them."
|
||||
)
|
||||
input_cap = max_input_tokens if max_input_tokens is not None else _STRUCTURED_DELTA_DEFAULT_MAX_INPUT_TOKENS
|
||||
doc_json, candidate, facts_body, input_truncated = _fit_structured_delta_prompt_parts(
|
||||
source_query=source_query,
|
||||
current_document_json=current_document_json,
|
||||
candidate_markdown=candidate_markdown,
|
||||
facts_block=facts_block,
|
||||
budget_hint=budget_hint,
|
||||
task_footer=task_footer,
|
||||
max_input_tokens=input_cap,
|
||||
)
|
||||
truncation_note = ""
|
||||
if input_truncated:
|
||||
truncation_note = (
|
||||
"\n\n*Note: Document, synthesis, or facts were truncated to fit the model "
|
||||
"context window. Prefer minimal, high-leverage operations.*"
|
||||
)
|
||||
|
||||
return (
|
||||
f"## Topic\n{source_query}\n\n"
|
||||
f"## CURRENT DOCUMENT (apply ops to this; reference section ids as listed)\n"
|
||||
f"```json\n{doc_json}\n```\n\n"
|
||||
f"## NEW INFORMATION SYNTHESIS (context for how new facts relate to the topic)\n"
|
||||
f"```markdown\n{candidate}\n```\n\n"
|
||||
f"## SUPPORTING FACTS (new since last refresh — integrate these)\n{facts_body}"
|
||||
f"{budget_hint}{truncation_note}\n\n"
|
||||
f"{task_footer}"
|
||||
)
|
||||
|
||||
|
||||
DELTA_SYSTEM_PROMPT = """You are performing a surgical delta update to an existing mental model document.
|
||||
|
||||
@@ -232,7 +232,6 @@ 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.
|
||||
@@ -248,9 +247,6 @@ 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
|
||||
@@ -264,8 +260,7 @@ def get_reflect_tools(
|
||||
if include_recall:
|
||||
tools.append(TOOL_RECALL)
|
||||
|
||||
if include_expand:
|
||||
tools.append(TOOL_EXPAND)
|
||||
tools.append(TOOL_EXPAND)
|
||||
|
||||
# Use directive-aware done tool if directives are present
|
||||
if directive_rules:
|
||||
|
||||
@@ -31,20 +31,8 @@ class LLMToolCallResult(BaseModel):
|
||||
content: str | None = Field(default=None, description="Text content if any")
|
||||
tool_calls: list[LLMToolCall] = Field(default_factory=list, description="Tool calls requested by the LLM")
|
||||
finish_reason: str | None = Field(default=None, description="Reason the LLM stopped: 'stop', 'tool_calls', etc.")
|
||||
input_tokens: int = Field(
|
||||
default=0,
|
||||
description="Input tokens used in this call (includes any cached prefix tokens reported by the provider)",
|
||||
)
|
||||
output_tokens: int = Field(
|
||||
default=0, description="Visible output tokens used in this call (excludes reasoning/thoughts)"
|
||||
)
|
||||
cached_tokens: int = Field(
|
||||
default=0, description="Cached prefix tokens, when reported by the provider. Subset of input_tokens."
|
||||
)
|
||||
thoughts_tokens: int = Field(
|
||||
default=0,
|
||||
description="Reasoning/thinking tokens. Billed at the output rate by some providers but not part of visible output.",
|
||||
)
|
||||
input_tokens: int = Field(default=0, description="Input tokens used in this call")
|
||||
output_tokens: int = Field(default=0, description="Output tokens used in this call")
|
||||
|
||||
|
||||
class ToolCallTrace(BaseModel):
|
||||
@@ -103,18 +91,8 @@ 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 visible output/completion tokens generated (excludes reasoning/thoughts)"
|
||||
)
|
||||
total_tokens: int = Field(default=0, description="Total tokens (input + output, excludes thoughts)")
|
||||
cached_tokens: int = Field(default=0, description="Cached/cache-read prompt tokens, when reported by the provider")
|
||||
thoughts_tokens: int = Field(
|
||||
default=0,
|
||||
description=(
|
||||
"Reasoning/thinking tokens generated by the model. Billed at the output rate by some providers "
|
||||
"(e.g. Gemini 2.5+ family) but not surfaced in the visible response."
|
||||
),
|
||||
)
|
||||
output_tokens: int = Field(default=0, description="Number of output/completion tokens generated")
|
||||
total_tokens: int = Field(default=0, description="Total tokens (input + output)")
|
||||
|
||||
def __add__(self, other: "TokenUsage") -> "TokenUsage":
|
||||
"""Allow aggregating token usage from multiple calls."""
|
||||
@@ -122,39 +100,9 @@ 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,
|
||||
thoughts_tokens=self.thoughts_tokens + other.thoughts_tokens,
|
||||
)
|
||||
|
||||
|
||||
class ExtractedFact(BaseModel):
|
||||
"""A single candidate fact produced by dry-run extraction (no resolution/links/persistence).
|
||||
|
||||
A deliberate subset of the persisted memory-unit shape — only the fields a fresh extraction
|
||||
yields. Storage/consolidation/curation fields (id, document_id, chunk_id, proof_count, state, …)
|
||||
are omitted because nothing is stored. Entities are raw, unresolved names.
|
||||
"""
|
||||
|
||||
text: str = Field(description="The extracted fact text.")
|
||||
fact_type: str = Field(description="Perspective classification: 'world' or 'experience'.")
|
||||
occurred_start: str | None = Field(default=None, description="ISO timestamp the fact's event started, if dated.")
|
||||
occurred_end: str | None = Field(default=None, description="ISO timestamp the fact's event ended, if dated.")
|
||||
entities: list[str] = Field(
|
||||
default_factory=list, description="Raw (unresolved) entity names mentioned in the fact."
|
||||
)
|
||||
|
||||
|
||||
class DryRunExtractionResult(BaseModel):
|
||||
"""Result of dry-run fact extraction: candidate facts plus aggregated LLM token usage."""
|
||||
|
||||
facts: list[ExtractedFact] = Field(
|
||||
default_factory=list, description="Candidate facts the retain step would extract."
|
||||
)
|
||||
usage: TokenUsage = Field(
|
||||
default_factory=TokenUsage, description="Aggregated token usage across the extraction LLM calls."
|
||||
)
|
||||
|
||||
|
||||
class DispositionTraits(BaseModel):
|
||||
"""
|
||||
Disposition traits for a memory bank.
|
||||
@@ -172,47 +120,6 @@ class DispositionTraits(BaseModel):
|
||||
model_config = ConfigDict(json_schema_extra={"example": {"skepticism": 3, "literalism": 3, "empathy": 3}})
|
||||
|
||||
|
||||
class RecallScores(BaseModel):
|
||||
"""Per-result recall scores from different stages of the pipeline.
|
||||
|
||||
``final`` is the value results are ranked by. The others are diagnostic and
|
||||
can be filtered on via the recall ``min_scores`` request parameter. ``semantic``
|
||||
and ``keyword`` are the raw per-strategy retrieval scores (``None`` when that
|
||||
strategy did not surface this result); ``reranker`` is the cross-encoder's
|
||||
normalized relevance.
|
||||
"""
|
||||
|
||||
final: float = Field(description="Final ranking score (combined reranker + recency/temporal/proof boosts)")
|
||||
reranker: float | None = Field(
|
||||
default=None,
|
||||
description="Cross-encoder relevance, normalized 0-1. None when the reranker is a passthrough (rrf/interleave modes).",
|
||||
)
|
||||
semantic: float | None = Field(
|
||||
default=None, description="Vector cosine similarity (0-1). None if this result was not surfaced semantically."
|
||||
)
|
||||
keyword: float | None = Field(
|
||||
default=None,
|
||||
description="Keyword/full-text (BM25) score (>= 0, unbounded). None if this result was not surfaced by keyword search.",
|
||||
)
|
||||
|
||||
|
||||
class MinScores(BaseModel):
|
||||
"""Optional per-stage score floors for recall (all inclusive, AND-ed).
|
||||
|
||||
``semantic`` and ``keyword`` are **retrieval-level** cutoffs pushed into the SQL
|
||||
arms (overriding the global ``semantic_min_similarity`` / ``bm25_min_score``
|
||||
config for this request), so they prune weak matches before fusion. ``reranker``
|
||||
and ``final`` are **post-query** filters applied to the scored results after
|
||||
reranking. Any field left None imposes no floor; all-None (the default) means
|
||||
no score filtering.
|
||||
"""
|
||||
|
||||
semantic: float | None = Field(default=None, description="Retrieval-level: minimum vector similarity (0-1).")
|
||||
keyword: float | None = Field(default=None, description="Retrieval-level: minimum keyword/full-text (BM25) score.")
|
||||
reranker: float | None = Field(default=None, description="Post-query: minimum normalized reranker score (0-1).")
|
||||
final: float | None = Field(default=None, description="Post-query: minimum final ranking score.")
|
||||
|
||||
|
||||
class MemoryFact(BaseModel):
|
||||
"""
|
||||
A single memory fact returned by search or think operations.
|
||||
@@ -243,7 +150,7 @@ class MemoryFact(BaseModel):
|
||||
|
||||
id: str = Field(description="Unique identifier for the memory fact")
|
||||
text: str = Field(description="The actual text content of the memory")
|
||||
fact_type: str = Field(description="Type of fact: 'world', 'experience', or 'observation'")
|
||||
fact_type: str = Field(description="Type of fact: 'world', 'experience', 'opinion', or 'observation'")
|
||||
entities: list[str] | None = Field(None, description="Entity names mentioned in this fact")
|
||||
context: str | None = Field(None, description="Additional context for the memory")
|
||||
occurred_start: str | None = Field(None, description="ISO format date when the event started occurring")
|
||||
@@ -272,10 +179,6 @@ class MemoryFact(BaseModel):
|
||||
None,
|
||||
description="IDs of source facts this observation was derived from (observation type only, when source_facts is enabled)",
|
||||
)
|
||||
scores: RecallScores | None = Field(
|
||||
None,
|
||||
description="Recall scores from each pipeline stage (final/reranker/semantic/keyword). Not returned for source facts.",
|
||||
)
|
||||
|
||||
|
||||
class ChunkInfo(BaseModel):
|
||||
@@ -374,8 +277,7 @@ class ReflectResult(BaseModel):
|
||||
],
|
||||
"experience": [],
|
||||
"opinion": [],
|
||||
"observation": [],
|
||||
"mental-models": [],
|
||||
"mental_models": [],
|
||||
"directives": [
|
||||
{
|
||||
"id": "directive-123",
|
||||
@@ -392,7 +294,7 @@ class ReflectResult(BaseModel):
|
||||
|
||||
text: str = Field(description="The formulated answer text")
|
||||
based_on: dict[str, Any] = Field(
|
||||
description="Facts used to formulate the answer, organized by type (world, experience, observation, mental-models, directives)"
|
||||
description="Facts used to formulate the answer, organized by type (world, experience, mental_models, directives)"
|
||||
)
|
||||
structured_output: dict[str, Any] | None = Field(
|
||||
default=None,
|
||||
|
||||
@@ -4,8 +4,8 @@ bank profile utilities for disposition and mission management.
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import TypedDict
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
@@ -105,18 +105,6 @@ 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."""
|
||||
|
||||
@@ -135,8 +123,8 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
|
||||
Returns:
|
||||
BankProfile with name, typed DispositionTraits, and mission
|
||||
"""
|
||||
result = await get_or_create_bank_profile(pool, bank_id)
|
||||
return result.profile
|
||||
profile, _ = await get_or_create_bank_profile(pool, bank_id)
|
||||
return profile
|
||||
|
||||
|
||||
async def get_bank_profile_if_exists(pool, bank_id: str) -> BankProfile | None:
|
||||
@@ -174,89 +162,70 @@ async def get_bank_profile_if_exists(pool, bank_id: str) -> BankProfile | None:
|
||||
)
|
||||
|
||||
|
||||
async def get_or_create_bank_profile(pool, bank_id: str) -> BankProfileResult:
|
||||
async def get_or_create_bank_profile(pool, bank_id: str) -> tuple[BankProfile, bool]:
|
||||
"""
|
||||
Get bank profile, auto-creating with defaults if it doesn't exist.
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
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.
|
||||
Returns:
|
||||
Tuple of (BankProfile, created) where created is True if the bank
|
||||
did not exist before this call.
|
||||
"""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
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,
|
||||
# Try to get existing bank
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT name, disposition, mission
|
||||
FROM {fq_table("banks")} WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
# 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,
|
||||
)
|
||||
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)
|
||||
|
||||
created = inserted is not None
|
||||
if created:
|
||||
# Fresh insert — create per-bank vector indexes (instant on empty bank)
|
||||
await create_bank_vector_indexes(conn, bank_id, str(internal_id), ops=ops)
|
||||
return (
|
||||
BankProfile(
|
||||
name=row["name"],
|
||||
disposition=DispositionTraits(**disposition_data),
|
||||
mission=row["mission"] or "",
|
||||
),
|
||||
False,
|
||||
)
|
||||
|
||||
return BankProfileResult(
|
||||
profile=BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission=""),
|
||||
created=created,
|
||||
)
|
||||
# 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,
|
||||
)
|
||||
|
||||
|
||||
async def update_bank_disposition(pool, bank_id: str, disposition: dict[str, int]) -> None:
|
||||
|
||||
@@ -8,7 +8,6 @@ import hashlib
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
from ...config import get_config
|
||||
from ..memory_engine import fq_table
|
||||
from .types import ChunkMetadata
|
||||
|
||||
@@ -89,11 +88,6 @@ 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 = []
|
||||
@@ -104,7 +98,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 if store_text else "")
|
||||
chunk_texts.append(chunk.chunk_text)
|
||||
chunk_indices.append(chunk.chunk_index)
|
||||
content_hashes.append(compute_chunk_hash(chunk.chunk_text))
|
||||
chunk_id_map[chunk.chunk_index] = chunk_id
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user