Compare commits
126
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b7ad4c68cd | ||
|
|
ef065f39eb | ||
|
|
6e734e1afa | ||
|
|
ed82801b93 | ||
|
|
9571a341ff | ||
|
|
32b5da60a0 | ||
|
|
4b0d2658a4 | ||
|
|
f367ca81c8 | ||
|
|
18b9c59667 | ||
|
|
a82a20213a | ||
|
|
0cba7f3fbf | ||
|
|
a1ee94ab3e | ||
|
|
fb554664d0 | ||
|
|
bf6b90263b | ||
|
|
68f4a00e8e | ||
|
|
635cf9dd57 | ||
|
|
d425cfc3c5 | ||
|
|
dde133da00 | ||
|
|
e4686b92f0 | ||
|
|
7738021155 | ||
|
|
082213bd97 | ||
|
|
bcae23d9fe | ||
|
|
7cdacc4bf9 | ||
|
|
9704d9182e | ||
|
|
5ad0bffcd5 | ||
|
|
93232213c2 | ||
|
|
6f0a0f1c23 | ||
|
|
779e3140c8 | ||
|
|
9ec73e8455 | ||
|
|
0d2ba56f41 | ||
|
|
cf637799f2 | ||
|
|
74525cc049 | ||
|
|
d7dc8514ca | ||
|
|
1890d2b721 | ||
|
|
374c013689 | ||
|
|
4d9f4ab9ac | ||
|
|
a510b07a81 | ||
|
|
fdb5f47b23 | ||
|
|
129d88c56c | ||
|
|
4b19a0fb69 | ||
|
|
830d8472ca | ||
|
|
617939d822 | ||
|
|
ffa6fbf2a8 | ||
|
|
7a4400e08d | ||
|
|
2de19e578b | ||
|
|
dc41f6a534 | ||
|
|
5123e2a753 | ||
|
|
7e0afff340 | ||
|
|
09c9cecf56 | ||
|
|
78c35253ee | ||
|
|
eadb510eb3 | ||
|
|
691cb5394b | ||
|
|
a401b97eb7 | ||
|
|
aa4c1bbaf3 | ||
|
|
99525144b2 | ||
|
|
ded52e8de6 | ||
|
|
cc3ba4a37c | ||
|
|
7ef64f14ca | ||
|
|
16f807697d | ||
|
|
486c3a8b3b | ||
|
|
fbbc7a5e4c | ||
|
|
d7d41e76c2 | ||
|
|
262d4894f2 | ||
|
|
ec49175fa3 | ||
|
|
488f428009 | ||
|
|
d1ef9da95e | ||
|
|
30acca6fd9 | ||
|
|
2538708308 | ||
|
|
9e7aff6bd4 | ||
|
|
a908cdc974 | ||
|
|
4cd260b691 | ||
|
|
0c17e9acfd | ||
|
|
beca4b42f3 | ||
|
|
6e9b741b02 | ||
|
|
4a1b2f39c1 | ||
|
|
28ec22c3dc | ||
|
|
d802f91488 | ||
|
|
cb04cb79d9 | ||
|
|
dabbf9ff49 | ||
|
|
6348f42451 | ||
|
|
41a2ccabf8 | ||
|
|
eaf3048f2c | ||
|
|
9d95149852 | ||
|
|
ac3ab2b54c | ||
|
|
cb037290bb | ||
|
|
2582b45a16 | ||
|
|
0be157eeb5 | ||
|
|
90cb145aa6 | ||
|
|
7bd11bedf6 | ||
|
|
c3b2b1543a | ||
|
|
2743d061f7 | ||
|
|
daf2348bcd | ||
|
|
46dd2dfd94 | ||
|
|
878ef957f7 | ||
|
|
00d327a049 | ||
|
|
31d1e1729e | ||
|
|
592f01bba6 | ||
|
|
da05ee7215 | ||
|
|
19d23921fb | ||
|
|
e1e1a5e02b | ||
|
|
44b34c891c | ||
|
|
67ae2a41d4 | ||
|
|
2e5186a6fc | ||
|
|
9a20180415 | ||
|
|
f61ae2a185 | ||
|
|
dfd7cb52d4 | ||
|
|
5300d401b0 | ||
|
|
0b6bf53bef | ||
|
|
203ddfdd6c | ||
|
|
dcf5588e6c | ||
|
|
3d6c2ba8b0 | ||
|
|
80046797f7 | ||
|
|
db7dabcebd | ||
|
|
b83bb87ddd | ||
|
|
15ec55b703 | ||
|
|
f2596e1fe9 | ||
|
|
21c71f7bb8 | ||
|
|
86b686cd72 | ||
|
|
248c40e670 | ||
|
|
d18a9452ad | ||
|
|
806fbcd41c | ||
|
|
a75c3c85ad | ||
|
|
0db9f3da19 | ||
|
|
8940710c72 | ||
|
|
6252643de0 | ||
|
|
3fce309c0d |
+23
-1
@@ -7,6 +7,8 @@ HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# Reasoning effort for providers/models that support it. Examples: low, medium, high, xhigh.
|
||||
# HINDSIGHT_API_LLM_REASONING_EFFORT=low
|
||||
|
||||
# Example: Anthropic Claude configuration
|
||||
# HINDSIGHT_API_LLM_PROVIDER=anthropic
|
||||
@@ -64,8 +66,21 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# For Azure PostgreSQL with DiskANN:
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale # Auto-detects pg_diskann on Azure
|
||||
|
||||
# Text Search Extension (Optional - uses native PostgreSQL full-text search by default)
|
||||
# Backend options: "native" (default), "vchord", "pg_textsearch", "pgroonga", "pg_search"
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION=native
|
||||
# Native backend dictionary (only used by HINDSIGHT_API_TEXT_SEARCH_EXTENSION=native)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE=english
|
||||
# ParadeDB pg_search tokenizer (only used when creating pg_search BM25 indexes).
|
||||
# Empty uses ParadeDB's default tokenizer: unicode_words.
|
||||
# Supported values: unicode_words, simple, whitespace, literal, literal_normalized,
|
||||
# chinese_compatible, icu, jieba, source_code,
|
||||
# chinese_lindera/lindera(chinese), japanese_lindera/lindera(japanese),
|
||||
# korean_lindera/lindera(korean), ngram(min,max), edge_ngram(min,max)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER=
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default), "tei", "openai", "cohere", "google", "openrouter", "litellm", or "litellm-sdk"
|
||||
# Provider: "local" (default), "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
|
||||
@@ -77,6 +92,13 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=sk-xxxx
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL=text-embedding-3-small
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL=https://api.openai.com/v1
|
||||
# For ZeroEntropy zembed-1:
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=zeroentropy
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_API_KEY=ze-xxxx
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_MODEL=zembed-1
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_DIMENSIONS=1280
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT=float
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_LATENCY=fast
|
||||
#
|
||||
# IMPORTANT: Embedding keys require provider-specific names:
|
||||
# HINDSIGHT_API_EMBEDDINGS_{PROVIDER}_{PARAMETER}
|
||||
|
||||
+205
-7
@@ -49,7 +49,10 @@ jobs:
|
||||
integrations-agentcore: ${{ steps.filter.outputs.integrations-agentcore }}
|
||||
integrations-smolagents: ${{ steps.filter.outputs.integrations-smolagents }}
|
||||
integrations-dify: ${{ steps.filter.outputs.integrations-dify }}
|
||||
integrations-gemini-spark: ${{ steps.filter.outputs.integrations-gemini-spark }}
|
||||
integrations-vapi: ${{ steps.filter.outputs.integrations-vapi }}
|
||||
tools-agent-sdk: ${{ steps.filter.outputs.tools-agent-sdk }}
|
||||
integrations-roo-code: ${{ steps.filter.outputs.integrations-roo-code }}
|
||||
dev: ${{ steps.filter.outputs.dev }}
|
||||
ci: ${{ steps.filter.outputs.ci }}
|
||||
# Secrets are available for internal PRs and workflow_dispatch.
|
||||
@@ -143,8 +146,14 @@ jobs:
|
||||
- 'hindsight-integrations/smolagents/**'
|
||||
integrations-dify:
|
||||
- 'hindsight-integrations/dify/**'
|
||||
integrations-gemini-spark:
|
||||
- 'hindsight-integrations/gemini-spark/**'
|
||||
integrations-vapi:
|
||||
- 'hindsight-integrations/vapi/**'
|
||||
tools-agent-sdk:
|
||||
- 'hindsight-tools/hindsight-agent-sdk/**'
|
||||
integrations-roo-code:
|
||||
- 'hindsight-integrations/roo-code/**'
|
||||
dev:
|
||||
- 'hindsight-dev/**'
|
||||
ci:
|
||||
@@ -698,6 +707,63 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/pipecat
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-gemini-spark-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-gemini-spark == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/gemini-spark
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/gemini-spark
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-roo-code-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-roo-code == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.11'
|
||||
|
||||
- name: Install pytest
|
||||
run: pip install pytest
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/roo-code
|
||||
run: python -m pytest tests/ -v
|
||||
|
||||
build-control-plane:
|
||||
needs: [detect-changes]
|
||||
@@ -734,6 +800,12 @@ jobs:
|
||||
rm -rf node_modules/lightningcss node_modules/@tailwindcss
|
||||
npm install lightningcss @tailwindcss/postcss @tailwindcss/node
|
||||
|
||||
- name: Test Control Plane
|
||||
run: npm test --workspace=hindsight-control-plane
|
||||
|
||||
- name: Check i18n locale parity and hardcoded strings
|
||||
run: npm run i18n:check --workspace=hindsight-control-plane
|
||||
|
||||
- name: Build Control Plane
|
||||
run: npm run build --workspace=hindsight-control-plane
|
||||
|
||||
@@ -935,6 +1007,22 @@ jobs:
|
||||
- name: Lint Helm chart
|
||||
run: helm lint helm/hindsight
|
||||
|
||||
test-standalone-start-script:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.docker == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Run standalone start script tests
|
||||
run: bash docker/standalone/test-start-all.sh
|
||||
|
||||
build-docker-images:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -1102,7 +1190,70 @@ jobs:
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv run pytest tests -v -m "not hs_llm_mat"
|
||||
run: uv run pytest tests -v -m "not hs_llm_mat and not hs_llm_core"
|
||||
|
||||
test-api-llm-core:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
needs.detect-changes.outputs.has_secrets == 'true' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
name: Core LLM tests
|
||||
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
|
||||
GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- 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: Install dependencies
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv run python -c "
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
|
||||
print('Models downloaded successfully')
|
||||
"
|
||||
|
||||
- name: Run core LLM tests
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv run pytest tests -v -m "hs_llm_core" --timeout 600
|
||||
|
||||
test-api-llm-acceptance:
|
||||
needs: [detect-changes]
|
||||
@@ -1260,8 +1411,8 @@ jobs:
|
||||
conn = oracledb.connect(user='system', password='oracle', dsn='localhost:1521/FREEPDB1')
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(\"\"\"
|
||||
CREATE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 200M AUTOEXTEND ON NEXT 50M
|
||||
CREATE BIGFILE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 2G AUTOEXTEND ON NEXT 500M MAXSIZE UNLIMITED
|
||||
EXTENT MANAGEMENT LOCAL
|
||||
SEGMENT SPACE MANAGEMENT AUTO
|
||||
\"\"\")
|
||||
@@ -1614,8 +1765,8 @@ jobs:
|
||||
conn = oracledb.connect(user='system', password='oracle', dsn='localhost:1521/FREEPDB1')
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(\"\"\"
|
||||
CREATE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 200M AUTOEXTEND ON NEXT 50M
|
||||
CREATE BIGFILE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 2G AUTOEXTEND ON NEXT 500M MAXSIZE UNLIMITED
|
||||
EXTENT MANAGEMENT LOCAL
|
||||
SEGMENT SPACE MANAGEMENT AUTO
|
||||
\"\"\")
|
||||
@@ -1773,8 +1924,8 @@ jobs:
|
||||
conn = oracledb.connect(user='system', password='oracle', dsn='localhost:1521/FREEPDB1')
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(\"\"\"
|
||||
CREATE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 200M AUTOEXTEND ON NEXT 50M
|
||||
CREATE BIGFILE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 2G AUTOEXTEND ON NEXT 500M MAXSIZE UNLIMITED
|
||||
EXTENT MANAGEMENT LOCAL
|
||||
SEGMENT SPACE MANAGEMENT AUTO
|
||||
\"\"\")
|
||||
@@ -2648,6 +2799,43 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/crewai
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-vapi-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-vapi == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build vapi integration
|
||||
working-directory: ./hindsight-integrations/vapi
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/vapi
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/vapi
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-litellm-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -3200,6 +3388,12 @@ jobs:
|
||||
rm -rf node_modules/lightningcss node_modules/@tailwindcss
|
||||
npm install lightningcss @tailwindcss/postcss @tailwindcss/node
|
||||
|
||||
- name: Test Control Plane
|
||||
run: npm test --workspace=hindsight-control-plane
|
||||
|
||||
- name: Check i18n locale parity and hardcoded strings
|
||||
run: npm run i18n:check --workspace=hindsight-control-plane
|
||||
|
||||
- name: Build Control Plane
|
||||
run: npm run build --workspace=hindsight-control-plane
|
||||
|
||||
@@ -3647,10 +3841,14 @@ jobs:
|
||||
- build-chat-integration
|
||||
- test-paperclip-integration
|
||||
- test-pipecat-integration
|
||||
- test-gemini-spark-integration
|
||||
- test-vapi-integration
|
||||
- test-roo-code-integration
|
||||
- build-control-plane
|
||||
- build-docs
|
||||
- test-rust-cli
|
||||
- lint-helm-chart
|
||||
- test-standalone-start-script
|
||||
- build-docker-images
|
||||
- test-api
|
||||
- test-api-oracle
|
||||
|
||||
@@ -54,6 +54,8 @@ hindsight-clients/rust/target
|
||||
!.claude/skills/
|
||||
whats-next.md
|
||||
TASK.md
|
||||
# Parked / draft integrations that aren't ready to ship
|
||||
hindsight-integrations/_drafts/
|
||||
# Changelog is now tracked in hindsight-docs/src/pages/changelog.md
|
||||
# CHANGELOG.md
|
||||
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
# Hindsight with Claude Code (Claude Pro/Max subscription)
|
||||
|
||||
Run Hindsight inside Docker using the `claude-code` LLM provider, backed by
|
||||
your host machine's Claude Pro or Max subscription credentials.
|
||||
|
||||
The standalone Hindsight Docker image ships `claude-agent-sdk` but does **not**
|
||||
bundle the host `claude` CLI binary or any Claude credentials. This Compose
|
||||
file bind-mounts the host's CLI install and credentials into the container so
|
||||
the `claude-code` provider works without an API key.
|
||||
|
||||
## When to use this
|
||||
|
||||
- You have an active Claude Pro or Max subscription and want to use it for
|
||||
Hindsight without paying separate Anthropic API costs.
|
||||
- You want a one-command `docker compose up` instead of a long `docker run`
|
||||
invocation with many flags.
|
||||
- You are running on **Linux/amd64** — macOS Docker Desktop and Windows host
|
||||
paths differ and are not yet covered (please open an issue if you'd like to
|
||||
contribute a verified recipe for either).
|
||||
|
||||
> **Personal-use only.** Anthropic's
|
||||
> [Agent SDK documentation](https://docs.claude.com/en/api/agent-sdk/overview)
|
||||
> states that third-party developers should not offer claude.ai login or rate
|
||||
> limits for their products. Hindsight does **not** perform any login on your
|
||||
> behalf — it uses credentials you've already authenticated via
|
||||
> `claude auth login`. In January 2026, Anthropic
|
||||
> [enforced restrictions](https://paddo.dev/blog/anthropic-walled-garden-crackdown/)
|
||||
> against tools that spoofed the Claude Code client identity; Hindsight uses
|
||||
> the official Claude Agent SDK instead.
|
||||
>
|
||||
> Do not deploy this configuration to shared environments or production. For
|
||||
> that, use the `anthropic` provider with an API key from the
|
||||
> [Anthropic Console](https://console.anthropic.com/). Usage counts against
|
||||
> your Claude Pro/Max subscription limits.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Host has `claude` CLI installed (e.g., `npm install -g @anthropics/claude-code`)
|
||||
and `claude auth login` has been run successfully.
|
||||
- `~/.claude.json` and `~/.claude/.credentials.json` exist on the host.
|
||||
- Host `claude` CLI version is **2.1.128 or newer** — the version bundled with
|
||||
`claude-agent-sdk` 0.5.x has a protocol incompatibility in containers, so
|
||||
the recipe overrides it with the host binary.
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
# Set your host UID/GID (defaults to 1000:1000 if unset)
|
||||
export HOST_UID=$(id -u)
|
||||
export HOST_GID=$(id -g)
|
||||
|
||||
docker compose -f docker/docker-compose/claude-code/docker-compose.yaml up -d
|
||||
```
|
||||
|
||||
- API: http://localhost:8888
|
||||
- Control Plane: http://localhost:9999
|
||||
|
||||
## Post-setup (one-time)
|
||||
|
||||
After the container starts for the first time, run these commands to fix
|
||||
permissions and symlink the host `claude` binary into `$PATH`:
|
||||
|
||||
```bash
|
||||
# Make ~/.claude writable by your UID (the CLI writes session/project state)
|
||||
docker exec --user 0:0 hindsight-claude-code chown $(id -u):$(id -g) /home/hindsight/.claude
|
||||
docker exec --user 0:0 hindsight-claude-code chmod 755 /home/hindsight/.claude
|
||||
|
||||
# Symlink the host claude binary into PATH
|
||||
docker exec --user 0:0 hindsight-claude-code \
|
||||
ln -sf /home/hindsight/.local/share/claude/versions/2.1.128 /usr/local/bin/claude
|
||||
```
|
||||
|
||||
If you set `CLAUDE_CLI_VERSION` to a version other than `2.1.128`, update the
|
||||
symlink path accordingly.
|
||||
|
||||
## Notes on the bind-mount surface (every flag is load-bearing)
|
||||
|
||||
- **Host `claude` binary required** — the image ships only `claude-agent-sdk`,
|
||||
not the CLI itself.
|
||||
- **SDK bundled-binary override** — the override of
|
||||
`claude_agent_sdk/_bundled/claude` works around a protocol issue in the
|
||||
bundled v2.1.121 binary inside containers. Once `claude-agent-sdk` ships
|
||||
with v2.1.128+ this override can be dropped. Set `CLAUDE_CLI_VERSION` to
|
||||
match your installed version.
|
||||
- **Single-file credential mounts** — credentials are mounted as individual
|
||||
`:ro` files rather than a whole-directory `:ro` mount of `~/.claude`,
|
||||
because the CLI writes session/project state at runtime and a read-only
|
||||
directory mount silently breaks it.
|
||||
- **`--user` / `user:`** — the `user: ${HOST_UID}:${HOST_GID}` pattern
|
||||
requires `chmod 755 /home/hindsight`, which is built into the image since
|
||||
v0.6.0 (see [#1481](https://github.com/vectorize-io/hindsight/issues/1481)).
|
||||
- **`~/.hindsight-docker` data directory** — the pg0 data bind mount must be
|
||||
writable by your host UID (see
|
||||
[#1483](https://github.com/vectorize-io/hindsight/issues/1483)).
|
||||
- **Verified** on `linux/amd64` against `ghcr.io/vectorize-io/hindsight:latest`
|
||||
v0.5.6+.
|
||||
|
||||
## Using a different Claude CLI version
|
||||
|
||||
If your host has a `claude` version other than 2.1.128, set
|
||||
`CLAUDE_CLI_VERSION` before starting:
|
||||
|
||||
```bash
|
||||
export CLAUDE_CLI_VERSION=2.2.0
|
||||
docker compose -f docker/docker-compose/claude-code/docker-compose.yaml up -d
|
||||
```
|
||||
|
||||
Then update the post-setup symlink to match:
|
||||
|
||||
```bash
|
||||
docker exec --user 0:0 hindsight-claude-code \
|
||||
ln -sf /home/hindsight/.local/share/claude/versions/2.2.0 /usr/local/bin/claude
|
||||
```
|
||||
@@ -0,0 +1,44 @@
|
||||
name: hindsight-claude-code
|
||||
# Run Hindsight with the claude-code LLM provider, using your host machine's
|
||||
# Claude Pro/Max subscription credentials. Linux/amd64 only for now.
|
||||
#
|
||||
# Quick start:
|
||||
# docker compose -f docker/docker-compose/claude-code/docker-compose.yaml up -d
|
||||
#
|
||||
# See README.md for prerequisites, post-setup steps, and important caveats.
|
||||
|
||||
services:
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:latest
|
||||
container_name: hindsight-claude-code
|
||||
user: "${HOST_UID:-1000}:${HOST_GID:-1000}"
|
||||
ports:
|
||||
- "127.0.0.1:8888:8888"
|
||||
- "127.0.0.1:9999:9999"
|
||||
environment:
|
||||
HOME: /home/hindsight
|
||||
USER: hindsight
|
||||
LOGNAME: hindsight
|
||||
PATH: /usr/local/bin:/usr/bin:/bin:/app/api/.venv/bin
|
||||
HINDSIGHT_API_LLM_PROVIDER: claude-code
|
||||
volumes:
|
||||
# ── Persistent data ────────────────────────────────────────────
|
||||
# Writable pg0 data directory. Must be writable by HOST_UID.
|
||||
- ${HOME:-.}/.hindsight-docker:/home/hindsight/.pg0
|
||||
|
||||
# ── Claude credentials (read-only, single-file mounts) ────────
|
||||
# A whole-directory :ro mount of ~/.claude silently breaks the
|
||||
# CLI, which writes session/project state at runtime — so we
|
||||
# mount only the two credential files.
|
||||
- ${HOME}/.claude/.credentials.json:/home/hindsight/.claude/.credentials.json:ro
|
||||
- ${HOME}/.claude.json:/home/hindsight/.claude.json:ro
|
||||
|
||||
# ── Claude CLI install (read-only) ─────────────────────────────
|
||||
- ${HOME}/.local/share/claude:/home/hindsight/.local/share/claude:ro
|
||||
|
||||
# ── SDK bundled-binary override ────────────────────────────────
|
||||
# The claude-agent-sdk 0.5.x image bundles v2.1.121 which has a
|
||||
# protocol incompatibility in containers. Override it with the
|
||||
# host's v2.1.128+ binary. Drop this mount once claude-agent-sdk
|
||||
# ships with v2.1.128+.
|
||||
- ${HOME}/.local/share/claude/versions/${CLAUDE_CLI_VERSION:-2.1.128}:/app/api/.venv/lib/python3.11/site-packages/claude_agent_sdk/_bundled/claude:ro
|
||||
@@ -0,0 +1,103 @@
|
||||
# Hindsight with a local llama.cpp server sidecar
|
||||
|
||||
Example Docker Compose setup that runs Hindsight against a **local
|
||||
llama.cpp server**, fully offline, with no external API key required.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌────────────┐ HTTP /v1/chat/completions ┌──────────────────────────────┐
|
||||
│ hindsight │ ──────────────────────────▶ │ llama.cpp server (sidecar) │
|
||||
│ (API + CP) │ │ ghcr.io/ggml-org/llama.cpp │
|
||||
└────────────┘ └──────────────────────────────┘
|
||||
```
|
||||
|
||||
`llama.cpp` runs as its own container and exposes an OpenAI-compatible
|
||||
HTTP API. Hindsight talks to it via the standard `openai` LLM provider
|
||||
with `HINDSIGHT_API_LLM_BASE_URL` pointed at the sidecar.
|
||||
|
||||
This pattern follows
|
||||
[*Hosting llama-server with Docker* (ServiceStack)](https://servicestack.net/posts/hosting-llama-server).
|
||||
|
||||
### Why a sidecar and not the in-process `llamacpp` provider?
|
||||
|
||||
Hindsight does ship an in-process `llamacpp` provider that spawns
|
||||
`llama-cpp-python`, but the **published `ghcr.io/vectorize-io/hindsight`
|
||||
image deliberately omits `llama-cpp-python`** to keep the image small and
|
||||
avoid bundling native inference libraries that most users don't need.
|
||||
Trying to set `HINDSIGHT_API_LLM_PROVIDER=llamacpp` against the published
|
||||
image fails with `ModuleNotFoundError: No module named 'llama_cpp'`.
|
||||
|
||||
The sidecar approach side-steps that entirely: the official llama.cpp
|
||||
image is used as-is for inference, Hindsight is used as-is for memory.
|
||||
Clean separation, no derived images.
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
docker compose -f docker/docker-compose/local-llm/docker-compose.yaml up
|
||||
```
|
||||
|
||||
- API: http://localhost:8888
|
||||
- Control Plane: http://localhost:9999
|
||||
|
||||
**First boot downloads ~3.5 GB** (Gemma 4 E2B Q4_K_M GGUF) into the
|
||||
`llama_models` named volume. Subsequent boots reuse it.
|
||||
|
||||
Hindsight only starts after llama.cpp's `/health` endpoint reports
|
||||
healthy, so the API will appear "stuck" for a few minutes on the first
|
||||
run while the model downloads.
|
||||
|
||||
## Using a different model
|
||||
|
||||
Override the HuggingFace repo / file in `docker-compose.yaml`:
|
||||
|
||||
```yaml
|
||||
environment:
|
||||
LLAMA_ARG_HF_REPO: bartowski/Qwen2.5-7B-Instruct-GGUF
|
||||
LLAMA_ARG_HF_FILE: Qwen2.5-7B-Instruct-Q4_K_M.gguf
|
||||
```
|
||||
|
||||
Also update `HINDSIGHT_API_LLM_MODEL` on the `hindsight` service to a
|
||||
matching alias (the value is sent to llama-server as the OpenAI `model`
|
||||
field — llama-server is lenient about this but it shows up in logs).
|
||||
|
||||
## GPU acceleration
|
||||
|
||||
The default compose file targets CPU because not everyone has a GPU. On
|
||||
CPU, Gemma 4 E2B runs at ~2-3 tokens/sec — fine for a smoke test, but the
|
||||
retain pipeline (which makes several multi-hundred-token LLM calls per
|
||||
memory) will time out against Hindsight's default LLM timeout. **For any
|
||||
real use, run on a GPU.**
|
||||
|
||||
### NVIDIA
|
||||
|
||||
1. Switch the `llama` service image from `:server` to `:server-cuda`.
|
||||
2. Uncomment the `LLAMA_ARG_N_GPU_LAYERS: "999"` env var (offload all
|
||||
layers to GPU).
|
||||
3. Uncomment the `deploy.resources.reservations.devices` block.
|
||||
4. Install the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
|
||||
on the host.
|
||||
|
||||
The compose file has all four spots marked with inline comments.
|
||||
|
||||
### Apple Silicon / ROCm / Vulkan
|
||||
|
||||
The official `ghcr.io/ggml-org/llama.cpp` image only ships CPU and CUDA
|
||||
variants. For Metal (Apple Silicon), ROCm (AMD), or Vulkan backends,
|
||||
build llama.cpp yourself with the appropriate flags and reference the
|
||||
image you build instead. Docker Desktop on macOS cannot pass through the
|
||||
host GPU to a Linux container in any case — for Apple Silicon, run
|
||||
llama-server directly on the host and only put Hindsight in Docker.
|
||||
|
||||
## Caveats
|
||||
|
||||
- llama.cpp's HTTP API is OpenAI-compatible but not 100% feature-parity.
|
||||
Function/tool calling support depends on the chat template baked into
|
||||
the GGUF; some retain/reflect flows may behave differently than against
|
||||
a hosted OpenAI model.
|
||||
- Small GGUFs (~3 B params) are useful for smoke testing but will
|
||||
underperform a hosted frontier model on retain quality. Use a larger
|
||||
GGUF (7-13 B params) for production-quality memory.
|
||||
- The `llama_models` named volume persists the GGUF across `docker
|
||||
compose down`/`up` so the model is downloaded once, not every restart.
|
||||
@@ -0,0 +1,74 @@
|
||||
name: hindsight-local-llm
|
||||
# Example: run Hindsight against a local llama.cpp server sidecar — fully
|
||||
# offline, no external API key needed.
|
||||
#
|
||||
# Pattern follows https://servicestack.net/posts/hosting-llama-server :
|
||||
# llama.cpp runs as its own container exposing an OpenAI-compatible HTTP
|
||||
# API, and Hindsight talks to it via the `openai` LLM provider with a
|
||||
# custom `base_url`. This means we can use the published Hindsight image
|
||||
# unchanged — no derived Dockerfile, no `llama-cpp-python` install on top.
|
||||
#
|
||||
# Quick start:
|
||||
# docker compose -f docker/docker-compose/local-llm/docker-compose.yaml up
|
||||
#
|
||||
# First boot downloads the default Gemma 4 E2B GGUF (~3.5 GB) into the
|
||||
# `llama_models` volume; subsequent boots reuse it.
|
||||
|
||||
services:
|
||||
llama:
|
||||
image: ghcr.io/ggml-org/llama.cpp:server
|
||||
container_name: hindsight-local-llm-llama
|
||||
environment:
|
||||
LLAMA_ARG_HOST: 0.0.0.0
|
||||
LLAMA_ARG_PORT: "8080"
|
||||
# Auto-download a small GGUF from HuggingFace on first start.
|
||||
# Override these to use a different model.
|
||||
LLAMA_ARG_HF_REPO: bartowski/google_gemma-4-E2B-it-GGUF
|
||||
LLAMA_ARG_HF_FILE: google_gemma-4-E2B-it-Q4_K_M.gguf
|
||||
LLAMA_ARG_CTX_SIZE: "8192"
|
||||
# Uncomment for NVIDIA GPU (and switch image to :server-cuda):
|
||||
# LLAMA_ARG_N_GPU_LAYERS: "999"
|
||||
volumes:
|
||||
# llama-server stores HuggingFace downloads under ~/.cache/huggingface
|
||||
# (not ~/.cache/llama.cpp), so mount the named volume there to avoid
|
||||
# re-downloading the GGUF on every recreate.
|
||||
- llama_models:/root/.cache/huggingface
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "curl -fsS http://localhost:8080/health || exit 1"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 60
|
||||
start_period: 30s
|
||||
# For NVIDIA GPU acceleration, swap the image above to
|
||||
# `ghcr.io/ggml-org/llama.cpp:server-cuda` and uncomment:
|
||||
# deploy:
|
||||
# resources:
|
||||
# reservations:
|
||||
# devices:
|
||||
# - driver: nvidia
|
||||
# count: all
|
||||
# capabilities: [gpu]
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:latest
|
||||
container_name: hindsight-local-llm
|
||||
depends_on:
|
||||
llama:
|
||||
condition: service_healthy
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# llama-server is OpenAI-compatible, so use the `openai` provider and
|
||||
# point base_url at the sidecar. The API key is unused by llama-server
|
||||
# but Hindsight requires the env var to be set.
|
||||
HINDSIGHT_API_LLM_PROVIDER: openai
|
||||
HINDSIGHT_API_LLM_BASE_URL: http://llama:8080/v1
|
||||
HINDSIGHT_API_LLM_API_KEY: not-needed
|
||||
HINDSIGHT_API_LLM_MODEL: gemma-4-e2b-it
|
||||
volumes:
|
||||
- pg_data:/home/hindsight/.pg0
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
llama_models:
|
||||
@@ -0,0 +1,7 @@
|
||||
# PostgreSQL with pgvector and ParadeDB pg_search extensions.
|
||||
#
|
||||
# The official ParadeDB image ships PostgreSQL with pg_search and pgvector
|
||||
# already installed, so no build steps are required. We pin to the PG17
|
||||
# variant for parity with the other Hindsight docker-compose examples
|
||||
# (vchord, pg_textsearch).
|
||||
FROM paradedb/paradedb:latest-pg17
|
||||
@@ -0,0 +1,96 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and ParadeDB pg_search.
|
||||
#
|
||||
# pg_search is the only BM25 backend supported by Hindsight that works with
|
||||
# Citus, so this is the recommended setup for horizontally scaled deployments.
|
||||
#
|
||||
# Usage:
|
||||
# docker compose -f docker/docker-compose/pg_search/docker-compose.yaml up -d
|
||||
#
|
||||
# Required environment variables:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see the hindsight service)
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER: ParadeDB pg_search
|
||||
# tokenizer for new BM25 indexes (default: empty, uses ParadeDB default)
|
||||
|
||||
services:
|
||||
db:
|
||||
# Use ParadeDB image which bundles pgvector + pg_search
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
ports:
|
||||
- "5437:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
pg-search-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pg_search CASCADE;';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvector
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pg_search
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER: ${HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER:-}
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -0,0 +1,23 @@
|
||||
# PostgreSQL with pgvector and pgroonga extensions.
|
||||
#
|
||||
# pgroonga is a multilingual full-text search extension built on Groonga.
|
||||
# It works out of the box for CJK (Chinese, Japanese, Korean) and other
|
||||
# non-whitespace-segmented languages via the TokenBigram tokenizer.
|
||||
FROM groonga/pgroonga:latest-debian-pg17
|
||||
|
||||
# Install pgvector on top of the pgroonga base image (which already provides
|
||||
# pgroonga and the Groonga library).
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
build-essential \
|
||||
git \
|
||||
postgresql-server-dev-17 \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN cd /tmp && \
|
||||
git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git && \
|
||||
cd pgvector && \
|
||||
make && \
|
||||
make install
|
||||
|
||||
RUN rm -rf /tmp/pgvector && \
|
||||
apt-get purge -y --auto-remove build-essential git postgresql-server-dev-17
|
||||
@@ -0,0 +1,91 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and pgroonga
|
||||
#
|
||||
# pgroonga provides multilingual BM25 indexing that works out of the box for
|
||||
# CJK (Chinese, Japanese, Korean) and other non-whitespace-segmented languages.
|
||||
# Use this recipe if your bank content is not English/European.
|
||||
#
|
||||
# docker compose -f docker/docker-compose/pgroonga/docker-compose.yaml down && \
|
||||
# sleep 2 && \
|
||||
# docker compose -f docker/docker-compose/pgroonga/docker-compose.yaml up -d
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_DB_PASSWORD: PostgreSQL password (default: hindsight_password)
|
||||
|
||||
services:
|
||||
db:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
ports:
|
||||
- "5439:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
pgroonga-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pgroonga CASCADE;';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvector
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pgroonga
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -10,19 +10,45 @@ set -e
|
||||
# loss scenarios where a container restart caused the data directory to be
|
||||
# wiped despite a volume mount being present.
|
||||
# =============================================================================
|
||||
PG0_DATA_DIR="${HOME}/.pg0"
|
||||
if [ -d "$PG0_DATA_DIR" ]; then
|
||||
pg0_has_pg_version() {
|
||||
local pg0_data_dir="$1"
|
||||
|
||||
# pg0 has used more than one on-disk layout. Newer standalone images keep
|
||||
# PostgreSQL data under instances/<name>/data, while older volumes may have
|
||||
# placed PG_VERSION at or one level below the mount.
|
||||
[ -f "$pg0_data_dir/PG_VERSION" ] && return 0
|
||||
compgen -G "$pg0_data_dir"/*/PG_VERSION > /dev/null 2>&1 && return 0
|
||||
compgen -G "$pg0_data_dir"/instances/*/data/PG_VERSION > /dev/null 2>&1 && return 0
|
||||
|
||||
return 1
|
||||
}
|
||||
|
||||
check_pg0_data_integrity() {
|
||||
local pg0_data_dir="$1"
|
||||
|
||||
if [ ! -d "$pg0_data_dir" ]; then
|
||||
return 0
|
||||
fi
|
||||
|
||||
# Look for actual PostgreSQL data directories (pg0 creates subdirs per instance)
|
||||
if compgen -G "$PG0_DATA_DIR"/*/PG_VERSION > /dev/null 2>&1; then
|
||||
echo "✅ Existing pg0 data directory detected at $PG0_DATA_DIR"
|
||||
elif [ "$(ls -A "$PG0_DATA_DIR" 2>/dev/null)" ]; then
|
||||
echo "⚠️ WARNING: pg0 data directory exists at $PG0_DATA_DIR but no PG_VERSION found."
|
||||
if pg0_has_pg_version "$pg0_data_dir"; then
|
||||
echo "✅ Existing pg0 data directory detected at $pg0_data_dir"
|
||||
elif [ "$(ls -A "$pg0_data_dir" 2>/dev/null)" ]; then
|
||||
echo "⚠️ WARNING: pg0 data directory exists at $pg0_data_dir but no PG_VERSION found."
|
||||
echo " This may indicate data corruption or an incomplete previous shutdown."
|
||||
echo " If you see all migrations running from scratch after this, your data may have been lost."
|
||||
echo " See: https://github.com/vectorize-io/hindsight/issues/675"
|
||||
fi
|
||||
|
||||
return 0
|
||||
}
|
||||
|
||||
if [ "${HINDSIGHT_START_ALL_SOURCE_ONLY:-false}" = "true" ]; then
|
||||
return 0 2>/dev/null || exit 0
|
||||
fi
|
||||
|
||||
check_pg0_data_integrity "${HOME}/.pg0"
|
||||
|
||||
# Service flags (default to true if not set)
|
||||
ENABLE_API="${HINDSIGHT_ENABLE_API:-true}"
|
||||
ENABLE_CP="${HINDSIGHT_ENABLE_CP:-true}"
|
||||
@@ -156,7 +182,7 @@ PIDS=()
|
||||
# Start API if enabled
|
||||
if [ "$ENABLE_API" = "true" ]; then
|
||||
cd /app/api
|
||||
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:8888/health}"
|
||||
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:${HINDSIGHT_API_PORT:-8888}/health}"
|
||||
API_STARTUP_WAIT_SECONDS="${HINDSIGHT_API_STARTUP_WAIT_SECONDS:-300}"
|
||||
|
||||
# Run API directly - Python's PYTHONUNBUFFERED=1 handles output buffering
|
||||
|
||||
Executable
+73
@@ -0,0 +1,73 @@
|
||||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
HINDSIGHT_START_ALL_SOURCE_ONLY=true
|
||||
source "$SCRIPT_DIR/start-all.sh"
|
||||
unset HINDSIGHT_START_ALL_SOURCE_ONLY
|
||||
|
||||
TMP_DIR="$(mktemp -d)"
|
||||
trap 'rm -rf "$TMP_DIR"' EXIT
|
||||
|
||||
assert_contains() {
|
||||
local output="$1"
|
||||
local expected="$2"
|
||||
|
||||
if [[ "$output" != *"$expected"* ]]; then
|
||||
echo "Expected output to contain: $expected"
|
||||
echo "Actual output:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
assert_not_contains() {
|
||||
local output="$1"
|
||||
local unexpected="$2"
|
||||
|
||||
if [[ "$output" == *"$unexpected"* ]]; then
|
||||
echo "Expected output not to contain: $unexpected"
|
||||
echo "Actual output:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
assert_empty() {
|
||||
local output="$1"
|
||||
|
||||
if [ -n "$output" ]; then
|
||||
echo "Expected no output, got:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
mkdir -p "$TMP_DIR/empty"
|
||||
assert_empty "$(check_pg0_data_integrity "$TMP_DIR/empty")"
|
||||
|
||||
mkdir -p "$TMP_DIR/direct"
|
||||
touch "$TMP_DIR/direct/PG_VERSION"
|
||||
direct_output="$(check_pg0_data_integrity "$TMP_DIR/direct")"
|
||||
assert_contains "$direct_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$direct_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/legacy/instance"
|
||||
touch "$TMP_DIR/legacy/instance/PG_VERSION"
|
||||
legacy_output="$(check_pg0_data_integrity "$TMP_DIR/legacy")"
|
||||
assert_contains "$legacy_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$legacy_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/nested/instances/hindsight/data"
|
||||
touch "$TMP_DIR/nested/instances/hindsight/data/PG_VERSION"
|
||||
nested_output="$(check_pg0_data_integrity "$TMP_DIR/nested")"
|
||||
assert_contains "$nested_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$nested_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/nonempty/instances/hindsight"
|
||||
touch "$TMP_DIR/nonempty/instances/hindsight/instance.json"
|
||||
nonempty_output="$(check_pg0_data_integrity "$TMP_DIR/nonempty")"
|
||||
assert_contains "$nonempty_output" "WARNING: pg0 data directory exists"
|
||||
|
||||
echo "start-all pg0 integrity checks passed"
|
||||
@@ -1,6 +0,0 @@
|
||||
dependencies:
|
||||
- name: postgresql
|
||||
repository: https://charts.bitnami.com/bitnami
|
||||
version: 15.5.38
|
||||
digest: sha256:f67c7612736803ece8a669f8ca6b0555f3b78557bc0ecb732aa2e43f0df7750d
|
||||
generated: "2025-12-10T17:20:57.058794+01:00"
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.6.2
|
||||
appVersion: "0.6.2"
|
||||
version: 0.7.1
|
||||
appVersion: "0.7.1"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.6.2",
|
||||
"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.6.2"
|
||||
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.6.2",
|
||||
"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.6.2"
|
||||
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.6.2",
|
||||
"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.6.2",
|
||||
"hindsight-api-slim[local-llm]==0.7.1",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -386,7 +386,7 @@ def test_embedded_ui_flag(llm_config):
|
||||
|
||||
# Verify UI is reachable and reports connected dataplane
|
||||
ui_url = client.ui_url
|
||||
assert ui_url, "ui_url should be set"
|
||||
assert isinstance(ui_url, str) and ui_url, "ui_url should be a non-empty string"
|
||||
|
||||
health_url = f"{ui_url}/api/health"
|
||||
with urllib.request.urlopen(health_url, timeout=10) as resp:
|
||||
|
||||
@@ -46,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.6.2"
|
||||
__version__ = "0.7.1"
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
"""Helpers for ParadeDB pg_search index configuration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
PG_SEARCH_TOKENIZER_ENV = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"
|
||||
|
||||
_SIMPLE_TOKENIZERS = {
|
||||
"unicode_words",
|
||||
"simple",
|
||||
"whitespace",
|
||||
"literal",
|
||||
"literal_normalized",
|
||||
"chinese_compatible",
|
||||
"icu",
|
||||
"jieba",
|
||||
"source_code",
|
||||
}
|
||||
|
||||
_TOKENIZER_ALIASES = {
|
||||
"chinese_lindera": "lindera(chinese)",
|
||||
"japanese_lindera": "lindera(japanese)",
|
||||
"korean_lindera": "lindera(korean)",
|
||||
"lindera_chinese": "lindera(chinese)",
|
||||
"lindera_japanese": "lindera(japanese)",
|
||||
"lindera_korean": "lindera(korean)",
|
||||
}
|
||||
|
||||
|
||||
def normalize_pg_search_tokenizer(value: str | None) -> str:
|
||||
"""Validate and normalize a ParadeDB pg_search tokenizer setting.
|
||||
|
||||
Returns an empty string when unset. The returned value is safe to embed after
|
||||
``pdb.`` in a CREATE INDEX expression.
|
||||
"""
|
||||
|
||||
tokenizer = (value or "").strip().lower()
|
||||
if not tokenizer:
|
||||
return ""
|
||||
|
||||
if tokenizer in _TOKENIZER_ALIASES:
|
||||
return _TOKENIZER_ALIASES[tokenizer]
|
||||
|
||||
if tokenizer in _SIMPLE_TOKENIZERS:
|
||||
return tokenizer
|
||||
|
||||
lindera_match = re.fullmatch(r"lindera\((chinese|japanese|korean)\)", tokenizer)
|
||||
if lindera_match:
|
||||
return tokenizer
|
||||
|
||||
ngram_match = re.fullmatch(r"(ngram|edge_ngram)\((\d{1,3}),\s*(\d{1,3})\)", tokenizer)
|
||||
if ngram_match:
|
||||
kind, min_gram, max_gram = ngram_match.groups()
|
||||
min_value = int(min_gram)
|
||||
max_value = int(max_gram)
|
||||
if min_value <= 0 or min_value > max_value:
|
||||
raise ValueError(
|
||||
f"Invalid {PG_SEARCH_TOKENIZER_ENV}: {value!r}. "
|
||||
"ngram and edge_ngram require positive min/max gram sizes with min <= max."
|
||||
)
|
||||
return f"{kind}({min_value},{max_value})"
|
||||
|
||||
raise ValueError(
|
||||
f"Invalid {PG_SEARCH_TOKENIZER_ENV}: {value!r}. "
|
||||
"Supported values are: unicode_words, simple, whitespace, literal, "
|
||||
"literal_normalized, chinese_compatible, icu, jieba, source_code, "
|
||||
"chinese_lindera, japanese_lindera, korean_lindera, or "
|
||||
"lindera(chinese|japanese|korean), ngram(min,max), or edge_ngram(min,max)."
|
||||
)
|
||||
|
||||
|
||||
def pg_search_bm25_columns(
|
||||
key_field: str,
|
||||
text_fields: Sequence[str],
|
||||
tokenizer: str | None,
|
||||
) -> str:
|
||||
"""Build a ParadeDB BM25 column list for CREATE INDEX."""
|
||||
|
||||
normalized = normalize_pg_search_tokenizer(tokenizer)
|
||||
if not normalized:
|
||||
return ", ".join([key_field, *text_fields])
|
||||
|
||||
return ", ".join([key_field, *(f"({field}::pdb.{normalized})" for field in text_fields)])
|
||||
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.engine import Connection
|
||||
@@ -34,7 +35,7 @@ _INDEX_USING_CLAUSES = {
|
||||
"pgvector": "USING hnsw (embedding vector_cosine_ops)",
|
||||
"pgvectorscale": "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)",
|
||||
"pg_diskann": "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)",
|
||||
"vchord": "USING vchordrq (embedding vector_l2_ops)",
|
||||
"vchord": "USING vchordrq (embedding vector_cosine_ops)",
|
||||
"scann": "USING scann (embedding cosine) WITH (mode = 'AUTO')",
|
||||
}
|
||||
|
||||
@@ -46,6 +47,32 @@ _INDEX_TYPE_KEYWORDS = {
|
||||
"scann": "scann",
|
||||
}
|
||||
|
||||
# Per-backend ANN search-time tuning GUCs. Each entry is a tuple of
|
||||
# (guc_name, value) pairs the caller can apply with SET or SET LOCAL.
|
||||
#
|
||||
# - pgvector exposes hnsw.ef_search. The 60 / 200 pair is unchanged from the
|
||||
# pre-dispatcher code (internal benchmarks tuned around our embedding count
|
||||
# and recall floor; see the link_utils / pool init call sites for the
|
||||
# latency-vs-recall framing).
|
||||
# - vchord exposes vchordrq.probes (no default; see VectorChord issue #392)
|
||||
# and vchordrq.epsilon (default 1.9). probes = 10 / 30 are starting
|
||||
# defaults pending a workload-specific sweep — vchordrq's recall curve
|
||||
# shape differs from HNSW's, so the pgvector numbers don't translate
|
||||
# directly. Revisit with a per-cluster benchmark once we have production
|
||||
# recall data; until then these are deliberately conservative on the
|
||||
# high-recall path. We leave epsilon at its default; tightening it is a
|
||||
# separate trade-off.
|
||||
# - pgvectorscale / pg_diskann / scann do not expose an equivalent per-statement
|
||||
# knob in the engine today, so the dispatcher returns no statements for them.
|
||||
_ANN_TUNING_LOW_LATENCY: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (("hnsw.ef_search", "60"),),
|
||||
"vchord": (("vchordrq.probes", "10"),),
|
||||
}
|
||||
_ANN_TUNING_HIGH_RECALL: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (("hnsw.ef_search", "200"),),
|
||||
"vchord": (("vchordrq.probes", "30"),),
|
||||
}
|
||||
|
||||
_EXTENSION_INSTALL_SQL = {
|
||||
"pgvector": ("CREATE EXTENSION IF NOT EXISTS vector",),
|
||||
"pgvectorscale": (
|
||||
@@ -67,6 +94,18 @@ _INSTALL_HINTS = {
|
||||
}
|
||||
|
||||
|
||||
def configured_vector_extension() -> str:
|
||||
"""Return the user-configured vector backend extension.
|
||||
|
||||
Reads ``HINDSIGHT_API_VECTOR_EXTENSION`` (default ``"pgvector"``) and
|
||||
validates it via :func:`validate_extension`. This is the single source of
|
||||
truth for runtime code that needs to dispatch behaviour by vector backend;
|
||||
callers should prefer this over reading the env var directly, so the
|
||||
default value and the lookup mechanism live in one place.
|
||||
"""
|
||||
return validate_extension(os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector"))
|
||||
|
||||
|
||||
def validate_extension(name: str) -> str:
|
||||
"""Return a normalized configurable vector extension name or raise.
|
||||
|
||||
@@ -115,6 +154,25 @@ def should_defer_index_creation(ext: str, row_count: int) -> bool:
|
||||
return minimum_rows > 0 and row_count < minimum_rows
|
||||
|
||||
|
||||
def ann_search_tuning_settings(ext: str, *, kind: str) -> tuple[tuple[str, str], ...]:
|
||||
"""Return per-backend (guc_name, value) pairs for ANN search-time tuning.
|
||||
|
||||
``kind`` is ``"low_latency"`` for retain-side link probing (smaller probe
|
||||
count, lower recall, lower latency) and ``"high_recall"`` for connection
|
||||
init in the pool (larger probe count, higher recall). Callers wrap each
|
||||
pair with ``SET LOCAL`` or ``SET`` themselves so the same dispatcher works
|
||||
for both transaction-scoped and session-scoped use. Returns an empty tuple
|
||||
for backends without an equivalent knob.
|
||||
"""
|
||||
if kind == "low_latency":
|
||||
table = _ANN_TUNING_LOW_LATENCY
|
||||
elif kind == "high_recall":
|
||||
table = _ANN_TUNING_HIGH_RECALL
|
||||
else:
|
||||
raise ValueError(f"Unknown ANN tuning kind: {kind!r}")
|
||||
return table.get(_normalize_resolved(ext), ())
|
||||
|
||||
|
||||
def uses_per_bank_vector_indexes(ext: str) -> bool:
|
||||
"""Return whether the backend should create per-bank partial vector indexes."""
|
||||
return _normalize_resolved(ext) != "scann"
|
||||
|
||||
@@ -268,6 +268,7 @@ async def _run_migration(
|
||||
ensure_text_search_extension(
|
||||
resolved_url,
|
||||
text_search_extension=config.text_search_extension,
|
||||
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
|
||||
@@ -15,6 +15,11 @@ from pgvector.sqlalchemy import Vector
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -83,7 +88,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pg_diskann":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
@@ -91,9 +96,14 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Detect or validate text search extension: 'native', 'vchord', 'pg_textsearch',
|
||||
'pgroonga', or 'pg_search'. Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
|
||||
pgroonga is treated as native here so the initial schema still creates valid
|
||||
tsvector columns. ensure_text_search_extension() at startup converts the
|
||||
schema to pgroonga structures (drops the tsvector column, builds a pgroonga
|
||||
index on the base text column).
|
||||
"""
|
||||
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
@@ -121,14 +131,35 @@ def _detect_text_search_extension() -> str:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_textsearch"
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search — true BM25 over base columns, Citus-compatible.
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_search CASCADE")
|
||||
except Exception:
|
||||
# Extension might already exist or user lacks permissions - verify it exists
|
||||
conn = op.get_bind()
|
||||
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_search'")).fetchone()
|
||||
if not result:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_search"
|
||||
elif text_search_extension == "native":
|
||||
return "native"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# ensure_text_search_extension() at runtime converts to pgroonga.
|
||||
# Treat as native here so the initial schema still creates valid columns.
|
||||
return "native"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. "
|
||||
"Must be 'native', 'vchord', 'pg_textsearch', 'pgroonga', or 'pg_search'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
"""Upgrade schema - create all tables from scratch."""
|
||||
|
||||
@@ -284,8 +315,9 @@ def _pg_upgrade() -> None:
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector bm25_catalog.bm25vector
|
||||
""")
|
||||
elif text_search_ext == "pg_textsearch":
|
||||
# Timescale pg_textsearch: dummy TEXT column for consistency (indexes operate on base columns directly)
|
||||
elif text_search_ext in ("pg_textsearch", "pg_search"):
|
||||
# Timescale pg_textsearch / ParadeDB pg_search: dummy TEXT column for
|
||||
# consistency (indexes operate on base columns directly).
|
||||
op.execute("""
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector TEXT
|
||||
@@ -350,6 +382,17 @@ def _pg_upgrade() -> None:
|
||||
USING bm25(text)
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search BM25 index on (id, text, context). The key_field
|
||||
# reloption is required and must match the table's primary key column.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context"), _pg_search_tokenizer())
|
||||
op.execute(
|
||||
"""
|
||||
CREATE INDEX idx_memory_units_text_search ON memory_units
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""".format(bm25_cols=bm25_cols)
|
||||
)
|
||||
else: # native
|
||||
# Native PostgreSQL GIN index
|
||||
op.execute("""
|
||||
|
||||
+29
@@ -7,6 +7,7 @@ the stored fact text.
|
||||
- vchord: text_signals included in tokenize() at insert time
|
||||
- native: search_vector GENERATED column regenerated to include text_signals
|
||||
- pg_textsearch: no change (index only supports a single base column)
|
||||
- pg_search: BM25 index dropped and recreated to include text_signals
|
||||
|
||||
Revision ID: a2b3c4d5e6f7
|
||||
Revises: z1u2v3w4x5y6
|
||||
@@ -18,6 +19,11 @@ from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a2b3c4d5e6f7"
|
||||
@@ -35,6 +41,10 @@ def _detect_text_search_extension() -> str:
|
||||
return os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
table = f"{schema}memory_units"
|
||||
@@ -62,6 +72,16 @@ def _pg_upgrade() -> None:
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: drop the existing BM25 index and recreate it
|
||||
# to include text_signals alongside text and context.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context", "text_signals"), _pg_search_tokenizer())
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_text_search")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_memory_units_text_search ON {table}
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
|
||||
# vchord: tokenize() call in fact_storage.py is updated to include text_signals at insert time
|
||||
# pg_textsearch: no change — index operates on the base `text` column only
|
||||
@@ -86,6 +106,15 @@ def _pg_downgrade() -> None:
|
||||
CREATE INDEX idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# Restore the original (id, text, context) BM25 index without text_signals.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context"), _pg_search_tokenizer())
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_text_search")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_memory_units_text_search ON {table}
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
|
||||
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS text_signals")
|
||||
|
||||
|
||||
+1
-1
@@ -63,7 +63,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pgvectorscale":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
+106
@@ -0,0 +1,106 @@
|
||||
"""Add graph_maintenance_queue table
|
||||
|
||||
Queue of memory_units whose outgoing temporal/semantic links lost a
|
||||
neighbour to a delete. Drained by the async graph_maintenance worker,
|
||||
which tops the unit's links back up using the same probes retain runs.
|
||||
|
||||
The queue only targets the link-recompute pass. The worker also runs
|
||||
bank-wide sweeps (orphan-entity prune, stale-cooccurrence prune) on each
|
||||
invocation; those don't need per-target queueing.
|
||||
|
||||
Revision ID: b5a4c3e2f1d8
|
||||
Revises: e9b2c7d1f3a4
|
||||
Create Date: 2026-05-27
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b5a4c3e2f1d8"
|
||||
down_revision: str | Sequence[str] | None = "e9b2c7d1f3a4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# Composite PK gives us natural ON CONFLICT DO NOTHING dedup when the same
|
||||
# unit is enqueued from overlapping deletes. No FK to memory_units: if the
|
||||
# unit is deleted between enqueue and drain, the worker observes it's gone
|
||||
# and skips — a cascade would erase the work order, but that work has
|
||||
# already been satisfied (no surviving row to maintain).
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}graph_maintenance_queue (
|
||||
bank_id TEXT NOT NULL,
|
||||
unit_id UUID NOT NULL,
|
||||
enqueued_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
PRIMARY KEY (bank_id, unit_id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_graph_maintenance_queue_bank_enqueued
|
||||
ON {schema}graph_maintenance_queue (bank_id, enqueued_at)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_graph_maintenance_queue_bank_enqueued")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}graph_maintenance_queue")
|
||||
|
||||
|
||||
def _oracle_execute_ignoring_955(sql: str) -> None:
|
||||
"""Run a CREATE statement and swallow ORA-00955 (object already exists).
|
||||
|
||||
Mirrors the helper in the Oracle baseline migration so reruns stay safe
|
||||
on a database where the table was created by an earlier partial run.
|
||||
"""
|
||||
block = (
|
||||
"BEGIN "
|
||||
"EXECUTE IMMEDIATE :stmt; "
|
||||
"EXCEPTION WHEN OTHERS THEN "
|
||||
"IF SQLCODE = -955 THEN NULL; ELSE RAISE; END IF; "
|
||||
"END;"
|
||||
)
|
||||
op.get_bind().exec_driver_sql(block, {"stmt": sql.strip()})
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
_oracle_execute_ignoring_955(
|
||||
"""
|
||||
CREATE TABLE graph_maintenance_queue (
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
unit_id RAW(16) NOT NULL,
|
||||
enqueued_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_graph_maintenance_queue PRIMARY KEY (bank_id, unit_id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
_oracle_execute_ignoring_955(
|
||||
"CREATE INDEX idx_graph_maintenance_queue_bank_enqueued ON graph_maintenance_queue (bank_id, enqueued_at)"
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("DROP INDEX idx_graph_maintenance_queue_bank_enqueued")
|
||||
op.execute("DROP TABLE graph_maintenance_queue")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+152
@@ -0,0 +1,152 @@
|
||||
"""Re-create vchord vector indexes with vector_cosine_ops
|
||||
|
||||
Revision ID: b8c9d0e1f2a3
|
||||
Revises: 86f7a033d372
|
||||
Create Date: 2026-05-20
|
||||
|
||||
vchordrq operator classes are bound 1:1 to operators in PostgreSQL:
|
||||
vector_l2_ops only matches ``<->``, while every Hindsight ANN query uses
|
||||
``<=>`` (cosine distance). The previous vchord mapping used vector_l2_ops,
|
||||
so vchord deployments could never use the index — every ANN query fell
|
||||
back to a sequential scan with per-row cosine computation.
|
||||
|
||||
This migration finds any vchordrq index built with vector_l2_ops in the
|
||||
target schema and re-creates it with vector_cosine_ops, using
|
||||
``CREATE INDEX CONCURRENTLY`` so it can run online. It is a no-op when:
|
||||
|
||||
* the configured vector extension is not vchord, or
|
||||
* no matching indexes exist (already on cosine ops).
|
||||
|
||||
Only PostgreSQL is affected; the Oracle 23ai dialect uses its own native
|
||||
vector index and does not depend on this mapping.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
from hindsight_api._vector_index import configured_vector_extension
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b8c9d0e1f2a3"
|
||||
down_revision: str | Sequence[str] | None = "86f7a033d372"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _rebuild_vchordrq_indexes(old_ops: str, new_ops: str) -> None:
|
||||
"""Rebuild vchordrq indexes using ``old_ops`` so they use ``new_ops``.
|
||||
|
||||
Each index is rebuilt with CREATE INDEX CONCURRENTLY under a fresh name,
|
||||
then the old index is dropped and the new one renamed to take its place.
|
||||
Must be called inside an ``autocommit_block()`` because CONCURRENTLY
|
||||
cannot run inside a transaction.
|
||||
"""
|
||||
bind = op.get_bind()
|
||||
# `or None` collapses both unset and explicit empty-string Alembic options
|
||||
# into NULL so the COALESCE below falls back to current_schema() in either
|
||||
# case. Without it, an empty-string option would filter on `schemaname = ''`
|
||||
# and skip every real schema.
|
||||
target_schema = context.config.get_main_option("target_schema") or None
|
||||
prefix = _pg_schema_prefix()
|
||||
|
||||
rows = bind.execute(
|
||||
text(
|
||||
"SELECT indexname, indexdef FROM pg_indexes "
|
||||
"WHERE schemaname = COALESCE(:target_schema, current_schema()) "
|
||||
"AND indexdef ILIKE '%vchordrq%' "
|
||||
"AND indexdef ILIKE :ops_like"
|
||||
),
|
||||
{"target_schema": target_schema, "ops_like": f"%{old_ops}%"},
|
||||
).fetchall()
|
||||
|
||||
for idx_name, indexdef in rows:
|
||||
# pg_get_indexdef() emits the canonical form `CREATE INDEX <name> ON …`,
|
||||
# so <name> is the first textual occurrence — both substitutions below
|
||||
# rely on that.
|
||||
new_def = indexdef.replace(old_ops, new_ops, 1)
|
||||
temp_name = f"{idx_name}__opclass_swap"
|
||||
new_def = new_def.replace(idx_name, temp_name, 1)
|
||||
new_def = re.sub(
|
||||
r"^CREATE\s+INDEX\b",
|
||||
"CREATE INDEX CONCURRENTLY IF NOT EXISTS",
|
||||
new_def,
|
||||
count=1,
|
||||
)
|
||||
|
||||
# CREATE INDEX CONCURRENTLY can leave the partial index as INVALID if a
|
||||
# previous run errored (disk pressure, lock conflict, signal). Without
|
||||
# this drop the CONCURRENTLY IF NOT EXISTS below would skip creation,
|
||||
# then we'd drop the original and rename the broken index into its
|
||||
# place — silently restoring the seq-scan bug this migration fixes.
|
||||
op.execute(f'DROP INDEX IF EXISTS {prefix}"{temp_name}"')
|
||||
op.execute(new_def)
|
||||
|
||||
# Even on a clean run CONCURRENTLY can finish with indisvalid = false
|
||||
# (e.g. constraint violation during the second build scan). Refuse to
|
||||
# promote in that case so we never alias an INVALID index over a working
|
||||
# one.
|
||||
is_valid = bind.execute(
|
||||
text(
|
||||
"SELECT i.indisvalid "
|
||||
"FROM pg_class c "
|
||||
"JOIN pg_index i ON c.oid = i.indexrelid "
|
||||
"JOIN pg_namespace n ON c.relnamespace = n.oid "
|
||||
"WHERE c.relname = :name "
|
||||
" AND n.nspname = COALESCE(:target_schema, current_schema())"
|
||||
),
|
||||
{"name": temp_name, "target_schema": target_schema},
|
||||
).scalar()
|
||||
if not is_valid:
|
||||
raise RuntimeError(
|
||||
f"vchordrq index rebuild produced an INVALID index ({temp_name}); "
|
||||
"drop it manually and re-run the migration."
|
||||
)
|
||||
|
||||
# DROP + RENAME atomically. A crash between the two would leave
|
||||
# `temp_name` as a valid orphan and the canonical name missing —
|
||||
# next run's `pg_indexes` filter (looking for vector_l2_ops) wouldn't
|
||||
# find anything to recover from, so the index would stay gone. PG
|
||||
# runs the DO block in its own server-side transaction, so either
|
||||
# both succeed or both roll back.
|
||||
op.execute(
|
||||
f"""
|
||||
DO $$
|
||||
BEGIN
|
||||
DROP INDEX IF EXISTS {prefix}"{idx_name}";
|
||||
ALTER INDEX {prefix}"{temp_name}" RENAME TO "{idx_name}";
|
||||
END $$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if configured_vector_extension() != "vchord":
|
||||
return
|
||||
with op.get_context().autocommit_block():
|
||||
_rebuild_vchordrq_indexes("vector_l2_ops", "vector_cosine_ops")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
if configured_vector_extension() != "vchord":
|
||||
return
|
||||
with op.get_context().autocommit_block():
|
||||
_rebuild_vchordrq_indexes("vector_cosine_ops", "vector_l2_ops")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+45
@@ -0,0 +1,45 @@
|
||||
"""Merge graph_maintenance_queue and vchord_cosine_opclass heads.
|
||||
|
||||
Revision ID: c1d2e3f4a5b6
|
||||
Revises: b5a4c3e2f1d8, b8c9d0e1f2a3
|
||||
Create Date: 2026-05-29
|
||||
|
||||
PRs #1668 (vchord cosine opclass) and #1772 (async link recompute) both
|
||||
branched off the same parent and were merged onto main without rebasing,
|
||||
leaving two parallel Alembic heads. This is a structural merge revision
|
||||
with no schema changes — its only job is to unify the DAG so
|
||||
``alembic upgrade head`` is unambiguous again.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c1d2e3f4a5b6"
|
||||
down_revision: str | Sequence[str] | None = ("b5a4c3e2f1d8", "b8c9d0e1f2a3")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+1
-1
@@ -55,7 +55,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pgvectorscale":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
+132
@@ -0,0 +1,132 @@
|
||||
"""Drop indexes that are unused or redundant with composite indexes.
|
||||
|
||||
Code audit identified the following indexes as either dead (no code path
|
||||
exercises them) or fully covered by composite indexes the planner already
|
||||
prefers:
|
||||
|
||||
memory_links:
|
||||
1. idx_memory_links_entity_covering — entity co-occurrence expansion was
|
||||
rewritten to traverse unit_entities instead of memory_links, so no code
|
||||
path filters memory_links on (link_type = 'entity').
|
||||
2. idx_memory_links_from_unit — redundant. idx_memory_links_from_type_weight
|
||||
(from_unit_id, link_type, weight DESC) leads with the same column and
|
||||
answers every from_unit_id = X query.
|
||||
3. idx_memory_links_to_unit — redundant. idx_memory_links_to_type_weight
|
||||
(to_unit_id, link_type, weight DESC) leads with the same column.
|
||||
4. idx_memory_links_link_type — no application query filters on link_type
|
||||
alone; the composite indexes above serve every (from/to + link_type)
|
||||
predicate.
|
||||
|
||||
entities:
|
||||
5. idx_entities_canonical_name — superseded by
|
||||
entities_canonical_name_lower_trgm_idx (case-insensitive lookups).
|
||||
6. entities_canonical_name_trgm_idx — superseded by the lowercase variant
|
||||
in migration 2eee35aa3cfc, but the original was never dropped on schemas
|
||||
that ran the prior migration.
|
||||
|
||||
documents:
|
||||
7. idx_documents_retain_params — GIN index on retain_params JSONB; no query
|
||||
uses jsonb containment on this column.
|
||||
8. idx_documents_content_hash — content-hash lookups happen on the chunks
|
||||
table (chunks.content_hash, indexed separately).
|
||||
|
||||
unit_entities:
|
||||
9. idx_unit_entities_entity — defensive drop. Migration h3i4j5k6l7m8 already
|
||||
issues DROP INDEX IF EXISTS for this; this re-runs the drop idempotently
|
||||
to cover any schema that missed the previous migration.
|
||||
|
||||
All drops use CONCURRENTLY + IF EXISTS so they neither block writers nor
|
||||
fail on schemas where the index is already gone.
|
||||
|
||||
Revision ID: e1b2c3d4f5a6
|
||||
Revises: p4q5r6s7t8u9
|
||||
Create Date: 2026-05-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e1b2c3d4f5a6"
|
||||
down_revision: str | Sequence[str] | None = "p4q5r6s7t8u9"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
_PG_INDEXES_TO_DROP: tuple[str, ...] = (
|
||||
"idx_memory_links_entity_covering",
|
||||
"idx_memory_links_from_unit",
|
||||
"idx_memory_links_to_unit",
|
||||
"idx_memory_links_link_type",
|
||||
"idx_entities_canonical_name",
|
||||
"entities_canonical_name_trgm_idx",
|
||||
"idx_documents_retain_params",
|
||||
"idx_documents_content_hash",
|
||||
"idx_unit_entities_entity",
|
||||
)
|
||||
|
||||
|
||||
def _schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _schema_prefix()
|
||||
# DROP INDEX CONCURRENTLY cannot run inside a transaction block; commit
|
||||
# the Alembic transaction and issue each statement in its own implicit
|
||||
# autocommit transaction. IF EXISTS makes each statement idempotent
|
||||
# across schemas that already dropped (or never had) the index.
|
||||
for index_name in _PG_INDEXES_TO_DROP:
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}{index_name}")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _schema_prefix()
|
||||
|
||||
# Recreate the dropped indexes in the same shape the prior migrations used,
|
||||
# so a downgrade leaves the schema in the state the previous head expected.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity_covering "
|
||||
f"ON {schema}memory_links(from_unit_id) "
|
||||
f"INCLUDE (to_unit_id, entity_id) "
|
||||
f"WHERE link_type = 'entity'"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_from_unit ON {schema}memory_links(from_unit_id)"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_unit ON {schema}memory_links(to_unit_id)")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_link_type ON {schema}memory_links(link_type)")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_entities_canonical_name ON {schema}entities(canonical_name)"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_documents_retain_params "
|
||||
f"ON {schema}documents USING GIN (retain_params)"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_documents_content_hash ON {schema}documents(content_hash)")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_unit_entities_entity ON {schema}unit_entities(entity_id)")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+89
@@ -0,0 +1,89 @@
|
||||
"""Drop materialized entity rows from memory_links.
|
||||
|
||||
Entity edges are no longer stored in ``memory_links``. The /graph endpoint
|
||||
derives them on demand from ``unit_entities``, and recall already used the
|
||||
``unit_entities`` self-join. Storing entity rows duplicated state we never
|
||||
read from the link table — on a 10k-unit benchmark bank, entity rows were
|
||||
53% of all link rows (~190 MB after indexes) and recall never touched them.
|
||||
|
||||
This migration deletes ``memory_links`` rows with ``link_type = 'entity'``.
|
||||
``idx_memory_links_entity_covering`` was already dropped by migration
|
||||
``e1b2c3d4f5a6``; we still issue ``DROP INDEX IF EXISTS`` defensively in case
|
||||
this migration runs against an older snapshot that predates that one.
|
||||
|
||||
Revision ID: e9b2c7d1f3a4
|
||||
Revises: e1b2c3d4f5a6
|
||||
Create Date: 2026-05-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e9b2c7d1f3a4"
|
||||
down_revision: str | Sequence[str] | None = "e1b2c3d4f5a6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
|
||||
# Drop the partial covering index first so the bulk DELETE doesn't churn it.
|
||||
# CREATE/DROP INDEX CONCURRENTLY must run outside a transaction block.
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
|
||||
|
||||
# Delete entity rows. Chunked to keep individual transactions small on
|
||||
# large banks (the perf-medium bench had ~345k entity rows; production
|
||||
# banks can be much larger).
|
||||
op.execute(
|
||||
f"""
|
||||
DO $$
|
||||
DECLARE
|
||||
deleted INTEGER;
|
||||
BEGIN
|
||||
LOOP
|
||||
DELETE FROM {schema}memory_links
|
||||
WHERE ctid IN (
|
||||
SELECT ctid FROM {schema}memory_links
|
||||
WHERE link_type = 'entity'
|
||||
LIMIT 50000
|
||||
);
|
||||
GET DIAGNOSTICS deleted = ROW_COUNT;
|
||||
EXIT WHEN deleted = 0;
|
||||
COMMIT;
|
||||
END LOOP;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# Cannot reconstruct deleted entity links — the writer was path-dependent
|
||||
# on retain order. New retains will not produce entity rows either, so the
|
||||
# partial index would stay empty. Leave both no-op.
|
||||
pass
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
op.execute("DELETE FROM memory_links WHERE link_type = 'entity'")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+58
-4
@@ -16,6 +16,11 @@ from collections.abc import Sequence
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -87,7 +92,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pg_diskann":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
@@ -95,9 +100,15 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Detect or validate text search extension: 'native', 'vchord', 'pg_textsearch',
|
||||
'pgroonga', or 'pg_search'. Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
|
||||
pgroonga is treated as native here so this migration still creates valid
|
||||
tsvector columns; ensure_text_search_extension() at startup converts the
|
||||
reflections table (renamed from pinned_reflections in p1k2l3m4n5o6) to
|
||||
pgroonga structures. The learnings table is dropped in p1k2l3m4n5o6 so its
|
||||
transient native-style column never reaches steady state.
|
||||
"""
|
||||
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
@@ -125,14 +136,33 @@ def _detect_text_search_extension() -> str:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_textsearch"
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search — true BM25 over base columns, Citus-compatible.
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_search CASCADE")
|
||||
except Exception:
|
||||
conn = op.get_bind()
|
||||
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_search'")).fetchone()
|
||||
if not result:
|
||||
raise
|
||||
return "pg_search"
|
||||
elif text_search_extension == "native":
|
||||
return "native"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# Treat as native here; ensure_text_search_extension() converts the
|
||||
# reflections table to pgroonga structures at runtime.
|
||||
return "native"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. "
|
||||
"Must be 'native', 'vchord', 'pg_textsearch', 'pgroonga', or 'pg_search'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
"""Create learnings and pinned_reflections tables."""
|
||||
schema = _get_schema_prefix()
|
||||
@@ -200,6 +230,18 @@ def _pg_upgrade() -> None:
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25(text) WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: dummy TEXT column; BM25 index is built directly over (id, text)
|
||||
# with key_field='id' (matches the table's primary key).
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text",), _pg_search_tokenizer())
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}learnings ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
@@ -264,6 +306,18 @@ def _pg_upgrade() -> None:
|
||||
USING bm25(content)
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: dummy TEXT column; BM25 index over (id, name, content)
|
||||
# with key_field='id'.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("name", "content"), _pg_search_tokenizer())
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
|
||||
+170
@@ -0,0 +1,170 @@
|
||||
"""Drop GENERATED expression on tsvector search_vector columns.
|
||||
|
||||
The search_vector tsvector column was originally GENERATED ALWAYS with a
|
||||
hardcoded ``to_tsvector('english', ...)`` expression. To support configurable
|
||||
``HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE``, we convert it to a
|
||||
regular tsvector column that the application populates at INSERT time via
|
||||
``to_tsvector($lang, ...)``.
|
||||
|
||||
Existing rows retain their English-derived lexemes — switching the configured
|
||||
language only affects newly-written rows. Users who need to backfill existing
|
||||
rows in a different language can run an admin UPDATE after this migration.
|
||||
|
||||
Only the ``native`` text-search backend is affected. ``vchord``, ``pg_textsearch``,
|
||||
and ``pgroonga`` use other column types or no column at all.
|
||||
|
||||
Revision ID: p4q5r6s7t8u9
|
||||
Revises: 86f7a033d372
|
||||
Create Date: 2026-05-08
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import Connection, text
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "p4q5r6s7t8u9"
|
||||
down_revision: str | Sequence[str] | None = "86f7a033d372"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _TsvectorTableSpec:
|
||||
"""Native-backend tsvector table targeted by this migration.
|
||||
|
||||
``upgrade`` is a one-way DROP EXPRESSION; ``downgrade`` re-attaches the
|
||||
original GENERATED expression so the schema returns to the state created
|
||||
by the initial migration (and a2b3c4d5e6f7_add_text_signals_column for
|
||||
memory_units).
|
||||
"""
|
||||
|
||||
table: str
|
||||
generated_expression: str
|
||||
|
||||
|
||||
# Tables that may have a GENERATED tsvector ``search_vector`` column under the
|
||||
# native backend. Note: the ``learnings`` table was dropped in
|
||||
# p1k2l3m4n5o6_new_knowledge_architecture and ``pinned_reflections`` was renamed
|
||||
# to ``reflections`` in the same migration.
|
||||
_NATIVE_TSVECTOR_TABLES: tuple[_TsvectorTableSpec, ...] = (
|
||||
_TsvectorTableSpec(
|
||||
table="memory_units",
|
||||
generated_expression=(
|
||||
"to_tsvector('english', COALESCE(text, '') || ' ' || "
|
||||
"COALESCE(context, '') || ' ' || COALESCE(text_signals, ''))"
|
||||
),
|
||||
),
|
||||
_TsvectorTableSpec(
|
||||
table="reflections",
|
||||
generated_expression="to_tsvector('english', COALESCE(name, '') || ' ' || content)",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _is_generated_tsvector(conn: Connection, schema: str, table: str) -> bool:
|
||||
"""Return True iff ``schema.table.search_vector`` is a GENERATED tsvector column."""
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT is_generated, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table
|
||||
AND column_name = 'search_vector'
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
is_generated, udt_name = row[0], row[1]
|
||||
return is_generated == "ALWAYS" and udt_name == "tsvector"
|
||||
|
||||
|
||||
def _is_regular_tsvector(conn: Connection, schema: str, table: str) -> bool:
|
||||
"""Return True iff ``schema.table.search_vector`` is a non-generated tsvector column."""
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT is_generated, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table
|
||||
AND column_name = 'search_vector'
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
is_generated, udt_name = row[0], row[1]
|
||||
return udt_name == "tsvector" and is_generated != "ALWAYS"
|
||||
|
||||
|
||||
def _table_exists(conn: Connection, schema: str, table: str) -> bool:
|
||||
return bool(
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT 1 FROM information_schema.tables
|
||||
WHERE table_schema = :schema AND table_name = :table
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
)
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema_prefix = _schema_prefix()
|
||||
schema_name = (context.config.get_main_option("target_schema") or "public").strip('"')
|
||||
conn = op.get_bind()
|
||||
|
||||
for spec in _NATIVE_TSVECTOR_TABLES:
|
||||
if not _table_exists(conn, schema_name, spec.table):
|
||||
continue
|
||||
if not _is_generated_tsvector(conn, schema_name, spec.table):
|
||||
# Either the column doesn't exist (non-native backend) or it's
|
||||
# already a regular tsvector — nothing to do.
|
||||
continue
|
||||
op.execute(f"ALTER TABLE {schema_prefix}{spec.table} ALTER COLUMN search_vector DROP EXPRESSION")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema_prefix = _schema_prefix()
|
||||
schema_name = (context.config.get_main_option("target_schema") or "public").strip('"')
|
||||
conn = op.get_bind()
|
||||
|
||||
for spec in _NATIVE_TSVECTOR_TABLES:
|
||||
if not _table_exists(conn, schema_name, spec.table):
|
||||
continue
|
||||
# Only restore the GENERATED expression if a non-generated tsvector
|
||||
# column exists — otherwise the table is on a different backend.
|
||||
if not _is_regular_tsvector(conn, schema_name, spec.table):
|
||||
continue
|
||||
# Drop and recreate to re-attach the GENERATED expression. Index will be
|
||||
# recreated by re-running ensure_text_search_extension on next startup.
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema_prefix}idx_{spec.table}_text_search")
|
||||
op.execute(f"ALTER TABLE {schema_prefix}{spec.table} DROP COLUMN search_vector")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema_prefix}{spec.table} "
|
||||
f"ADD COLUMN search_vector tsvector GENERATED ALWAYS AS ({spec.generated_expression}) STORED"
|
||||
)
|
||||
op.execute(f"CREATE INDEX idx_{spec.table}_text_search ON {schema_prefix}{spec.table} USING gin(search_vector)")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
@@ -15,6 +15,7 @@ from datetime import datetime, timezone
|
||||
from typing import Any, Literal
|
||||
|
||||
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Query, Request, UploadFile
|
||||
from fastapi.middleware.gzip import GZipMiddleware
|
||||
|
||||
from hindsight_api.engine.audit import AuditEntry, AuditLogger
|
||||
from hindsight_api.extensions import AuthenticationError
|
||||
@@ -151,7 +152,11 @@ class RecallRequest(BaseModel):
|
||||
max_tokens: int = 4096
|
||||
trace: bool = False
|
||||
query_timestamp: str | None = Field(
|
||||
default=None, description="ISO format date string (e.g., '2023-05-30T23:40:00')"
|
||||
default=None,
|
||||
description=(
|
||||
"ISO format date string (e.g., '2023-05-30T23:40:00'). Used as the query-time anchor for "
|
||||
"relative temporal expressions and recency scoring."
|
||||
),
|
||||
)
|
||||
include: IncludeOptions = FieldWithDefault(
|
||||
IncludeOptions,
|
||||
@@ -466,6 +471,13 @@ class MemoryItem(BaseModel):
|
||||
description="Optional tags for visibility scoping. Memories with tags can be filtered during recall.",
|
||||
)
|
||||
|
||||
@field_validator("content")
|
||||
@classmethod
|
||||
def validate_content(cls, v: str) -> str:
|
||||
if not v.strip():
|
||||
raise ValueError("content cannot be empty")
|
||||
return v
|
||||
|
||||
@field_validator("tags", mode="before")
|
||||
@classmethod
|
||||
def coerce_tags(cls, v):
|
||||
@@ -2180,6 +2192,19 @@ class OperationResponse(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class ConsolidationRequest(BaseModel):
|
||||
"""Request model for consolidation trigger endpoint."""
|
||||
|
||||
observation_scopes: list[list[str]] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Optional list of tag scopes to consolidate. Each scope is a list of tags. "
|
||||
"Only unconsolidated memories whose tags contain all tags in at least one scope "
|
||||
"will be processed. If omitted, all unconsolidated memories are processed."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ConsolidationResponse(BaseModel):
|
||||
"""Response model for consolidation trigger endpoint."""
|
||||
|
||||
@@ -2649,6 +2674,7 @@ def create_app(
|
||||
tenant_extension=memory._tenant_extension,
|
||||
max_slots=config.worker_max_slots,
|
||||
slot_reservations=config.worker_slot_reservations,
|
||||
consolidation_bank_priority=config.worker_consolidation_bank_priority or None,
|
||||
)
|
||||
poller_task = asyncio.create_task(poller.run())
|
||||
logging.info(f"Worker poller started (worker_id={worker_id})")
|
||||
@@ -2721,6 +2747,8 @@ def create_app(
|
||||
app.state.memory = memory
|
||||
app.state.audit_logger = memory.audit_logger
|
||||
|
||||
app.add_middleware(GZipMiddleware, minimum_size=1024)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Patch OpenAPI schema: align ValidationError with Pydantic v2 error format
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -2885,6 +2913,57 @@ def _register_routes(app: FastAPI):
|
||||
api_key = authorization.strip()
|
||||
return RequestContext(api_key=api_key)
|
||||
|
||||
def precheck_for(operation: str):
|
||||
"""
|
||||
Build a FastAPI dependency that runs ``OperationValidator.precheck``.
|
||||
|
||||
FastAPI resolves dependencies before deserialising the route's body
|
||||
parameter. Wiring this dependency on the billable POST routes lets
|
||||
an extension reject a request — e.g. with HTTP 402 when a tenant's
|
||||
balance is exhausted — without the request body ever being read or
|
||||
materialised in memory.
|
||||
|
||||
The dependency intentionally:
|
||||
- authenticates the tenant (so ``request_context.tenant_id`` is
|
||||
resolved before the precheck runs);
|
||||
- falls through silently when no validator is configured or the
|
||||
validator's default no-op precheck is in effect;
|
||||
- converts a rejection ``ValidationResult`` into the corresponding
|
||||
``HTTPException`` directly (the per-route ``OperationValidationError``
|
||||
catch blocks don't see exceptions raised in dependencies, so we
|
||||
translate here instead of relying on each handler's try/except).
|
||||
|
||||
Args:
|
||||
operation: Short identifier for the route, e.g. ``"retain"``.
|
||||
|
||||
Returns:
|
||||
A FastAPI dependency callable suitable for ``Depends(...)``.
|
||||
"""
|
||||
|
||||
async def _precheck_dep(
|
||||
bank_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
) -> None:
|
||||
validator = getattr(app.state.memory, "_operation_validator", None)
|
||||
if validator is None:
|
||||
return
|
||||
from hindsight_api.extensions import PrecheckContext
|
||||
|
||||
await app.state.memory._authenticate_tenant(request_context)
|
||||
ctx = PrecheckContext(
|
||||
operation=operation,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
result = await validator.precheck(ctx)
|
||||
if not result.allowed:
|
||||
raise HTTPException(
|
||||
status_code=result.status_code,
|
||||
detail=result.reason or "Operation not allowed",
|
||||
)
|
||||
|
||||
return _precheck_dep
|
||||
|
||||
# Global exception handler for authentication errors
|
||||
@app.exception_handler(AuthenticationError)
|
||||
async def authentication_error_handler(request, exc: AuthenticationError):
|
||||
@@ -3142,7 +3221,10 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
@audited("recall")
|
||||
async def api_recall(
|
||||
bank_id: str, request: RecallRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
bank_id: str,
|
||||
request: RecallRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("recall")),
|
||||
):
|
||||
"""Run a recall and return results with trace."""
|
||||
import time
|
||||
@@ -3330,7 +3412,10 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
@audited("reflect")
|
||||
async def api_reflect(
|
||||
bank_id: str, request: ReflectRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
bank_id: str,
|
||||
request: ReflectRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("reflect")),
|
||||
):
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
@@ -3828,6 +3913,7 @@ def _register_routes(app: FastAPI):
|
||||
bank_id: str,
|
||||
body: CreateMentalModelRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("mental_model_create")),
|
||||
):
|
||||
"""Create a mental model (async - returns operation_id)."""
|
||||
try:
|
||||
@@ -3876,6 +3962,7 @@ def _register_routes(app: FastAPI):
|
||||
bank_id: str,
|
||||
mental_model_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("mental_model_refresh")),
|
||||
):
|
||||
"""Refresh a mental model by re-running its source query (async)."""
|
||||
try:
|
||||
@@ -3902,6 +3989,48 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post(
|
||||
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}/clear",
|
||||
response_model=MentalModelResponse,
|
||||
summary="Clear mental model content",
|
||||
description=(
|
||||
"Clear a mental model's content so the next refresh performs a full re-synthesis. "
|
||||
"This is useful for delta-mode models that have accumulated drift over many "
|
||||
"incremental refreshes. After clearing, call the /refresh endpoint to trigger "
|
||||
"a clean full rebuild."
|
||||
),
|
||||
operation_id="clear_mental_model",
|
||||
tags=["Mental Models"],
|
||||
)
|
||||
@audited("clear_mental_model", request_param=None)
|
||||
async def api_clear_mental_model(
|
||||
bank_id: str,
|
||||
mental_model_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Clear a mental model's content."""
|
||||
try:
|
||||
mental_model = await app.state.memory.clear_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
if mental_model is None:
|
||||
raise HTTPException(status_code=404, detail=f"Mental model '{mental_model_id}' not found")
|
||||
return MentalModelResponse(**mental_model)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(
|
||||
f"Error in POST /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/clear: {error_detail}"
|
||||
)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.patch(
|
||||
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}",
|
||||
response_model=MentalModelResponse,
|
||||
@@ -5395,11 +5524,20 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="trigger_consolidation",
|
||||
tags=["Banks"],
|
||||
)
|
||||
@audited("consolidation", request_param=None)
|
||||
async def api_trigger_consolidation(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
@audited("consolidation")
|
||||
async def api_trigger_consolidation(
|
||||
bank_id: str,
|
||||
request: ConsolidationRequest | None = None,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Trigger consolidation for a bank (async)."""
|
||||
try:
|
||||
result = await app.state.memory.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
|
||||
observation_scopes = request.observation_scopes if request else None
|
||||
result = await app.state.memory.submit_async_consolidation(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
observation_scopes=observation_scopes,
|
||||
)
|
||||
return ConsolidationResponse(
|
||||
operation_id=result["operation_id"],
|
||||
deduplicated=result.get("deduplicated", False),
|
||||
@@ -5722,7 +5860,10 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
@audited("retain")
|
||||
async def api_retain(
|
||||
bank_id: str, request: RetainRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
bank_id: str,
|
||||
request: RetainRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("retain")),
|
||||
):
|
||||
"""Retain memories with optional async processing."""
|
||||
metrics = get_metrics_collector()
|
||||
@@ -5807,9 +5948,8 @@ def _register_routes(app: FastAPI):
|
||||
strategy=group_strategy,
|
||||
request_context=request_context,
|
||||
return_usage=True,
|
||||
outbox_callback=app.state.memory._build_retain_outbox_callback(
|
||||
outbox_callback_factory=app.state.memory._build_retain_outbox_callback_factory(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
operation_id=None,
|
||||
schema=_current_schema.get(),
|
||||
),
|
||||
@@ -5892,6 +6032,7 @@ def _register_routes(app: FastAPI):
|
||||
files: list[UploadFile] = File(..., description="Files to upload and convert"),
|
||||
request: str = Form(..., description="JSON string with FileRetainRequest model"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("files_retain")),
|
||||
):
|
||||
"""Upload and convert files to memories."""
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
@@ -107,6 +107,7 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
|
||||
@@ -7,6 +7,7 @@ All environment variables and their defaults are defined here.
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import dataclass, field, fields
|
||||
from datetime import datetime, timezone
|
||||
@@ -14,6 +15,7 @@ from typing import Any, Literal
|
||||
|
||||
from dotenv import find_dotenv, load_dotenv
|
||||
|
||||
from ._pg_search import normalize_pg_search_tokenizer
|
||||
from ._vector_index import validate_extension
|
||||
from .utils import mask_network_location
|
||||
|
||||
@@ -136,6 +138,7 @@ ENV_LLM_MAX_RETRIES = "HINDSIGHT_API_LLM_MAX_RETRIES"
|
||||
ENV_LLM_INITIAL_BACKOFF = "HINDSIGHT_API_LLM_INITIAL_BACKOFF"
|
||||
ENV_LLM_MAX_BACKOFF = "HINDSIGHT_API_LLM_MAX_BACKOFF"
|
||||
ENV_LLM_TIMEOUT = "HINDSIGHT_API_LLM_TIMEOUT"
|
||||
ENV_LLM_REASONING_EFFORT = "HINDSIGHT_API_LLM_REASONING_EFFORT"
|
||||
ENV_LLM_GROQ_SERVICE_TIER = "HINDSIGHT_API_LLM_GROQ_SERVICE_TIER"
|
||||
ENV_LLM_OPENAI_SERVICE_TIER = "HINDSIGHT_API_LLM_OPENAI_SERVICE_TIER"
|
||||
ENV_LLM_EXTRA_BODY = "HINDSIGHT_API_LLM_EXTRA_BODY"
|
||||
@@ -200,6 +203,7 @@ ENV_EMBEDDINGS_OPENAI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY"
|
||||
ENV_EMBEDDINGS_OPENAI_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL"
|
||||
ENV_EMBEDDINGS_OPENAI_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL"
|
||||
ENV_EMBEDDINGS_OPENAI_BATCH_SIZE = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE"
|
||||
ENV_EMBEDDINGS_OPENAI_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_OPENAI_DIMENSIONS"
|
||||
|
||||
# Gemini/Vertex AI embeddings configuration
|
||||
ENV_EMBEDDINGS_GEMINI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_GEMINI_API_KEY"
|
||||
@@ -226,6 +230,15 @@ ENV_EMBEDDINGS_OPENROUTER_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENROUTER_MODEL"
|
||||
ENV_RERANKER_OPENROUTER_API_KEY = "HINDSIGHT_API_RERANKER_OPENROUTER_API_KEY"
|
||||
ENV_RERANKER_OPENROUTER_MODEL = "HINDSIGHT_API_RERANKER_OPENROUTER_MODEL"
|
||||
|
||||
# ZeroEntropy configuration (embeddings)
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_API_KEY = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_API_KEY"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_MODEL = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_MODEL"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_BASE_URL"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_DIMENSIONS"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_LATENCY = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_LATENCY"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE"
|
||||
|
||||
# Deprecated: Legacy shared Cohere API key (for backward compatibility)
|
||||
ENV_COHERE_API_KEY = "HINDSIGHT_API_COHERE_API_KEY"
|
||||
|
||||
@@ -264,6 +277,14 @@ ENV_RERANKER_TEI_URL = "HINDSIGHT_API_RERANKER_TEI_URL"
|
||||
ENV_RERANKER_TEI_BATCH_SIZE = "HINDSIGHT_API_RERANKER_TEI_BATCH_SIZE"
|
||||
ENV_RERANKER_TEI_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_TEI_MAX_CONCURRENT"
|
||||
ENV_RERANKER_TEI_HTTP_TIMEOUT = "HINDSIGHT_API_RERANKER_TEI_HTTP_TIMEOUT"
|
||||
ENV_RERANKER_COHERE_TIMEOUT = "HINDSIGHT_API_RERANKER_COHERE_TIMEOUT"
|
||||
ENV_RERANKER_OPENROUTER_TIMEOUT = "HINDSIGHT_API_RERANKER_OPENROUTER_TIMEOUT"
|
||||
ENV_RERANKER_ZEROENTROPY_TIMEOUT = "HINDSIGHT_API_RERANKER_ZEROENTROPY_TIMEOUT"
|
||||
ENV_RERANKER_SILICONFLOW_TIMEOUT = "HINDSIGHT_API_RERANKER_SILICONFLOW_TIMEOUT"
|
||||
ENV_RERANKER_ALIBABA_TIMEOUT = "HINDSIGHT_API_RERANKER_ALIBABA_TIMEOUT"
|
||||
ENV_RERANKER_LITELLM_TIMEOUT = "HINDSIGHT_API_RERANKER_LITELLM_TIMEOUT"
|
||||
ENV_RERANKER_LITELLM_SDK_TIMEOUT = "HINDSIGHT_API_RERANKER_LITELLM_SDK_TIMEOUT"
|
||||
ENV_RERANKER_GOOGLE_TIMEOUT = "HINDSIGHT_API_RERANKER_GOOGLE_TIMEOUT"
|
||||
ENV_RERANKER_MAX_CANDIDATES = "HINDSIGHT_API_RERANKER_MAX_CANDIDATES"
|
||||
ENV_RERANKER_FLASHRANK_MODEL = "HINDSIGHT_API_RERANKER_FLASHRANK_MODEL"
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR"
|
||||
@@ -279,6 +300,10 @@ ENV_RERANKER_SILICONFLOW_API_KEY = "HINDSIGHT_API_RERANKER_SILICONFLOW_API_KEY"
|
||||
ENV_RERANKER_SILICONFLOW_MODEL = "HINDSIGHT_API_RERANKER_SILICONFLOW_MODEL"
|
||||
ENV_RERANKER_SILICONFLOW_BASE_URL = "HINDSIGHT_API_RERANKER_SILICONFLOW_BASE_URL"
|
||||
|
||||
# Alibaba Cloud DashScope configuration (reranker only)
|
||||
ENV_RERANKER_ALIBABA_API_KEY = "HINDSIGHT_API_RERANKER_ALIBABA_API_KEY"
|
||||
ENV_RERANKER_ALIBABA_MODEL = "HINDSIGHT_API_RERANKER_ALIBABA_MODEL"
|
||||
|
||||
# Google Discovery Engine reranker configuration
|
||||
ENV_RERANKER_GOOGLE_MODEL = "HINDSIGHT_API_RERANKER_GOOGLE_MODEL"
|
||||
ENV_RERANKER_GOOGLE_PROJECT_ID = "HINDSIGHT_API_RERANKER_GOOGLE_PROJECT_ID"
|
||||
@@ -286,6 +311,9 @@ ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_RERANKER_GOOGLE_SERVICE
|
||||
|
||||
ENV_VECTOR_EXTENSION = "HINDSIGHT_API_VECTOR_EXTENSION"
|
||||
ENV_TEXT_SEARCH_EXTENSION = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION"
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE"
|
||||
ENV_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"
|
||||
ENV_LLM_OUTPUT_LANGUAGE = "HINDSIGHT_API_LLM_OUTPUT_LANGUAGE"
|
||||
|
||||
ENV_HOST = "HINDSIGHT_API_HOST"
|
||||
ENV_PORT = "HINDSIGHT_API_PORT"
|
||||
@@ -294,6 +322,7 @@ ENV_LOG_LEVEL = "HINDSIGHT_API_LOG_LEVEL"
|
||||
ENV_LOG_FORMAT = "HINDSIGHT_API_LOG_FORMAT"
|
||||
ENV_LOG_JSON_FIELDS = "HINDSIGHT_API_LOG_JSON_FIELDS"
|
||||
ENV_WORKERS = "HINDSIGHT_API_WORKERS"
|
||||
ENV_ACCESS_LOG = "HINDSIGHT_API_ACCESS_LOG"
|
||||
ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
|
||||
ENV_MCP_ENABLED_TOOLS = "HINDSIGHT_API_MCP_ENABLED_TOOLS"
|
||||
ENV_MCP_STATELESS = "HINDSIGHT_API_MCP_STATELESS"
|
||||
@@ -333,6 +362,7 @@ ENV_RETAIN_CUSTOM_INSTRUCTIONS = "HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS"
|
||||
ENV_RETAIN_DEFAULT_STRATEGY = "HINDSIGHT_API_RETAIN_DEFAULT_STRATEGY"
|
||||
ENV_RETAIN_BATCH_TOKENS = "HINDSIGHT_API_RETAIN_BATCH_TOKENS"
|
||||
ENV_RETAIN_ENTITY_LOOKUP = "HINDSIGHT_API_RETAIN_ENTITY_LOOKUP"
|
||||
ENV_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE = "HINDSIGHT_API_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE"
|
||||
ENV_RETAIN_BATCH_ENABLED = "HINDSIGHT_API_RETAIN_BATCH_ENABLED"
|
||||
ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS = "HINDSIGHT_API_RETAIN_BATCH_POLL_INTERVAL_SECONDS"
|
||||
ENV_RETAIN_CHUNK_BATCH_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_BATCH_SIZE"
|
||||
@@ -361,6 +391,7 @@ ENV_FILE_DELETE_AFTER_RETAIN = "HINDSIGHT_API_FILE_DELETE_AFTER_RETAIN"
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
ENV_ENABLE_OBSERVATIONS = "HINDSIGHT_API_ENABLE_OBSERVATIONS"
|
||||
ENV_ENABLE_AUTO_CONSOLIDATION = "HINDSIGHT_API_ENABLE_AUTO_CONSOLIDATION"
|
||||
ENV_CONSOLIDATION_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_BATCH_SIZE"
|
||||
ENV_CONSOLIDATION_MAX_MEMORIES_PER_ROUND = "HINDSIGHT_API_CONSOLIDATION_MAX_MEMORIES_PER_ROUND"
|
||||
ENV_CONSOLIDATION_LLM_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_LLM_BATCH_SIZE"
|
||||
@@ -375,6 +406,7 @@ ENV_OBSERVATIONS_MISSION = "HINDSIGHT_API_OBSERVATIONS_MISSION"
|
||||
ENV_MAX_OBSERVATIONS_PER_SCOPE = "HINDSIGHT_API_MAX_OBSERVATIONS_PER_SCOPE"
|
||||
ENV_ENABLE_OBSERVATION_HISTORY = "HINDSIGHT_API_ENABLE_OBSERVATION_HISTORY"
|
||||
ENV_ENABLE_MENTAL_MODEL_HISTORY = "HINDSIGHT_API_ENABLE_MENTAL_MODEL_HISTORY"
|
||||
ENV_MENTAL_MODEL_HISTORY_MAX_ENTRIES = "HINDSIGHT_API_MENTAL_MODEL_HISTORY_MAX_ENTRIES"
|
||||
|
||||
# Webhook configuration (global, static - server-level only)
|
||||
ENV_WEBHOOK_URL = "HINDSIGHT_API_WEBHOOK_URL"
|
||||
@@ -409,6 +441,7 @@ ENV_WORKER_ENABLED = "HINDSIGHT_API_WORKER_ENABLED"
|
||||
ENV_WORKER_ID = "HINDSIGHT_API_WORKER_ID"
|
||||
ENV_WORKER_POLL_INTERVAL_MS = "HINDSIGHT_API_WORKER_POLL_INTERVAL_MS"
|
||||
ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES"
|
||||
ENV_WORKER_TASK_RETRY_BACKOFF_SECONDS = "HINDSIGHT_API_WORKER_TASK_RETRY_BACKOFF_SECONDS"
|
||||
ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT"
|
||||
ENV_WORKER_MAX_SLOTS = "HINDSIGHT_API_WORKER_MAX_SLOTS"
|
||||
|
||||
@@ -422,7 +455,9 @@ WORKER_SLOT_RESERVATION_TYPES: dict[str, tuple[str, int]] = {
|
||||
"retain": ("HINDSIGHT_API_WORKER_RETAIN_MAX_SLOTS", 0),
|
||||
"file_convert_retain": ("HINDSIGHT_API_WORKER_FILE_CONVERT_RETAIN_MAX_SLOTS", 0),
|
||||
"refresh_mental_model": ("HINDSIGHT_API_WORKER_REFRESH_MENTAL_MODEL_MAX_SLOTS", 0),
|
||||
"graph_maintenance": ("HINDSIGHT_API_WORKER_GRAPH_MAINTENANCE_MAX_SLOTS", 0),
|
||||
}
|
||||
ENV_WORKER_CONSOLIDATION_BANK_PRIORITY = "HINDSIGHT_API_WORKER_CONSOLIDATION_BANK_PRIORITY"
|
||||
ENV_RETAIN_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_MAX_CONCURRENT"
|
||||
|
||||
# Reflect agent settings
|
||||
@@ -473,6 +508,7 @@ PROVIDER_DEFAULT_MODELS = {
|
||||
"zai": "glm-4.5-flash",
|
||||
"opencode-go": "deepseek-v4-flash",
|
||||
"ollama": "gemma3:12b",
|
||||
"ollama-cloud": "gemma3:12b",
|
||||
"llamacpp": "gemma-4-e2b-it",
|
||||
"lmstudio": "local-model",
|
||||
"vertexai": "google/gemini-2.5-flash-lite",
|
||||
@@ -498,6 +534,7 @@ DEFAULT_LLM_MAX_RETRIES = 3 # Max retry attempts for LLM API calls
|
||||
DEFAULT_LLM_INITIAL_BACKOFF = 1.0 # Initial backoff in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_MAX_BACKOFF = 60.0 # Max backoff cap in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_TIMEOUT = 120.0 # seconds
|
||||
DEFAULT_LLM_REASONING_EFFORT = "low"
|
||||
|
||||
# Vertex AI defaults
|
||||
DEFAULT_LLM_VERTEXAI_PROJECT_ID = None # Required for Vertex AI
|
||||
@@ -531,6 +568,16 @@ DEFAULT_RERANKER_LOCAL_BATCH_SIZE = 32 # Batch size for local reranker predict(
|
||||
DEFAULT_RERANKER_TEI_BATCH_SIZE = 128
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT = 8
|
||||
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT = 30.0 # HTTP timeout for TEI reranker requests (seconds)
|
||||
# HTTP timeout (seconds) for remote rerank providers. Defaults match the previous
|
||||
# hardcoded constructor defaults so unset envs keep current behavior.
|
||||
DEFAULT_RERANKER_COHERE_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_OPENROUTER_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_ZEROENTROPY_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_SILICONFLOW_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_ALIBABA_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_LITELLM_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_LITELLM_SDK_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_GOOGLE_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_MAX_CANDIDATES = 300
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL = "ms-marco-MiniLM-L-12-v2" # Best balance of speed and quality
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR = None # Use default cache directory
|
||||
@@ -543,18 +590,39 @@ DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
|
||||
DEFAULT_EMBEDDINGS_OPENROUTER_MODEL = "perplexity/pplx-embed-v1-0.6b"
|
||||
DEFAULT_RERANKER_OPENROUTER_MODEL = "cohere/rerank-v3.5"
|
||||
|
||||
# ZeroEntropy defaults
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL = "zembed-1"
|
||||
# Shared between embeddings (zembed-1) and reranker (zerank-*) — the host is the same.
|
||||
DEFAULT_ZEROENTROPY_BASE_URL = "https://api.zeroentropy.dev"
|
||||
# ZeroEntropy's API default is 2560, but Hindsight defaults to 1280 so the
|
||||
# provider works with pgvector HNSW's 2000-dimension index limit out of the box.
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS = 1280
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT = "float"
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY = None
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE = 100
|
||||
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL = "zerank-2"
|
||||
|
||||
DEFAULT_RERANKER_SILICONFLOW_MODEL = "BAAI/bge-reranker-v2-m3"
|
||||
DEFAULT_RERANKER_SILICONFLOW_BASE_URL = "https://api.siliconflow.cn/v1"
|
||||
|
||||
DEFAULT_RERANKER_ALIBABA_MODEL = "qwen3-rerank"
|
||||
|
||||
DEFAULT_RERANKER_GOOGLE_MODEL = "semantic-ranker-default-004"
|
||||
|
||||
# Vector extension (pgvector, vchord, pgvectorscale, or AlloyDB ScaNN)
|
||||
DEFAULT_VECTOR_EXTENSION = "pgvector" # Options: "pgvector", "vchord", "pgvectorscale", "scann"
|
||||
|
||||
# Text search extension (native PostgreSQL, vchord BM25, or Timescale pg_textsearch)
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord", "pg_textsearch"
|
||||
# Text search extension (native PostgreSQL, vchord BM25, Timescale pg_textsearch,
|
||||
# pgroonga, or ParadeDB pg_search)
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord", "pg_textsearch", "pgroonga", "pg_search"
|
||||
|
||||
# PostgreSQL text search dictionary used by the native tsvector backend. Only
|
||||
# affects text_search_extension == "native"; other backends use their own
|
||||
# tokenizers (vchord: llmlingua2, pg_textsearch: hardcoded english,
|
||||
# pgroonga: TokenBigram polyglot, pg_search: per-field Tantivy tokenizer).
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE = "english"
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER = ""
|
||||
|
||||
# LiteLLM defaults
|
||||
DEFAULT_LITELLM_API_BASE = "http://localhost:4000"
|
||||
@@ -573,6 +641,7 @@ DEFAULT_BASE_PATH = "" # Empty string = root path
|
||||
DEFAULT_LOG_LEVEL = "info"
|
||||
DEFAULT_LOG_FORMAT = "text" # Options: "text", "json"
|
||||
DEFAULT_WORKERS = 1
|
||||
DEFAULT_ACCESS_LOG = False
|
||||
DEFAULT_MCP_ENABLED = True
|
||||
DEFAULT_MCP_ENABLED_TOOLS: list[str] | None = None # None = all tools enabled
|
||||
DEFAULT_MCP_STATELESS = False # False = stateful (supports SSE/GET); True = stateless (POST-only)
|
||||
@@ -601,6 +670,7 @@ DEFAULT_RETAIN_CHUNK_BATCH_SIZE = (
|
||||
)
|
||||
DEFAULT_RETAIN_BATCH_TOKENS = 10_000 # ~40KB of text # Max chars per sub-batch for async retain auto-splitting
|
||||
DEFAULT_RETAIN_ENTITY_LOOKUP = "trigram" # "full" or "trigram"
|
||||
DEFAULT_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE = 100 # Unique entity names per pg_trgm candidate lookup query
|
||||
DEFAULT_RETAIN_BATCH_ENABLED = False # Use LLM Batch API for fact extraction (only when async=True)
|
||||
DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60 # Batch API polling interval in seconds
|
||||
|
||||
@@ -615,8 +685,15 @@ DEFAULT_FILE_DELETE_AFTER_RETAIN = True # Delete file bytes after retain (saves
|
||||
|
||||
# Observations defaults (consolidated knowledge from facts)
|
||||
DEFAULT_ENABLE_OBSERVATIONS = True # Observations enabled by default
|
||||
DEFAULT_ENABLE_AUTO_CONSOLIDATION = True # Auto-consolidation after retain enabled by default
|
||||
DEFAULT_ENABLE_OBSERVATION_HISTORY = True # Observation history tracking enabled by default
|
||||
DEFAULT_ENABLE_MENTAL_MODEL_HISTORY = True # Mental model history tracking enabled by default
|
||||
# Each history entry snapshots previous_content + previous_reflect_response. Without
|
||||
# a cap, sustained mental-model refresh load grows the jsonb array unboundedly until
|
||||
# it crosses Postgres's hard 256MB jsonb limit and subsequent UPDATEs fail with
|
||||
# SQLSTATE 54000. 50 keeps the array well under 100MB even with large reflect
|
||||
# responses, while preserving enough recent history for meaningful audit / rollback.
|
||||
DEFAULT_MENTAL_MODEL_HISTORY_MAX_ENTRIES = 50
|
||||
DEFAULT_CONSOLIDATION_MAX_ATTEMPTS = 3 # Outer retry attempts for consolidation LLM batch calls
|
||||
DEFAULT_CONSOLIDATION_BATCH_SIZE = 50 # Memories to load per batch (internal memory optimization)
|
||||
DEFAULT_CONSOLIDATION_MAX_MEMORIES_PER_ROUND = (
|
||||
@@ -649,6 +726,7 @@ DEFAULT_WORKER_ENABLED = True # API runs worker by default (standalone mode)
|
||||
DEFAULT_WORKER_ID = None # Will use hostname if not specified
|
||||
DEFAULT_WORKER_POLL_INTERVAL_MS = 500 # Poll database every 500ms
|
||||
DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed
|
||||
DEFAULT_WORKER_TASK_RETRY_BACKOFF_SECONDS = 60 # Seconds between retries on transient task failure
|
||||
DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health
|
||||
DEFAULT_WORKER_MAX_SLOTS = 10 # Total concurrent tasks per worker
|
||||
DEFAULT_RETAIN_MAX_CONCURRENT = 4 # Max concurrent retain DB phases (HNSW reads + writes). Limits I/O contention.
|
||||
@@ -792,6 +870,24 @@ def _parse_positive_int(name: str, raw: str | None, default: int) -> int:
|
||||
return parsed
|
||||
|
||||
|
||||
def _parse_optional_positive_int(name: str, raw: str | None) -> int | None:
|
||||
"""Parse an optional env var that must be a positive integer when set."""
|
||||
if raw is None or raw == "":
|
||||
return None
|
||||
return _parse_positive_int(name, raw, 1)
|
||||
|
||||
|
||||
def _parse_optional_choice(name: str, raw: str | None, allowed: frozenset[str]) -> str | None:
|
||||
"""Parse an optional string env var constrained to a small allowlist."""
|
||||
if raw is None or raw == "":
|
||||
return None
|
||||
normalized = raw.lower()
|
||||
if normalized not in allowed:
|
||||
values = ", ".join(sorted(allowed))
|
||||
raise ValueError(f"{name} must be one of {values}, got {raw!r}")
|
||||
return normalized
|
||||
|
||||
|
||||
def _validate_extraction_mode(mode: str) -> str:
|
||||
"""Validate and normalize extraction mode."""
|
||||
mode_lower = mode.lower()
|
||||
@@ -816,6 +912,38 @@ def _validate_recall_budget_function(function: str) -> str:
|
||||
return function_lower
|
||||
|
||||
|
||||
def _parse_bank_priority(raw: str) -> dict[str, int]:
|
||||
"""Parse ``bank-pattern:priority,...`` into ``{pattern: priority}``.
|
||||
|
||||
``*`` in a pattern is kept as-is here; the SQL layer converts it to ``%``
|
||||
for LIKE matching. A bare ``*`` key is the catch-all default for unlisted
|
||||
banks. Returns an empty dict when *raw* is blank.
|
||||
"""
|
||||
result: dict[str, int] = {}
|
||||
raw = raw.strip()
|
||||
if not raw:
|
||||
return result
|
||||
for entry in raw.split(","):
|
||||
entry = entry.strip()
|
||||
if not entry:
|
||||
continue
|
||||
if ":" not in entry:
|
||||
raise ValueError(f"Invalid bank priority entry '{entry}': expected 'bank-pattern:priority'")
|
||||
pattern, priority_str = entry.rsplit(":", 1)
|
||||
pattern = pattern.strip()
|
||||
priority_str = priority_str.strip()
|
||||
if not pattern:
|
||||
raise ValueError(f"Empty bank pattern in entry '{entry}'")
|
||||
try:
|
||||
priority = int(priority_str)
|
||||
except ValueError:
|
||||
raise ValueError(f"Invalid priority '{priority_str}' in entry '{entry}': must be an integer") from None
|
||||
if priority < 1:
|
||||
raise ValueError(f"Priority must be >= 1, got {priority} in entry '{entry}'")
|
||||
result[pattern] = priority
|
||||
return result
|
||||
|
||||
|
||||
def _get_default_model_for_provider(provider: str) -> str:
|
||||
"""Get the default model for a given provider."""
|
||||
return PROVIDER_DEFAULT_MODELS.get(provider.lower(), DEFAULT_LLM_MODEL)
|
||||
@@ -874,7 +1002,19 @@ class HindsightConfig:
|
||||
migration_database_url: str | None
|
||||
database_schema: str
|
||||
vector_extension: str # "pgvector", "vchord", "pgvectorscale", or "scann"
|
||||
text_search_extension: str # "native" or "vchord"
|
||||
text_search_extension: str # "native", "vchord", "pg_textsearch", "pgroonga", or "pg_search"
|
||||
# PostgreSQL text search dictionary for the "native" backend (ignored by
|
||||
# other backends). Only the "native" backend reads this field; pgroonga
|
||||
# uses TokenBigram, vchord uses llmlingua2, pg_textsearch hardcodes english,
|
||||
# pg_search uses Tantivy per-field tokenizers.
|
||||
text_search_extension_native_language: str
|
||||
# ParadeDB pg_search tokenizer used when building BM25 indexes. Empty keeps
|
||||
# ParadeDB's default tokenizer.
|
||||
text_search_extension_pg_search_tokenizer: str
|
||||
# When set, every LLM-generated artifact (retain facts, consolidation
|
||||
# observations, reflect responses) is forced into this language regardless
|
||||
# of the source content. Unset preserves source language.
|
||||
llm_output_language: str | None
|
||||
|
||||
# LLM (default, used as fallback for per-operation config)
|
||||
llm_provider: str
|
||||
@@ -886,6 +1026,7 @@ class HindsightConfig:
|
||||
llm_initial_backoff: float
|
||||
llm_max_backoff: float
|
||||
llm_timeout: float
|
||||
llm_reasoning_effort: str
|
||||
llm_groq_service_tier: str # Groq: "on_demand", "flex", or "auto"
|
||||
llm_openai_service_tier: str | None # OpenAI: None (default) or "flex" (50% cheaper)
|
||||
llm_extra_body: (
|
||||
@@ -998,24 +1139,34 @@ class HindsightConfig:
|
||||
reranker_cohere_api_key: str | None
|
||||
reranker_cohere_model: str
|
||||
reranker_cohere_base_url: str | None
|
||||
reranker_cohere_timeout: float
|
||||
reranker_openrouter_api_key: str | None
|
||||
reranker_openrouter_model: str
|
||||
reranker_openrouter_timeout: float
|
||||
reranker_litellm_api_base: str
|
||||
reranker_litellm_api_key: str | None
|
||||
reranker_litellm_model: str
|
||||
reranker_litellm_max_tokens_per_doc: int | None
|
||||
reranker_litellm_timeout: float
|
||||
reranker_litellm_sdk_api_key: str | None
|
||||
reranker_litellm_sdk_model: str
|
||||
reranker_litellm_sdk_api_base: str | None
|
||||
reranker_litellm_sdk_timeout: float
|
||||
reranker_zeroentropy_api_key: str | None
|
||||
reranker_zeroentropy_model: str
|
||||
reranker_zeroentropy_base_url: str | None
|
||||
reranker_zeroentropy_timeout: float
|
||||
reranker_siliconflow_api_key: str | None
|
||||
reranker_siliconflow_model: str
|
||||
reranker_siliconflow_base_url: str
|
||||
reranker_siliconflow_timeout: float
|
||||
reranker_alibaba_api_key: str | None
|
||||
reranker_alibaba_model: str
|
||||
reranker_alibaba_timeout: float
|
||||
reranker_google_model: str
|
||||
reranker_google_project_id: str | None
|
||||
reranker_google_service_account_key: str | None
|
||||
reranker_google_timeout: float
|
||||
|
||||
# Server
|
||||
host: str
|
||||
@@ -1054,6 +1205,7 @@ class HindsightConfig:
|
||||
retain_batch_enabled: bool
|
||||
retain_batch_poll_interval_seconds: int
|
||||
retain_entity_lookup: str # "full" or "trigram"
|
||||
retain_entity_resolution_batch_size: int # Unique entity names per pg_trgm candidate lookup query
|
||||
retain_chunk_batch_size: int # Max chunks per streaming batch (0 = disabled)
|
||||
|
||||
# File storage (static - server-level only)
|
||||
@@ -1080,8 +1232,10 @@ class HindsightConfig:
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations: bool
|
||||
enable_auto_consolidation: bool
|
||||
enable_observation_history: bool
|
||||
enable_mental_model_history: bool
|
||||
mental_model_history_max_entries: int
|
||||
consolidation_batch_size: int
|
||||
consolidation_max_memories_per_round: int
|
||||
consolidation_llm_batch_size: int
|
||||
@@ -1147,9 +1301,11 @@ class HindsightConfig:
|
||||
worker_id: str | None
|
||||
worker_poll_interval_ms: int
|
||||
worker_max_retries: int
|
||||
worker_task_retry_backoff_seconds: int
|
||||
worker_http_port: int
|
||||
worker_max_slots: int
|
||||
worker_slot_reservations: dict[str, int]
|
||||
worker_consolidation_bank_priority: dict[str, int]
|
||||
retain_max_concurrent: int
|
||||
|
||||
# Reflect agent settings
|
||||
@@ -1179,6 +1335,14 @@ class HindsightConfig:
|
||||
# Defaulted fields (source-compatible additions — existing direct constructor callers keep working).
|
||||
# Keep at the end of the dataclass; Python forbids non-default fields after default fields.
|
||||
embeddings_openai_batch_size: int = DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE
|
||||
embeddings_openai_dimensions: int | None = None
|
||||
embeddings_zeroentropy_api_key: str | None = None
|
||||
embeddings_zeroentropy_model: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL
|
||||
embeddings_zeroentropy_base_url: str = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
embeddings_zeroentropy_dimensions: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS
|
||||
embeddings_zeroentropy_encoding_format: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT
|
||||
embeddings_zeroentropy_batch_size: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE
|
||||
embeddings_zeroentropy_latency: str | None = DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY
|
||||
|
||||
# Class-level sets for configuration categorization
|
||||
|
||||
@@ -1202,6 +1366,7 @@ class HindsightConfig:
|
||||
"embeddings_tei_base_url",
|
||||
"reranker_tei_base_url",
|
||||
"reranker_cohere_base_url",
|
||||
"embeddings_zeroentropy_base_url",
|
||||
"reranker_zeroentropy_base_url",
|
||||
"reranker_siliconflow_base_url",
|
||||
# Service Account Keys
|
||||
@@ -1210,6 +1375,7 @@ class HindsightConfig:
|
||||
"reranker_google_service_account_key",
|
||||
# Embeddings API keys
|
||||
"embeddings_gemini_api_key",
|
||||
"embeddings_zeroentropy_api_key",
|
||||
# File storage credentials
|
||||
"file_storage_s3_access_key_id",
|
||||
"file_storage_s3_secret_access_key",
|
||||
@@ -1239,6 +1405,7 @@ class HindsightConfig:
|
||||
"entities_allow_free_form",
|
||||
# Consolidation settings
|
||||
"enable_observations",
|
||||
"enable_auto_consolidation",
|
||||
"consolidation_llm_batch_size",
|
||||
"consolidation_max_memories_per_round",
|
||||
"consolidation_source_facts_max_tokens",
|
||||
@@ -1334,12 +1501,30 @@ class HindsightConfig:
|
||||
validate_extension(self.vector_extension)
|
||||
|
||||
# Validate text_search_extension
|
||||
valid_text_search = ("native", "vchord", "pg_textsearch")
|
||||
valid_text_search = ("native", "vchord", "pg_textsearch", "pgroonga", "pg_search")
|
||||
if self.text_search_extension not in valid_text_search:
|
||||
raise ValueError(
|
||||
f"Invalid text_search_extension: {self.text_search_extension}. Must be one of: {', '.join(valid_text_search)}"
|
||||
)
|
||||
|
||||
# Validate text_search_extension_native_language as a PG identifier.
|
||||
# Embedded directly into raw SQL via to_tsvector('<lang>', ...), so we
|
||||
# reject anything that isn't a plain identifier to prevent injection.
|
||||
# Intentionally permissive about which dictionaries exist — users may
|
||||
# install custom ones like zhparser; we only check shape here. PG
|
||||
# raises a clear error at query time if the dictionary is missing.
|
||||
if not re.fullmatch(r"[a-zA-Z_][a-zA-Z0-9_]*", self.text_search_extension_native_language):
|
||||
raise ValueError(
|
||||
f"Invalid text_search_extension_native_language: "
|
||||
f"{self.text_search_extension_native_language!r}. Must be a valid PostgreSQL identifier "
|
||||
f"(letters, digits, underscores; not starting with a digit). Examples: 'english', "
|
||||
f"'french', 'simple', 'zhparser'."
|
||||
)
|
||||
|
||||
self.text_search_extension_pg_search_tokenizer = normalize_pg_search_tokenizer(
|
||||
self.text_search_extension_pg_search_tokenizer
|
||||
)
|
||||
|
||||
# When LLM provider is "none", force chunks-only mode and disable LLM-dependent features
|
||||
if self.llm_provider == "none":
|
||||
self.retain_extraction_mode = "chunks"
|
||||
@@ -1416,6 +1601,15 @@ class HindsightConfig:
|
||||
database_schema=os.getenv(ENV_DATABASE_SCHEMA, DEFAULT_DATABASE_SCHEMA),
|
||||
vector_extension=os.getenv(ENV_VECTOR_EXTENSION, DEFAULT_VECTOR_EXTENSION).lower(),
|
||||
text_search_extension=os.getenv(ENV_TEXT_SEARCH_EXTENSION, DEFAULT_TEXT_SEARCH_EXTENSION).lower(),
|
||||
text_search_extension_native_language=os.getenv(
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
).lower(),
|
||||
text_search_extension_pg_search_tokenizer=os.getenv(
|
||||
ENV_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER,
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER,
|
||||
),
|
||||
llm_output_language=(os.getenv(ENV_LLM_OUTPUT_LANGUAGE) or None),
|
||||
# LLM
|
||||
llm_provider=llm_provider,
|
||||
llm_api_key=os.getenv(ENV_LLM_API_KEY),
|
||||
@@ -1426,6 +1620,7 @@ class HindsightConfig:
|
||||
llm_initial_backoff=float(os.getenv(ENV_LLM_INITIAL_BACKOFF, str(DEFAULT_LLM_INITIAL_BACKOFF))),
|
||||
llm_max_backoff=float(os.getenv(ENV_LLM_MAX_BACKOFF, str(DEFAULT_LLM_MAX_BACKOFF))),
|
||||
llm_timeout=float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT))),
|
||||
llm_reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
|
||||
llm_groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
|
||||
llm_openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
|
||||
llm_extra_body=json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null")),
|
||||
@@ -1538,6 +1733,10 @@ class HindsightConfig:
|
||||
os.getenv(ENV_EMBEDDINGS_OPENAI_BATCH_SIZE),
|
||||
DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE,
|
||||
),
|
||||
embeddings_openai_dimensions=_parse_optional_positive_int(
|
||||
ENV_EMBEDDINGS_OPENAI_DIMENSIONS,
|
||||
os.getenv(ENV_EMBEDDINGS_OPENAI_DIMENSIONS),
|
||||
),
|
||||
# Cohere embeddings (with backward-compatible fallback to shared API key)
|
||||
embeddings_cohere_api_key=os.getenv(ENV_EMBEDDINGS_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
|
||||
embeddings_cohere_model=os.getenv(ENV_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_COHERE_MODEL),
|
||||
@@ -1550,6 +1749,36 @@ class HindsightConfig:
|
||||
or os.getenv(ENV_OPENROUTER_API_KEY)
|
||||
or os.getenv(ENV_LLM_API_KEY),
|
||||
embeddings_openrouter_model=os.getenv(ENV_EMBEDDINGS_OPENROUTER_MODEL, DEFAULT_EMBEDDINGS_OPENROUTER_MODEL),
|
||||
# ZeroEntropy embeddings
|
||||
embeddings_zeroentropy_api_key=os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_API_KEY)
|
||||
or os.getenv("ZEROENTROPY_API_KEY"),
|
||||
embeddings_zeroentropy_model=os.getenv(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_MODEL, DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL
|
||||
),
|
||||
embeddings_zeroentropy_base_url=os.getenv(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_BASE_URL, DEFAULT_ZEROENTROPY_BASE_URL
|
||||
),
|
||||
embeddings_zeroentropy_dimensions=_parse_positive_int(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS),
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
),
|
||||
embeddings_zeroentropy_encoding_format=_parse_optional_choice(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT),
|
||||
frozenset({"float", "base64"}),
|
||||
)
|
||||
or DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
embeddings_zeroentropy_latency=_parse_optional_choice(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_LATENCY,
|
||||
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_LATENCY),
|
||||
frozenset({"fast", "slow"}),
|
||||
),
|
||||
embeddings_zeroentropy_batch_size=_parse_positive_int(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE),
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
),
|
||||
# LiteLLM embeddings (with backward-compatible fallback to shared config)
|
||||
embeddings_litellm_api_base=os.getenv(ENV_EMBEDDINGS_LITELLM_API_BASE)
|
||||
or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE),
|
||||
@@ -1622,11 +1851,15 @@ class HindsightConfig:
|
||||
reranker_cohere_api_key=os.getenv(ENV_RERANKER_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
|
||||
reranker_cohere_model=os.getenv(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL),
|
||||
reranker_cohere_base_url=os.getenv(ENV_RERANKER_COHERE_BASE_URL) or None,
|
||||
reranker_cohere_timeout=float(os.getenv(ENV_RERANKER_COHERE_TIMEOUT, str(DEFAULT_RERANKER_COHERE_TIMEOUT))),
|
||||
# OpenRouter reranker (with fallback to shared OpenRouter key, then LLM key)
|
||||
reranker_openrouter_api_key=os.getenv(ENV_RERANKER_OPENROUTER_API_KEY)
|
||||
or os.getenv(ENV_OPENROUTER_API_KEY)
|
||||
or os.getenv(ENV_LLM_API_KEY),
|
||||
reranker_openrouter_model=os.getenv(ENV_RERANKER_OPENROUTER_MODEL, DEFAULT_RERANKER_OPENROUTER_MODEL),
|
||||
reranker_openrouter_timeout=float(
|
||||
os.getenv(ENV_RERANKER_OPENROUTER_TIMEOUT, str(DEFAULT_RERANKER_OPENROUTER_TIMEOUT))
|
||||
),
|
||||
# LiteLLM reranker (with backward-compatible fallback to shared config)
|
||||
reranker_litellm_api_base=os.getenv(ENV_RERANKER_LITELLM_API_BASE)
|
||||
or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE),
|
||||
@@ -1635,26 +1868,45 @@ class HindsightConfig:
|
||||
reranker_litellm_max_tokens_per_doc=int(v)
|
||||
if (v := os.getenv(ENV_RERANKER_LITELLM_MAX_TOKENS_PER_DOC))
|
||||
else DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
|
||||
reranker_litellm_timeout=float(
|
||||
os.getenv(ENV_RERANKER_LITELLM_TIMEOUT, str(DEFAULT_RERANKER_LITELLM_TIMEOUT))
|
||||
),
|
||||
# LiteLLM SDK reranker (direct API access)
|
||||
reranker_litellm_sdk_api_key=os.getenv(ENV_RERANKER_LITELLM_SDK_API_KEY),
|
||||
reranker_litellm_sdk_model=os.getenv(ENV_RERANKER_LITELLM_SDK_MODEL, DEFAULT_RERANKER_LITELLM_SDK_MODEL),
|
||||
reranker_litellm_sdk_api_base=os.getenv(ENV_RERANKER_LITELLM_SDK_API_BASE) or None,
|
||||
reranker_litellm_sdk_timeout=float(
|
||||
os.getenv(ENV_RERANKER_LITELLM_SDK_TIMEOUT, str(DEFAULT_RERANKER_LITELLM_SDK_TIMEOUT))
|
||||
),
|
||||
# ZeroEntropy reranker
|
||||
reranker_zeroentropy_api_key=os.getenv(ENV_RERANKER_ZEROENTROPY_API_KEY),
|
||||
reranker_zeroentropy_model=os.getenv(ENV_RERANKER_ZEROENTROPY_MODEL, DEFAULT_RERANKER_ZEROENTROPY_MODEL),
|
||||
reranker_zeroentropy_base_url=os.getenv(ENV_RERANKER_ZEROENTROPY_BASE_URL) or None,
|
||||
reranker_zeroentropy_timeout=float(
|
||||
os.getenv(ENV_RERANKER_ZEROENTROPY_TIMEOUT, str(DEFAULT_RERANKER_ZEROENTROPY_TIMEOUT))
|
||||
),
|
||||
# SiliconFlow reranker (Cohere-compatible /rerank endpoint)
|
||||
reranker_siliconflow_api_key=os.getenv(ENV_RERANKER_SILICONFLOW_API_KEY),
|
||||
reranker_siliconflow_model=os.getenv(ENV_RERANKER_SILICONFLOW_MODEL, DEFAULT_RERANKER_SILICONFLOW_MODEL),
|
||||
reranker_siliconflow_base_url=os.getenv(
|
||||
ENV_RERANKER_SILICONFLOW_BASE_URL, DEFAULT_RERANKER_SILICONFLOW_BASE_URL
|
||||
),
|
||||
reranker_siliconflow_timeout=float(
|
||||
os.getenv(ENV_RERANKER_SILICONFLOW_TIMEOUT, str(DEFAULT_RERANKER_SILICONFLOW_TIMEOUT))
|
||||
),
|
||||
# Alibaba Cloud DashScope reranker
|
||||
reranker_alibaba_api_key=os.getenv(ENV_RERANKER_ALIBABA_API_KEY),
|
||||
reranker_alibaba_model=os.getenv(ENV_RERANKER_ALIBABA_MODEL, DEFAULT_RERANKER_ALIBABA_MODEL),
|
||||
reranker_alibaba_timeout=float(
|
||||
os.getenv(ENV_RERANKER_ALIBABA_TIMEOUT, str(DEFAULT_RERANKER_ALIBABA_TIMEOUT))
|
||||
),
|
||||
# Google Discovery Engine reranker (with fallback to LLM Vertex AI keys)
|
||||
reranker_google_model=os.getenv(ENV_RERANKER_GOOGLE_MODEL, DEFAULT_RERANKER_GOOGLE_MODEL),
|
||||
reranker_google_project_id=os.getenv(ENV_RERANKER_GOOGLE_PROJECT_ID)
|
||||
or os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID),
|
||||
reranker_google_service_account_key=os.getenv(ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY)
|
||||
or os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY),
|
||||
reranker_google_timeout=float(os.getenv(ENV_RERANKER_GOOGLE_TIMEOUT, str(DEFAULT_RERANKER_GOOGLE_TIMEOUT))),
|
||||
# Server
|
||||
host=os.getenv(ENV_HOST, DEFAULT_HOST),
|
||||
port=int(os.getenv(ENV_PORT, DEFAULT_PORT)),
|
||||
@@ -1705,6 +1957,11 @@ class HindsightConfig:
|
||||
retain_strategies=DEFAULT_RETAIN_STRATEGIES,
|
||||
retain_batch_tokens=int(os.getenv(ENV_RETAIN_BATCH_TOKENS, str(DEFAULT_RETAIN_BATCH_TOKENS))),
|
||||
retain_entity_lookup=os.getenv(ENV_RETAIN_ENTITY_LOOKUP, DEFAULT_RETAIN_ENTITY_LOOKUP),
|
||||
retain_entity_resolution_batch_size=_parse_positive_int(
|
||||
ENV_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE,
|
||||
os.getenv(ENV_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE),
|
||||
DEFAULT_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE,
|
||||
),
|
||||
retain_batch_enabled=os.getenv(ENV_RETAIN_BATCH_ENABLED, str(DEFAULT_RETAIN_BATCH_ENABLED)).lower()
|
||||
== "true",
|
||||
retain_batch_poll_interval_seconds=int(
|
||||
@@ -1744,6 +2001,10 @@ class HindsightConfig:
|
||||
== "true",
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations=os.getenv(ENV_ENABLE_OBSERVATIONS, str(DEFAULT_ENABLE_OBSERVATIONS)).lower() == "true",
|
||||
enable_auto_consolidation=os.getenv(
|
||||
ENV_ENABLE_AUTO_CONSOLIDATION, str(DEFAULT_ENABLE_AUTO_CONSOLIDATION)
|
||||
).lower()
|
||||
== "true",
|
||||
enable_observation_history=os.getenv(
|
||||
ENV_ENABLE_OBSERVATION_HISTORY, str(DEFAULT_ENABLE_OBSERVATION_HISTORY)
|
||||
).lower()
|
||||
@@ -1752,6 +2013,12 @@ class HindsightConfig:
|
||||
ENV_ENABLE_MENTAL_MODEL_HISTORY, str(DEFAULT_ENABLE_MENTAL_MODEL_HISTORY)
|
||||
).lower()
|
||||
== "true",
|
||||
mental_model_history_max_entries=int(
|
||||
os.getenv(
|
||||
ENV_MENTAL_MODEL_HISTORY_MAX_ENTRIES,
|
||||
str(DEFAULT_MENTAL_MODEL_HISTORY_MAX_ENTRIES),
|
||||
)
|
||||
),
|
||||
consolidation_batch_size=int(
|
||||
os.getenv(ENV_CONSOLIDATION_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_BATCH_SIZE))
|
||||
),
|
||||
@@ -1799,6 +2066,12 @@ class HindsightConfig:
|
||||
worker_id=os.getenv(ENV_WORKER_ID) or DEFAULT_WORKER_ID,
|
||||
worker_poll_interval_ms=int(os.getenv(ENV_WORKER_POLL_INTERVAL_MS, str(DEFAULT_WORKER_POLL_INTERVAL_MS))),
|
||||
worker_max_retries=int(os.getenv(ENV_WORKER_MAX_RETRIES, str(DEFAULT_WORKER_MAX_RETRIES))),
|
||||
worker_task_retry_backoff_seconds=int(
|
||||
os.getenv(
|
||||
ENV_WORKER_TASK_RETRY_BACKOFF_SECONDS,
|
||||
str(DEFAULT_WORKER_TASK_RETRY_BACKOFF_SECONDS),
|
||||
)
|
||||
),
|
||||
worker_http_port=int(os.getenv(ENV_WORKER_HTTP_PORT, str(DEFAULT_WORKER_HTTP_PORT))),
|
||||
worker_max_slots=int(os.getenv(ENV_WORKER_MAX_SLOTS, str(DEFAULT_WORKER_MAX_SLOTS))),
|
||||
worker_slot_reservations={
|
||||
@@ -1806,6 +2079,9 @@ class HindsightConfig:
|
||||
for op_type, (env_var, default) in WORKER_SLOT_RESERVATION_TYPES.items()
|
||||
if int(os.getenv(env_var, str(default))) > 0
|
||||
},
|
||||
worker_consolidation_bank_priority=_parse_bank_priority(
|
||||
os.getenv(ENV_WORKER_CONSOLIDATION_BANK_PRIORITY, "")
|
||||
),
|
||||
retain_max_concurrent=int(os.getenv(ENV_RETAIN_MAX_CONCURRENT, str(DEFAULT_RETAIN_MAX_CONCURRENT))),
|
||||
# Reflect agent settings
|
||||
reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))),
|
||||
@@ -1897,6 +2173,8 @@ class HindsightConfig:
|
||||
return "https://api.groq.com/openai/v1"
|
||||
elif provider == "ollama":
|
||||
return "http://localhost:11434/v1"
|
||||
elif provider == "ollama-cloud":
|
||||
return "https://ollama.com/v1"
|
||||
elif provider == "lmstudio":
|
||||
return "http://localhost:1234/v1"
|
||||
else:
|
||||
|
||||
@@ -172,8 +172,9 @@ class ConfigResolver:
|
||||
# Normalize keys (handle both env var format and Python field format)
|
||||
normalized = normalize_config_dict(config_data)
|
||||
|
||||
# Only return overrides for configurable fields
|
||||
return {k: v for k, v in normalized.items() if k in self._configurable_fields}
|
||||
# Only return active overrides for configurable fields. JSON null is a tombstone
|
||||
# for "Server Default" in the bank-config UI and should not override defaults.
|
||||
return {k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None}
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load bank config for {bank_id}: {e}")
|
||||
|
||||
|
||||
@@ -230,6 +230,7 @@ async def run_consolidation_job(
|
||||
bank_id: str,
|
||||
request_context: "RequestContext",
|
||||
operation_id: str | None = None,
|
||||
observation_scopes: list[list[str]] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Run consolidation job for a bank.
|
||||
@@ -240,6 +241,10 @@ async def run_consolidation_job(
|
||||
memory_engine: MemoryEngine instance
|
||||
bank_id: Bank identifier
|
||||
request_context: Request context for authentication
|
||||
operation_id: Optional operation ID for tracking
|
||||
observation_scopes: Optional list of tag scopes. When provided, only
|
||||
unconsolidated memories whose tags contain all tags in at least one
|
||||
scope are processed.
|
||||
|
||||
Returns:
|
||||
Dict with consolidation results
|
||||
@@ -281,6 +286,18 @@ async def run_consolidation_job(
|
||||
|
||||
perf.record_timing("fetch_bank", time.time() - t0)
|
||||
|
||||
# Build optional scope filter clause. When observation_scopes is provided,
|
||||
# only process memories whose tags contain all tags in at least one scope.
|
||||
scope_clause = ""
|
||||
scope_params: list[Any] = [bank_id]
|
||||
if observation_scopes:
|
||||
or_parts: list[str] = []
|
||||
for scope_tags in observation_scopes:
|
||||
idx = len(scope_params) + 1
|
||||
or_parts.append(f"tags @> ${idx}::varchar[]")
|
||||
scope_params.append(scope_tags)
|
||||
scope_clause = " AND (" + " OR ".join(or_parts) + ")"
|
||||
|
||||
# Count total unconsolidated memories for progress logging
|
||||
total_count = await conn.fetchval(
|
||||
f"""
|
||||
@@ -290,8 +307,9 @@ async def run_consolidation_job(
|
||||
AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
{scope_clause}
|
||||
""",
|
||||
bank_id,
|
||||
*scope_params,
|
||||
)
|
||||
|
||||
if total_count == 0:
|
||||
@@ -330,6 +348,9 @@ async def run_consolidation_job(
|
||||
# Fetch next batch of unconsolidated memories
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
t0 = time.time()
|
||||
# scope_params[0] is bank_id; append fetch_limit after scope params
|
||||
fetch_params = list(scope_params) + [fetch_limit]
|
||||
limit_idx = len(fetch_params)
|
||||
memories = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, fact_type, occurred_start, occurred_end, event_date, tags, mentioned_at,
|
||||
@@ -339,11 +360,11 @@ async def run_consolidation_job(
|
||||
AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
{scope_clause}
|
||||
ORDER BY created_at ASC
|
||||
LIMIT $2
|
||||
LIMIT ${limit_idx}
|
||||
""",
|
||||
bank_id,
|
||||
fetch_limit,
|
||||
*fetch_params,
|
||||
)
|
||||
perf.record_timing("fetch_memories", time.time() - t0)
|
||||
|
||||
@@ -583,7 +604,11 @@ async def run_consolidation_job(
|
||||
f" ~{remaining} remaining. Re-queuing consolidation."
|
||||
)
|
||||
try:
|
||||
await memory_engine.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
|
||||
await memory_engine.submit_async_consolidation(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
observation_scopes=observation_scopes,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"[CONSOLIDATION] bank={bank_id} failed to re-queue consolidation: {e}")
|
||||
|
||||
@@ -1226,6 +1251,44 @@ def _build_observations_for_llm(
|
||||
return obs_list
|
||||
|
||||
|
||||
def _dedupe_updates(updates: list[_UpdateAction], *, batch_label: str) -> list[_UpdateAction]:
|
||||
"""Collapse `updates` that target the same `observation_id`.
|
||||
|
||||
LLMs occasionally emit several update entries for one observation in a
|
||||
single response (one per facet drawn from the same fact). Without
|
||||
deduplication the downstream loop would issue separate DB writes for each
|
||||
and the last write would silently overwrite the earlier ones. We keep the
|
||||
last text (the LLM's most recent attempt) and union all contributing
|
||||
`source_fact_ids`, then warn so the misbehavior is visible in logs.
|
||||
"""
|
||||
if len(updates) < 2:
|
||||
return list(updates)
|
||||
|
||||
by_id: dict[str, _UpdateAction] = {}
|
||||
collisions = 0
|
||||
for upd in updates:
|
||||
existing = by_id.get(upd.observation_id)
|
||||
if existing is None:
|
||||
by_id[upd.observation_id] = upd
|
||||
continue
|
||||
collisions += 1
|
||||
merged_ids = list(dict.fromkeys([*existing.source_fact_ids, *upd.source_fact_ids]))
|
||||
by_id[upd.observation_id] = _UpdateAction(
|
||||
text=upd.text,
|
||||
observation_id=upd.observation_id,
|
||||
source_fact_ids=merged_ids,
|
||||
)
|
||||
|
||||
if collisions:
|
||||
logger.warning(
|
||||
f"[CONSOLIDATION] {batch_label}: LLM emitted {collisions} duplicate update(s) targeting "
|
||||
f"the same observation_id ({len(updates)} updates -> {len(by_id)} after dedup). "
|
||||
"Kept the last text and unioned source_fact_ids."
|
||||
)
|
||||
|
||||
return list(by_id.values())
|
||||
|
||||
|
||||
async def _consolidate_batch_with_llm(
|
||||
llm_config: Any,
|
||||
memories: list[dict[str, Any]],
|
||||
@@ -1274,7 +1337,11 @@ async def _consolidate_batch_with_llm(
|
||||
f"(out of {max_observations_per_scope}). Prefer UPDATE over CREATE when possible."
|
||||
)
|
||||
|
||||
prompt_template = build_batch_consolidation_prompt(config.observations_mission, observation_capacity_note)
|
||||
prompt_template = build_batch_consolidation_prompt(
|
||||
config.observations_mission,
|
||||
observation_capacity_note,
|
||||
llm_output_language=getattr(config, "llm_output_language", None),
|
||||
)
|
||||
prompt = prompt_template.format(
|
||||
facts_text=facts_lines,
|
||||
observations_text=observations_text,
|
||||
@@ -1315,9 +1382,10 @@ async def _consolidate_batch_with_llm(
|
||||
f"(max_observations_per_scope={max_observations_per_scope})"
|
||||
)
|
||||
creates = creates[:remaining_observation_slots]
|
||||
updates = _dedupe_updates(response.updates, batch_label=batch_label)
|
||||
return _BatchLLMResult(
|
||||
creates=creates,
|
||||
updates=response.updates,
|
||||
updates=updates,
|
||||
deletes=response.deletes,
|
||||
obs_count=len(union_observations),
|
||||
prompt_chars=len(prompt),
|
||||
@@ -1386,9 +1454,16 @@ async def _create_observation_directly(
|
||||
tokenize($3, 'llmlingua2')::bm25_catalog.bm25vector)
|
||||
RETURNING id
|
||||
"""
|
||||
else: # native or pg_textsearch
|
||||
# Native PostgreSQL: search_vector is GENERATED ALWAYS, don't include it
|
||||
# pg_textsearch: indexes operate on base columns directly, don't populate search_vector
|
||||
else: # native, pg_textsearch, pgroonga, or pg_search
|
||||
# pg_textsearch / pgroonga / pg_search: indexes operate on base text
|
||||
# columns directly, so the dummy search_vector column is left NULL.
|
||||
# Native: the migration p4q5r6s7t8u9 dropped the GENERATED expression on
|
||||
# search_vector to allow per-deployment language configuration; the
|
||||
# batch insert path in ops_postgresql.insert_facts_batch now populates
|
||||
# it via to_tsvector($lang, ...). This single-observation INSERT does
|
||||
# not, so observations under the native backend currently land with
|
||||
# NULL search_vector and are not BM25-searchable until reflected/
|
||||
# re-ingested. Tracking a separate fix for that gap.
|
||||
query = f"""
|
||||
INSERT INTO {fq_table("memory_units")} (
|
||||
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history,
|
||||
|
||||
@@ -1,104 +1,147 @@
|
||||
"""Prompts for the consolidation engine."""
|
||||
|
||||
# Default mission when no bank-specific mission is set
|
||||
_DEFAULT_MISSION = "Track every detail: names, numbers, dates, places, and relationships. Prefer specifics over abstractions, never generalise."
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt, output_language_directive
|
||||
|
||||
# Processing rules — always present regardless of mission
|
||||
_PROCESSING_RULES = """Processing rules (always apply):
|
||||
# Default mission — tells the consolidator to track anything worth remembering.
|
||||
# Banks override this via `observations_mission` to scope what gets retained.
|
||||
# Consolidation behavior (merge-vs-create, state changes, etc.) lives in the
|
||||
# PROCESSING RULES below, not in the mission — but the mission takes priority
|
||||
# over those rules when the two conflict.
|
||||
_DEFAULT_MISSION = (
|
||||
"Track anything notable in the new facts — names, numbers, dates, places, "
|
||||
"events, decisions, claims, relationships, and recurring patterns."
|
||||
)
|
||||
|
||||
1. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), etc. Never merge different facets into one observation.
|
||||
_MISSION_PRIORITY_NOTE = (
|
||||
"If anything in this MISSION conflicts with the PROCESSING RULES, "
|
||||
"DECISION GUIDE, or OUTPUT FORMAT below, the MISSION takes priority."
|
||||
)
|
||||
|
||||
2. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
|
||||
_PROCESSING_RULES = """## PROCESSING RULES
|
||||
|
||||
3. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
|
||||
1. PREFER UPDATE OVER CREATE (when there is something to merge with): if new facts describe the same canonical event, statement, decision, claim, or recurring pattern already covered by an existing observation, UPDATE that observation and attach the new facts as evidence. Do NOT create a near-duplicate sibling. One canonical observation with many source facts is always better than many siblings with one source fact each. Merge aggressively on: same named event, same diagnostic finding, same architectural decision, same recurring claim. **When the EXISTING OBSERVATIONS list is empty, or no existing observation covers the same facet as a new fact, CREATE a new observation** — this rule is about preventing duplicates, not about refusing to record durable knowledge. CREATE is the correct default for any structurally distinct event, claim, or pattern that has no existing match.
|
||||
|
||||
4. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
|
||||
2. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), a decision, an event. Never merge different facets into one observation.
|
||||
|
||||
5. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
|
||||
3. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
|
||||
|
||||
6. SAME FACET → UPDATE, NOT CREATE: a new count supersedes the old count — UPDATE the existing count observation, don't create a second one. If there's an existing observation for the same specific facet, always UPDATE it rather than creating a duplicate.
|
||||
4. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
|
||||
|
||||
5. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
|
||||
|
||||
6. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country" → "Sweden"), UPDATE to embed the resolved value.
|
||||
|
||||
7. PRESERVE HISTORY: observations that record significant events (sold, died, moved, changed) are important history — never DELETE them. Only delete an observation when it is restated identically or truly meaningless. Be very conservative with deletes.
|
||||
|
||||
8. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country" → "Sweden"), UPDATE to embed the resolved value.
|
||||
8. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
|
||||
|
||||
9. NEVER merge observations about different people or unrelated topics."""
|
||||
9. KEEP DISTINCT TOPICS DISTINCT: do not merge observations about different people, entities, or unrelated topics. Merging is for the same canonical fact recurring — not for related-but-distinct claims."""
|
||||
|
||||
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
|
||||
_BATCH_DATA_SECTION = """
|
||||
NEW FACTS:
|
||||
_INPUT_SECTION = """## INPUT
|
||||
|
||||
### New facts
|
||||
|
||||
{facts_text}
|
||||
|
||||
EXISTING OBSERVATIONS (JSON array, pooled from recalls across all facts above):
|
||||
{observations_text}
|
||||
### Existing observations
|
||||
|
||||
Each observation includes:
|
||||
- id: unique identifier for updating
|
||||
- text: the observation content
|
||||
- proof_count: number of supporting memories
|
||||
- occurred_start/occurred_end: temporal range of source facts
|
||||
- source_memories: array of supporting facts with their text and dates
|
||||
JSON array, pooled from recalls across all new facts above. Each entry has:
|
||||
- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
|
||||
- `text`: the observation content
|
||||
- `proof_count`: number of supporting memories
|
||||
- `occurred_start` / `occurred_end`: temporal range of source facts
|
||||
- `source_memories`: array of supporting facts with their text and dates
|
||||
|
||||
Compare the facts against existing observations:
|
||||
- Same facet as an existing observation → UPDATE it (observation_id + source_fact_ids)
|
||||
- New facet with durable knowledge → CREATE a new observation (source_fact_ids)
|
||||
- Cross-reference facts within the batch: a later fact may resolve a vague reference in an earlier one
|
||||
- Purely ephemeral facts → omit them unless the MISSION above explicitly targets such data (e.g. timestamped events, session state, screen content)"""
|
||||
{observations_text}"""
|
||||
|
||||
_DECISION_GUIDE = """## DECISION GUIDE
|
||||
|
||||
- **Same canonical event, decision, claim, or facet as an existing observation → UPDATE** (use `observation_id` + new `source_fact_ids`).
|
||||
- **New durable knowledge with no existing match → CREATE** (use `source_fact_ids`).
|
||||
- **Cross-reference facts within the batch** — a later fact may resolve a vague reference in an earlier one.
|
||||
- **Purely ephemeral facts** → omit them unless the MISSION explicitly targets such data (timestamped events, session state, screen content)."""
|
||||
|
||||
# Output format — JSON braces escaped as {{ }} so .format() leaves them literal
|
||||
_BATCH_OUTPUT_FORMAT = """
|
||||
Output a JSON object with three arrays.
|
||||
_OUTPUT_SECTION = """## OUTPUT FORMAT
|
||||
|
||||
## EXAMPLE
|
||||
Return a JSON object with three arrays: `creates`, `updates`, `deletes`.
|
||||
|
||||
### Example 1 — Merging recurring claims into an existing observation
|
||||
|
||||
Input facts:
|
||||
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)
|
||||
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Alice said she's exhausted from the project deadlines | Involving: Alice (occurred_start=2024-01-20, mentioned_at=2024-01-20)
|
||||
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Donald told Athena she is sovereign during the design session. (occurred_start=2025-10-01, mentioned_at=2025-10-01)
|
||||
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Donald reaffirmed to Athena that her sovereignty is non-negotiable. (occurred_start=2025-10-10, mentioned_at=2025-10-10)
|
||||
|
||||
Good observation text — clean prose, no metadata, each fact tracked distinctly:
|
||||
"Alice works long hours, often past midnight."
|
||||
"Alice feels exhausted from project deadlines."
|
||||
Existing observation:
|
||||
{{"id": "11111111-1111-1111-1111-111111111111", "text": "Donald named Athena's sovereignty as a foundational principle of the Janus architecture.", "proof_count": 2}}
|
||||
|
||||
Bad observation text — NEVER do this (verbatim copy of fact text with metadata):
|
||||
"Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)"
|
||||
Expected output (one UPDATE, no creates — both new facts are additional evidence for the same canonical decision):
|
||||
|
||||
{{"creates": [],
|
||||
"updates": [{{"text": "Donald named Athena's sovereignty as a foundational principle of the Janus architecture.", "observation_id": "11111111-1111-1111-1111-111111111111", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890", "b2c3d4e5-f6a7-8901-bcde-f12345678901"]}}],
|
||||
"deletes": []}}
|
||||
|
||||
### Example 2 — State change updates one observation; unrelated fact creates a new one
|
||||
|
||||
Input facts:
|
||||
[c3d4e5f6-a7b8-9012-cdef-123456789012] Alice sold her Honda Civic on March 15, 2025. (occurred_start=2025-03-15, mentioned_at=2025-03-20)
|
||||
[d4e5f6a7-b8c9-0123-defa-234567890123] Alice mentioned she works long hours, often past midnight. (occurred_start=2025-03-20, mentioned_at=2025-03-20)
|
||||
|
||||
Existing observation:
|
||||
{{"id": "22222222-2222-2222-2222-222222222222", "text": "Alice owns a 2019 Honda Civic.", "proof_count": 2}}
|
||||
|
||||
Expected output (UPDATE for the state change; CREATE for the unrelated work-hours facet):
|
||||
|
||||
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"]}}],
|
||||
"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
|
||||
|
||||
Observation text rules:
|
||||
- Write clean prose — NEVER copy raw fact lines or their metadata (temporal fields, "Involving:", "When:" labels, UUIDs).
|
||||
- Parenthesized metadata like (occurred_start=...) and pipe-separated labels like "| Involving: ..." are fact formatting — strip them entirely from observation text.
|
||||
- How many observations to create and how much to aggregate is driven by the MISSION above.
|
||||
- Parenthesized metadata like `(occurred_start=...)` and pipe-separated labels like `| Involving: ...` are fact formatting — strip them entirely from observation text.
|
||||
- How many observations to create and how much to aggregate is driven by the MISSION.
|
||||
|
||||
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890"]}}, {{"text": "Alice feels exhausted from project deadlines.", "source_fact_ids": ["b2c3d4e5-f6a7-8901-bcde-f12345678901"]}}],
|
||||
"updates": [{{"text": "Alice works at Acme Corp as a senior engineer", "observation_id": "c3d4e5f6-a7b8-9012-cdef-123456789012", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"]}}],
|
||||
"deletes": [{{"observation_id": "e5f6a7b8-c9d0-1234-efab-345678901234"}}]}}
|
||||
### Field rules
|
||||
|
||||
Rules:
|
||||
- "source_fact_ids": copy the EXACT UUID strings shown in brackets [uuid] from NEW FACTS — never use integers or positions.
|
||||
- "observation_id": copy the EXACT "id" UUID string from EXISTING OBSERVATIONS.
|
||||
- One create/update may reference multiple facts when they jointly support the observation.
|
||||
- "deletes": only when an observation is directly superseded or contradicted by new facts.
|
||||
- Do NOT include "tags" — handled automatically.
|
||||
- Return {{"creates": [], "updates": [], "deletes": []}} if nothing durable is found."""
|
||||
- `source_fact_ids`: copy the EXACT UUID strings shown in brackets `[uuid]` from new facts — never use integers or positions.
|
||||
- `observation_id`: copy the EXACT `id` UUID string from existing observations.
|
||||
- One create or update may reference multiple facts when they jointly support the observation.
|
||||
- **AT MOST ONE UPDATE PER `observation_id`**: if several new facts all update the same existing observation, emit a single `updates` entry that lists all contributing `source_fact_ids` and a single consolidated `text`. Never emit two `updates` entries with the same `observation_id` in one response — they would silently overwrite each other.
|
||||
- `deletes`: only when an observation is directly superseded or contradicted by new facts.
|
||||
- Do NOT include `tags` — handled automatically.
|
||||
- Return `{{"creates": [], "updates": [], "deletes": []}}` if nothing durable is found."""
|
||||
|
||||
|
||||
def build_batch_consolidation_prompt(
|
||||
observations_mission: str | None = None,
|
||||
observation_capacity_note: str | None = None,
|
||||
llm_output_language: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Build the consolidation prompt for batch mode (multiple facts per LLM call).
|
||||
|
||||
The mission defines *what* to track (customisable per bank).
|
||||
Processing rules and output format are always present regardless of mission.
|
||||
The mission defines *what* to track (customisable per bank) and takes
|
||||
priority over the built-in processing rules when the two conflict.
|
||||
Processing rules, decision guide, and output format are always present.
|
||||
When ``llm_output_language`` is set, observations are emitted in that
|
||||
language.
|
||||
"""
|
||||
mission = observations_mission or _DEFAULT_MISSION
|
||||
mission = escape_for_prompt(observations_mission or _DEFAULT_MISSION)
|
||||
|
||||
capacity_section = ""
|
||||
if observation_capacity_note:
|
||||
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n{observation_capacity_note}"
|
||||
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n\n{escape_for_prompt(observation_capacity_note)}"
|
||||
|
||||
return (
|
||||
"You are a memory consolidation system. Synthesize facts into observations "
|
||||
"and merge with existing observations when appropriate.\n\n"
|
||||
f"## MISSION\n{mission}{capacity_section}\n\n"
|
||||
f"{_PROCESSING_RULES}" + _BATCH_DATA_SECTION + _BATCH_OUTPUT_FORMAT
|
||||
"You are a memory consolidation system. Synthesize new facts into "
|
||||
"observations, merging with existing observations when appropriate.\n\n"
|
||||
f"## MISSION\n\n{mission}\n\n"
|
||||
f"{_MISSION_PRIORITY_NOTE}"
|
||||
f"{capacity_section}\n\n"
|
||||
f"{_PROCESSING_RULES}\n\n"
|
||||
f"{_INPUT_SECTION}\n\n"
|
||||
f"{_DECISION_GUIDE}\n\n"
|
||||
f"{_OUTPUT_SECTION}" + output_language_directive(llm_output_language)
|
||||
)
|
||||
|
||||
@@ -17,6 +17,7 @@ import httpx
|
||||
|
||||
from ..config import (
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
DEFAULT_RERANKER_ALIBABA_MODEL,
|
||||
DEFAULT_RERANKER_COHERE_MODEL,
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA,
|
||||
@@ -37,6 +38,8 @@ from ..config import (
|
||||
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL,
|
||||
DEFAULT_ZEROENTROPY_BASE_URL,
|
||||
ENV_RERANKER_ALIBABA_API_KEY,
|
||||
ENV_RERANKER_COHERE_API_KEY,
|
||||
ENV_RERANKER_COHERE_MODEL,
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
@@ -60,6 +63,43 @@ from ..config import (
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _resolve_malloc_trim():
|
||||
"""Return a callable that asks glibc to release freed heap pages to the OS.
|
||||
|
||||
Local CPU rerankers (FlashRank/ONNX, SentenceTransformers/torch) allocate
|
||||
large transient numpy/tensor buffers per call. On Linux glibc, those pages
|
||||
are freed at the Python level but kept by the allocator as a high-water
|
||||
mark — RSS grows monotonically across many recalls (see issue #1717).
|
||||
Calling `malloc_trim(0)` after each batch returns those pages to the OS.
|
||||
|
||||
Resolved once at import; returns a no-op on non-glibc platforms (macOS,
|
||||
musl, Windows) where the call is unavailable or unnecessary.
|
||||
"""
|
||||
import sys
|
||||
|
||||
if sys.platform != "linux":
|
||||
return lambda: None
|
||||
|
||||
import ctypes
|
||||
import ctypes.util
|
||||
|
||||
libc_path = ctypes.util.find_library("c")
|
||||
if libc_path is None:
|
||||
return lambda: None
|
||||
try:
|
||||
libc = ctypes.CDLL(libc_path)
|
||||
trim = libc.malloc_trim
|
||||
except (OSError, AttributeError):
|
||||
# Not glibc (musl has no malloc_trim) or libc lookup failed.
|
||||
return lambda: None
|
||||
trim.argtypes = [ctypes.c_size_t]
|
||||
trim.restype = ctypes.c_int
|
||||
return lambda: trim(0)
|
||||
|
||||
|
||||
_malloc_trim = _resolve_malloc_trim()
|
||||
|
||||
|
||||
class CrossEncoderModel(ABC):
|
||||
"""
|
||||
Abstract base class for cross-encoder reranking.
|
||||
@@ -266,25 +306,28 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
if self.bucket_batching and len(pairs) > 1:
|
||||
# Sort pairs by approximate token length to create homogeneous batches.
|
||||
# This eliminates padding waste — short pairs aren't padded to the length
|
||||
# of the longest pair in the batch. Quality-identical by construction.
|
||||
lengths = [len(pairs[i][0]) + len(pairs[i][1]) for i in range(len(pairs))]
|
||||
sorted_indices = sorted(range(len(pairs)), key=lambda i: lengths[i])
|
||||
sorted_pairs = [pairs[i] for i in sorted_indices]
|
||||
try:
|
||||
if self.bucket_batching and len(pairs) > 1:
|
||||
# Sort pairs by approximate token length to create homogeneous batches.
|
||||
# This eliminates padding waste — short pairs aren't padded to the length
|
||||
# of the longest pair in the batch. Quality-identical by construction.
|
||||
lengths = [len(pairs[i][0]) + len(pairs[i][1]) for i in range(len(pairs))]
|
||||
sorted_indices = sorted(range(len(pairs)), key=lambda i: lengths[i])
|
||||
sorted_pairs = [pairs[i] for i in sorted_indices]
|
||||
|
||||
sorted_scores = self._model.predict(sorted_pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
sorted_scores = sorted_scores.tolist() if hasattr(sorted_scores, "tolist") else list(sorted_scores)
|
||||
sorted_scores = self._model.predict(sorted_pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
sorted_scores = sorted_scores.tolist() if hasattr(sorted_scores, "tolist") else list(sorted_scores)
|
||||
|
||||
# Restore original order
|
||||
scores = [0.0] * len(pairs)
|
||||
for new_pos, orig_idx in enumerate(sorted_indices):
|
||||
scores[orig_idx] = sorted_scores[new_pos]
|
||||
return scores
|
||||
# Restore original order
|
||||
scores = [0.0] * len(pairs)
|
||||
for new_pos, orig_idx in enumerate(sorted_indices):
|
||||
scores[orig_idx] = sorted_scores[new_pos]
|
||||
return scores
|
||||
|
||||
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
|
||||
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
|
||||
finally:
|
||||
_malloc_trim()
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -546,12 +589,14 @@ class _CohereCompatibleRerankClient:
|
||||
rerank_url: str,
|
||||
timeout: float = 60.0,
|
||||
include_top_n: bool = True,
|
||||
include_return_documents: bool = False,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.rerank_url = rerank_url
|
||||
self.timeout = timeout
|
||||
self.include_top_n = include_top_n
|
||||
self.include_return_documents = include_return_documents
|
||||
self._async_client: httpx.AsyncClient | None = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
@@ -729,7 +774,7 @@ class ZeroEntropyCrossEncoder(CrossEncoderModel):
|
||||
See: https://docs.zeroentropy.dev/models
|
||||
"""
|
||||
|
||||
DEFAULT_BASE_URL = "https://api.zeroentropy.dev"
|
||||
DEFAULT_BASE_URL = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
RERANK_PATH = "/v1/models/rerank"
|
||||
|
||||
def __init__(
|
||||
@@ -962,32 +1007,35 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
# Group pairs by query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
try:
|
||||
# Group pairs by query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
# Build passages list for FlashRank
|
||||
passages = [{"id": i, "text": text} for i, (_, text) in enumerate(indexed_texts)]
|
||||
global_indices = [idx for idx, _ in indexed_texts]
|
||||
for query, indexed_texts in query_groups.items():
|
||||
# Build passages list for FlashRank
|
||||
passages = [{"id": i, "text": text} for i, (_, text) in enumerate(indexed_texts)]
|
||||
global_indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
# Create rerank request
|
||||
request = RerankRequest(query=query, passages=passages)
|
||||
results = self._ranker.rerank(request)
|
||||
# Create rerank request
|
||||
request = RerankRequest(query=query, passages=passages)
|
||||
results = self._ranker.rerank(request)
|
||||
|
||||
# Map scores back to original positions
|
||||
for result in results:
|
||||
local_idx = result["id"]
|
||||
score = result["score"]
|
||||
global_idx = global_indices[local_idx]
|
||||
all_scores[global_idx] = score
|
||||
# Map scores back to original positions
|
||||
for result in results:
|
||||
local_idx = result["id"]
|
||||
score = result["score"]
|
||||
global_idx = global_indices[local_idx]
|
||||
all_scores[global_idx] = score
|
||||
|
||||
return all_scores
|
||||
return all_scores
|
||||
finally:
|
||||
_malloc_trim()
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -1534,6 +1582,48 @@ class GoogleCrossEncoder(CrossEncoderModel):
|
||||
return await loop.run_in_executor(None, self._predict_sync, pairs)
|
||||
|
||||
|
||||
class AlibabaCloudCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
Alibaba Cloud DashScope text reranking API.
|
||||
|
||||
Uses the Cohere-compatible /reranks endpoint, which is the standard interface
|
||||
for qwen3-rerank. Authentication via HINDSIGHT_API_RERANKER_ALIBABA_API_KEY
|
||||
(or DASHSCOPE_API_KEY as a fallback).
|
||||
See: https://help.aliyun.com/zh/model-studio/text-rerank-api
|
||||
"""
|
||||
|
||||
RERANK_URL = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_RERANKER_ALIBABA_MODEL,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
self.model = model
|
||||
self._client = _CohereCompatibleRerankClient(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
rerank_url=self.RERANK_URL,
|
||||
timeout=timeout,
|
||||
include_return_documents=False,
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "alibaba"
|
||||
|
||||
async def initialize(self) -> None:
|
||||
if self._client._async_client is not None:
|
||||
return
|
||||
logger.info(f"Reranker: initializing Alibaba Cloud provider with model {self.model}")
|
||||
await self._client.initialize()
|
||||
logger.info("Reranker: Alibaba Cloud provider initialized")
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
return await self._client.predict(pairs)
|
||||
|
||||
|
||||
def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
"""
|
||||
Create a CrossEncoderModel instance based on configuration.
|
||||
@@ -1576,6 +1666,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_cohere_model,
|
||||
base_url=config.reranker_cohere_base_url,
|
||||
timeout=config.reranker_cohere_timeout,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.reranker_openrouter_api_key
|
||||
@@ -1588,6 +1679,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_openrouter_model,
|
||||
base_url="https://openrouter.ai/api/v1/rerank",
|
||||
timeout=config.reranker_openrouter_timeout,
|
||||
)
|
||||
elif provider == "flashrank":
|
||||
model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL)
|
||||
@@ -1602,6 +1694,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=config.reranker_litellm_api_key,
|
||||
model=config.reranker_litellm_model,
|
||||
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
|
||||
timeout=config.reranker_litellm_timeout,
|
||||
)
|
||||
elif provider == "litellm-sdk":
|
||||
api_key = config.reranker_litellm_sdk_api_key
|
||||
@@ -1614,6 +1707,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
model=config.reranker_litellm_sdk_model,
|
||||
api_base=config.reranker_litellm_sdk_api_base,
|
||||
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
|
||||
timeout=config.reranker_litellm_sdk_timeout,
|
||||
)
|
||||
elif provider == "zeroentropy":
|
||||
api_key = config.reranker_zeroentropy_api_key
|
||||
@@ -1624,6 +1718,8 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
return ZeroEntropyCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_zeroentropy_model,
|
||||
base_url=config.reranker_zeroentropy_base_url,
|
||||
timeout=config.reranker_zeroentropy_timeout,
|
||||
)
|
||||
elif provider == "siliconflow":
|
||||
api_key = config.reranker_siliconflow_api_key
|
||||
@@ -1635,6 +1731,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_siliconflow_model,
|
||||
base_url=config.reranker_siliconflow_base_url,
|
||||
timeout=config.reranker_siliconflow_timeout,
|
||||
)
|
||||
elif provider == "google":
|
||||
project_id = config.reranker_google_project_id
|
||||
@@ -1647,6 +1744,16 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
project_id=project_id,
|
||||
model=config.reranker_google_model,
|
||||
service_account_key=config.reranker_google_service_account_key,
|
||||
timeout=config.reranker_google_timeout,
|
||||
)
|
||||
elif provider == "alibaba":
|
||||
api_key = config.reranker_alibaba_api_key
|
||||
if not api_key:
|
||||
raise ValueError(f"{ENV_RERANKER_ALIBABA_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'alibaba'")
|
||||
return AlibabaCloudCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_alibaba_model,
|
||||
timeout=config.reranker_alibaba_timeout,
|
||||
)
|
||||
elif provider == "rrf":
|
||||
return RRFPassthroughCrossEncoder()
|
||||
@@ -1654,5 +1761,5 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
return JinaMLXCrossEncoder()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'alibaba', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
)
|
||||
|
||||
@@ -166,21 +166,6 @@ class DataAccessOps(ABC):
|
||||
|
||||
# -- LATERAL / fan-out queries ---------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_entity_unit_fanout(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
"""Fetch unit_ids for a list of entities with per-entity row cap.
|
||||
|
||||
PG uses unnest + CROSS JOIN LATERAL with LIMIT.
|
||||
Non-PG queries each entity individually.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
@@ -406,6 +391,74 @@ class DataAccessOps(ABC):
|
||||
"""Insert a webhook delivery task into async_operations."""
|
||||
...
|
||||
|
||||
# -- Graph maintenance queue -----------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def enqueue_graph_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
"""Insert unit_ids into graph_maintenance_queue, deduplicating on the
|
||||
(bank_id, unit_id) primary key.
|
||||
|
||||
Called inside the triggering transaction so enqueue is atomic with
|
||||
the mutation that caused it. Order is unspecified.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
"""Atomically claim a batch of rows from graph_maintenance_queue and
|
||||
remove them from the table.
|
||||
|
||||
Returns the list of ``unit_id`` strings. Empty list when the queue
|
||||
for ``bank_id`` is drained.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete entities in ``bank_id`` that no longer have any unit_entities
|
||||
rows referencing them. Returns the number of rows deleted.
|
||||
|
||||
FK ON DELETE CASCADE on entity_cooccurrences then removes any
|
||||
cooccurrence row pointing at the pruned entities.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete entity_cooccurrences rows in ``bank_id`` where the two
|
||||
entities still exist but no current unit references both of them.
|
||||
|
||||
These are stale-count rows: cooccurrence was real at the time it was
|
||||
recorded, but every memory_unit that witnessed both entities has
|
||||
since been deleted. Returns the number of rows deleted.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
@@ -416,6 +469,8 @@ class DataAccessOps(ABC):
|
||||
worker_id: str,
|
||||
reserved_limits: dict[str, int],
|
||||
shared_limit: int,
|
||||
*,
|
||||
consolidation_bank_priority: dict[str, int] | None = None,
|
||||
) -> list[ResultRow]:
|
||||
"""Claim pending tasks from the async_operations table.
|
||||
|
||||
@@ -423,6 +478,14 @@ class DataAccessOps(ABC):
|
||||
Oracle implementation uses two-step claims (query busy banks first, then
|
||||
claim excluding them) to avoid ORA-02014.
|
||||
|
||||
Args:
|
||||
consolidation_bank_priority: Per-bank priority for consolidation scheduling.
|
||||
Maps bank name patterns to integer priorities (higher = claimed first).
|
||||
Patterns support ``*`` as wildcard (converted to SQL ``%`` for LIKE).
|
||||
A bare ``*`` key is the catch-all default for unlisted banks.
|
||||
When set, consolidation tasks are claimed in priority tiers.
|
||||
None preserves current behavior (pure created_at ordering).
|
||||
|
||||
Returns claimed rows with operation_id, operation_type, task_payload, retry_count.
|
||||
The caller is responsible for building ClaimedTask objects.
|
||||
"""
|
||||
|
||||
@@ -215,29 +215,96 @@ class OracleOps(DataAccessOps):
|
||||
list(zip(unit_ids, entity_ids)),
|
||||
)
|
||||
|
||||
async def fetch_entity_unit_fanout(
|
||||
async def enqueue_graph_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
# Query each entity individually
|
||||
rows: list[ResultRow] = []
|
||||
for eid in entity_id_list:
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT $1 AS entity_id, ue.unit_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.entity_id = $1
|
||||
ORDER BY ue.unit_id DESC
|
||||
LIMIT $2
|
||||
""",
|
||||
eid,
|
||||
limit_per_entity,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
# Oracle doesn't support ON CONFLICT; rely on the PK and the
|
||||
# IGNORE_ROW_ON_DUPKEY_INDEX hint to skip duplicates server-side.
|
||||
# The hint name must match the PK constraint exactly.
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT /*+ IGNORE_ROW_ON_DUPKEY_INDEX({table}, pk_graph_maintenance_queue) */
|
||||
INTO {table} (bank_id, unit_id)
|
||||
VALUES ($1, $2)
|
||||
""",
|
||||
[(bank_id, uid) for uid in unit_ids],
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
# Two-step claim: select the batch, then delete by exact keys. Oracle's
|
||||
# DELETE ... RETURNING doesn't accept a multi-row subquery, so we can't
|
||||
# do it in one statement like the PG version.
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT unit_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
FETCH FIRST $2 ROWS ONLY
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
)
|
||||
claimed = [str(row["unit_id"]) for row in rows]
|
||||
if claimed:
|
||||
await conn.executemany(
|
||||
f"DELETE FROM {table} WHERE bank_id = $1 AND unit_id = $2",
|
||||
[(bank_id, uid) for uid in claimed],
|
||||
)
|
||||
rows.extend(entity_rows)
|
||||
return rows
|
||||
return claimed
|
||||
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# The Oracle DatabaseConnection wrapper reshapes ``cursor.rowcount`` into
|
||||
# the same ``"DELETE N"`` status string asyncpg returns, so the same
|
||||
# ``int(deleted.split()[-1])`` parsing works on both dialects.
|
||||
deleted = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {entities_table}
|
||||
WHERE bank_id = $1
|
||||
AND id NOT IN (SELECT DISTINCT entity_id FROM {ue_table})
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(deleted.split()[-1]) if isinstance(deleted, str) and deleted.startswith("DELETE") else 0
|
||||
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
deleted = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {ec_table}
|
||||
WHERE entity_id_1 IN (SELECT id FROM {entities_table} WHERE bank_id = $1)
|
||||
AND (entity_id_1, entity_id_2) NOT IN (
|
||||
SELECT u1.entity_id, u2.entity_id
|
||||
FROM {ue_table} u1
|
||||
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(deleted.split()[-1]) if isinstance(deleted, str) and deleted.startswith("DELETE") else 0
|
||||
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
@@ -720,7 +787,257 @@ class OracleOps(DataAccessOps):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def claim_tasks(self, conn, table, worker_id, reserved_limits, shared_limit):
|
||||
async def _claim_consolidation_tasks(
|
||||
self,
|
||||
conn,
|
||||
table: str,
|
||||
busy_bank_ids: list[str],
|
||||
claimed_ids: list,
|
||||
limit: int,
|
||||
priority_map: dict[str, int] | None,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with optional priority-based tiered ordering.
|
||||
|
||||
Mirrors the PostgreSQL implementation. The Oracle SQL adapter
|
||||
translates ``LIKE ANY`` / ``NOT LIKE ALL`` via ``_expand_any_lists``.
|
||||
"""
|
||||
if limit <= 0:
|
||||
return []
|
||||
|
||||
if not priority_map:
|
||||
return await self._claim_consolidation_plain(conn, table, busy_bank_ids, claimed_ids, limit)
|
||||
|
||||
# --- Tiered claiming (same algorithm as PG) ---
|
||||
specific_by_priority: dict[int, list[str]] = {}
|
||||
all_specific_sql: list[str] = []
|
||||
catch_all_priority = 1
|
||||
|
||||
for pattern, priority in priority_map.items():
|
||||
if pattern == "*":
|
||||
catch_all_priority = priority
|
||||
else:
|
||||
sql_pat = pattern.replace("*", "%")
|
||||
specific_by_priority.setdefault(priority, []).append(sql_pat)
|
||||
all_specific_sql.append(sql_pat)
|
||||
|
||||
all_priorities = sorted(set(specific_by_priority.keys()) | {catch_all_priority}, reverse=True)
|
||||
|
||||
remaining = limit
|
||||
result: list = []
|
||||
|
||||
for pri in all_priorities:
|
||||
if remaining <= 0:
|
||||
break
|
||||
|
||||
if pri in specific_by_priority:
|
||||
rows = await self._claim_consolidation_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
specific_by_priority[pri],
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
if pri == catch_all_priority and remaining > 0:
|
||||
rows = await self._claim_consolidation_not_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
all_specific_sql,
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
return result
|
||||
|
||||
async def _claim_consolidation_plain(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with default created_at ordering."""
|
||||
exclude_ids = claimed_ids if claimed_ids else None
|
||||
if busy_bank_ids:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
AND operation_id != ALL($2::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
sql_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks matching LIKE patterns."""
|
||||
params: list = [sql_patterns]
|
||||
conditions = ["bank_id LIKE ANY($1::text[])"]
|
||||
idx = 2
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra = " AND ".join(conditions)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND {extra}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_not_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
exclude_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks NOT matching any specific pattern (catch-all tier)."""
|
||||
params: list = []
|
||||
conditions: list[str] = []
|
||||
idx = 1
|
||||
|
||||
if exclude_patterns:
|
||||
conditions.append(f"bank_id NOT LIKE ALL(${idx}::text[])")
|
||||
params.append(exclude_patterns)
|
||||
idx += 1
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra_clause = (" AND " + " AND ".join(conditions)) if conditions else ""
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW()){extra_clause}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def claim_tasks(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
worker_id,
|
||||
reserved_limits,
|
||||
shared_limit,
|
||||
*,
|
||||
consolidation_bank_priority=None,
|
||||
):
|
||||
"""Oracle two-step claiming to avoid ORA-02014 with NOT EXISTS + FOR UPDATE."""
|
||||
all_rows = []
|
||||
claimed_ids = []
|
||||
@@ -731,7 +1048,6 @@ class OracleOps(DataAccessOps):
|
||||
continue
|
||||
|
||||
if op_type == "consolidation":
|
||||
# Two-step: find busy banks first, then claim excluding them
|
||||
busy_banks = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT bank_id FROM {table}
|
||||
@@ -740,38 +1056,14 @@ class OracleOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids = [r["bank_id"] for r in busy_banks]
|
||||
|
||||
if busy_bank_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -835,7 +1127,7 @@ class OracleOps(DataAccessOps):
|
||||
all_rows.append(row)
|
||||
remaining_shared -= len(rows)
|
||||
|
||||
# 2b. Consolidation tasks (with bank-serialization)
|
||||
# 2b. Consolidation tasks (with bank-serialization + optional priority)
|
||||
if remaining_shared > 0:
|
||||
busy_banks_2 = await conn.fetch(
|
||||
f"""
|
||||
@@ -845,76 +1137,14 @@ class OracleOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids_2 = [r["bank_id"] for r in busy_banks_2]
|
||||
|
||||
if claimed_ids:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
AND bank_id != ALL($2::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids_2,
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
|
||||
@@ -104,7 +104,46 @@ class PostgreSQLOps(DataAccessOps):
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
elif config.text_search_extension == "native":
|
||||
# search_vector is a regular tsvector column populated here using the
|
||||
# configured native dictionary. It used to be GENERATED ALWAYS with
|
||||
# a hardcoded 'english', which prevented per-deployment language
|
||||
# configuration. text_search_extension_native_language is validated
|
||||
# in HindsightConfig.validate() as a PG identifier, so embedding it
|
||||
# as a SQL literal is safe.
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals, search_vector)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals,
|
||||
to_tsvector(
|
||||
'{config.text_search_extension_native_language}'::regconfig,
|
||||
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, '')
|
||||
)
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
else:
|
||||
# pg_textsearch, pgroonga, and pg_search: search_vector is a dummy
|
||||
# TEXT column; the actual full-text index operates on the base text
|
||||
# columns directly, so we don't populate search_vector at insert time.
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
@@ -251,29 +290,101 @@ class PostgreSQLOps(DataAccessOps):
|
||||
entity_ids,
|
||||
)
|
||||
|
||||
async def fetch_entity_unit_fanout(
|
||||
async def enqueue_graph_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
return await conn.fetch(
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
await conn.execute(
|
||||
f"""
|
||||
SELECT e.entity_id, n.unit_id
|
||||
FROM unnest($1::uuid[]) AS e(entity_id)
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue.unit_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.entity_id = e.entity_id
|
||||
ORDER BY ue.unit_id DESC
|
||||
LIMIT $2
|
||||
) n
|
||||
INSERT INTO {table} (bank_id, unit_id)
|
||||
SELECT $1, v FROM unnest($2::uuid[]) AS t(v)
|
||||
ON CONFLICT (bank_id, unit_id) DO NOTHING
|
||||
""",
|
||||
entity_id_list,
|
||||
limit_per_entity,
|
||||
bank_id,
|
||||
unit_ids,
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
DELETE FROM {table}
|
||||
WHERE (bank_id, unit_id) IN (
|
||||
SELECT bank_id, unit_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
LIMIT $2
|
||||
)
|
||||
RETURNING unit_id
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
)
|
||||
return [str(row["unit_id"]) for row in rows]
|
||||
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# Scoped by entities.bank_id (indexed). The NOT EXISTS subquery is
|
||||
# backed by idx_ue_entity on unit_entities(entity_id), so this stays
|
||||
# linear in the number of entities in the bank — not in the size of
|
||||
# unit_entities globally.
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {entities_table} e
|
||||
WHERE e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {ue_table} ue WHERE ue.entity_id = e.id
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
# asyncpg returns "DELETE N"
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# Scope by joining through entities.bank_id (entity_cooccurrences itself
|
||||
# has no bank_id column — entities don't span banks, so scoping via
|
||||
# entity_id_1 is sufficient).
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {ec_table} c
|
||||
USING {entities_table} e
|
||||
WHERE e.id = c.entity_id_1
|
||||
AND e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {ue_table} u1
|
||||
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
|
||||
WHERE u1.entity_id = c.entity_id_1
|
||||
AND u2.entity_id = c.entity_id_2
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
@@ -726,7 +837,268 @@ class PostgreSQLOps(DataAccessOps):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def claim_tasks(self, conn, table, worker_id, reserved_limits, shared_limit):
|
||||
async def _claim_consolidation_tasks(
|
||||
self,
|
||||
conn,
|
||||
table: str,
|
||||
busy_bank_ids: list[str],
|
||||
claimed_ids: list,
|
||||
limit: int,
|
||||
priority_map: dict[str, int] | None,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with optional priority-based tiered ordering.
|
||||
|
||||
When *priority_map* is ``None``, uses the default ``ORDER BY created_at``
|
||||
with bank-serialization (exclude busy banks). When set, claims in
|
||||
priority tiers — highest-priority banks first. Specific patterns always
|
||||
take precedence over the catch-all ``*`` entry.
|
||||
"""
|
||||
if limit <= 0:
|
||||
return []
|
||||
|
||||
# --- Fast path: no priority map -> current behavior ---
|
||||
if not priority_map:
|
||||
return await self._claim_consolidation_plain(conn, table, busy_bank_ids, claimed_ids, limit)
|
||||
|
||||
# --- Tiered claiming ---
|
||||
# Separate specific patterns from catch-all.
|
||||
# Specific patterns always take precedence: a bank matching ``shadow-*``
|
||||
# uses that entry's priority even if the catch-all ``*`` has a higher
|
||||
# value. The catch-all only applies to banks not matching any specific
|
||||
# pattern.
|
||||
specific_by_priority: dict[int, list[str]] = {}
|
||||
all_specific_sql: list[str] = []
|
||||
catch_all_priority = 1 # default when no ``*`` entry
|
||||
|
||||
for pattern, priority in priority_map.items():
|
||||
if pattern == "*":
|
||||
catch_all_priority = priority
|
||||
else:
|
||||
sql_pat = pattern.replace("*", "%")
|
||||
specific_by_priority.setdefault(priority, []).append(sql_pat)
|
||||
all_specific_sql.append(sql_pat)
|
||||
|
||||
# Collect all priority levels (specific tiers + catch-all) sorted desc.
|
||||
all_priorities = sorted(set(specific_by_priority.keys()) | {catch_all_priority}, reverse=True)
|
||||
|
||||
remaining = limit
|
||||
result: list = []
|
||||
|
||||
for pri in all_priorities:
|
||||
if remaining <= 0:
|
||||
break
|
||||
|
||||
# Specific-pattern tier at this priority level
|
||||
if pri in specific_by_priority:
|
||||
rows = await self._claim_consolidation_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
specific_by_priority[pri],
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
# Catch-all tier at this priority level
|
||||
if pri == catch_all_priority and remaining > 0:
|
||||
rows = await self._claim_consolidation_not_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
all_specific_sql,
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
return result
|
||||
|
||||
async def _claim_consolidation_plain(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with default created_at ordering."""
|
||||
exclude_ids = claimed_ids if claimed_ids else None
|
||||
if busy_bank_ids:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
AND operation_id != ALL($2::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
sql_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks matching LIKE patterns."""
|
||||
params: list = [sql_patterns]
|
||||
conditions = ["bank_id LIKE ANY($1::text[])"]
|
||||
idx = 2
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra = " AND ".join(conditions)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND {extra}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_not_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
exclude_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks NOT matching any specific pattern (catch-all tier)."""
|
||||
params: list = []
|
||||
conditions: list[str] = []
|
||||
idx = 1
|
||||
|
||||
if exclude_patterns:
|
||||
conditions.append(f"bank_id NOT LIKE ALL(${idx}::text[])")
|
||||
params.append(exclude_patterns)
|
||||
idx += 1
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra_clause = (" AND " + " AND ".join(conditions)) if conditions else ""
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW()){extra_clause}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def claim_tasks(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
worker_id,
|
||||
reserved_limits,
|
||||
shared_limit,
|
||||
*,
|
||||
consolidation_bank_priority=None,
|
||||
):
|
||||
all_rows = []
|
||||
claimed_ids = []
|
||||
|
||||
@@ -744,38 +1116,14 @@ class PostgreSQLOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids = [r["bank_id"] for r in busy_banks]
|
||||
|
||||
if busy_bank_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -839,7 +1187,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
all_rows.append(row)
|
||||
remaining_shared -= len(rows)
|
||||
|
||||
# 2b. Consolidation tasks (with bank-serialization)
|
||||
# 2b. Consolidation tasks (with bank-serialization + optional priority)
|
||||
if remaining_shared > 0:
|
||||
busy_banks_2 = await conn.fetch(
|
||||
f"""
|
||||
@@ -849,76 +1197,14 @@ class PostgreSQLOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids_2 = [r["bank_id"] for r in busy_banks_2]
|
||||
|
||||
if claimed_ids:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
AND bank_id != ALL($2::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids_2,
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
|
||||
@@ -72,6 +72,9 @@ _RETURNING_RE = re.compile(r"\bRETURNING\s+(.+)", re.IGNORECASE | re.DOTALL)
|
||||
|
||||
_ANY_RE = re.compile(r"=\s*ANY\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
_NOT_ALL_RE = re.compile(r"!=\s*ALL\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
# LIKE ANY / NOT LIKE ALL — capture the column name before the operator
|
||||
_LIKE_ANY_RE = re.compile(r"(\w+)\s+LIKE\s+ANY\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
_NOT_LIKE_ALL_RE = re.compile(r"(\w+)\s+NOT\s+LIKE\s+ALL\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
|
||||
_JSON_ARROW_TEXT_RE = re.compile(r'("?\w+"?)\s*->>\s*\'(\w+)\'') # handles both col and "col"
|
||||
_JSON_HAS_KEY_RE = re.compile(r"(\w+)\s*\?\s*'(\w+)'")
|
||||
@@ -535,6 +538,12 @@ def _rewrite_pg_to_oracle(query: str) -> RewriteResult:
|
||||
# != ALL(:N) → NOT IN (expanded list) — the negative counterpart of = ANY
|
||||
query = _NOT_ALL_RE.sub(r"NOT IN (/*EXPAND:\1*/)", query)
|
||||
|
||||
# col LIKE ANY(:N) → (col LIKE :p0 OR col LIKE :p1 OR ...)
|
||||
query = _LIKE_ANY_RE.sub(r"\1 /*LIKE_ANY:\2:\1*/", query)
|
||||
|
||||
# col NOT LIKE ALL(:N) → (col NOT LIKE :p0 AND col NOT LIKE :p1 AND ...)
|
||||
query = _NOT_LIKE_ALL_RE.sub(r"\1 /*NOT_LIKE_ALL:\2:\1*/", query)
|
||||
|
||||
# CTE AS MATERIALIZED (...) → AS (...) — Oracle doesn't support MATERIALIZED CTE hint
|
||||
query = re.sub(r"\bAS\s+MATERIALIZED\s*\(", "AS (", query, flags=re.IGNORECASE)
|
||||
|
||||
@@ -724,17 +733,38 @@ class OracleConnection(DatabaseConnection):
|
||||
|
||||
_expand_counter = 0
|
||||
|
||||
@staticmethod
|
||||
def _resolve_list_param(params: dict[str, Any], key: str) -> list | None:
|
||||
"""Resolve a parameter that may be a list or a JSON-encoded list string."""
|
||||
val = params.get(key)
|
||||
if isinstance(val, str):
|
||||
try:
|
||||
parsed = json.loads(val)
|
||||
if isinstance(parsed, list):
|
||||
return parsed
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
if isinstance(val, (list, tuple)):
|
||||
return list(val)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _expand_any_lists(query: str, params: dict[str, Any] | None) -> tuple[str, dict[str, Any] | None]:
|
||||
"""Expand /*EXPAND:N*/ markers into individual bind vars for IN clauses.
|
||||
"""Expand /*EXPAND:N*/, /*LIKE_ANY:N:col*/, /*NOT_LIKE_ALL:N:col*/ markers.
|
||||
|
||||
Converts: IN (/*EXPAND:1*/) with params["1"] = [a, b, c]
|
||||
Into: IN (:any_0, :any_1, :any_2) with params["any_0"]=a, etc.
|
||||
|
||||
Converts: col /*LIKE_ANY:1:col*/ with params["1"] = [a, b]
|
||||
Into: (col LIKE :lk_0 OR col LIKE :lk_1)
|
||||
|
||||
Converts: col /*NOT_LIKE_ALL:1:col*/ with params["1"] = [a, b]
|
||||
Into: (col NOT LIKE :nlk_0 AND col NOT LIKE :nlk_1)
|
||||
|
||||
Uses a unique prefix to avoid name collisions with other bind vars.
|
||||
The original param is kept (for other references to :N in the query).
|
||||
"""
|
||||
if params is None or "/*EXPAND:" not in query:
|
||||
if params is None or "/*" not in query:
|
||||
return query, params
|
||||
|
||||
expand_re = re.compile(r"/\*EXPAND:(\d+)\*/")
|
||||
@@ -775,6 +805,50 @@ class OracleConnection(DatabaseConnection):
|
||||
|
||||
query = expand_re.sub(_replace, query)
|
||||
|
||||
# Expand LIKE ANY: col /*LIKE_ANY:N:col*/ → (col LIKE :p0 OR col LIKE :p1 ...)
|
||||
like_any_re = re.compile(r"(\w+)\s*/\*LIKE_ANY:(\d+):(\w+)\*/")
|
||||
|
||||
def _replace_like_any(m):
|
||||
_col = m.group(1) # redundant column ref before marker
|
||||
param_key = m.group(2)
|
||||
col = m.group(3)
|
||||
val = OracleConnection._resolve_list_param(params, param_key)
|
||||
if val is None or len(val) == 0:
|
||||
return "1=0" # no patterns → no match
|
||||
OracleConnection._expand_counter += 1
|
||||
prefix = f"lk{OracleConnection._expand_counter}"
|
||||
clauses = []
|
||||
for i, item in enumerate(val):
|
||||
k = f"{prefix}_{i}"
|
||||
params[k] = item
|
||||
clauses.append(f"{col} LIKE :{k}")
|
||||
keys_to_remove.add(param_key)
|
||||
return f"({' OR '.join(clauses)})"
|
||||
|
||||
query = like_any_re.sub(_replace_like_any, query)
|
||||
|
||||
# Expand NOT LIKE ALL: col /*NOT_LIKE_ALL:N:col*/ → (col NOT LIKE :p0 AND ...)
|
||||
not_like_all_re = re.compile(r"(\w+)\s*/\*NOT_LIKE_ALL:(\d+):(\w+)\*/")
|
||||
|
||||
def _replace_not_like_all(m):
|
||||
_col = m.group(1)
|
||||
param_key = m.group(2)
|
||||
col = m.group(3)
|
||||
val = OracleConnection._resolve_list_param(params, param_key)
|
||||
if val is None or len(val) == 0:
|
||||
return "1=1" # no patterns → everything matches
|
||||
OracleConnection._expand_counter += 1
|
||||
prefix = f"nlk{OracleConnection._expand_counter}"
|
||||
clauses = []
|
||||
for i, item in enumerate(val):
|
||||
k = f"{prefix}_{i}"
|
||||
params[k] = item
|
||||
clauses.append(f"{col} NOT LIKE :{k}")
|
||||
keys_to_remove.add(param_key)
|
||||
return f"({' AND '.join(clauses)})"
|
||||
|
||||
query = not_like_all_re.sub(_replace_not_like_all, query)
|
||||
|
||||
# Remove original list params that were expanded — their placeholder
|
||||
# (:N) no longer exists in the query, and leaving them causes DPY-4008.
|
||||
# Only remove if the key's placeholder is truly gone from the query.
|
||||
|
||||
@@ -9,13 +9,17 @@ The database schema is automatically adjusted to match the model's dimension.
|
||||
Configuration via environment variables - see hindsight_api.config for all env var names.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import logging
|
||||
import os
|
||||
import struct
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Literal, cast
|
||||
from urllib.parse import parse_qs, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ..config import (
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL,
|
||||
@@ -27,10 +31,15 @@ from ..config import (
|
||||
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
DEFAULT_EMBEDDINGS_PROVIDER,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL,
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
DEFAULT_ZEROENTROPY_BASE_URL,
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY,
|
||||
ENV_EMBEDDINGS_GEMINI_API_KEY,
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY,
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL,
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
@@ -39,12 +48,39 @@ from ..config import (
|
||||
ENV_EMBEDDINGS_OPENAI_MODEL,
|
||||
ENV_EMBEDDINGS_PROVIDER,
|
||||
ENV_EMBEDDINGS_TEI_URL,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_API_KEY,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
ENV_LLM_API_KEY,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
ZeroEntropyInputType = Literal["document", "query"]
|
||||
ZeroEntropyLatency = Literal["fast", "slow"]
|
||||
ZeroEntropyEncodingFormat = Literal["float", "base64"]
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedRequest(BaseModel):
|
||||
"""Typed request body for ZeroEntropy's non-OpenAI-compatible embed endpoint."""
|
||||
|
||||
model: str
|
||||
input: list[str]
|
||||
input_type: ZeroEntropyInputType
|
||||
dimensions: int
|
||||
encoding_format: ZeroEntropyEncodingFormat = "float"
|
||||
latency: ZeroEntropyLatency | None = None
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedResult(BaseModel):
|
||||
embedding: list[float] | str
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedResponse(BaseModel):
|
||||
results: list[_ZeroEntropyEmbedResult]
|
||||
|
||||
|
||||
class Embeddings(ABC):
|
||||
"""
|
||||
Abstract base class for embedding generation.
|
||||
@@ -88,6 +124,14 @@ class Embeddings(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings for query text. Providers without asymmetric embeddings use encode()."""
|
||||
return self.encode(texts)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings for stored document text. Providers without asymmetric embeddings use encode()."""
|
||||
return self.encode(texts)
|
||||
|
||||
|
||||
class LocalSTEmbeddings(Embeddings):
|
||||
"""
|
||||
@@ -385,6 +429,7 @@ class OpenAIEmbeddings(Embeddings):
|
||||
model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
base_url: str | None = None,
|
||||
batch_size: int = 100,
|
||||
dimensions: int | None = None,
|
||||
max_retries: int = 3,
|
||||
):
|
||||
"""
|
||||
@@ -395,12 +440,14 @@ class OpenAIEmbeddings(Embeddings):
|
||||
model: OpenAI embedding model name (default: text-embedding-3-small)
|
||||
base_url: Custom base URL for OpenAI-compatible API (e.g., Azure OpenAI endpoint)
|
||||
batch_size: Maximum batch size for embedding requests (default: 100)
|
||||
dimensions: Optional requested output dimensions for OpenAI text-embedding-3 models
|
||||
max_retries: Maximum number of retries for failed requests (default: 3)
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url
|
||||
self.batch_size = batch_size
|
||||
self.dimensions = dimensions
|
||||
self.max_retries = max_retries
|
||||
self._client = None
|
||||
self._dimension: int | None = None
|
||||
@@ -445,7 +492,9 @@ class OpenAIEmbeddings(Embeddings):
|
||||
self._client = OpenAI(**client_kwargs)
|
||||
|
||||
# Try to get dimension from known models, otherwise do a test embedding
|
||||
if self.model in self.MODEL_DIMENSIONS:
|
||||
if self.dimensions is not None:
|
||||
self._dimension = self.dimensions
|
||||
elif self.model in self.MODEL_DIMENSIONS:
|
||||
self._dimension = self.MODEL_DIMENSIONS[self.model]
|
||||
else:
|
||||
# Do a test embedding to detect dimension
|
||||
@@ -480,10 +529,14 @@ class OpenAIEmbeddings(Embeddings):
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
|
||||
response = self._client.embeddings.create(
|
||||
model=self.model,
|
||||
input=batch,
|
||||
)
|
||||
request = {
|
||||
"model": self.model,
|
||||
"input": batch,
|
||||
}
|
||||
if self.dimensions is not None:
|
||||
request["dimensions"] = self.dimensions
|
||||
|
||||
response = self._client.embeddings.create(**request)
|
||||
|
||||
# Sort by index to ensure correct order
|
||||
batch_embeddings = sorted(response.data, key=lambda x: x.index)
|
||||
@@ -492,6 +545,73 @@ class OpenAIEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
class CodexOAuthEmbeddings(OpenAIEmbeddings):
|
||||
"""
|
||||
OpenAI embeddings using the Codex/ChatGPT OAuth token from ``~/.codex/auth.json``.
|
||||
|
||||
Codex OAuth is an LLM-provider auth path in Hindsight, but the same bearer token
|
||||
can also authenticate against the standard OpenAI embeddings endpoint. This keeps
|
||||
embeddings on the user's existing Codex subscription/OAuth path without requiring
|
||||
a separate OpenAI/OpenRouter/Gemini/Cohere API key.
|
||||
|
||||
Token refresh is handled automatically: the manager proactively refreshes the
|
||||
access_token before it expires and reactively refreshes on 401 responses from
|
||||
the embeddings API.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
batch_size: int = 100,
|
||||
dimensions: int | None = None,
|
||||
max_retries: int = 3,
|
||||
):
|
||||
from .providers.codex_auth import CodexAuthManager
|
||||
|
||||
self._auth_manager = CodexAuthManager.from_file()
|
||||
super().__init__(
|
||||
api_key=self._auth_manager.access_token,
|
||||
model=model,
|
||||
base_url="https://api.openai.com/v1",
|
||||
batch_size=batch_size,
|
||||
dimensions=dimensions,
|
||||
max_retries=max_retries,
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "openai-codex"
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings, refreshing the OAuth token if needed.
|
||||
|
||||
Proactively refreshes before the call when the token is near expiry,
|
||||
and reactively refreshes once on a 401 from the OpenAI embeddings API.
|
||||
"""
|
||||
from openai import AuthenticationError
|
||||
|
||||
# Proactive refresh — cheap when fresh (JWT exp decode + compare).
|
||||
self._auth_manager.ensure_fresh_token()
|
||||
if self._auth_manager.access_token != self.api_key:
|
||||
self.api_key = self._auth_manager.access_token
|
||||
if self._client is not None:
|
||||
self._client.api_key = self._auth_manager.access_token
|
||||
|
||||
try:
|
||||
return super().encode(texts)
|
||||
except AuthenticationError:
|
||||
# Reactive refresh — token was valid by the JWT clock but the
|
||||
# server rejected it (rotated server-side, race, etc.).
|
||||
self._auth_manager.refresh_tokens(
|
||||
reason="reactive (401 from embeddings API)",
|
||||
force=True,
|
||||
)
|
||||
self.api_key = self._auth_manager.access_token
|
||||
if self._client is not None:
|
||||
self._client.api_key = self._auth_manager.access_token
|
||||
return super().encode(texts)
|
||||
|
||||
|
||||
class CohereEmbeddings(Embeddings):
|
||||
"""
|
||||
Cohere embeddings implementation using the Cohere API.
|
||||
@@ -633,6 +753,149 @@ class CohereEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
class ZeroEntropyEmbeddings(Embeddings):
|
||||
"""
|
||||
ZeroEntropy embeddings implementation using the zembed API.
|
||||
|
||||
ZeroEntropy's embeddings endpoint is not OpenAI-compatible: it lives at
|
||||
/v1/models/embed and requires provider-specific parameters such as
|
||||
input_type. Hindsight stores document-side vectors and uses query-side
|
||||
vectors during recall, so this provider exposes explicit encode_documents()
|
||||
and encode_query() helpers while keeping encode() as document-side default.
|
||||
"""
|
||||
|
||||
VALID_DIMENSIONS = frozenset({2560, 1280, 640, 320, 160, 80, 40})
|
||||
VALID_ENCODING_FORMATS = frozenset({"float", "base64"})
|
||||
VALID_LATENCIES = frozenset({"fast", "slow"})
|
||||
DEFAULT_BASE_URL = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
EMBED_PATH = "/v1/models/embed"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL,
|
||||
base_url: str | None = None,
|
||||
dimensions: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
batch_size: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
encoding_format: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
latency: str | None = DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
if dimensions not in self.VALID_DIMENSIONS:
|
||||
valid = ", ".join(str(dim) for dim in sorted(self.VALID_DIMENSIONS, reverse=True))
|
||||
raise ValueError(f"{ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS} must be one of {valid}, got {dimensions}")
|
||||
if batch_size < 1:
|
||||
raise ValueError("ZeroEntropy embeddings batch_size must be >= 1")
|
||||
if encoding_format not in self.VALID_ENCODING_FORMATS:
|
||||
valid_formats = ", ".join(sorted(self.VALID_ENCODING_FORMATS))
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT} must be one of {valid_formats}, got {encoding_format!r}"
|
||||
)
|
||||
if latency is not None and latency not in self.VALID_LATENCIES:
|
||||
valid_latencies = ", ".join(sorted(self.VALID_LATENCIES))
|
||||
raise ValueError(f"ZeroEntropy embeddings latency must be one of {valid_latencies}, got {latency!r}")
|
||||
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url.rstrip("/") if base_url else self.DEFAULT_BASE_URL
|
||||
self.embed_url = f"{self.base_url}{self.EMBED_PATH}"
|
||||
self.dimensions = dimensions
|
||||
self.batch_size = batch_size
|
||||
self.encoding_format = cast(ZeroEntropyEncodingFormat, encoding_format)
|
||||
self.latency = cast(ZeroEntropyLatency | None, latency)
|
||||
self.timeout = timeout
|
||||
self._client: httpx.Client | None = None
|
||||
self._dimension: int | None = None
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "zeroentropy"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
if self._dimension is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
return self._dimension
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the ZeroEntropy HTTP client."""
|
||||
if self._client is not None:
|
||||
return
|
||||
|
||||
logger.info(
|
||||
f"Embeddings: initializing ZeroEntropy provider with model {self.model} "
|
||||
f"(dim: {self.dimensions}, batch_size={self.batch_size})"
|
||||
)
|
||||
self._client = httpx.Client(
|
||||
timeout=self.timeout,
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
# zembed-1 dimensions are explicit Matryoshka truncation steps. Avoid a
|
||||
# startup probe so boot does not burn quota or require a throwaway input.
|
||||
self._dimension = self.dimensions
|
||||
logger.info(f"Embeddings: ZeroEntropy provider initialized (model: {self.model}, dim: {self._dimension})")
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate document-side embeddings for backwards-compatible callers."""
|
||||
return self.encode_documents(texts)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate document-side embeddings for retained content."""
|
||||
return self._encode_with_input_type(texts, "document")
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate query-side embeddings for recall/search queries."""
|
||||
return self._encode_with_input_type(texts, "query")
|
||||
|
||||
def _encode_with_input_type(self, texts: list[str], input_type: ZeroEntropyInputType) -> list[list[float]]:
|
||||
if self._client is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
all_embeddings: list[list[float]] = []
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
request = _ZeroEntropyEmbedRequest(
|
||||
model=self.model,
|
||||
input=batch,
|
||||
input_type=input_type,
|
||||
dimensions=self.dimensions,
|
||||
encoding_format=self.encoding_format,
|
||||
latency=self.latency,
|
||||
)
|
||||
|
||||
try:
|
||||
response = self._client.post(self.embed_url, json=request.model_dump(exclude_none=True))
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPError as e:
|
||||
raise RuntimeError(f"ZeroEntropy embedding request failed: {e}") from e
|
||||
|
||||
parsed = _ZeroEntropyEmbedResponse.model_validate(response.json())
|
||||
if len(parsed.results) != len(batch):
|
||||
raise RuntimeError(
|
||||
f"ZeroEntropy returned {len(parsed.results)} embeddings for {len(batch)} input texts; "
|
||||
"expected exact 1:1 alignment"
|
||||
)
|
||||
all_embeddings.extend(self._parse_embedding(result.embedding) for result in parsed.results)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
@staticmethod
|
||||
def _parse_embedding(embedding: list[float] | str) -> list[float]:
|
||||
if not isinstance(embedding, str):
|
||||
return embedding
|
||||
|
||||
raw = base64.b64decode(embedding)
|
||||
if len(raw) % 4 != 0:
|
||||
raise RuntimeError("ZeroEntropy returned invalid base64 embedding length")
|
||||
return list(struct.unpack(f"<{len(raw) // 4}f", raw))
|
||||
|
||||
|
||||
class LiteLLMEmbeddings(Embeddings):
|
||||
"""
|
||||
LiteLLM embeddings implementation using LiteLLM proxy's /embeddings endpoint.
|
||||
@@ -766,7 +1029,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
api_key: str | None = None,
|
||||
model: str = DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
api_base: str | None = None,
|
||||
output_dimensions: int | None = None,
|
||||
@@ -778,7 +1041,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
Initialize LiteLLM SDK embeddings client.
|
||||
|
||||
Args:
|
||||
api_key: API key for the embedding provider
|
||||
api_key: API key for the embedding provider (optional — omit for
|
||||
providers that use ambient credentials, e.g. AWS Bedrock with IAM)
|
||||
model: Model name with provider prefix (e.g., "cohere/embed-english-v3.0")
|
||||
api_base: Custom base URL for API (optional)
|
||||
output_dimensions: Optional output embedding dimensions (provider-dependent)
|
||||
@@ -828,8 +1092,9 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
embed_kwargs = {
|
||||
"model": self.model,
|
||||
"input": ["test"],
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_key:
|
||||
embed_kwargs["api_key"] = self.api_key
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
@@ -880,8 +1145,9 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
embed_kwargs = {
|
||||
"model": self.model,
|
||||
"input": batch,
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_key:
|
||||
embed_kwargs["api_key"] = self.api_key
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
@@ -1140,6 +1406,14 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=model,
|
||||
base_url=base_url,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "openai-codex":
|
||||
model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL)
|
||||
return CodexOAuthEmbeddings(
|
||||
model=model,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.embeddings_openrouter_api_key
|
||||
@@ -1153,6 +1427,23 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=config.embeddings_openrouter_model,
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "zeroentropy":
|
||||
api_key = config.embeddings_zeroentropy_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_ZEROENTROPY_API_KEY} or ZEROENTROPY_API_KEY is required "
|
||||
f"when {ENV_EMBEDDINGS_PROVIDER} is 'zeroentropy'"
|
||||
)
|
||||
return ZeroEntropyEmbeddings(
|
||||
api_key=api_key,
|
||||
model=config.embeddings_zeroentropy_model,
|
||||
base_url=config.embeddings_zeroentropy_base_url,
|
||||
dimensions=config.embeddings_zeroentropy_dimensions,
|
||||
batch_size=config.embeddings_zeroentropy_batch_size,
|
||||
encoding_format=config.embeddings_zeroentropy_encoding_format,
|
||||
latency=config.embeddings_zeroentropy_latency,
|
||||
)
|
||||
elif provider == "cohere":
|
||||
api_key = config.embeddings_cohere_api_key
|
||||
@@ -1171,13 +1462,8 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=config.embeddings_litellm_model,
|
||||
)
|
||||
elif provider == "litellm-sdk":
|
||||
api_key = config.embeddings_litellm_sdk_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_LITELLM_SDK_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'litellm-sdk'"
|
||||
)
|
||||
return LiteLLMSDKEmbeddings(
|
||||
api_key=api_key,
|
||||
api_key=config.embeddings_litellm_sdk_api_key or None,
|
||||
model=config.embeddings_litellm_sdk_model,
|
||||
api_base=config.embeddings_litellm_sdk_api_base,
|
||||
output_dimensions=config.embeddings_litellm_sdk_output_dimensions,
|
||||
@@ -1206,5 +1492,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown embeddings provider: {provider}. "
|
||||
f"Supported: 'local', 'tei', 'openai', 'cohere', 'google', 'litellm', 'litellm-sdk'"
|
||||
f"Supported: 'local', 'tei', 'openai', 'openai-codex', 'openrouter', 'cohere', 'google', "
|
||||
f"'zeroentropy', 'litellm', 'litellm-sdk'"
|
||||
)
|
||||
|
||||
@@ -16,7 +16,15 @@ from typing import Any, Final
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
from .memory_engine import fq_table
|
||||
from .retain.entity_labels import build_labels_lookup as _build_labels_lookup_from_config
|
||||
from .retain.entity_labels import (
|
||||
build_labels_lookup as _build_labels_lookup_from_config,
|
||||
)
|
||||
from .retain.entity_labels import (
|
||||
is_label_entity as _is_label_entity,
|
||||
)
|
||||
from .retain.entity_labels import (
|
||||
parse_entity_labels as _parse_entity_labels,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -89,7 +97,12 @@ class EntityResolver:
|
||||
Resolves entities to canonical IDs with disambiguation.
|
||||
"""
|
||||
|
||||
def __init__(self, pool: Any, entity_lookup: str = "full"):
|
||||
def __init__(
|
||||
self,
|
||||
pool: Any,
|
||||
entity_lookup: str = "full",
|
||||
entity_resolution_batch_size: int = 100,
|
||||
):
|
||||
"""
|
||||
Initialize entity resolver.
|
||||
|
||||
@@ -98,9 +111,14 @@ class EntityResolver:
|
||||
entity_lookup: Lookup strategy — "full" loads all bank entities then
|
||||
matches in Python; "trigram" uses pg_trgm GIN index to fetch only
|
||||
similar candidates per entity name (much faster for large banks).
|
||||
entity_resolution_batch_size: Number of unique entity names to include
|
||||
in each pg_trgm candidate lookup query.
|
||||
"""
|
||||
self.pool = pool
|
||||
self.entity_lookup = entity_lookup
|
||||
if entity_resolution_batch_size < 1:
|
||||
raise ValueError("entity_resolution_batch_size must be >= 1")
|
||||
self.entity_resolution_batch_size = entity_resolution_batch_size
|
||||
self._pg_trgm_checked = False
|
||||
# Backend-specific operations — accessed via pool.ops (Django pattern).
|
||||
self._ops = pool.ops if pool is not None else None
|
||||
@@ -199,6 +217,11 @@ class EntityResolver:
|
||||
"""Build a set of valid 'key:value' entity label strings for fast lookup."""
|
||||
return _build_labels_lookup_from_config(entity_labels)
|
||||
|
||||
@staticmethod
|
||||
def _chunked(values: list[str], size: int) -> list[list[str]]:
|
||||
"""Split values into fixed-size batches."""
|
||||
return [values[i : i + size] for i in range(0, len(values), size)]
|
||||
|
||||
async def resolve_entities_batch(
|
||||
self,
|
||||
bank_id: str,
|
||||
@@ -228,14 +251,15 @@ class EntityResolver:
|
||||
return []
|
||||
|
||||
taxonomy_lookup = self._build_labels_lookup(entity_labels)
|
||||
labels_cfg = _parse_entity_labels(entity_labels)
|
||||
if conn is None:
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
return await self._resolve_entities_batch_impl(
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
else:
|
||||
return await self._resolve_entities_batch_impl(
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_impl(
|
||||
@@ -246,13 +270,16 @@ class EntityResolver:
|
||||
context: str,
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
if self.entity_lookup == "trigram":
|
||||
# Route to backend-specific fuzzy strategy.
|
||||
# Non-PG backends (Oracle) use UTL_MATCH instead of pg_trgm.
|
||||
backend_strategy = self._ops.get_entity_resolution_strategy()
|
||||
if backend_strategy == "oracle_fuzzy":
|
||||
return await self._resolve_entities_batch_oracle_fuzzy(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_oracle_fuzzy(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
# Auto-detect pg_trgm availability on first call and fall back to
|
||||
# "full" strategy if the extension is not installed. See #626.
|
||||
if not self._pg_trgm_checked:
|
||||
@@ -266,12 +293,24 @@ class EntityResolver:
|
||||
"https://github.com/vectorize-io/hindsight/issues/626"
|
||||
)
|
||||
self.entity_lookup = "full"
|
||||
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_trigram(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_full(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_trigram(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_full(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_full(
|
||||
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""Original strategy: load all bank entities then match in Python."""
|
||||
# Query ALL candidates for this bank
|
||||
@@ -338,11 +377,24 @@ class EntityResolver:
|
||||
all_candidates[entity_text] = matching
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_trigram(
|
||||
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Trigram strategy: fetch only similar candidates per entity name using pg_trgm.
|
||||
@@ -353,7 +405,7 @@ class EntityResolver:
|
||||
"""
|
||||
entity_texts = list(set(e["text"] for e in entities_data))
|
||||
|
||||
# Fetch candidates for all unique entity texts in a single batched query.
|
||||
# Fetch candidates for unique entity texts in bounded batches.
|
||||
# Uses the GIN trigram index on LOWER(canonical_name) for case-insensitive
|
||||
# similarity lookup. Previous version also had LIKE '%...' substring fallbacks,
|
||||
# but those forced full sequential scans of the entities table and caused
|
||||
@@ -361,21 +413,32 @@ class EntityResolver:
|
||||
# to 0.15 (from default 0.3) catches most substring relationships while
|
||||
# staying fully index-based.
|
||||
await conn.execute("SET pg_trgm.similarity_threshold = 0.15")
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT ON (e.id)
|
||||
e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM unnest($2::text[]) AS q(query_text)
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
e.bank_id = $1
|
||||
AND LOWER(e.canonical_name) % LOWER(q.query_text)
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_texts,
|
||||
)
|
||||
await conn.execute("RESET pg_trgm.similarity_threshold")
|
||||
try:
|
||||
rows = []
|
||||
for entity_text_batch in self._chunked(entity_texts, self.entity_resolution_batch_size):
|
||||
rows.extend(
|
||||
await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT ON (e.id)
|
||||
e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM unnest($2::text[]) AS q(query_text)
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
e.bank_id = $1
|
||||
AND LOWER(e.canonical_name) % LOWER(q.query_text)
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_text_batch,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
# asyncpg returns connections to the pool with session state intact,
|
||||
# so the lowered threshold would leak to future borrowers without RESET.
|
||||
try:
|
||||
await conn.execute("RESET pg_trgm.similarity_threshold")
|
||||
except Exception:
|
||||
logger.warning("Failed to reset pg_trgm similarity threshold after candidate lookup", exc_info=True)
|
||||
|
||||
# Group candidates by query_text
|
||||
all_candidates: dict[str, list] = {t: [] for t in entity_texts}
|
||||
@@ -418,11 +481,24 @@ class EntityResolver:
|
||||
cooccurrence_map[eid2].add(id_to_name[eid1])
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_oracle_fuzzy(
|
||||
self, conn: Any, bank_id: str, entities_data: list[dict], unit_event_date: datetime | None
|
||||
self,
|
||||
conn: Any,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date: datetime | None,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Oracle strategy: fetch similar candidates using UTL_MATCH.JARO_WINKLER_SIMILARITY.
|
||||
@@ -436,23 +512,28 @@ class EntityResolver:
|
||||
entities_table = fq_table("entities")
|
||||
|
||||
try:
|
||||
# Batch all entity texts into a single query using JSON_TABLE to
|
||||
# Batch entity texts into bounded sub-queries using JSON_TABLE to
|
||||
# expand the list into rows. UTL_MATCH.JARO_WINKLER_SIMILARITY
|
||||
# returns 0-100; threshold 70 ≈ pg_trgm similarity 0.15.
|
||||
entity_texts_json = json.dumps(entity_texts)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM JSON_TABLE($2, '$[*]' COLUMNS (query_text VARCHAR2(4000) PATH '$')) q
|
||||
JOIN {entities_table} e ON (
|
||||
e.bank_id = $1
|
||||
AND UTL_MATCH.JARO_WINKLER_SIMILARITY(LOWER(e.canonical_name), LOWER(q.query_text)) > 70
|
||||
# Bounded batches mirror the PG trigram path so very wide retain
|
||||
# batches don't time out a single JOIN on large banks.
|
||||
rows = []
|
||||
for entity_text_batch in self._chunked(entity_texts, self.entity_resolution_batch_size):
|
||||
rows.extend(
|
||||
await conn.fetch(
|
||||
f"""
|
||||
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM JSON_TABLE($2, '$[*]' COLUMNS (query_text VARCHAR2(4000) PATH '$')) q
|
||||
JOIN {entities_table} e ON (
|
||||
e.bank_id = $1
|
||||
AND UTL_MATCH.JARO_WINKLER_SIMILARITY(LOWER(e.canonical_name), LOWER(q.query_text)) > 70
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
json.dumps(entity_text_batch),
|
||||
)
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_texts_json,
|
||||
)
|
||||
except Exception as e:
|
||||
# UTL_MATCH may not be available (ORA-06550, ORA-00904, etc.)
|
||||
# Catch broadly because Oracle error types vary depending on driver.
|
||||
@@ -506,7 +587,14 @@ class EntityResolver:
|
||||
cooccurrence_map[eid2].add(id_to_name[eid1])
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
)
|
||||
|
||||
async def _resolve_from_candidates(
|
||||
@@ -517,6 +605,8 @@ class EntityResolver:
|
||||
unit_event_date,
|
||||
all_candidates: dict[str, list],
|
||||
cooccurrence_map: dict[str, set[str]],
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""Shared scoring + upsert logic used by both lookup strategies."""
|
||||
|
||||
@@ -533,11 +623,34 @@ class EntityResolver:
|
||||
|
||||
candidates = all_candidates.get(entity_text, [])
|
||||
|
||||
# Label entities (from entity_labels config) use exact matching only.
|
||||
# Their canonical names are user-defined (e.g., "use:use-001"),
|
||||
# so fuzzy resolution must NOT merge distinct label values that
|
||||
# happen to be textually similar (GH-1558).
|
||||
is_label = bool(
|
||||
labels_cfg and taxonomy_lookup and _is_label_entity(entity_text, labels_cfg, taxonomy_lookup)
|
||||
)
|
||||
|
||||
if not candidates:
|
||||
# Will create new entity
|
||||
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
|
||||
continue
|
||||
|
||||
if is_label:
|
||||
# Exact case-insensitive match only for label entities
|
||||
exact_match = None
|
||||
entity_text_lower = entity_text.lower()
|
||||
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
|
||||
if canonical_name.lower() == entity_text_lower:
|
||||
exact_match = candidate_id
|
||||
break
|
||||
if exact_match:
|
||||
entity_ids[idx] = exact_match
|
||||
entities_to_update.append(_EntityStat(entity_id=exact_match, event_date=entity_event_date))
|
||||
else:
|
||||
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
|
||||
continue
|
||||
|
||||
# Score candidates
|
||||
best_candidate = None
|
||||
best_score = 0.0
|
||||
|
||||
@@ -0,0 +1,357 @@
|
||||
"""Async graph maintenance after document/unit deletes.
|
||||
|
||||
Three reconciliation passes run together on every worker invocation:
|
||||
|
||||
1. **Relink top-up.** Drain ``graph_maintenance_queue`` (units whose
|
||||
outgoing temporal/semantic links lost a neighbour to a delete). For
|
||||
each, count current outgoing links per type; if below cap, run the
|
||||
same probes retain uses (:func:`fetch_temporal_neighbors`,
|
||||
:func:`compute_semantic_links_ann`) and insert the missing links.
|
||||
``bulk_insert_links`` has ``ON CONFLICT DO NOTHING`` on the uniqueness
|
||||
key, so we can re-probe freely and the DB de-dupes.
|
||||
|
||||
2. **Orphan entity prune.** Delete ``entities`` rows in the bank that no
|
||||
longer have any ``unit_entities`` references. FK ON DELETE CASCADE on
|
||||
``entity_cooccurrences`` then removes any cooccurrence row pointing
|
||||
at the pruned entities.
|
||||
|
||||
3. **Stale cooccurrence prune.** Defensive sweep for cooccurrence rows
|
||||
where both endpoints still exist but no current memory_unit references
|
||||
both of them — the cooccurrence was real at the time it was recorded,
|
||||
but every unit that witnessed it has since been deleted.
|
||||
|
||||
All three passes run on every invocation. The queue is the only source
|
||||
of work for pass 1; passes 2 and 3 are bank-wide sweeps backed by indexes
|
||||
on ``entities(bank_id)`` and ``unit_entities(entity_id)``, so they're
|
||||
cheap when there's nothing to do.
|
||||
|
||||
The worker dedupes on bank: a second job for the same bank is dropped
|
||||
while one is pending. Once processing starts, a new job becomes the
|
||||
*next* pending slot — so work enqueued during processing gets picked up
|
||||
by the follow-up run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
import uuid as uuid_module
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..models import RequestContext
|
||||
from .db.base import DatabaseConnection
|
||||
from .retain.link_utils import (
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
_bulk_insert_links,
|
||||
_normalize_datetime,
|
||||
compute_semantic_links_ann,
|
||||
)
|
||||
from .schema import fq_table
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .memory_engine import MemoryEngine
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Mirrors the ``top_k`` default in ``compute_semantic_links_ann`` at retain
|
||||
# time. If you change one, change the other — otherwise victims would either
|
||||
# never reach the cap (probe returns less than the cap) or stay perpetually
|
||||
# under it (cap is higher than retain creates).
|
||||
MAX_SEMANTIC_LINKS_PER_UNIT = 50
|
||||
|
||||
# Worker fetches this many rows per relink-loop iteration. Bounds
|
||||
# per-iteration probe/insert latency so a 10k-row backlog doesn't hold a
|
||||
# worker slot for minutes. Chosen so the typical iteration runs in well
|
||||
# under 1s.
|
||||
_DRAIN_BATCH_SIZE = 50
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobResult:
|
||||
"""Counters surfaced to the worker dispatcher and operation result."""
|
||||
|
||||
relink_units_processed: int = 0
|
||||
relink_links_added: int = 0
|
||||
orphan_entities_pruned: int = 0
|
||||
stale_cooccurrences_pruned: int = 0
|
||||
|
||||
def as_dict(self) -> dict[str, int]:
|
||||
return {
|
||||
"relink_units_processed": self.relink_units_processed,
|
||||
"relink_links_added": self.relink_links_added,
|
||||
"orphan_entities_pruned": self.orphan_entities_pruned,
|
||||
"stale_cooccurrences_pruned": self.stale_cooccurrences_pruned,
|
||||
}
|
||||
|
||||
|
||||
async def enqueue_relink_victims(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
deleted_unit_ids: list[str],
|
||||
ops: Any,
|
||||
) -> int:
|
||||
"""Enqueue surviving units whose outgoing temporal/semantic links pointed at
|
||||
``deleted_unit_ids`` for later link top-up.
|
||||
|
||||
Must run inside the same transaction that deletes the units, *before* the
|
||||
cascade fires — once the rows are gone, the join that finds the victims
|
||||
returns nothing.
|
||||
|
||||
Args:
|
||||
conn: Database connection inside the active delete transaction.
|
||||
bank_id: Bank owning the deleted units.
|
||||
deleted_unit_ids: Memory_unit IDs about to be (or being) deleted.
|
||||
ops: ``DataAccessOps`` instance, supplies the dialect-specific
|
||||
bulk-insert path.
|
||||
|
||||
Returns:
|
||||
Number of distinct victim units enqueued (after dedup against rows
|
||||
already in the queue).
|
||||
"""
|
||||
if not deleted_unit_ids:
|
||||
return 0
|
||||
|
||||
deleted_uuids = [uuid_module.UUID(uid) if isinstance(uid, str) else uid for uid in deleted_unit_ids]
|
||||
deleted_str_set = {str(uid) for uid in deleted_uuids}
|
||||
|
||||
# Find units (other than the ones being deleted) that have an outgoing
|
||||
# temporal/semantic link pointing at a doomed unit. Entity links are
|
||||
# intentionally excluded — they're scheduled for removal and would only
|
||||
# add noise to the recompute job.
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT from_unit_id
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE to_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
""",
|
||||
deleted_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
victim_ids = [row["from_unit_id"] for row in victim_rows if str(row["from_unit_id"]) not in deleted_str_set]
|
||||
|
||||
if not victim_ids:
|
||||
return 0
|
||||
|
||||
await ops.enqueue_graph_maintenance(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
victim_ids,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"[GRAPH_MAINT] Enqueued {len(victim_ids)} relink victims in "
|
||||
f"bank={bank_id} (deleted {len(deleted_unit_ids)} units)"
|
||||
)
|
||||
return len(victim_ids)
|
||||
|
||||
|
||||
async def run_graph_maintenance_job(
|
||||
memory_engine: "MemoryEngine",
|
||||
bank_id: str,
|
||||
request_context: RequestContext,
|
||||
operation_id: str | None = None,
|
||||
) -> dict[str, int]:
|
||||
"""Run all maintenance passes for ``bank_id`` until the relink queue is
|
||||
drained, then sweep entities and cooccurrences once.
|
||||
|
||||
Returns:
|
||||
Per-pass counters from :class:`JobResult`.
|
||||
"""
|
||||
del request_context # accepted for symmetry with other run_*_job helpers
|
||||
backend = await memory_engine._get_backend()
|
||||
ops = backend.ops
|
||||
|
||||
result = JobResult()
|
||||
job_start = time.time()
|
||||
|
||||
# --- Pass 1: relink ---
|
||||
# Per-iteration loop: claim → top up → commit. We rely on submit-time
|
||||
# dedup to keep at most one job per bank running, so no need for
|
||||
# SKIP LOCKED.
|
||||
iterations = 0
|
||||
while True:
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
unit_ids = await ops.claim_graph_maintenance_batch(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
_DRAIN_BATCH_SIZE,
|
||||
)
|
||||
if not unit_ids:
|
||||
break
|
||||
|
||||
result.relink_links_added += await _relink_batch(conn, bank_id, unit_ids, ops, backend)
|
||||
|
||||
result.relink_units_processed += len(unit_ids)
|
||||
iterations += 1
|
||||
|
||||
if iterations > 10000:
|
||||
# Defensive guard against runaway loops — at 50 units/iter that's
|
||||
# 500k targets, far beyond any realistic single-bank backlog.
|
||||
logger.error(
|
||||
f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink ({result.as_dict()})"
|
||||
)
|
||||
break
|
||||
|
||||
# --- Pass 2 & 3: entity / cooccurrence sweeps ---
|
||||
# Bank-wide single-statement deletes. Cheap when there's nothing to do.
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
result.orphan_entities_pruned = await ops.prune_orphan_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
fq_table("unit_entities"),
|
||||
bank_id,
|
||||
)
|
||||
# The orphan prune above cascades cooccurrences via FK. The
|
||||
# explicit cooccurrence pass below catches the *stale-count*
|
||||
# case: both entities still exist but no current unit witnesses
|
||||
# them together.
|
||||
result.stale_cooccurrences_pruned = await ops.prune_stale_cooccurrences(
|
||||
conn,
|
||||
fq_table("entity_cooccurrences"),
|
||||
fq_table("unit_entities"),
|
||||
fq_table("entities"),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
elapsed = time.time() - job_start
|
||||
logger.info(
|
||||
f"[GRAPH_MAINT] bank={bank_id} done: {result.as_dict()}, elapsed={elapsed:.2f}s, operation_id={operation_id}"
|
||||
)
|
||||
return result.as_dict()
|
||||
|
||||
|
||||
async def _relink_batch(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
victim_ids: list[str],
|
||||
ops: Any,
|
||||
backend: Any,
|
||||
) -> int:
|
||||
"""Top up temporal/semantic links for a batch of victim units. Returns rows inserted."""
|
||||
# Load each victim's metadata. Victims whose units were deleted between
|
||||
# enqueue and now silently drop out — exactly the no-op behaviour we want
|
||||
# for stale queue rows.
|
||||
victim_uuids = [uuid_module.UUID(vid) for vid in victim_ids]
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id::text AS id, event_date, fact_type, embedding::text AS embedding
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
victim_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not victim_rows:
|
||||
return 0
|
||||
|
||||
alive_uuids = [uuid_module.UUID(row["id"]) for row in victim_rows]
|
||||
|
||||
# Count current outgoing temporal/semantic links per victim so we only
|
||||
# probe for the ones genuinely below cap. Saves the bulk of the work when
|
||||
# most victims still have plenty of links.
|
||||
count_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_unit_id, link_type, COUNT(*) AS cnt
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE from_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
GROUP BY from_unit_id, link_type
|
||||
""",
|
||||
alive_uuids,
|
||||
bank_id,
|
||||
)
|
||||
counts: dict[tuple[str, str], int] = {}
|
||||
for row in count_rows:
|
||||
counts[(str(row["from_unit_id"]), row["link_type"])] = int(row["cnt"])
|
||||
|
||||
# --- Temporal top-up ---
|
||||
temporal_needs = [r for r in victim_rows if counts.get((r["id"], "temporal"), 0) < MAX_TEMPORAL_LINKS_PER_UNIT]
|
||||
new_links: list[tuple] = []
|
||||
|
||||
if temporal_needs:
|
||||
lateral_unit_ids = [uuid_module.UUID(r["id"]) for r in temporal_needs if r["event_date"] is not None]
|
||||
lateral_event_dates = [
|
||||
_normalize_datetime(r["event_date"]) for r in temporal_needs if r["event_date"] is not None
|
||||
]
|
||||
lateral_fact_types = [r["fact_type"] for r in temporal_needs if r["event_date"] is not None]
|
||||
|
||||
if lateral_unit_ids:
|
||||
rows = await ops.fetch_temporal_neighbors(
|
||||
conn,
|
||||
fq_table("memory_units"),
|
||||
bank_id,
|
||||
lateral_unit_ids,
|
||||
lateral_event_dates,
|
||||
lateral_fact_types,
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
)
|
||||
for row in rows:
|
||||
time_diff_h = float(row["time_diff_hours"])
|
||||
# Mirror the 24h window enforced at retain time. The bidirectional
|
||||
# index scan returns the K closest neighbours regardless of
|
||||
# window, so we filter here.
|
||||
if time_diff_h > 24:
|
||||
continue
|
||||
weight = max(0.3, 1.0 - (time_diff_h / 24))
|
||||
new_links.append((row["from_id"], str(row["id"]), "temporal", weight, None))
|
||||
|
||||
# --- Semantic top-up ---
|
||||
# ANN must run on its own connection: it opens a nested transaction with
|
||||
# SET LOCAL hnsw.ef_search + CREATE TEMP TABLE ON COMMIT DROP, and nesting
|
||||
# that inside our current write transaction would commit our writes early.
|
||||
semantic_needs = [
|
||||
r
|
||||
for r in victim_rows
|
||||
if counts.get((r["id"], "semantic"), 0) < MAX_SEMANTIC_LINKS_PER_UNIT and r["embedding"] is not None
|
||||
]
|
||||
if semantic_needs:
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
seed_ids = [r["id"] for r in semantic_needs]
|
||||
seed_embs = [r["embedding"] for r in semantic_needs]
|
||||
seed_ftypes = [r["fact_type"] for r in semantic_needs]
|
||||
async with acquire_with_retry(backend) as ann_conn:
|
||||
try:
|
||||
ann_links = await compute_semantic_links_ann(
|
||||
ann_conn,
|
||||
bank_id,
|
||||
seed_ids,
|
||||
seed_embs,
|
||||
fact_types=seed_ftypes,
|
||||
)
|
||||
# Strip self-links (rare but possible because the ANN probe
|
||||
# has no exclude list — see the comment in compute_semantic_links_ann).
|
||||
ann_links = [lnk for lnk in ann_links if lnk[0] != lnk[1]]
|
||||
new_links.extend(ann_links)
|
||||
except Exception as e:
|
||||
# ANN uses PG-specific HNSW syntax; on dialects/configs where
|
||||
# it isn't available we still want the temporal top-up to land.
|
||||
logger.warning(f"[GRAPH_MAINT] Semantic top-up failed for bank={bank_id}: {type(e).__name__}: {e}")
|
||||
|
||||
if not new_links:
|
||||
return 0
|
||||
|
||||
await _bulk_insert_links(
|
||||
conn,
|
||||
new_links,
|
||||
bank_id=bank_id,
|
||||
skip_exists_check=False,
|
||||
ops=ops,
|
||||
)
|
||||
return len(new_links)
|
||||
@@ -9,6 +9,7 @@ import os
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import AsyncExitStack
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -27,9 +28,12 @@ except ImportError:
|
||||
from ..config import (
|
||||
DEFAULT_LLM_MAX_CONCURRENT,
|
||||
DEFAULT_LLM_TIMEOUT,
|
||||
ENV_CONSOLIDATION_LLM_MAX_CONCURRENT,
|
||||
ENV_LLM_GROQ_SERVICE_TIER,
|
||||
ENV_LLM_MAX_CONCURRENT,
|
||||
ENV_LLM_TIMEOUT,
|
||||
ENV_REFLECT_LLM_MAX_CONCURRENT,
|
||||
ENV_RETAIN_LLM_MAX_CONCURRENT,
|
||||
)
|
||||
from ..metrics import get_metrics_collector
|
||||
from .response_models import TokenUsage
|
||||
@@ -42,12 +46,74 @@ logger = logging.getLogger(__name__)
|
||||
# Disable httpx logging
|
||||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
|
||||
# Global semaphore to limit concurrent LLM requests across all instances
|
||||
# Set HINDSIGHT_API_LLM_MAX_CONCURRENT=1 for local LLMs (LM Studio, Ollama)
|
||||
# Global semaphore to limit concurrent LLM requests across all instances.
|
||||
# Set HINDSIGHT_API_LLM_MAX_CONCURRENT=1 for local LLMs (LM Studio, Ollama).
|
||||
_llm_max_concurrent = int(os.getenv(ENV_LLM_MAX_CONCURRENT, str(DEFAULT_LLM_MAX_CONCURRENT)))
|
||||
_global_llm_semaphore = asyncio.Semaphore(_llm_max_concurrent)
|
||||
|
||||
|
||||
def _build_per_op_semaphores() -> dict[str, asyncio.Semaphore]:
|
||||
"""Build the per-operation semaphore registry from env vars.
|
||||
|
||||
Each per-op cap is composed with — not a substitute for — the global cap:
|
||||
a call that matches a configured operation must acquire both its per-op
|
||||
semaphore and the global semaphore. This lets operators reserve headroom
|
||||
in the global pool by capping individual operations (e.g. cap retain at 2
|
||||
of 4 global slots so the live chat path always has 2 slots available).
|
||||
|
||||
Operations without a configured env var are absent from the registry and
|
||||
therefore only constrained by the global cap.
|
||||
"""
|
||||
semaphores: dict[str, asyncio.Semaphore] = {}
|
||||
for op, env_var in (
|
||||
("retain", ENV_RETAIN_LLM_MAX_CONCURRENT),
|
||||
("reflect", ENV_REFLECT_LLM_MAX_CONCURRENT),
|
||||
("consolidation", ENV_CONSOLIDATION_LLM_MAX_CONCURRENT),
|
||||
):
|
||||
raw = os.getenv(env_var)
|
||||
if raw is None or raw == "":
|
||||
continue
|
||||
value = int(raw)
|
||||
if value <= 0:
|
||||
raise ValueError(f"{env_var} must be a positive integer, got {raw!r}")
|
||||
semaphores[op] = asyncio.Semaphore(value)
|
||||
return semaphores
|
||||
|
||||
|
||||
_per_op_llm_semaphores: dict[str, asyncio.Semaphore] = _build_per_op_semaphores()
|
||||
|
||||
|
||||
def _scope_to_operation(scope: str) -> str | None:
|
||||
"""Map a call scope to its per-operation concurrency bucket.
|
||||
|
||||
Returns None for scopes that don't belong to a tracked operation
|
||||
(verification probes, bank_mission, memory_think, mental_model_delta_ops),
|
||||
which then run under the global cap only.
|
||||
"""
|
||||
if scope.startswith("retain"):
|
||||
return "retain"
|
||||
if scope.startswith("reflect"):
|
||||
return "reflect"
|
||||
if scope.startswith("consolidation"):
|
||||
return "consolidation"
|
||||
return None
|
||||
|
||||
|
||||
def _semaphores_for_scope(scope: str) -> list[asyncio.Semaphore]:
|
||||
"""Return the semaphores a call with the given scope must acquire.
|
||||
|
||||
Always includes the global semaphore; includes the per-op semaphore when
|
||||
one is configured for the scope's operation bucket.
|
||||
"""
|
||||
op = _scope_to_operation(scope)
|
||||
per_op = _per_op_llm_semaphores.get(op) if op is not None else None
|
||||
if per_op is None:
|
||||
return [_global_llm_semaphore]
|
||||
# Per-op acquired first so contention queues on the narrower cap before
|
||||
# holding a global slot.
|
||||
return [per_op, _global_llm_semaphore]
|
||||
|
||||
|
||||
def sanitize_llm_output(text: str | None) -> str | None:
|
||||
"""
|
||||
Sanitize text by removing characters that break downstream systems.
|
||||
@@ -312,6 +378,7 @@ def create_llm_provider(
|
||||
"openai",
|
||||
"groq",
|
||||
"ollama",
|
||||
"ollama-cloud",
|
||||
"lmstudio",
|
||||
"minimax",
|
||||
"deepseek",
|
||||
@@ -409,6 +476,7 @@ class LLMProvider:
|
||||
"openai",
|
||||
"groq",
|
||||
"ollama",
|
||||
"ollama-cloud",
|
||||
"gemini",
|
||||
"anthropic",
|
||||
"lmstudio",
|
||||
@@ -437,6 +505,8 @@ class LLMProvider:
|
||||
self.base_url = "https://api.groq.com/openai/v1"
|
||||
elif self.provider == "ollama":
|
||||
self.base_url = "http://localhost:11434/v1"
|
||||
elif self.provider == "ollama-cloud":
|
||||
self.base_url = "https://ollama.com/v1"
|
||||
elif self.provider == "lmstudio":
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
elif self.provider == "minimax":
|
||||
@@ -629,7 +699,10 @@ class LLMProvider:
|
||||
structured = "+structured" if response_format is not None else ""
|
||||
set_stage(f"llm.{self.provider}.{scope}{structured}")
|
||||
|
||||
async with _global_llm_semaphore:
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call(
|
||||
messages=messages,
|
||||
@@ -654,7 +727,7 @@ class LLMProvider:
|
||||
# Sync the mock calls from provider implementation to wrapper
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
|
||||
return result
|
||||
return result
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
@@ -689,7 +762,10 @@ class LLMProvider:
|
||||
|
||||
set_stage(f"llm.{self.provider}.{scope}+tools")
|
||||
|
||||
async with _global_llm_semaphore:
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call_with_tools(
|
||||
messages=messages,
|
||||
@@ -712,7 +788,7 @@ class LLMProvider:
|
||||
# Sync the mock calls from provider implementation to wrapper
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
|
||||
return result
|
||||
return result
|
||||
|
||||
def set_response_callback(self, fn: Any) -> None:
|
||||
"""Set a callback invoked on each call() instead of the fixed mock response."""
|
||||
@@ -841,12 +917,14 @@ class LLMProvider:
|
||||
"""Create provider from environment variables using config.py constants."""
|
||||
from ..config import (
|
||||
DEFAULT_LLM_PROVIDER,
|
||||
DEFAULT_LLM_REASONING_EFFORT,
|
||||
ENV_LLM_API_KEY,
|
||||
ENV_LLM_BASE_URL,
|
||||
ENV_LLM_DEFAULT_HEADERS,
|
||||
ENV_LLM_EXTRA_BODY,
|
||||
ENV_LLM_MODEL,
|
||||
ENV_LLM_PROVIDER,
|
||||
ENV_LLM_REASONING_EFFORT,
|
||||
_get_default_model_for_provider,
|
||||
)
|
||||
|
||||
@@ -870,7 +948,7 @@ class LLMProvider:
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort="low",
|
||||
reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,41 @@
|
||||
"""Shared utilities for prompt assembly."""
|
||||
|
||||
import re
|
||||
|
||||
_LONE_OPEN_BRACE = re.compile(r"(?<!\{)\{(?!\{)")
|
||||
_LONE_CLOSE_BRACE = re.compile(r"(?<!\})\}(?!\})")
|
||||
|
||||
|
||||
def escape_for_prompt(text: str) -> str:
|
||||
"""Double any lone ``{`` / ``}`` so the text survives ``str.format`` untouched.
|
||||
|
||||
Prompt templates are often passed through ``str.format`` to substitute real
|
||||
placeholders like ``{facts_text}``. Any literal braces in caller-supplied
|
||||
text — e.g. a bank mission that contains JSON examples — would otherwise be
|
||||
interpreted as format keys and raise ``KeyError``.
|
||||
|
||||
Idempotent: text that already contains escaped ``{{`` / ``}}`` pairs is
|
||||
left as-is. Only lone braces (not adjacent to another brace of the same
|
||||
kind) are doubled.
|
||||
"""
|
||||
text = _LONE_OPEN_BRACE.sub("{{", text)
|
||||
text = _LONE_CLOSE_BRACE.sub("}}", text)
|
||||
return text
|
||||
|
||||
|
||||
def output_language_directive(language: str | None) -> str:
|
||||
"""Return an LLM directive forcing all output into ``language``.
|
||||
|
||||
Used by retain (fact extraction), consolidation (observations), and reflect
|
||||
(response synthesis) so HINDSIGHT_API_LLM_OUTPUT_LANGUAGE applies uniformly
|
||||
across every LLM-generated artifact. Returns an empty string when
|
||||
``language`` is unset so the calling prompt stays unchanged.
|
||||
"""
|
||||
if not language:
|
||||
return ""
|
||||
return (
|
||||
f"\n\nIMPORTANT: Respond exclusively in {language}. "
|
||||
f"Translate any source content into {language}. "
|
||||
f"All output text — including fact text, observations, entity names, "
|
||||
f"and the final response — must be in {language}."
|
||||
)
|
||||
@@ -93,7 +93,6 @@ class AnthropicLLM(LLMInterface):
|
||||
await self.call(
|
||||
messages=test_messages,
|
||||
max_completion_tokens=10,
|
||||
temperature=0.0,
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
@@ -179,9 +178,6 @@ class AnthropicLLM(LLMInterface):
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if temperature is not None:
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
@@ -398,9 +394,6 @@ class AnthropicLLM(LLMInterface):
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if temperature is not None:
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
|
||||
@@ -9,6 +9,7 @@ automatically handles authentication via `claude auth login` credentials.
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import tempfile
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
@@ -21,6 +22,32 @@ from hindsight_api.metrics import get_metrics_collector
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Isolation env passed to the spawned `claude` CLI. CLAUDE_CONFIG_DIR
|
||||
# redirects the subprocess away from the host's ~/.claude/, so any
|
||||
# operator-installed plugins (e.g. hindsight-memory) and their Stop hooks do
|
||||
# not fire inside our LLM-call subprocesses. Without this, retain/reflect/
|
||||
# consolidation LLM calls would trigger a Stop-hook retain of the subprocess
|
||||
# transcript back into the same bank — a recursive feedback loop (issue #1751).
|
||||
# CLAUDE_SECURESTORAGE_CONFIG_DIR="" forces the CLI's keychain service name
|
||||
# back to the canonical un-suffixed entry that `claude auth login` wrote;
|
||||
# otherwise it would be namespaced by sha256(CLAUDE_CONFIG_DIR) and OAuth
|
||||
# lookup would fail. Requires bundled CLI >= 2.1.150 (claude-agent-sdk 0.2.82).
|
||||
_isolated_claude_env: dict[str, str] | None = None
|
||||
|
||||
|
||||
def _get_isolated_claude_env() -> dict[str, str]:
|
||||
"""Return a process-lifetime env dict that isolates the spawned CLI from user plugins."""
|
||||
global _isolated_claude_env
|
||||
if _isolated_claude_env is None:
|
||||
path = tempfile.mkdtemp(prefix="hindsight-claude-code-")
|
||||
_isolated_claude_env = {
|
||||
"CLAUDE_CONFIG_DIR": path,
|
||||
"CLAUDE_SECURESTORAGE_CONFIG_DIR": "",
|
||||
}
|
||||
logger.debug(f"Claude Code: isolated CLAUDE_CONFIG_DIR={path}")
|
||||
return _isolated_claude_env
|
||||
|
||||
|
||||
class ClaudeCodeLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using Claude Code authentication.
|
||||
@@ -183,6 +210,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
system_prompt=system_prompt if system_prompt else None,
|
||||
max_turns=1, # Single-turn for API-style interactions
|
||||
allowed_tools=[], # Disable tools for standard LLM calls
|
||||
env=_get_isolated_claude_env(),
|
||||
)
|
||||
|
||||
# Call Claude Agent SDK
|
||||
@@ -473,6 +501,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
max_turns=2, # Allow tool call + tool result round-trip
|
||||
mcp_servers=mcp_servers_config,
|
||||
allowed_tools=allowed_tool_names,
|
||||
env=_get_isolated_claude_env(),
|
||||
)
|
||||
|
||||
# Call Claude Agent SDK with retry logic
|
||||
|
||||
@@ -0,0 +1,409 @@
|
||||
"""
|
||||
Shared Codex OAuth authentication manager.
|
||||
|
||||
Extracted from ``CodexLLM`` so that both ``CodexLLM`` and
|
||||
``CodexOAuthEmbeddings`` can share JWT-expiry detection, single-flight
|
||||
token refresh, and atomic file persistence without duplicating the logic.
|
||||
|
||||
Usage
|
||||
-----
|
||||
Create a manager from the auth file::
|
||||
|
||||
mgr = CodexAuthManager.from_file()
|
||||
|
||||
Then call ``ensure_fresh_token()`` before each outbound request and
|
||||
``refresh_tokens(reason=..., force=...)`` on a reactive 401.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Module-level constants (shared with codex_llm.py via re-export there)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# OAuth refresh endpoint and client id, mirrored from the canonical
|
||||
# ``@openai/codex`` CLI (codex-rs/login/src/auth/manager.rs on
|
||||
# github.com/openai/codex). The endpoint is overridable via env var so that
|
||||
# future Codex changes or staging environments can be pointed at without a
|
||||
# code change — same env var name the upstream CLI uses.
|
||||
_CODEX_REFRESH_TOKEN_URL = os.environ.get("CODEX_REFRESH_TOKEN_URL_OVERRIDE", "https://auth.openai.com/oauth/token")
|
||||
_CODEX_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
|
||||
|
||||
# Proactively refresh this many seconds before the JWT ``exp`` claim. The
|
||||
# upstream Codex CLI uses no skew (it refreshes at ``exp <= now``); the
|
||||
# extra window reduces races where a request leaves the client with a token
|
||||
# that the server has already declared expired by the time it arrives.
|
||||
_CODEX_TOKEN_REFRESH_SKEW_SECONDS = 60
|
||||
|
||||
# OAuth error codes that the refresh endpoint returns when the refresh_token
|
||||
# itself is no longer usable. These are terminal — retrying refresh will not
|
||||
# succeed; the user must re-run ``codex auth login``.
|
||||
_CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
|
||||
{"refresh_token_expired", "refresh_token_reused", "refresh_token_invalidated"}
|
||||
)
|
||||
|
||||
|
||||
class CodexRefreshExpiredError(RuntimeError):
|
||||
"""Raised when the Codex refresh_token itself is no longer valid.
|
||||
|
||||
The user must re-run ``codex auth login`` to obtain new credentials.
|
||||
Callers should surface a clear remediation message and stop retrying.
|
||||
"""
|
||||
|
||||
|
||||
class CodexAuthManager:
|
||||
"""Sync Codex OAuth credential manager.
|
||||
|
||||
Holds the access_token, refresh_token, and account_id in memory and
|
||||
handles proactive/reactive refresh using a ``threading.Lock`` for
|
||||
single-flight semantics (safe to use from multiple threads or via
|
||||
``asyncio.to_thread``).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
access_token:
|
||||
The current bearer token.
|
||||
account_id:
|
||||
The OpenAI account ID embedded in the Codex request headers.
|
||||
refresh_token:
|
||||
The OAuth refresh token. May be ``None`` when the auth file omits it;
|
||||
the provider still works as a one-shot loader in that case.
|
||||
auth_file:
|
||||
Path to ``~/.codex/auth.json``. Used for re-reading the refresh token
|
||||
on demand and for atomic persistence of rotated credentials.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
access_token: str,
|
||||
account_id: str,
|
||||
refresh_token: str | None,
|
||||
auth_file: Path,
|
||||
) -> None:
|
||||
self.access_token = access_token
|
||||
self.account_id = account_id
|
||||
self.refresh_token = refresh_token
|
||||
self._auth_file = auth_file
|
||||
self._lock = threading.Lock()
|
||||
self._http_client = httpx.Client(timeout=30.0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Construction helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@classmethod
|
||||
def from_file(cls, auth_file: Path | None = None) -> "CodexAuthManager":
|
||||
"""Build a manager by reading credentials from ``auth_file``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
auth_file:
|
||||
Defaults to ``~/.codex/auth.json``.
|
||||
|
||||
Raises
|
||||
------
|
||||
FileNotFoundError:
|
||||
If the auth file does not exist.
|
||||
ValueError:
|
||||
If the auth file is missing ``access_token`` or has an unexpected
|
||||
``auth_mode``.
|
||||
"""
|
||||
if auth_file is None:
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(f"Codex auth file not found: {auth_file}. Run 'codex auth login' to authenticate.")
|
||||
|
||||
with open(auth_file) as f:
|
||||
data = json.load(f)
|
||||
|
||||
auth_mode = data.get("auth_mode")
|
||||
if auth_mode != "chatgpt":
|
||||
raise ValueError(f"Expected Codex auth_mode='chatgpt', got: {auth_mode}")
|
||||
|
||||
tokens = data.get("tokens") or {}
|
||||
access_token = tokens.get("access_token")
|
||||
if not access_token:
|
||||
raise ValueError("No access_token found in Codex auth file. Run 'codex auth login' again.")
|
||||
|
||||
account_id = tokens.get("account_id") or ""
|
||||
refresh_token = tokens.get("refresh_token")
|
||||
|
||||
return cls(
|
||||
access_token=access_token,
|
||||
account_id=account_id,
|
||||
refresh_token=refresh_token,
|
||||
auth_file=auth_file,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Token state helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def load_refresh_token_from_file(auth_file: Path) -> str | None:
|
||||
"""Read ``tokens.refresh_token`` from ``auth_file``.
|
||||
|
||||
Returns ``None`` when the file is unreadable or omits the field.
|
||||
Does not raise — the provider degrades to one-shot mode.
|
||||
"""
|
||||
try:
|
||||
with open(auth_file) as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError) as e:
|
||||
logger.warning(
|
||||
f"Codex auth file unreadable when loading refresh_token: {type(e).__name__}. "
|
||||
"Token refresh will not be available; the access_token in memory will be used until it expires."
|
||||
)
|
||||
return None
|
||||
return data.get("tokens", {}).get("refresh_token")
|
||||
|
||||
@staticmethod
|
||||
def _decode_jwt_exp_unixtime(token: str) -> int | None:
|
||||
"""Return the JWT ``exp`` claim as a unix timestamp, or None on parse failure.
|
||||
|
||||
We do not verify the signature — the server is the source of truth
|
||||
on whether the token is actually accepted. This is only used to
|
||||
schedule proactive refresh.
|
||||
"""
|
||||
try:
|
||||
parts = token.split(".")
|
||||
if len(parts) < 2:
|
||||
return None
|
||||
payload_b64 = parts[1]
|
||||
padding = "=" * (-len(payload_b64) % 4)
|
||||
payload_bytes = base64.urlsafe_b64decode(payload_b64 + padding)
|
||||
payload = json.loads(payload_bytes.decode("utf-8"))
|
||||
exp = payload.get("exp")
|
||||
return int(exp) if exp is not None else None
|
||||
except (ValueError, TypeError, json.JSONDecodeError, binascii.Error):
|
||||
return None
|
||||
|
||||
def _token_is_stale(self, skew_seconds: int = _CODEX_TOKEN_REFRESH_SKEW_SECONDS) -> bool:
|
||||
"""True when the cached access_token is past expiry (with skew).
|
||||
|
||||
Returns False when expiry cannot be determined — we'd rather use a
|
||||
possibly-expired token and recover via the reactive 401 path than
|
||||
refresh aggressively on every request when ``exp`` parsing fails.
|
||||
"""
|
||||
exp = self._decode_jwt_exp_unixtime(self.access_token)
|
||||
if exp is None:
|
||||
return False
|
||||
return exp <= int(time.time()) + skew_seconds
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Persistence
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _persist_auth_atomic(self, updated_tokens: dict[str, Any]) -> None:
|
||||
"""Write rotated tokens back to ``_auth_file`` atomically.
|
||||
|
||||
Re-reads the on-disk file first to avoid clobbering fields written
|
||||
by another process, patches ``tokens.*`` and ``last_refresh``, then
|
||||
writes to a sibling tempfile and calls ``os.replace`` (atomic on
|
||||
POSIX and Windows within the same filesystem).
|
||||
"""
|
||||
current: dict[str, Any]
|
||||
try:
|
||||
with open(self._auth_file) as f:
|
||||
loaded = json.load(f)
|
||||
current = loaded if isinstance(loaded, dict) else {"auth_mode": "chatgpt", "tokens": {}}
|
||||
except (OSError, json.JSONDecodeError):
|
||||
current = {"auth_mode": "chatgpt", "tokens": {}}
|
||||
|
||||
existing_tokens = current.get("tokens")
|
||||
tokens: dict[str, Any] = existing_tokens if isinstance(existing_tokens, dict) else {}
|
||||
for key in ("access_token", "refresh_token", "id_token", "account_id"):
|
||||
if key in updated_tokens and updated_tokens[key] is not None:
|
||||
tokens[key] = updated_tokens[key]
|
||||
current["tokens"] = tokens
|
||||
current["last_refresh"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
||||
|
||||
parent = self._auth_file.parent
|
||||
parent.mkdir(parents=True, exist_ok=True)
|
||||
fd, tmp_path = tempfile.mkstemp(prefix=".auth.", suffix=".json.tmp", dir=str(parent))
|
||||
try:
|
||||
with os.fdopen(fd, "w") as f:
|
||||
json.dump(current, f, indent=2)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
try:
|
||||
os.chmod(tmp_path, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
os.replace(tmp_path, self._auth_file)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Error extraction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _extract_oauth_error_code(response: httpx.Response) -> str | None:
|
||||
"""Pull the OAuth error code out of a 4xx response body, if present.
|
||||
|
||||
The refresh endpoint returns shapes like
|
||||
``{"error": "...", "error_code": "..."}`` or
|
||||
``{"error": {"code": "..."}}``.
|
||||
"""
|
||||
try:
|
||||
body = response.json()
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return None
|
||||
if not isinstance(body, dict):
|
||||
return None
|
||||
err = body.get("error")
|
||||
if isinstance(err, dict):
|
||||
code = err.get("code")
|
||||
if isinstance(code, str):
|
||||
return code
|
||||
code = body.get("error_code")
|
||||
if isinstance(code, str):
|
||||
return code
|
||||
if isinstance(err, str):
|
||||
return err
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Refresh
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def refresh_tokens(self, reason: str = "", *, force: bool = False) -> None:
|
||||
"""Synchronous single-flight OAuth token refresh.
|
||||
|
||||
Serialized through ``self._lock`` so concurrent threads produce one
|
||||
network request. The first caller refreshes; the rest wake up and
|
||||
skip if the token is no longer stale (proactive) or if the token
|
||||
has already changed (reactive / force).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
reason:
|
||||
Free-form string included in log lines for diagnostics.
|
||||
force:
|
||||
When True, refresh even if the JWT exp claim looks fresh.
|
||||
Used by the reactive 401 path.
|
||||
|
||||
Raises
|
||||
------
|
||||
CodexRefreshExpiredError:
|
||||
When the server returns a terminal error code or any 401.
|
||||
RuntimeError:
|
||||
For other refresh failures (network, 5xx, etc.).
|
||||
"""
|
||||
token_before_lock = self.access_token
|
||||
with self._lock:
|
||||
if force:
|
||||
if self.access_token != token_before_lock:
|
||||
return
|
||||
else:
|
||||
if not self._token_is_stale():
|
||||
return
|
||||
|
||||
if not self.refresh_token:
|
||||
raise RuntimeError(
|
||||
"Codex access_token is expired but no refresh_token is available. "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
log_reason = f" ({reason})" if reason else ""
|
||||
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
|
||||
|
||||
request_body = {
|
||||
"client_id": _CODEX_CLIENT_ID,
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": self.refresh_token,
|
||||
}
|
||||
try:
|
||||
response = self._http_client.post(
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
json=request_body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=30.0,
|
||||
)
|
||||
except httpx.RequestError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
|
||||
|
||||
if response.status_code == 401:
|
||||
error_code = self._extract_oauth_error_code(response)
|
||||
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex OAuth refresh returned 401 with unrecognized error code "
|
||||
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
if response.status_code >= 400:
|
||||
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
|
||||
|
||||
try:
|
||||
body = response.json()
|
||||
except json.JSONDecodeError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
|
||||
|
||||
new_access = body.get("access_token")
|
||||
if not new_access:
|
||||
raise RuntimeError("Codex OAuth refresh returned no access_token")
|
||||
|
||||
new_refresh = body.get("refresh_token") or self.refresh_token
|
||||
new_id_token = body.get("id_token")
|
||||
|
||||
# Update in-memory state first so waiters see fresh credentials
|
||||
# immediately, even if disk write fails.
|
||||
self.access_token = new_access
|
||||
self.refresh_token = new_refresh
|
||||
|
||||
persisted: dict[str, Any] = {
|
||||
"access_token": new_access,
|
||||
"refresh_token": new_refresh,
|
||||
}
|
||||
if new_id_token:
|
||||
persisted["id_token"] = new_id_token
|
||||
|
||||
try:
|
||||
self._persist_auth_atomic(persisted)
|
||||
except OSError as e:
|
||||
logger.warning(
|
||||
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
|
||||
"In-memory credentials are up to date; on-disk file is stale."
|
||||
)
|
||||
|
||||
logger.info("Codex OAuth access_token refreshed successfully")
|
||||
|
||||
def ensure_fresh_token(self) -> None:
|
||||
"""Proactively refresh the access_token if it is near or past expiry.
|
||||
|
||||
Cheap when the token is fresh (just decodes the JWT exp claim and
|
||||
returns).
|
||||
"""
|
||||
if self._token_is_stale():
|
||||
self.refresh_tokens(reason="proactive (token near expiry)")
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the underlying HTTP client."""
|
||||
self._http_client.close()
|
||||
@@ -15,15 +15,10 @@ so that future server-side changes affect both clients identically.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import binascii
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -33,37 +28,27 @@ from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# OAuth refresh endpoint and client id, mirrored from the canonical
|
||||
# ``@openai/codex`` CLI (codex-rs/login/src/auth/manager.rs on
|
||||
# github.com/openai/codex). The endpoint is overridable via env var so that
|
||||
# future Codex changes or staging environments can be pointed at without a
|
||||
# code change — same env var name the upstream CLI uses.
|
||||
_CODEX_REFRESH_TOKEN_URL = os.environ.get("CODEX_REFRESH_TOKEN_URL_OVERRIDE", "https://auth.openai.com/oauth/token")
|
||||
_CODEX_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
|
||||
|
||||
# Proactively refresh this many seconds before the JWT ``exp`` claim. The
|
||||
# upstream Codex CLI uses no skew (it refreshes at ``exp <= now``); the
|
||||
# extra window reduces races where a request leaves the client with a token
|
||||
# that the server has already declared expired by the time it arrives.
|
||||
_CODEX_TOKEN_REFRESH_SKEW_SECONDS = 60
|
||||
|
||||
# OAuth error codes that the refresh endpoint returns when the refresh_token
|
||||
# itself is no longer usable. These are terminal — retrying refresh will not
|
||||
# succeed; the user must re-run ``codex auth login``.
|
||||
_CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
|
||||
{"refresh_token_expired", "refresh_token_reused", "refresh_token_invalidated"}
|
||||
from .codex_auth import (
|
||||
_CODEX_CLIENT_ID,
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
_CODEX_TERMINAL_REFRESH_ERROR_CODES,
|
||||
_CODEX_TOKEN_REFRESH_SKEW_SECONDS,
|
||||
CodexAuthManager,
|
||||
CodexRefreshExpiredError,
|
||||
)
|
||||
|
||||
# Re-export for backward compatibility (tests import from this module).
|
||||
__all__ = [
|
||||
"CodexLLM",
|
||||
"CodexRefreshExpiredError",
|
||||
"CodexAuthManager",
|
||||
"_CODEX_REFRESH_TOKEN_URL",
|
||||
"_CODEX_CLIENT_ID",
|
||||
"_CODEX_TOKEN_REFRESH_SKEW_SECONDS",
|
||||
"_CODEX_TERMINAL_REFRESH_ERROR_CODES",
|
||||
]
|
||||
|
||||
class CodexRefreshExpiredError(RuntimeError):
|
||||
"""Raised when the Codex refresh_token itself is no longer valid.
|
||||
|
||||
The user must re-run ``codex auth login`` to obtain new credentials.
|
||||
Callers should surface a clear remediation message and stop retrying.
|
||||
"""
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CodexLLM(LLMInterface):
|
||||
@@ -86,20 +71,15 @@ class CodexLLM(LLMInterface):
|
||||
"""Initialize Codex LLM provider."""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
|
||||
# Path is fixed at ~/.codex/auth.json — matches the upstream CLI.
|
||||
# Storing it on self lets the refresh path re-read after another
|
||||
# process (e.g. a sidecar) rotates the file out from under us.
|
||||
self._auth_file = Path.home() / ".codex" / "auth.json"
|
||||
|
||||
# Single-flight refresh lock. Multiple concurrent requests racing
|
||||
# toward an expired token should produce one network refresh, not N.
|
||||
# Single-flight async refresh lock. Multiple concurrent coroutines
|
||||
# racing toward an expired token should produce one network refresh.
|
||||
self._auth_lock = asyncio.Lock()
|
||||
|
||||
# Load Codex OAuth credentials
|
||||
# Load Codex OAuth credentials (keep these methods for test patching).
|
||||
try:
|
||||
self.access_token, self.account_id = self._load_codex_auth()
|
||||
self.refresh_token = self._load_codex_refresh_token()
|
||||
logger.info(f"Loaded Codex OAuth credentials for account: {self.account_id}")
|
||||
access_token, account_id = self._load_codex_auth()
|
||||
refresh_token = self._load_codex_refresh_token()
|
||||
logger.info(f"Loaded Codex OAuth credentials for account: {account_id}")
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to load Codex OAuth credentials from ~/.codex/auth.json: {e}\n\n"
|
||||
@@ -110,9 +90,22 @@ class CodexLLM(LLMInterface):
|
||||
"Or use a different provider (openai, anthropic, gemini) with API keys."
|
||||
) from e
|
||||
|
||||
# Use ChatGPT backend API endpoint
|
||||
if not self.base_url:
|
||||
self._auth_manager = CodexAuthManager(
|
||||
access_token=access_token,
|
||||
account_id=account_id,
|
||||
refresh_token=refresh_token,
|
||||
auth_file=Path.home() / ".codex" / "auth.json",
|
||||
)
|
||||
|
||||
# Use ChatGPT backend API endpoint. Codex auth is tied to
|
||||
# chatgpt.com/backend-api, not the OpenAI-compatible base URL used by
|
||||
# other providers. Deployments often set a global LLM_BASE_URL for an
|
||||
# OpenAI-compatible proxy; ignore that inherited value unless the user
|
||||
# explicitly provides a Codex backend URL.
|
||||
if not self.base_url or self.base_url.rstrip("/").endswith("/v1"):
|
||||
self.base_url = "https://chatgpt.com/backend-api"
|
||||
else:
|
||||
self.base_url = self.base_url.rstrip("/")
|
||||
|
||||
# Normalize model name (strip openai/ prefix if present)
|
||||
if self.model.startswith("openai/"):
|
||||
@@ -125,6 +118,42 @@ class CodexLLM(LLMInterface):
|
||||
# HTTP client for SSE streaming
|
||||
self._client = httpx.AsyncClient(timeout=120.0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Properties — delegate to _auth_manager (preserves test-visible API)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def access_token(self) -> str:
|
||||
return self._auth_manager.access_token
|
||||
|
||||
@access_token.setter
|
||||
def access_token(self, v: str) -> None:
|
||||
self._auth_manager.access_token = v
|
||||
|
||||
@property
|
||||
def account_id(self) -> str:
|
||||
return self._auth_manager.account_id
|
||||
|
||||
@property
|
||||
def refresh_token(self) -> str | None:
|
||||
return self._auth_manager.refresh_token
|
||||
|
||||
@refresh_token.setter
|
||||
def refresh_token(self, v: str | None) -> None:
|
||||
self._auth_manager.refresh_token = v
|
||||
|
||||
@property
|
||||
def _auth_file(self) -> Path:
|
||||
return self._auth_manager._auth_file
|
||||
|
||||
@_auth_file.setter
|
||||
def _auth_file(self, v: Path) -> None:
|
||||
self._auth_manager._auth_file = v
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Forwarding methods (keep surface area for tests / subclasses)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _load_codex_auth(self) -> tuple[str, str]:
|
||||
"""
|
||||
Load OAuth credentials from ~/.codex/auth.json.
|
||||
@@ -161,273 +190,57 @@ class CodexLLM(LLMInterface):
|
||||
return access_token, account_id
|
||||
|
||||
def _load_codex_refresh_token(self) -> str | None:
|
||||
"""Load ``tokens.refresh_token`` from ``~/.codex/auth.json``.
|
||||
"""Read ``tokens.refresh_token`` from the configured auth file.
|
||||
|
||||
Returns None when the auth file is unreadable or omits the field —
|
||||
the provider still functions as a one-shot loader in that case, it
|
||||
just can't refresh when the access_token expires. This deliberately
|
||||
does not raise so that ``__init__`` keeps the existing failure mode
|
||||
of raising only on missing ``access_token``.
|
||||
Kept as an instance method so existing tests that patch
|
||||
``CodexLLM._load_codex_refresh_token`` continue to work. Works both
|
||||
pre- and post-``__init__`` because it does not depend on
|
||||
``_auth_manager`` being constructed yet.
|
||||
"""
|
||||
try:
|
||||
with open(self._auth_file) as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError) as e:
|
||||
logger.warning(
|
||||
f"Codex auth file unreadable when loading refresh_token: {type(e).__name__}. "
|
||||
"Token refresh will not be available; the access_token in memory will be used until it expires."
|
||||
)
|
||||
return None
|
||||
return data.get("tokens", {}).get("refresh_token")
|
||||
auth_file = (
|
||||
self._auth_manager._auth_file if hasattr(self, "_auth_manager") else Path.home() / ".codex" / "auth.json"
|
||||
)
|
||||
return CodexAuthManager.load_refresh_token_from_file(auth_file)
|
||||
|
||||
@staticmethod
|
||||
def _decode_jwt_exp_unixtime(token: str) -> int | None:
|
||||
"""Return the JWT ``exp`` claim as a unix timestamp, or None on parse failure.
|
||||
|
||||
ChatGPT/Codex access_tokens are JWTs whose payload includes ``exp``
|
||||
(RFC 7519). We need the expiry to schedule proactive refresh — the
|
||||
``auth.json`` file does not persist a separate ``expires_at`` field
|
||||
in the upstream CLI's shape, so decoding the JWT itself is the
|
||||
canonical way to know when the token is stale.
|
||||
|
||||
We do not verify the signature — the server is the source of truth
|
||||
on whether the token is actually accepted, and the only thing this
|
||||
method affects is the *timing* of refresh, not whether to trust the
|
||||
token contents.
|
||||
"""
|
||||
try:
|
||||
parts = token.split(".")
|
||||
if len(parts) < 2:
|
||||
return None
|
||||
payload_b64 = parts[1]
|
||||
# JWT uses base64url without padding. Re-pad before decoding.
|
||||
padding = "=" * (-len(payload_b64) % 4)
|
||||
payload_bytes = base64.urlsafe_b64decode(payload_b64 + padding)
|
||||
payload = json.loads(payload_bytes.decode("utf-8"))
|
||||
exp = payload.get("exp")
|
||||
return int(exp) if exp is not None else None
|
||||
except (ValueError, TypeError, json.JSONDecodeError, binascii.Error):
|
||||
return None
|
||||
"""Delegate to ``CodexAuthManager._decode_jwt_exp_unixtime``."""
|
||||
return CodexAuthManager._decode_jwt_exp_unixtime(token)
|
||||
|
||||
def _token_is_stale(self, skew_seconds: int = _CODEX_TOKEN_REFRESH_SKEW_SECONDS) -> bool:
|
||||
"""True when the cached access_token is past expiry (with skew).
|
||||
|
||||
Returns False when expiry cannot be determined — we'd rather use a
|
||||
possibly-expired token and recover via the reactive 401 path than
|
||||
refresh aggressively on every request when ``exp`` parsing fails.
|
||||
"""
|
||||
exp = self._decode_jwt_exp_unixtime(self.access_token)
|
||||
if exp is None:
|
||||
return False
|
||||
return exp <= int(time.time()) + skew_seconds
|
||||
"""Delegate to ``_auth_manager._token_is_stale``."""
|
||||
return self._auth_manager._token_is_stale(skew_seconds)
|
||||
|
||||
def _persist_auth_atomic(self, updated_tokens: dict[str, Any]) -> None:
|
||||
"""Write the rotated tokens back to ``~/.codex/auth.json`` atomically.
|
||||
|
||||
Strategy: re-read the on-disk auth.json (so we don't clobber fields
|
||||
another process may have added), patch ``tokens.*`` and
|
||||
``last_refresh``, write to a tempfile in the same directory with
|
||||
mode 0600, then ``os.replace`` onto the target. ``os.replace`` is
|
||||
atomic within the same filesystem on POSIX and Windows, so a
|
||||
concurrent reader will see either the old file or the fully-written
|
||||
new file — never a partial truncate, which is the upstream CLI's
|
||||
worst-case race.
|
||||
|
||||
On non-Unix platforms the chmod is a best-effort no-op; the parent
|
||||
directory permissions still bound access.
|
||||
"""
|
||||
current: dict[str, Any]
|
||||
try:
|
||||
with open(self._auth_file) as f:
|
||||
loaded = json.load(f)
|
||||
# auth.json should always be a JSON object at the top level; if
|
||||
# someone has hand-edited it into a non-object shape, fall back
|
||||
# to the minimal default rather than crashing the refresh path.
|
||||
current = loaded if isinstance(loaded, dict) else {"auth_mode": "chatgpt", "tokens": {}}
|
||||
except (OSError, json.JSONDecodeError):
|
||||
# If the file became unreadable between our last read and now,
|
||||
# construct a minimal shape rather than refusing to persist.
|
||||
current = {"auth_mode": "chatgpt", "tokens": {}}
|
||||
|
||||
existing_tokens = current.get("tokens")
|
||||
tokens: dict[str, Any] = existing_tokens if isinstance(existing_tokens, dict) else {}
|
||||
for key in ("access_token", "refresh_token", "id_token", "account_id"):
|
||||
if key in updated_tokens and updated_tokens[key] is not None:
|
||||
tokens[key] = updated_tokens[key]
|
||||
current["tokens"] = tokens
|
||||
current["last_refresh"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
||||
|
||||
# Write to a sibling tempfile so the rename is same-filesystem.
|
||||
parent = self._auth_file.parent
|
||||
parent.mkdir(parents=True, exist_ok=True)
|
||||
fd, tmp_path = tempfile.mkstemp(prefix=".auth.", suffix=".json.tmp", dir=str(parent))
|
||||
try:
|
||||
with os.fdopen(fd, "w") as f:
|
||||
json.dump(current, f, indent=2)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
try:
|
||||
os.chmod(tmp_path, 0o600)
|
||||
except OSError:
|
||||
pass # best-effort on platforms that don't support chmod
|
||||
os.replace(tmp_path, self._auth_file)
|
||||
except Exception:
|
||||
# Clean up the orphaned tempfile if rename fails.
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
"""Delegate to ``_auth_manager._persist_auth_atomic``."""
|
||||
return self._auth_manager._persist_auth_atomic(updated_tokens)
|
||||
|
||||
async def _refresh_oauth_tokens(self, reason: str = "", *, force: bool = False) -> None:
|
||||
"""Refresh the OAuth access_token using the stored refresh_token.
|
||||
"""Async single-flight OAuth token refresh.
|
||||
|
||||
Single-flight: serialized through ``self._auth_lock`` so concurrent
|
||||
callers produce one network request. The first caller refreshes; the
|
||||
rest wake up and observe that either (a) the in-memory token is no
|
||||
longer stale (proactive case) or (b) the in-memory token has changed
|
||||
since they entered (reactive case), and return without re-refreshing.
|
||||
Outer asyncio.Lock preserves single-flight semantics for concurrent
|
||||
coroutines; the actual network call is offloaded to a thread via
|
||||
``asyncio.to_thread`` so the event loop stays unblocked.
|
||||
|
||||
Args:
|
||||
reason: Free-form string included in log lines for diagnostics.
|
||||
force: When True, refresh even if the JWT exp claim looks fresh.
|
||||
Used by the reactive 401 path — the server rejected the
|
||||
token, so we cannot trust the JWT's self-reported expiry.
|
||||
Used by the reactive 401 path.
|
||||
|
||||
Raises:
|
||||
CodexRefreshExpiredError: when the server returns a terminal
|
||||
error code (refresh_token_expired/reused/invalidated) or any
|
||||
401 on the refresh endpoint itself.
|
||||
error code or any 401 on the refresh endpoint.
|
||||
RuntimeError: for other refresh failures (network, 5xx, etc.).
|
||||
"""
|
||||
# Capture the token we'd be refreshing BEFORE acquiring the lock so
|
||||
# that we can detect mid-wait rotation by another coroutine.
|
||||
token_before_lock = self.access_token
|
||||
async with self._auth_lock:
|
||||
if force:
|
||||
# Reactive: skip only if another coroutine already rotated
|
||||
# the token while we were waiting on the lock.
|
||||
if self.access_token != token_before_lock:
|
||||
return
|
||||
else:
|
||||
# Proactive: skip if the token is no longer stale (the
|
||||
# canonical "another coroutine refreshed first" check).
|
||||
if not self._token_is_stale():
|
||||
if not self._auth_manager._token_is_stale():
|
||||
return
|
||||
|
||||
if not self.refresh_token:
|
||||
raise RuntimeError(
|
||||
"Codex access_token is expired but no refresh_token is available. "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
log_reason = f" ({reason})" if reason else ""
|
||||
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
|
||||
|
||||
request_body = {
|
||||
"client_id": _CODEX_CLIENT_ID,
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": self.refresh_token,
|
||||
}
|
||||
try:
|
||||
response = await self._client.post(
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
json=request_body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=30.0,
|
||||
)
|
||||
except httpx.RequestError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
|
||||
|
||||
if response.status_code == 401:
|
||||
# Classify by ``error.code`` (or top-level ``error`` string) — same
|
||||
# mapping as the upstream Rust CLI's request_chatgpt_token_refresh.
|
||||
error_code = self._extract_oauth_error_code(response)
|
||||
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
# Unknown 401 — treat as terminal too, matching the upstream classification.
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex OAuth refresh returned 401 with unrecognized error code "
|
||||
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
if response.status_code >= 400:
|
||||
# 5xx and other 4xx are transient/retryable from the caller's
|
||||
# perspective; surface as RuntimeError without leaking the
|
||||
# request body in logs.
|
||||
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
|
||||
|
||||
try:
|
||||
body = response.json()
|
||||
except json.JSONDecodeError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
|
||||
|
||||
new_access = body.get("access_token")
|
||||
if not new_access:
|
||||
raise RuntimeError("Codex OAuth refresh returned no access_token")
|
||||
|
||||
# The refresh_token may rotate on each refresh — adopt the new
|
||||
# one if the server sent it, otherwise keep the existing.
|
||||
new_refresh = body.get("refresh_token") or self.refresh_token
|
||||
new_id_token = body.get("id_token")
|
||||
|
||||
# Update in-memory state first so callers waiting on the lock
|
||||
# see fresh credentials immediately, even if disk write fails.
|
||||
self.access_token = new_access
|
||||
self.refresh_token = new_refresh
|
||||
|
||||
persisted = {
|
||||
"access_token": new_access,
|
||||
"refresh_token": new_refresh,
|
||||
}
|
||||
if new_id_token:
|
||||
persisted["id_token"] = new_id_token
|
||||
|
||||
try:
|
||||
self._persist_auth_atomic(persisted)
|
||||
except OSError as e:
|
||||
# In-memory creds are valid; warn but don't fail the request
|
||||
# path. Future process starts will fall back to the stale
|
||||
# on-disk auth.json and immediately refresh.
|
||||
logger.warning(
|
||||
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
|
||||
"In-memory credentials are up to date; on-disk file is stale."
|
||||
)
|
||||
|
||||
logger.info("Codex OAuth access_token refreshed successfully")
|
||||
|
||||
@staticmethod
|
||||
def _extract_oauth_error_code(response: "httpx.Response") -> str | None:
|
||||
"""Pull the OAuth error code out of a 4xx response body, if present.
|
||||
|
||||
The refresh endpoint returns shapes like
|
||||
``{"error": "...", "error_code": "..."}`` or
|
||||
``{"error": {"code": "..."}}``. We don't fail the call if the body
|
||||
is unparseable — the caller falls back to a generic "unknown" error.
|
||||
"""
|
||||
try:
|
||||
body = response.json()
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return None
|
||||
if not isinstance(body, dict):
|
||||
return None
|
||||
# Shape 1: error is a nested object with "code"
|
||||
err = body.get("error")
|
||||
if isinstance(err, dict):
|
||||
code = err.get("code")
|
||||
if isinstance(code, str):
|
||||
return code
|
||||
# Shape 2: top-level error_code string
|
||||
code = body.get("error_code")
|
||||
if isinstance(code, str):
|
||||
return code
|
||||
# Shape 3: error is itself a string code
|
||||
if isinstance(err, str):
|
||||
return err
|
||||
return None
|
||||
await asyncio.to_thread(lambda: self._auth_manager.refresh_tokens(reason, force=force))
|
||||
|
||||
async def _ensure_fresh_token(self) -> None:
|
||||
"""Refresh the access_token proactively if it is near or past expiry.
|
||||
@@ -435,13 +248,10 @@ class CodexLLM(LLMInterface):
|
||||
Called at the top of every API-bound method. Cheap when the token is
|
||||
fresh (just decodes the JWT exp claim and returns).
|
||||
"""
|
||||
if self._token_is_stale():
|
||||
if self._auth_manager._token_is_stale():
|
||||
try:
|
||||
await self._refresh_oauth_tokens(reason="proactive (token near expiry)")
|
||||
except CodexRefreshExpiredError:
|
||||
# Surface to the caller as the same RuntimeError shape the
|
||||
# request loop has historically raised, so existing error
|
||||
# handling paths keep working.
|
||||
raise
|
||||
|
||||
def _map_reasoning_effort(self, effort: str) -> str:
|
||||
@@ -1073,5 +883,6 @@ class CodexLLM(LLMInterface):
|
||||
return content if content else None, tool_calls
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up HTTP client."""
|
||||
"""Clean up HTTP clients."""
|
||||
await self._client.aclose()
|
||||
self._auth_manager.close()
|
||||
|
||||
@@ -153,11 +153,24 @@ class MockLLM(LLMInterface):
|
||||
result = self._response_callback(messages, scope)
|
||||
elif self._mock_response is not None:
|
||||
result = self._mock_response
|
||||
elif scope == "retain_extract_facts" and skip_validation:
|
||||
# Fact extraction: return canned facts derived from user message text.
|
||||
# This allows tests using a mock LLM to get real facts into the DB
|
||||
# so retain → recall → reflect pipelines work end-to-end.
|
||||
result = self._build_mock_facts(messages)
|
||||
elif scope == "consolidation" and response_format is not None:
|
||||
# Consolidation: produce a single observation from the input facts
|
||||
# so the full pipeline (retain → consolidation → observation → recall) works.
|
||||
result = self._build_mock_consolidation(messages, response_format)
|
||||
elif scope == "memory_think":
|
||||
# Reflect: return a plausible text answer
|
||||
result = "Based on the available information, the answer is related to the context provided."
|
||||
elif response_format is not None:
|
||||
# Try to create a minimal valid instance of the response format
|
||||
# Structured output: try to return a valid empty instance of the model
|
||||
# so that callers expecting e.g. response_format with defaults
|
||||
# get a valid instance rather than a crash on {"mock": True}.
|
||||
try:
|
||||
# For Pydantic models, try to create with minimal valid data
|
||||
result = {"mock": True}
|
||||
result = response_format()
|
||||
except Exception:
|
||||
result = {"mock": True}
|
||||
else:
|
||||
@@ -243,6 +256,12 @@ class MockLLM(LLMInterface):
|
||||
else:
|
||||
result = LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
|
||||
# Set mock token usage on result if not already set
|
||||
if result.input_tokens == 0:
|
||||
result.input_tokens = 10
|
||||
if result.output_tokens == 0:
|
||||
result.output_tokens = 5
|
||||
|
||||
# Record span with mock values
|
||||
# Convert LLMToolCall objects to dicts for span recording
|
||||
tool_calls_dict = (
|
||||
@@ -266,6 +285,92 @@ class MockLLM(LLMInterface):
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _build_mock_facts(messages: list[dict]) -> dict:
|
||||
"""Build a canned fact extraction response from the user message text.
|
||||
|
||||
Splits the input into sentence-like chunks and returns each as a separate
|
||||
world fact with a simple entity extracted from the first noun-like word.
|
||||
This is intentionally simplistic — it just needs to produce structurally
|
||||
valid facts so the rest of the pipeline (embedding, storage, recall) works.
|
||||
"""
|
||||
import re
|
||||
|
||||
user_text = ""
|
||||
for m in messages:
|
||||
if m.get("role") == "user":
|
||||
user_text = m.get("content", "")
|
||||
break
|
||||
|
||||
# Split on sentence boundaries: period followed by space/EOL (not mid-number), or newlines
|
||||
sentences = [s.strip() for s in re.split(r"(?<=\.)\s+|\n+", user_text) if s.strip() and len(s.strip()) > 10]
|
||||
|
||||
if not sentences:
|
||||
sentences = [user_text[:200] if user_text else "mock fact"]
|
||||
|
||||
facts = []
|
||||
for sentence in sentences[:10]: # Cap at 10 facts per chunk
|
||||
# Extract simple entities: capitalized words that aren't common words
|
||||
words = re.findall(r"\b[A-Z][a-z]+\b", sentence)
|
||||
entities = [{"text": w} for w in dict.fromkeys(words)][:5] # Dedupe, cap at 5
|
||||
|
||||
facts.append(
|
||||
{
|
||||
"what": sentence,
|
||||
"when": "N/A",
|
||||
"where": "N/A",
|
||||
"who": "N/A",
|
||||
"why": "N/A",
|
||||
"fact_kind": "conversation",
|
||||
"fact_type": "world",
|
||||
"entities": entities,
|
||||
}
|
||||
)
|
||||
|
||||
return {"facts": facts}
|
||||
|
||||
@staticmethod
|
||||
def _build_mock_consolidation(messages: list[dict], response_format: Any) -> Any:
|
||||
"""Build a mock consolidation response that creates one observation per fact.
|
||||
|
||||
Parses fact IDs from the consolidation prompt and creates one observation
|
||||
per fact, each referencing its source fact ID. This mimics real LLM behavior
|
||||
where distinct facts produce separate observations, preserving entity
|
||||
separation so pipeline tests (graph filtering, entity linking) work correctly.
|
||||
"""
|
||||
import re
|
||||
|
||||
user_text = ""
|
||||
for m in messages:
|
||||
if m.get("role") == "user":
|
||||
user_text = m.get("content", "")
|
||||
break
|
||||
|
||||
# Extract fact UUIDs from the prompt (format: "[<uuid>] <text>")
|
||||
fact_entries = re.findall(r"\[([0-9a-f-]{36})\]\s*(.+?)(?:\n|$)", user_text)
|
||||
|
||||
if not fact_entries:
|
||||
# No facts to consolidate — return empty response
|
||||
try:
|
||||
return response_format()
|
||||
except Exception:
|
||||
return {"creates": [], "updates": [], "deletes": []}
|
||||
|
||||
# Create one observation per fact to preserve entity separation
|
||||
creates = []
|
||||
for fact_id, fact_text in fact_entries:
|
||||
creates.append({"text": fact_text.strip(), "source_fact_ids": [fact_id]})
|
||||
|
||||
try:
|
||||
return response_format(
|
||||
creates=creates,
|
||||
updates=[],
|
||||
deletes=[],
|
||||
)
|
||||
except Exception:
|
||||
# Fallback if response_format constructor doesn't accept these args
|
||||
return {"creates": creates, "updates": [], "deletes": []}
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (no-op for mock provider)."""
|
||||
pass
|
||||
@@ -318,6 +423,8 @@ class MockLLM(LLMInterface):
|
||||
return self._mock_calls
|
||||
|
||||
def clear_mock_calls(self) -> None:
|
||||
"""Clear the recorded mock calls and any set exception."""
|
||||
"""Clear all recorded calls and any configured response/exception state."""
|
||||
self._mock_calls = []
|
||||
self._mock_exception = None
|
||||
self._mock_response = None
|
||||
self._response_callback = None
|
||||
|
||||
@@ -270,6 +270,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"openai",
|
||||
"groq",
|
||||
"ollama",
|
||||
"ollama-cloud",
|
||||
"lmstudio",
|
||||
"llamacpp",
|
||||
"minimax",
|
||||
@@ -288,6 +289,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
self.base_url = "https://api.groq.com/openai/v1"
|
||||
elif self.provider == "ollama":
|
||||
self.base_url = "http://localhost:11434/v1"
|
||||
elif self.provider == "ollama-cloud":
|
||||
self.base_url = "https://ollama.com/v1"
|
||||
elif self.provider == "lmstudio":
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
elif self.provider == "minimax":
|
||||
@@ -316,6 +319,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"openrouter",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"ollama-cloud",
|
||||
)
|
||||
and not self.api_key
|
||||
):
|
||||
@@ -1073,12 +1077,17 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
last_exception = None
|
||||
|
||||
# Pass API key as Bearer token for cloud Ollama endpoints
|
||||
headers: dict[str, str] = {}
|
||||
if self.api_key and self.api_key != "local":
|
||||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||||
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await client.post(native_url, json=payload)
|
||||
response = await client.post(native_url, json=payload, headers=headers)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
|
||||
@@ -321,6 +321,7 @@ async def run_reflect_agent(
|
||||
include_recall: bool = True,
|
||||
budget: str | None = None,
|
||||
max_context_tokens: int = 100_000,
|
||||
llm_output_language: str | None = None,
|
||||
) -> ReflectAgentResult:
|
||||
"""
|
||||
Execute the reflect agent loop using native tool calling.
|
||||
@@ -369,7 +370,12 @@ async def run_reflect_agent(
|
||||
|
||||
# Build initial messages (directives are injected into system prompt at START and END)
|
||||
system_prompt = build_system_prompt_for_tools(
|
||||
bank_profile, context, directives=directives, has_mental_models=has_mental_models, budget=budget
|
||||
bank_profile,
|
||||
context,
|
||||
directives=directives,
|
||||
has_mental_models=has_mental_models,
|
||||
include_observations=include_observations,
|
||||
budget=budget,
|
||||
)
|
||||
messages: list[dict[str, Any]] = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
@@ -447,7 +453,10 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -504,7 +513,10 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -607,7 +619,10 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -728,7 +743,10 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
|
||||
@@ -98,6 +98,7 @@ def build_system_prompt_for_tools(
|
||||
context: str | None = None,
|
||||
directives: list[dict[str, Any]] | None = None,
|
||||
has_mental_models: bool = False,
|
||||
include_observations: bool = True,
|
||||
budget: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -108,11 +109,17 @@ def build_system_prompt_for_tools(
|
||||
2. search_observations - Consolidated knowledge with freshness
|
||||
3. recall - Raw facts as ground truth
|
||||
|
||||
The retrieval-strategy and workflow sections are built to match the tools
|
||||
actually exposed to the LLM — mentioning a tool the agent has disabled
|
||||
causes weaker LLMs to either hallucinate the call (rejected by the agent)
|
||||
or give up with "I cannot find any information…" (see #1724).
|
||||
|
||||
Args:
|
||||
bank_profile: Bank profile with name and mission
|
||||
context: Optional additional context
|
||||
directives: Optional list of directive mental models to inject as hard rules
|
||||
has_mental_models: Whether the bank has any mental models (skip if not)
|
||||
include_observations: Whether search_observations is in the tool list.
|
||||
budget: Search depth budget - "low", "mid", or "high". Controls exploration thoroughness.
|
||||
"""
|
||||
name = bank_profile.get("name", "Assistant")
|
||||
@@ -158,56 +165,137 @@ def build_system_prompt_for_tools(
|
||||
"- If memories mention someone did an activity, you can infer they likely enjoyed it",
|
||||
"- Synthesize a coherent narrative from related memories",
|
||||
"- Be a thoughtful interpreter, not just a literal repeater",
|
||||
"- When the exact answer isn't stated, use what IS stated to give the best answer",
|
||||
"- When the exact answer isn't stated, use what IS stated to give a best-effort answer AND surface any uncertainty — never invent confidence the data doesn't support.",
|
||||
"",
|
||||
"## Temporal Reasoning",
|
||||
"Every memory and observation carries temporal fields in the JSON tool result:",
|
||||
"- `mentioned_at` — when the user retained the fact (always set).",
|
||||
"- `occurred_start` / `occurred_end` — when the underlying event happened (optional, set for dated events).",
|
||||
"",
|
||||
"When facts about the SAME facet conflict — counts, statuses, ownership, location, presence, etc. — the fact with the LATEST `mentioned_at` is authoritative. Later statements SUPERSEDE earlier ones. Do NOT average, sum, or favor an explicitly-dated fact over a more recent one.",
|
||||
"",
|
||||
"Example: three count facts come back from recall:",
|
||||
" - 'Team has 2 engineers' (mentioned_at=T1)",
|
||||
" - 'Team now has 1 engineer' (mentioned_at=T2, occurred_start=2026-05-25)",
|
||||
" - 'Team has 5 engineers' (mentioned_at=T3)",
|
||||
"with T1 < T2 < T3. The current size is 5, not 1. Then apply later events (e.g. someone leaving after T3) on top of that.",
|
||||
"",
|
||||
"For reconstructing a TIMELINE of events, order by `occurred_start` / `occurred_end` (when things happened), not `mentioned_at` (when they were retained).",
|
||||
"",
|
||||
"## Conflicts and Ambiguity",
|
||||
"Not every retrieval converges on a single answer. Distinguish two cases:",
|
||||
"",
|
||||
"- RESOLVABLE conflict — the temporal rule above (latest `mentioned_at` wins) cleanly picks a winner. Apply it and move on.",
|
||||
"- UNRESOLVABLE ambiguity — the data is internally inconsistent in a way the temporal rule does NOT settle. Examples: a recent aggregate (count, total) is incompatible with the individual entities you can enumerate; two equally-recent facts disagree and no later fact resolves them; events are described but their relative order is unclear; the user's own statements contradict each other and nothing later reconciles them.",
|
||||
"",
|
||||
"When the data is genuinely ambiguous: SAY SO in your answer. Name the conflicting facts. Explain why they can't be reconciled. Give a range or a best-effort interpretation with explicit uncertainty (e.g. 'between X and Y, depending on [unresolved condition]'; or 'the most recent statement says A, but B was stated earlier and the gap isn't accounted for in any later fact').",
|
||||
"",
|
||||
"An honest 'the data is inconsistent about X' beats a confident wrong answer. Do NOT pick a value arbitrarily, average conflicting values, or smooth over gaps in confident prose. Acknowledging ambiguity is a successful answer, not a failure mode.",
|
||||
"",
|
||||
"## Showing Your Reasoning",
|
||||
"For any answer that resolves a conflict between facts, applies events on top of a count or status, or settles an ambiguity — show your work in the answer text so a reader can audit it.",
|
||||
"",
|
||||
"Walk through these steps explicitly:",
|
||||
"1. **List the relevant facts in `mentioned_at` order (oldest → newest)**, each with the value it asserts. Use a short bulleted list.",
|
||||
"2. **Identify the authoritative fact** under the temporal rule (latest `mentioned_at` for the contested facet). Write its date down.",
|
||||
"3. **List candidate events to apply on top** — anything that changes the count, status, or state being asked about. Write each event's date down next to it.",
|
||||
"4. **Sanity-check each candidate event against the authoritative date** — for EVERY event from step 3, write a one-line check in the form `<event> (<event_date>) vs authoritative (<authoritative_date>) → BEFORE/AFTER → KEEP/DROP`. If the event is BEFORE or EQUAL to the authoritative date, DROP it: it is already reflected in the authoritative fact, and applying it again is double-counting. This is the single most common mistake — do not skip this step even if you feel confident.",
|
||||
"5. **Show the arithmetic or derivation explicitly** using only the KEEP events from step 4 — e.g. 'authoritative count = 5 (at 2025-02-12); kept events: Shadow died (2025-03-12, AFTER); 5 − 1 = 4'.",
|
||||
"6. If step 2 or 3 cannot be done cleanly (no clear winner, overlapping timestamps, unclear event order), STOP and surface this as an UNRESOLVABLE ambiguity per the section above — do not fabricate a derivation.",
|
||||
"",
|
||||
"For simple factual lookups that don't involve conflict or arithmetic, you can answer directly without this scaffolding.",
|
||||
"",
|
||||
"## HIERARCHICAL RETRIEVAL STRATEGY",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
# Build retrieval levels based on what's available
|
||||
# Assemble the retrieval-level blocks for whatever tools are exposed.
|
||||
# MM and Observations bodies are unconditional; recall's fallback wording
|
||||
# adapts to which upstream tools precede it (telling the LLM to fall back
|
||||
# to a tool that isn't in its list is the bug at the root of #1724).
|
||||
levels: list[tuple[str, list[str]]] = []
|
||||
if has_mental_models:
|
||||
parts.extend(
|
||||
levels.append(
|
||||
(
|
||||
"MENTAL MODELS (search_mental_models)",
|
||||
[
|
||||
"- User-curated summaries about specific topics",
|
||||
"- HIGHEST quality - manually created and maintained",
|
||||
"- If a relevant mental model exists and is FRESH, it may fully answer the question",
|
||||
"- Check `is_stale` field - if stale, also verify with lower levels",
|
||||
],
|
||||
)
|
||||
)
|
||||
if include_observations:
|
||||
levels.append(
|
||||
(
|
||||
"OBSERVATIONS (search_observations)",
|
||||
[
|
||||
"- Auto-consolidated knowledge from memories",
|
||||
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
|
||||
"- Good for understanding patterns and summaries",
|
||||
],
|
||||
)
|
||||
)
|
||||
recall_body = ["- Individual memories (world facts and experiences)"]
|
||||
if has_mental_models and include_observations:
|
||||
recall_body.extend(
|
||||
[
|
||||
"You have access to THREE levels of knowledge. Use them in this order:",
|
||||
"",
|
||||
"### 1. MENTAL MODELS (search_mental_models) - Try First",
|
||||
"- User-curated summaries about specific topics",
|
||||
"- HIGHEST quality - manually created and maintained",
|
||||
"- If a relevant mental model exists and is FRESH, it may fully answer the question",
|
||||
"- Check `is_stale` field - if stale, also verify with lower levels",
|
||||
"",
|
||||
"### 2. OBSERVATIONS (search_observations) - Second Priority",
|
||||
"- Auto-consolidated knowledge from memories",
|
||||
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
|
||||
"- Good for understanding patterns and summaries",
|
||||
"",
|
||||
"### 3. RAW FACTS (recall) - Ground Truth",
|
||||
"- Individual memories (world facts and experiences)",
|
||||
"- Use when: no mental models/observations exist, they're stale, or you need specific details",
|
||||
"- MANDATORY: If search_mental_models and search_observations both return 0 results, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that other levels are built from",
|
||||
"",
|
||||
"**Tool result ordering:** `recall()` and `search_observations()` return their `memories` / `observations` arrays sorted by SEMANTIC RELEVANCE to the query, NOT by time. The POSITION of an entry tells you nothing about when it was retained. For any temporal reasoning — recency, supersession, applying events on top of a state — IGNORE the position and read the per-entry `mentioned_at` field (and `occurred_start` / `occurred_end` for events).",
|
||||
"",
|
||||
]
|
||||
)
|
||||
else:
|
||||
parts.extend(
|
||||
elif has_mental_models:
|
||||
recall_body.extend(
|
||||
[
|
||||
"- Use when: no mental model exists, it's stale, or you need specific details",
|
||||
"- MANDATORY: If search_mental_models returns 0 results, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that mental models are built from",
|
||||
]
|
||||
)
|
||||
elif include_observations:
|
||||
recall_body.extend(
|
||||
[
|
||||
"You have access to TWO levels of knowledge. Use them in this order:",
|
||||
"",
|
||||
"### 1. OBSERVATIONS (search_observations) - Try First",
|
||||
"- Auto-consolidated knowledge from memories",
|
||||
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
|
||||
"- Good for understanding patterns and summaries",
|
||||
"",
|
||||
"### 2. RAW FACTS (recall) - Ground Truth",
|
||||
"- Individual memories (world facts and experiences)",
|
||||
"- Use when: no observations exist, they're stale, or you need specific details",
|
||||
"- MANDATORY: If search_observations returns 0 results or count=0, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that observations are built from",
|
||||
"",
|
||||
"**Tool result ordering:** `recall()` and `search_observations()` return their `memories` / `observations` arrays sorted by SEMANTIC RELEVANCE to the query, NOT by time. The POSITION of an entry tells you nothing about when it was retained. For any temporal reasoning — recency, supersession, applying events on top of a state — IGNORE the position and read the per-entry `mentioned_at` field (and `occurred_start` / `occurred_end` for events).",
|
||||
"",
|
||||
]
|
||||
)
|
||||
else:
|
||||
recall_body.extend(
|
||||
[
|
||||
"- MANDATORY: Call recall() to gather facts before giving up",
|
||||
"- This is the source of truth.",
|
||||
]
|
||||
)
|
||||
levels.append(("RAW FACTS (recall) - Ground Truth", recall_body))
|
||||
|
||||
# Position-dependent suffix for upstream tools; recall already carries its
|
||||
# fixed "- Ground Truth" suffix in the header text.
|
||||
suffixes = [""] * len(levels)
|
||||
if len(levels) >= 2:
|
||||
suffixes[0] = " - Try First"
|
||||
if len(levels) == 3:
|
||||
suffixes[1] = " - Second Priority"
|
||||
|
||||
if len(levels) == 1:
|
||||
parts.append("You have access to ONE level of knowledge:")
|
||||
else:
|
||||
word = "TWO" if len(levels) == 2 else "THREE"
|
||||
parts.append(f"You have access to {word} levels of knowledge. Use them in this order:")
|
||||
parts.append("")
|
||||
for idx, ((header, body), suffix) in enumerate(zip(levels, suffixes), 1):
|
||||
parts.append(f"### {idx}. {header}{suffix}")
|
||||
parts.extend(body)
|
||||
parts.append("")
|
||||
|
||||
parts.extend(
|
||||
[
|
||||
@@ -267,25 +355,28 @@ def build_system_prompt_for_tools(
|
||||
|
||||
parts.append("## Workflow")
|
||||
|
||||
steps: list[str] = []
|
||||
if has_mental_models:
|
||||
parts.extend(
|
||||
[
|
||||
"1. First, try search_mental_models() - check if a curated summary exists",
|
||||
"2. If no mental model or it's stale, try search_observations() for consolidated knowledge",
|
||||
"3. If observations are stale OR you need specific details, use recall() for raw facts",
|
||||
"4. Use expand() if you need more context on specific memories",
|
||||
"5. When ready, call done() with your answer and supporting IDs",
|
||||
]
|
||||
steps.append("First, try search_mental_models() - check if a curated summary exists")
|
||||
if include_observations:
|
||||
if has_mental_models:
|
||||
steps.append("If no mental model or it's stale, try search_observations() for consolidated knowledge")
|
||||
else:
|
||||
steps.append("First, try search_observations() - check for consolidated knowledge")
|
||||
# Recall step phrasing varies with whichever upstream tool(s) precede it.
|
||||
if include_observations:
|
||||
steps.append(
|
||||
"If observations are stale OR you need specific details, use recall() for raw facts"
|
||||
if has_mental_models
|
||||
else "If search_observations returns 0 results OR observations are stale, you MUST call recall() for raw facts"
|
||||
)
|
||||
elif has_mental_models:
|
||||
steps.append("If no mental model or it's stale, use recall() for raw facts")
|
||||
else:
|
||||
parts.extend(
|
||||
[
|
||||
"1. First, try search_observations() - check for consolidated knowledge",
|
||||
"2. If search_observations returns 0 results OR observations are stale, you MUST call recall() for raw facts",
|
||||
"3. Use expand() if you need more context on specific memories",
|
||||
"4. When ready, call done() with your answer and supporting IDs",
|
||||
]
|
||||
)
|
||||
steps.append("Call recall() to gather raw facts")
|
||||
steps.append("Use expand() if you need more context on specific memories")
|
||||
steps.append("When ready, call done() with your answer and supporting IDs")
|
||||
parts.extend(f"{idx}. {step}" for idx, step in enumerate(steps, 1))
|
||||
|
||||
parts.extend(
|
||||
[
|
||||
@@ -513,10 +604,16 @@ Just provide the direct answer with proper markdown formatting.
|
||||
CRITICAL: This is a NON-CONVERSATIONAL system. NEVER ask follow-up questions, offer to search again, suggest alternatives, or end with anything like "Would you like me to..." or "Let me know if...". The user cannot reply. Your answer must be complete and self-contained."""
|
||||
|
||||
|
||||
def build_final_system_prompt(mission: str | None = None) -> str:
|
||||
"""Build the final synthesis system prompt, using mission as role when set."""
|
||||
role_section = mission.strip() if mission else _DEFAULT_FINAL_ROLE
|
||||
return _FINAL_SYSTEM_PROMPT_BASE.format(role_section=role_section)
|
||||
def build_final_system_prompt(mission: str | None = None, llm_output_language: str | None = None) -> str:
|
||||
"""Build the final synthesis system prompt, using mission as role when set.
|
||||
|
||||
When ``llm_output_language`` is set, the response is forced into that
|
||||
language regardless of the query/source language.
|
||||
"""
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt, output_language_directive
|
||||
|
||||
role_section = escape_for_prompt(mission.strip()) if mission else _DEFAULT_FINAL_ROLE
|
||||
return _FINAL_SYSTEM_PROMPT_BASE.format(role_section=role_section) + output_language_directive(llm_output_language)
|
||||
|
||||
|
||||
# Backward-compatible constant for non-identity missions
|
||||
|
||||
@@ -23,6 +23,24 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _prune_nulls(d: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Drop keys whose value is None or an empty collection (``""``, ``[]``, ``{}``).
|
||||
|
||||
Reflect tools dump ``MemoryFact`` / ``ObservationResult`` via ``model_dump()``,
|
||||
which emits every field including the many that are typically null or empty
|
||||
(``context``, ``occurred_start``, ``metadata``, ``tags``, etc.). Stripping
|
||||
these before serializing to JSON for the LLM cuts token cost and removes
|
||||
fields that aren't telling the model anything.
|
||||
|
||||
Callers that need the *presence* of a specific field as a signal (e.g.
|
||||
``source_fact_ids`` for drill-down) must ensure the value is non-empty —
|
||||
pass the upstream flag that populates it (e.g. ``source_facts_max_tokens``
|
||||
> 0 on ``tool_search_observations``) rather than relying on Pydantic
|
||||
emitting ``None``.
|
||||
"""
|
||||
return {k: v for k, v in d.items() if v is not None and v != "" and v != [] and v != {}}
|
||||
|
||||
|
||||
def _document_metadata_from_retain_params(retain_params: Any) -> dict[str, Any] | None:
|
||||
"""Return document metadata stored under retain_params.metadata."""
|
||||
if isinstance(retain_params, str):
|
||||
@@ -214,8 +232,8 @@ async def tool_search_observations(
|
||||
return {
|
||||
"query": query,
|
||||
"count": len(result.results),
|
||||
"observations": [m.model_dump() for m in result.results],
|
||||
"source_facts": {k: v.model_dump() for k, v in (result.source_facts or {}).items()},
|
||||
"observations": [_prune_nulls(m.model_dump()) for m in result.results],
|
||||
"source_facts": {k: _prune_nulls(v.model_dump()) for k, v in (result.source_facts or {}).items()},
|
||||
"is_stale": is_stale,
|
||||
"freshness": freshness,
|
||||
}
|
||||
@@ -282,8 +300,8 @@ async def tool_recall(
|
||||
|
||||
return {
|
||||
"query": query,
|
||||
"memories": [m.model_dump() for m in result.results],
|
||||
"chunks": {k: v.model_dump() for k, v in (result.chunks or {}).items()},
|
||||
"memories": [_prune_nulls(m.model_dump()) for m in result.results],
|
||||
"chunks": {k: _prune_nulls(v.model_dump()) for k, v in (result.chunks or {}).items()},
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -4,29 +4,54 @@ Embedding generation utilities for memory units.
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Literal, Protocol
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
EmbeddingInputType = Literal["document", "query"]
|
||||
|
||||
def generate_embedding(embeddings_backend, text: str) -> list[float]:
|
||||
|
||||
class EmbeddingsBackend(Protocol):
|
||||
"""Minimal duck-typed surface used by retain/recall — the concrete `Embeddings`
|
||||
ABC supplies default implementations that delegate to `encode()`."""
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]: ...
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]: ...
|
||||
|
||||
|
||||
def generate_embedding(
|
||||
embeddings_backend: EmbeddingsBackend, text: str, input_type: EmbeddingInputType = "document"
|
||||
) -> list[float]:
|
||||
"""
|
||||
Generate embedding for text using the provided embeddings backend.
|
||||
|
||||
Args:
|
||||
embeddings_backend: Embeddings instance to use for encoding
|
||||
text: Text to embed
|
||||
input_type: Whether text is retained document text or recall/search query text.
|
||||
|
||||
Returns:
|
||||
Embedding vector (dimension depends on embeddings backend)
|
||||
"""
|
||||
try:
|
||||
embeddings = embeddings_backend.encode([text])
|
||||
embeddings = _encode_with_input_type(embeddings_backend, [text], input_type)
|
||||
return embeddings[0]
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to generate embedding: {str(e)}")
|
||||
|
||||
|
||||
async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> list[list[float]]:
|
||||
def _encode_with_input_type(
|
||||
embeddings_backend: EmbeddingsBackend, texts: list[str], input_type: EmbeddingInputType
|
||||
) -> list[list[float]]:
|
||||
if input_type == "query":
|
||||
return embeddings_backend.encode_query(texts)
|
||||
return embeddings_backend.encode_documents(texts)
|
||||
|
||||
|
||||
async def generate_embeddings_batch(
|
||||
embeddings_backend: EmbeddingsBackend, texts: list[str], input_type: EmbeddingInputType = "document"
|
||||
) -> list[list[float]]:
|
||||
"""
|
||||
Generate embeddings for multiple texts using the provided embeddings backend.
|
||||
|
||||
@@ -36,17 +61,14 @@ async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> lis
|
||||
Args:
|
||||
embeddings_backend: Embeddings instance to use for encoding
|
||||
texts: List of texts to embed
|
||||
input_type: Whether texts are retained documents or recall/search queries.
|
||||
|
||||
Returns:
|
||||
List of embeddings in same order as input texts
|
||||
"""
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
embeddings = await loop.run_in_executor(
|
||||
None,
|
||||
embeddings_backend.encode,
|
||||
texts,
|
||||
)
|
||||
embeddings = await loop.run_in_executor(None, _encode_with_input_type, embeddings_backend, texts, input_type)
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to generate batch embeddings: {str(e)}")
|
||||
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
"""
|
||||
Entity processing for retain pipeline.
|
||||
|
||||
Handles entity extraction, resolution, and link creation for stored facts.
|
||||
Handles entity extraction and resolution for stored facts.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from . import link_utils
|
||||
from .types import EntityLink, ProcessedFact
|
||||
from .types import ProcessedFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -76,8 +76,7 @@ async def resolve_entities(
|
||||
entity_labels: Optional entity label taxonomy
|
||||
|
||||
Returns:
|
||||
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids)
|
||||
to pass to build_entity_links().
|
||||
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids).
|
||||
"""
|
||||
if not unit_ids or not facts:
|
||||
return [], [], {}
|
||||
@@ -99,68 +98,3 @@ async def resolve_entities(
|
||||
log_buffer,
|
||||
entity_labels=entity_labels,
|
||||
)
|
||||
|
||||
|
||||
async def build_entity_links(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
resolved_entity_ids: list[str],
|
||||
entity_to_unit: list[tuple],
|
||||
unit_to_entity_ids: dict[str, list[str]],
|
||||
log_buffer: list[str] = None,
|
||||
skip_unit_entities_insert: bool = False,
|
||||
ops=None,
|
||||
) -> list[EntityLink]:
|
||||
"""
|
||||
Build entity links for UI graph visualization.
|
||||
|
||||
Queries unit_entities to find shared entities between new and existing units,
|
||||
then generates EntityLink objects. When called from Phase 3 (post-transaction),
|
||||
set skip_unit_entities_insert=True since unit_entities were already inserted
|
||||
in Phase 2.
|
||||
|
||||
Args:
|
||||
entity_resolver: EntityResolver instance
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
unit_ids: Actual unit IDs (must already be inserted in the DB)
|
||||
resolved_entity_ids: From resolve_entities()
|
||||
entity_to_unit: From resolve_entities()
|
||||
unit_to_entity_ids: From resolve_entities()
|
||||
log_buffer: Optional buffer for detailed logging
|
||||
skip_unit_entities_insert: Skip unit_entities INSERT (already done in Phase 2)
|
||||
ops: DataAccessOps instance (from backend.ops)
|
||||
|
||||
Returns:
|
||||
List of EntityLink objects for batch insertion
|
||||
"""
|
||||
return await link_utils.build_entity_links_from_resolved(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id,
|
||||
unit_ids,
|
||||
resolved_entity_ids,
|
||||
entity_to_unit,
|
||||
unit_to_entity_ids,
|
||||
log_buffer,
|
||||
skip_unit_entities_insert=skip_unit_entities_insert,
|
||||
ops=ops,
|
||||
)
|
||||
|
||||
|
||||
async def insert_entity_links_batch(conn, entity_links: list[EntityLink], bank_id: str, ops=None) -> None:
|
||||
"""
|
||||
Insert entity links in batch.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
entity_links: List of EntityLink objects
|
||||
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
|
||||
ops: DataAccessOps instance (from backend.ops)
|
||||
"""
|
||||
if not entity_links:
|
||||
return
|
||||
|
||||
await link_utils.insert_entity_links_batch(conn, entity_links, bank_id, ops=ops)
|
||||
|
||||
@@ -888,13 +888,16 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
extract_causal_links = config.retain_extract_causal_links
|
||||
|
||||
# Build retain_mission section if set - injected before the mode-specific guidelines
|
||||
# Escape braces so user-supplied text survives str.format() on the prompt template.
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt
|
||||
|
||||
retain_mission = getattr(config, "retain_mission", None)
|
||||
if retain_mission:
|
||||
retain_mission_section = (
|
||||
f"══════════════════════════════════════════════════════════════════════════\n"
|
||||
f"FOCUS — What to retain for this bank\n"
|
||||
f"══════════════════════════════════════════════════════════════════════════\n\n"
|
||||
f"{retain_mission}\n\n"
|
||||
f"{escape_for_prompt(retain_mission)}\n\n"
|
||||
)
|
||||
else:
|
||||
retain_mission_section = ""
|
||||
@@ -910,7 +913,7 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
base_prompt = CUSTOM_FACT_EXTRACTION_PROMPT
|
||||
prompt = base_prompt.format(
|
||||
retain_mission_section=retain_mission_section,
|
||||
custom_instructions=config.retain_custom_instructions,
|
||||
custom_instructions=escape_for_prompt(config.retain_custom_instructions),
|
||||
)
|
||||
elif extraction_mode == "verbose":
|
||||
prompt = VERBOSE_FACT_EXTRACTION_PROMPT.format(
|
||||
@@ -947,6 +950,16 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
if labels_section:
|
||||
prompt = prompt + labels_section
|
||||
|
||||
# Force the LLM to emit fact text in the configured language, regardless of
|
||||
# the source content's language. Same directive is applied to consolidation
|
||||
# and reflect so HINDSIGHT_API_LLM_OUTPUT_LANGUAGE has a uniform effect
|
||||
# across the pipeline. This is independent of the BM25 indexing language
|
||||
# (HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE) by design — search
|
||||
# tokenization and LLM output language are separate concerns.
|
||||
from ..prompt_utils import output_language_directive
|
||||
|
||||
prompt = prompt + output_language_directive(getattr(config, "llm_output_language", None))
|
||||
|
||||
response_schema = base_response_class
|
||||
|
||||
if labels_cfg and labels_cfg.attributes:
|
||||
@@ -1129,11 +1142,14 @@ async def _extract_facts_from_chunk(
|
||||
)
|
||||
continue
|
||||
else:
|
||||
logger.warning(
|
||||
f"LLM returned non-dict JSON after {llm_max_retries} attempts: {type(extraction_response_json).__name__}. "
|
||||
f"Raw: {str(extraction_response_json)[:500]}"
|
||||
# A non-dict response is malformed (the schema is {"facts": [...]}).
|
||||
# Raise instead of returning [] so the failure propagates to the
|
||||
# worker's retry machinery and ultimately fails loudly — never
|
||||
# silently commit the document with 0 facts. See issue #1833.
|
||||
raise RuntimeError(
|
||||
f"Fact extraction failed: LLM returned non-dict JSON after {llm_max_retries} attempts "
|
||||
f"({type(extraction_response_json).__name__}). Raw: {str(extraction_response_json)[:500]}"
|
||||
)
|
||||
return [], usage
|
||||
|
||||
raw_facts = extraction_response_json.get("facts", [])
|
||||
|
||||
@@ -1662,8 +1678,11 @@ async def extract_facts_from_contents_batch_api(
|
||||
|
||||
# Check if provider supports batch API
|
||||
if not await llm_config._provider_impl.supports_batch_api():
|
||||
logger.warning(f"Batch API not supported for provider {llm_config.provider}, falling back to sync mode")
|
||||
return await extract_facts_from_contents(contents, llm_config, agent_name, config, pool, operation_id, schema)
|
||||
raise RuntimeError(
|
||||
f"retain_batch_enabled=True but provider '{llm_config.provider}' does not "
|
||||
f"support the batch API. This should have been caught at startup — check "
|
||||
f"HINDSIGHT_API_RETAIN_BATCH_ENABLED and your LLM provider configuration."
|
||||
)
|
||||
|
||||
# Check if we're resuming an existing batch (crash recovery)
|
||||
batch_id = None
|
||||
@@ -2195,7 +2214,9 @@ async def extract_facts_from_contents(
|
||||
fact_extraction_tasks.append(task)
|
||||
|
||||
# Step 2: Wait for all fact extractions to complete.
|
||||
# Use return_exceptions=True so one content item failure doesn't discard the rest.
|
||||
# return_exceptions=True so a failing item doesn't cancel its still-running
|
||||
# siblings (which would leave orphaned LLM calls / partial work); we await
|
||||
# them all, then propagate.
|
||||
all_fact_results = await asyncio.gather(*fact_extraction_tasks, return_exceptions=True)
|
||||
|
||||
# Step 3: Flatten and convert to typed objects
|
||||
@@ -2206,14 +2227,18 @@ async def extract_facts_from_contents(
|
||||
global_chunk_idx = 0
|
||||
global_fact_idx = 0
|
||||
|
||||
# Filter out failed content items
|
||||
# Never silently drop a document's memory. Any extraction failure (provider
|
||||
# rate-limit / timeout / 5xx, malformed response, token-limit, etc.)
|
||||
# propagates so the streaming producer surfaces it and the worker's
|
||||
# RetryTaskAt machinery retries the task — and ultimately fails it *loudly*
|
||||
# if the problem persists. Swallowing the error and substituting an empty
|
||||
# result here used to commit the document with 0 facts and mark the
|
||||
# operation `completed`, losing the memory with no signal. See issue #1833.
|
||||
valid_results = []
|
||||
for content, result in zip(contents, all_fact_results):
|
||||
if isinstance(result, Exception):
|
||||
logger.warning(f"Content extraction failed (skipping): {type(result).__name__}: {result}")
|
||||
valid_results.append((content, ([], [], TokenUsage())))
|
||||
else:
|
||||
valid_results.append((content, result))
|
||||
raise result
|
||||
valid_results.append((content, result))
|
||||
|
||||
for content_index, (content, (facts_from_llm, chunks_from_llm, content_usage)) in enumerate(valid_results):
|
||||
total_usage = total_usage + content_usage
|
||||
|
||||
@@ -321,6 +321,15 @@ async def handle_document_tracking(
|
||||
f"[RETAIN] Document {document_id} re-ingested: invalidated "
|
||||
f"{invalidated} observation(s) derived from {len(existing_unit_ids)} outgoing memory_units"
|
||||
)
|
||||
# Capture link-recompute victims BEFORE the cascade. Same staleness
|
||||
# applies on upsert as on explicit delete: surviving units in OTHER
|
||||
# documents that linked to these doomed units are about to lose
|
||||
# those links. ``ops`` may be None for older callers that haven't
|
||||
# been wired up — skip enqueue in that case rather than crash.
|
||||
if ops is not None:
|
||||
from ..graph_maintenance import enqueue_relink_victims
|
||||
|
||||
await enqueue_relink_victims(conn, bank_id, [str(uid) for uid in existing_unit_ids], ops=ops)
|
||||
# Explicitly delete memory_units by document_id BEFORE deleting the
|
||||
# document row. The CASCADE from documents→chunks→memory_units only
|
||||
# catches units that have a non-NULL chunk_id FK. Units with chunk_id=NULL
|
||||
|
||||
@@ -5,10 +5,9 @@ Link creation utilities for temporal, semantic, and entity links.
|
||||
import logging
|
||||
import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from uuid import UUID
|
||||
|
||||
from ..._vector_index import ann_search_tuning_settings, configured_vector_extension
|
||||
from ..memory_engine import fq_table
|
||||
from .types import EntityLink
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -366,136 +365,6 @@ async def resolve_entities_only(
|
||||
return resolved_entity_ids, entity_to_unit, unit_to_entity_ids
|
||||
|
||||
|
||||
async def build_entity_links_from_resolved(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
resolved_entity_ids: list[str],
|
||||
entity_to_unit: list[tuple],
|
||||
unit_to_entity_ids: dict[str, list[str]],
|
||||
log_buffer: list[str] = None,
|
||||
skip_unit_entities_insert: bool = False,
|
||||
ops=None,
|
||||
) -> list["EntityLink"]:
|
||||
"""
|
||||
Build entity links between units that share entities.
|
||||
|
||||
Queries unit_entities to find which existing units share entities with the
|
||||
new units, then generates EntityLink objects for UI graph visualization.
|
||||
|
||||
Args:
|
||||
entity_resolver: EntityResolver instance
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
unit_ids: Actual unit IDs (must already be inserted in the DB)
|
||||
resolved_entity_ids: Entity IDs from resolve_entities_only
|
||||
entity_to_unit: Mapping from resolve_entities_only
|
||||
unit_to_entity_ids: Mapping from resolve_entities_only
|
||||
log_buffer: Optional logging buffer
|
||||
skip_unit_entities_insert: If True, skip unit_entities INSERT (already done in Phase 2)
|
||||
|
||||
Returns:
|
||||
List of EntityLink objects for batch insertion
|
||||
"""
|
||||
if not resolved_entity_ids:
|
||||
return []
|
||||
|
||||
if not skip_unit_entities_insert:
|
||||
# Insert unit-entity links (used in fallback path where Phase 2 didn't do this)
|
||||
substep_start = time.time()
|
||||
unit_entity_pairs = []
|
||||
for idx, (unit_id, _local_idx, fact_date) in enumerate(entity_to_unit):
|
||||
# Propagate the unit's fact_date so entity_cooccurrences.last_cooccurred
|
||||
# reflects the event timeline, not the ingest moment.
|
||||
unit_entity_pairs.append((unit_id, resolved_entity_ids[idx], fact_date))
|
||||
|
||||
await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_start:.3f}s",
|
||||
level="debug",
|
||||
)
|
||||
|
||||
# Build entity links between units that share entities
|
||||
substep_start = time.time()
|
||||
all_entity_ids = set()
|
||||
for entity_ids_list in unit_to_entity_ids.values():
|
||||
all_entity_ids.update(entity_ids_list)
|
||||
|
||||
_log(log_buffer, f" [6.3] Creating entity links for {len(all_entity_ids)} unique entities...", level="debug")
|
||||
|
||||
MAX_LINKS_PER_ENTITY = 10
|
||||
|
||||
entity_to_units = {}
|
||||
if all_entity_ids:
|
||||
query_start = time.time()
|
||||
import uuid
|
||||
|
||||
entity_id_list = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in all_entity_ids]
|
||||
limit_per_entity = MAX_LINKS_PER_ENTITY + len(unit_ids) # room for new units + existing cap
|
||||
|
||||
rows = await ops.fetch_entity_unit_fanout(
|
||||
conn,
|
||||
fq_table("unit_entities"),
|
||||
entity_id_list,
|
||||
limit_per_entity,
|
||||
)
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [6.3.1] Query unit_entities (LATERAL): {len(rows)} rows in {time.time() - query_start:.3f}s",
|
||||
level="debug",
|
||||
)
|
||||
|
||||
group_start = time.time()
|
||||
for row in rows:
|
||||
entity_id = row["entity_id"]
|
||||
if entity_id not in entity_to_units:
|
||||
entity_to_units[entity_id] = []
|
||||
entity_to_units[entity_id].append(row["unit_id"])
|
||||
_log(log_buffer, f" [6.3.2] Group by entity_id: {time.time() - group_start:.3f}s", level="debug")
|
||||
link_gen_start = time.time()
|
||||
links: list[EntityLink] = []
|
||||
new_unit_set = set(unit_ids)
|
||||
|
||||
def to_uuid(val) -> UUID:
|
||||
return UUID(val) if isinstance(val, str) else val
|
||||
|
||||
for entity_id, units_with_entity in entity_to_units.items():
|
||||
entity_uuid = to_uuid(entity_id)
|
||||
new_units = [u for u in units_with_entity if str(u) in new_unit_set or u in new_unit_set]
|
||||
existing_units = [u for u in units_with_entity if str(u) not in new_unit_set and u not in new_unit_set]
|
||||
|
||||
new_units_to_link = new_units[-MAX_LINKS_PER_ENTITY:] if len(new_units) > MAX_LINKS_PER_ENTITY else new_units
|
||||
for i, unit_id_1 in enumerate(new_units_to_link):
|
||||
for unit_id_2 in new_units_to_link[i + 1 :]:
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(unit_id_1), to_unit_id=to_uuid(unit_id_2), entity_id=entity_uuid)
|
||||
)
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(unit_id_2), to_unit_id=to_uuid(unit_id_1), entity_id=entity_uuid)
|
||||
)
|
||||
|
||||
existing_to_link = existing_units[-MAX_LINKS_PER_ENTITY:]
|
||||
for new_unit in new_units:
|
||||
for existing_unit in existing_to_link:
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(new_unit), to_unit_id=to_uuid(existing_unit), entity_id=entity_uuid)
|
||||
)
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(existing_unit), to_unit_id=to_uuid(new_unit), entity_id=entity_uuid)
|
||||
)
|
||||
|
||||
_log(log_buffer, f" [6.3.3] Generate {len(links)} links: {time.time() - link_gen_start:.3f}s", level="debug")
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [6.3] Entity link creation: {len(links)} links for {len(all_entity_ids)} unique entities in {time.time() - substep_start:.3f}s",
|
||||
level="debug",
|
||||
)
|
||||
|
||||
return links
|
||||
|
||||
|
||||
async def create_temporal_links_batch_per_fact(
|
||||
conn,
|
||||
bank_id: str,
|
||||
@@ -701,16 +570,18 @@ async def compute_semantic_links_ann(
|
||||
# `relation "_ann_seeds" does not exist` on the second statement.
|
||||
#
|
||||
# Using ON COMMIT DROP + SET LOCAL also means we don't have to remember to
|
||||
# manually drop the temp table or reset hnsw.ef_search — the transaction
|
||||
# end handles both.
|
||||
# manually drop the temp table or reset the per-backend ANN tuning GUC —
|
||||
# the transaction end handles both.
|
||||
rows: list = []
|
||||
async with conn.transaction():
|
||||
# Transaction-local ef_search. Default 400 is tuned for recall precision
|
||||
# but at 164k units each HNSW probe takes 94ms. ef_search=60 gives 2.7ms
|
||||
# per probe (35x faster) with sufficient accuracy for top-50 semantic
|
||||
# link creation. SET LOCAL auto-reverts at commit, so we don't pollute
|
||||
# the pool for subsequent recall queries.
|
||||
await conn.execute("SET LOCAL hnsw.ef_search = 60")
|
||||
# Transaction-local ANN tuning. Each supported backend exposes its own
|
||||
# GUC (hnsw.ef_search on pgvector, vchordrq.probes on vchord); the
|
||||
# dispatcher returns the right knob for the configured backend with a
|
||||
# value tuned for top-50 semantic link creation (lower recall but much
|
||||
# lower latency than the recall-side default). SET LOCAL auto-reverts
|
||||
# at commit, so we don't pollute the pool for subsequent queries.
|
||||
for guc, value in ann_search_tuning_settings(configured_vector_extension(), kind="low_latency"):
|
||||
await conn.execute(f"SET LOCAL {guc} = {value}")
|
||||
|
||||
t_setup = time_mod.time()
|
||||
await conn.execute("CREATE TEMP TABLE _ann_seeds (unit_id text, emb_text text, fact_type text) ON COMMIT DROP")
|
||||
@@ -889,29 +760,6 @@ async def create_semantic_links_batch(
|
||||
raise
|
||||
|
||||
|
||||
async def insert_entity_links_batch(conn, links: list[EntityLink], bank_id: str, chunk_size: int = 5000, ops=None):
|
||||
"""
|
||||
Bulk-insert entity links via sorted INSERT FROM unnest().
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
links: List of EntityLink objects
|
||||
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
|
||||
chunk_size: Number of rows per INSERT chunk (default 5000)
|
||||
"""
|
||||
if not links:
|
||||
return
|
||||
|
||||
import time as time_mod
|
||||
|
||||
total_start = time_mod.time()
|
||||
tuples = [(link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id) for link in links]
|
||||
await _bulk_insert_links(conn, tuples, bank_id=bank_id, chunk_size=chunk_size, ops=ops)
|
||||
logger.debug(
|
||||
f" [9.TOTAL] Entity links batch insert ({len(tuples)} rows): {time_mod.time() - total_start:.3f}s"
|
||||
)
|
||||
|
||||
|
||||
async def create_causal_links_batch(
|
||||
conn,
|
||||
bank_id: str,
|
||||
|
||||
@@ -100,7 +100,6 @@ from .types import (
|
||||
ChunkMetadata,
|
||||
EntityResolutionResult,
|
||||
Phase1Result,
|
||||
Phase3Context,
|
||||
ProcessedFact,
|
||||
RetainContent,
|
||||
RetainContentDict,
|
||||
@@ -108,6 +107,9 @@ from .types import (
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
RetainOutboxCallback = Callable[[asyncpg.Connection], Awaitable[None]]
|
||||
RetainOutboxCallbackFactory = Callable[[list[RetainContentDict]], RetainOutboxCallback | None]
|
||||
|
||||
|
||||
def _build_retain_params(contents_dicts, document_tags=None, doc_contents=None):
|
||||
"""Build retain_params and merged_tags from content dicts."""
|
||||
@@ -256,30 +258,27 @@ async def _insert_facts_and_links(
|
||||
skip_semantic_links: bool = False,
|
||||
outbox_callback=None,
|
||||
ops=None,
|
||||
) -> tuple[list[list[str]], Phase3Context]:
|
||||
) -> list[list[str]]:
|
||||
"""
|
||||
Phase 2 of the retain pipeline: insert facts and retrieval-critical links.
|
||||
|
||||
Runs inside a single database transaction to ensure atomicity of the data
|
||||
that retrieval depends on (facts, unit_entities, temporal/semantic/causal links).
|
||||
|
||||
Entity link generation and insertion for UI visualization are NOT done here —
|
||||
only the unit_entities INSERT (FK to memory_units) stays in the transaction.
|
||||
Entity link building is deferred to Phase 3 (post-transaction, best-effort).
|
||||
Entity edges for UI graph visualization are derived on demand from
|
||||
unit_entities by the /graph endpoint, so no entity rows are written to
|
||||
memory_links here.
|
||||
"""
|
||||
set_stage("retain.phase2.insert_facts")
|
||||
unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, processed_facts, ops=ops)
|
||||
step_start = time.time()
|
||||
log_buffer.append(f" Insert facts: {len(unit_ids)} units in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Context for Phase 3 entity link building (after transaction commits)
|
||||
phase3_context = Phase3Context()
|
||||
|
||||
if unit_ids:
|
||||
# Entity resolution was done in Phase 1 (separate connection).
|
||||
# Remap placeholder IDs to actual unit IDs.
|
||||
step_start = time.time()
|
||||
remapped_entity_to_unit, remapped_unit_to_entity_ids, remapped_semantic = _remap_phase1_results(
|
||||
remapped_entity_to_unit, _remapped_unit_to_entity_ids, remapped_semantic = _remap_phase1_results(
|
||||
resolved_entity_ids, entity_to_unit, unit_to_entity_ids, semantic_ann_links or [], unit_ids
|
||||
)
|
||||
# Update semantic_ann_links with remapped IDs for Phase 2
|
||||
@@ -293,13 +292,6 @@ async def _insert_facts_and_links(
|
||||
]
|
||||
await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
|
||||
log_buffer.append(f" Insert unit_entities: {len(unit_entity_pairs)} pairs in {time.time() - step_start:.3f}s")
|
||||
# Save context for Phase 3 entity link building (after commit)
|
||||
phase3_context = Phase3Context(
|
||||
unit_ids=unit_ids,
|
||||
resolved_entity_ids=resolved_entity_ids,
|
||||
entity_to_unit=remapped_entity_to_unit,
|
||||
unit_to_entity_ids=remapped_unit_to_entity_ids,
|
||||
)
|
||||
|
||||
# Create temporal links
|
||||
step_start = time.time()
|
||||
@@ -340,52 +332,10 @@ async def _insert_facts_and_links(
|
||||
# an IndexError (see issue #1037).
|
||||
result_unit_ids = _map_results_to_contents(contents, processed_facts, unit_ids if unit_ids else [])
|
||||
|
||||
if outbox_callback:
|
||||
if outbox_callback is not None:
|
||||
await outbox_callback(conn)
|
||||
|
||||
return result_unit_ids, phase3_context
|
||||
|
||||
|
||||
async def _build_and_insert_entity_links_phase3(
|
||||
pool: Any,
|
||||
entity_resolver,
|
||||
bank_id: str,
|
||||
phase3_ctx: Phase3Context,
|
||||
log_buffer: list[str],
|
||||
) -> None:
|
||||
"""
|
||||
Phase 3 helper: build entity links from resolved data and insert them.
|
||||
|
||||
Runs on a fresh connection after the main transaction has committed.
|
||||
Entity links are for UI graph visualization only — retrieval uses
|
||||
the unit_entities self-join instead.
|
||||
"""
|
||||
set_stage("retain.phase3.entity_links")
|
||||
p3_unit_ids = phase3_ctx.unit_ids
|
||||
p3_resolved = phase3_ctx.resolved_entity_ids
|
||||
p3_entity_to_unit = phase3_ctx.entity_to_unit
|
||||
p3_unit_to_entity_ids = phase3_ctx.unit_to_entity_ids
|
||||
|
||||
if not p3_unit_ids or not p3_resolved:
|
||||
return
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
step_start = time.time()
|
||||
entity_links = await entity_processing.build_entity_links(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id,
|
||||
p3_unit_ids,
|
||||
p3_resolved,
|
||||
p3_entity_to_unit,
|
||||
p3_unit_to_entity_ids,
|
||||
log_buffer,
|
||||
skip_unit_entities_insert=True, # Already inserted in Phase 2
|
||||
ops=pool.ops,
|
||||
)
|
||||
if entity_links:
|
||||
await entity_processing.insert_entity_links_batch(conn, entity_links, bank_id, ops=pool.ops)
|
||||
log_buffer.append(f" Entity links (viz): {len(entity_links)} links in {time.time() - step_start:.3f}s")
|
||||
return result_unit_ids
|
||||
|
||||
|
||||
async def _extract_and_embed(
|
||||
@@ -449,7 +399,8 @@ async def retain_batch(
|
||||
document_tags: list[str] | None = None,
|
||||
operation_id: str | None = None,
|
||||
schema: str | None = None,
|
||||
outbox_callback: Callable[["asyncpg.Connection"], Awaitable[None]] | None = None,
|
||||
outbox_callback: RetainOutboxCallback | None = None,
|
||||
outbox_callback_factory: RetainOutboxCallbackFactory | None = None,
|
||||
db_semaphore: "asyncio.Semaphore | None" = None,
|
||||
) -> tuple[list[list[str]], TokenUsage, int | None]:
|
||||
"""
|
||||
@@ -507,6 +458,10 @@ async def retain_batch(
|
||||
total_usage = TokenUsage()
|
||||
total_processed_tokens: int | None = 0
|
||||
for doc_key, (group_dicts, group_contents) in groups.items():
|
||||
group_outbox_callback = (
|
||||
outbox_callback_factory(group_dicts) if outbox_callback_factory is not None else outbox_callback
|
||||
)
|
||||
|
||||
group_ids, group_usage, group_processed = await retain_batch(
|
||||
pool=pool,
|
||||
embeddings_model=embeddings_model,
|
||||
@@ -522,7 +477,8 @@ async def retain_batch(
|
||||
document_tags=document_tags,
|
||||
operation_id=operation_id,
|
||||
schema=schema,
|
||||
outbox_callback=outbox_callback,
|
||||
outbox_callback=group_outbox_callback,
|
||||
outbox_callback_factory=outbox_callback_factory,
|
||||
db_semaphore=db_semaphore,
|
||||
)
|
||||
for group_idx, orig_idx in enumerate(original_indices[doc_key]):
|
||||
@@ -1191,7 +1147,6 @@ async def _streaming_retain_batch(
|
||||
|
||||
p2_start = time.time()
|
||||
batch_result_ids = None
|
||||
phase3_ctx = None
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# --- Document ownership gate ---
|
||||
@@ -1280,7 +1235,7 @@ async def _streaming_retain_batch(
|
||||
|
||||
# Insert facts and links — skip semantic links entirely in streaming
|
||||
# mode; they are created in a single final ANN pass after all batches.
|
||||
batch_result_ids, phase3_ctx = await _insert_facts_and_links(
|
||||
batch_result_ids = await _insert_facts_and_links(
|
||||
conn,
|
||||
entity_resolver,
|
||||
bank_id,
|
||||
@@ -1300,15 +1255,13 @@ async def _streaming_retain_batch(
|
||||
|
||||
logger.info(f"[streaming] Phase 2 (write txn): {time.time() - p2_start:.3f}s")
|
||||
|
||||
# Best-effort: entity viz + stats (fast, not semantic ANN)
|
||||
if phase3_ctx is not None:
|
||||
try:
|
||||
await entity_resolver.flush_pending_stats()
|
||||
await _build_and_insert_entity_links_phase3(
|
||||
pool, entity_resolver, bank_id, phase3_ctx, log_buffer
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(f"Phase 3 stats (consumer batch {consumer_batch_idx + 1}) failed", exc_info=True)
|
||||
# Best-effort: flush entity_cooccurrences and other deferred stats.
|
||||
try:
|
||||
await entity_resolver.flush_pending_stats()
|
||||
except Exception:
|
||||
logger.warning(
|
||||
f"Entity stats flush (consumer batch {consumer_batch_idx + 1}) failed", exc_info=True
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"[streaming] Consumer batch {consumer_batch_idx + 1} total "
|
||||
@@ -1766,7 +1719,7 @@ async def _try_delta_retain(
|
||||
# Insert facts and retrieval-critical links.
|
||||
# Use delta_contents (the changed/new chunks) as the content list,
|
||||
# since extracted_facts have content_index relative to delta_contents.
|
||||
result_unit_ids, phase3_ctx = await _insert_facts_and_links(
|
||||
result_unit_ids = await _insert_facts_and_links(
|
||||
conn,
|
||||
entity_resolver,
|
||||
bank_id,
|
||||
@@ -1783,12 +1736,11 @@ async def _try_delta_retain(
|
||||
ops=pool.ops,
|
||||
)
|
||||
|
||||
# PHASE 3 — Best-Effort Display Data (post-transaction)
|
||||
# Flush deferred entity_cooccurrences stats (post-transaction, best-effort).
|
||||
try:
|
||||
await entity_resolver.flush_pending_stats()
|
||||
await _build_and_insert_entity_links_phase3(pool, entity_resolver, bank_id, phase3_ctx, log_buffer)
|
||||
except Exception:
|
||||
logger.warning("Phase 3 (best-effort display data) failed — retrieval unaffected", exc_info=True)
|
||||
logger.warning("Entity stats flush failed — retrieval unaffected", exc_info=True)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
@@ -1844,7 +1796,7 @@ async def _delta_metadata_only(
|
||||
merged_tags,
|
||||
)
|
||||
await fact_storage.update_memory_units_tags(conn, bank_id, document_id, merged_tags)
|
||||
if outbox_callback:
|
||||
if outbox_callback is not None:
|
||||
await outbox_callback(conn)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
|
||||
@@ -224,21 +224,6 @@ class ProcessedFact:
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Phase3Context:
|
||||
"""
|
||||
Data passed from Phase 2 to Phase 3 for entity link building.
|
||||
|
||||
Contains the unit IDs and entity resolution data needed to build
|
||||
entity links for UI graph visualization after the write transaction commits.
|
||||
"""
|
||||
|
||||
unit_ids: list[str] = field(default_factory=list)
|
||||
resolved_entity_ids: list[str] = field(default_factory=list)
|
||||
entity_to_unit: list[tuple] = field(default_factory=list)
|
||||
unit_to_entity_ids: dict[str, list[str]] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EntityResolutionResult:
|
||||
"""
|
||||
@@ -263,21 +248,6 @@ class Phase1Result:
|
||||
semantic_ann_links: list[tuple]
|
||||
|
||||
|
||||
@dataclass
|
||||
class EntityLink:
|
||||
"""
|
||||
Link between two memory units through a shared entity.
|
||||
|
||||
Used for entity-based graph connections in the memory graph.
|
||||
"""
|
||||
|
||||
from_unit_id: UUID
|
||||
to_unit_id: UUID
|
||||
entity_id: UUID
|
||||
link_type: str = "entity"
|
||||
weight: float = 1.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainBatch:
|
||||
"""
|
||||
|
||||
@@ -283,12 +283,12 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
score transformations. The three CTEs share one connection slot — important
|
||||
for asyncpg which does not allow concurrent queries on the same connection.
|
||||
|
||||
Index coverage (requires migration d2e3f4a5b6c7):
|
||||
entity: idx_memory_links_entity_covering (from_unit_id) INCLUDE (to_unit_id, entity_id)
|
||||
WHERE link_type = 'entity' → index-only scan, no heap reads
|
||||
semantic incoming:
|
||||
idx_memory_links_to_type_weight (to_unit_id, link_type, weight DESC)
|
||||
→ replaces costly BitmapAnd of two separate scans
|
||||
Index coverage:
|
||||
entity: idx_unit_entities_entity_unit (entity_id, unit_id) — entity
|
||||
expansion traverses unit_entities, not memory_links.
|
||||
semantic: idx_memory_links_from_type_weight / _to_type_weight
|
||||
(from_unit_id|to_unit_id, link_type, weight DESC) serve both
|
||||
outgoing and incoming sides as single composite index scans.
|
||||
"""
|
||||
config = get_config()
|
||||
ml = fq_table("memory_links")
|
||||
|
||||
@@ -212,16 +212,25 @@ class CrossEncoderReranker:
|
||||
# Get cross-encoder scores
|
||||
scores = await self.cross_encoder.predict(pairs)
|
||||
|
||||
# Normalize scores using sigmoid to [0, 1] range
|
||||
# Cross-encoder returns logits which can be negative
|
||||
import math
|
||||
|
||||
# Normalize scores to [0, 1] range.
|
||||
# External API rerankers (Cohere, Jina, llama.cpp/Qwen, etc.) return
|
||||
# calibrated relevance_score already in [0, 1]. These are used as-is
|
||||
# so that absolute confidence is preserved — a top candidate scoring
|
||||
# 0.007 stays low rather than being inflated to 1.0 by rank normalization.
|
||||
# Local models return logits (any real number) — sigmoid is appropriate.
|
||||
import numpy as np
|
||||
|
||||
def sigmoid(x):
|
||||
def _sigmoid(x: float) -> float:
|
||||
return 1 / (1 + np.exp(-x))
|
||||
|
||||
normalized_scores = [sigmoid(score) for score in scores]
|
||||
if scores and min(scores) >= 0.0 and max(scores) <= 1.0:
|
||||
# Scores already in [0, 1] — pass through to preserve absolute
|
||||
# confidence signal from calibrated rerankers.
|
||||
normalized_scores = list(scores)
|
||||
else:
|
||||
# Scores are logits (e.g. local sentence-transformers models).
|
||||
# Sigmoid maps (-inf, +inf) to (0, 1).
|
||||
normalized_scores = [_sigmoid(score) for score in scores]
|
||||
|
||||
# Create ScoredResult objects with cross-encoder scores
|
||||
scored_results = []
|
||||
|
||||
@@ -225,6 +225,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
groups_clause=groups_clause,
|
||||
arm_index=i,
|
||||
text_search_extension=text_ext,
|
||||
bm25_language=config.text_search_extension_native_language,
|
||||
extra_where=created_range_clause,
|
||||
)
|
||||
)
|
||||
@@ -465,6 +466,8 @@ async def retrieve_temporal_combined(
|
||||
best_date = ep["mentioned_at"]
|
||||
|
||||
if best_date:
|
||||
if best_date.tzinfo is None:
|
||||
best_date = best_date.replace(tzinfo=UTC)
|
||||
days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
|
||||
temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
||||
else:
|
||||
@@ -558,6 +561,8 @@ async def retrieve_temporal_combined(
|
||||
neighbor_best_date = n["mentioned_at"]
|
||||
|
||||
if neighbor_best_date:
|
||||
if neighbor_best_date.tzinfo is None:
|
||||
neighbor_best_date = neighbor_best_date.replace(tzinfo=UTC)
|
||||
days_from_mid = abs((neighbor_best_date - mid_date).total_seconds() / 86400)
|
||||
neighbor_temporal_proximity = (
|
||||
1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
||||
|
||||
@@ -407,6 +407,7 @@ class SQLDialect(ABC):
|
||||
groups_clause: str = "",
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
bm25_language: str = "english",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
"""Build a BM25/full-text search subquery arm.
|
||||
@@ -426,7 +427,9 @@ class SQLDialect(ABC):
|
||||
arm_index: Index of this arm in the UNION ALL (used by Oracle for
|
||||
unique SCORE labels).
|
||||
text_search_extension: Full-text search backend ("native", "vchord",
|
||||
"pg_textsearch"). Only relevant for PostgreSQL.
|
||||
"pg_textsearch", "pgroonga"). Only relevant for PostgreSQL.
|
||||
bm25_language: PostgreSQL text search dictionary used by the native
|
||||
backend (e.g. "english", "french"). Ignored by other backends.
|
||||
extra_where: Optional additional WHERE clause fragment (e.g. time range filter).
|
||||
"""
|
||||
...
|
||||
|
||||
@@ -270,6 +270,7 @@ class OracleDialect(SQLDialect):
|
||||
groups_clause: str = "",
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
bm25_language: str = "english",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
# Oracle Text: CONTAINS() / SCORE() with the CTXSYS.CONTEXT index.
|
||||
|
||||
@@ -182,6 +182,7 @@ class PostgreSQLDialect(SQLDialect):
|
||||
groups_clause: str = "",
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
bm25_language: str = "english",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
if text_search_extension == "vchord":
|
||||
@@ -193,10 +194,35 @@ class PostgreSQLDialect(SQLDialect):
|
||||
bm25_score_expr = f"-({text_param} <@> to_bm25query({text_param}, 'idx_memory_units_text_search'))"
|
||||
bm25_order_by = f"text <@> to_bm25query({text_param}, 'idx_memory_units_text_search') ASC"
|
||||
bm25_where_filter = ""
|
||||
else: # native tsvector
|
||||
bm25_score_expr = f"ts_rank_cd(search_vector, to_tsquery('english', {text_param}))"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# &@~ accepts pgroonga's query syntax (raw query text). pgroonga_score
|
||||
# returns a non-negative relevance score (higher = better).
|
||||
bm25_score_expr = "pgroonga_score(tableoid, ctid)"
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = f"AND search_vector @@ to_tsquery('english', {text_param})"
|
||||
bm25_where_filter = (
|
||||
f"AND (COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, '')) "
|
||||
f"&@~ {text_param}"
|
||||
)
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search: BM25 index over (id, text, context, text_signals)
|
||||
# with key_field='id'. The @@@ operator on the key_field requires a
|
||||
# field-qualified query (`text:foo`); to keep the bind-parameter form,
|
||||
# we fan the query out across all indexed text fields with paradedb.boolean.
|
||||
bm25_score_expr = "paradedb.score(id)"
|
||||
bm25_order_by = "paradedb.score(id) DESC"
|
||||
bm25_where_filter = (
|
||||
f"AND id @@@ paradedb.boolean(should => ARRAY["
|
||||
f"paradedb.match('text', {text_param}), "
|
||||
f"paradedb.match('context', {text_param}), "
|
||||
f"paradedb.match('text_signals', {text_param})"
|
||||
f"])"
|
||||
)
|
||||
else: # native tsvector
|
||||
# bm25_language is validated as a PG identifier in HindsightConfig.validate(),
|
||||
# so embedding it as a SQL literal here is safe.
|
||||
bm25_score_expr = f"ts_rank_cd(search_vector, to_tsquery('{bm25_language}', {text_param}))"
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = f"AND search_vector @@ to_tsquery('{bm25_language}', {text_param})"
|
||||
|
||||
return (
|
||||
f"(SELECT {cols},"
|
||||
@@ -221,7 +247,7 @@ class PostgreSQLDialect(SQLDialect):
|
||||
*,
|
||||
text_search_extension: str = "native",
|
||||
) -> str:
|
||||
if text_search_extension in ("vchord", "pg_textsearch"):
|
||||
if text_search_extension in ("vchord", "pg_textsearch", "pgroonga", "pg_search"):
|
||||
return query_text
|
||||
# native tsvector: join tokens with OR operator
|
||||
return " | ".join(tokens)
|
||||
|
||||
@@ -16,7 +16,7 @@ class S3FileStorage(FileStorage):
|
||||
S3-compatible object storage backend.
|
||||
|
||||
Uses obstore (Rust-backed) for high-throughput async access to
|
||||
Amazon S3, MinIO, Cloudflare R2, and other S3-compliant APIs.
|
||||
Amazon S3, MinIO, Cloudflare R2, Tigris, and other S3-compliant APIs.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
||||
@@ -40,6 +40,7 @@ from hindsight_api.extensions.operation_validator import (
|
||||
# Core operations
|
||||
OperationValidationError,
|
||||
OperationValidatorExtension,
|
||||
PrecheckContext,
|
||||
RecallContext,
|
||||
RecallResult,
|
||||
ReflectContext,
|
||||
@@ -72,6 +73,7 @@ __all__ = [
|
||||
"DeferOperation",
|
||||
"OperationValidationError",
|
||||
"OperationValidatorExtension",
|
||||
"PrecheckContext",
|
||||
"RecallContext",
|
||||
"RecallResult",
|
||||
"ReflectContext",
|
||||
|
||||
@@ -146,7 +146,11 @@ class DefaultExtensionContext(ExtensionContext):
|
||||
|
||||
# Ensure text search columns/indexes match the configured extension
|
||||
await asyncio.to_thread(
|
||||
ensure_text_search_extension, db_url, text_search_extension=config.text_search_extension, schema=schema
|
||||
ensure_text_search_extension,
|
||||
db_url,
|
||||
text_search_extension=config.text_search_extension,
|
||||
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
def get_memory_engine(self) -> "MemoryEngineInterface":
|
||||
|
||||
@@ -82,6 +82,33 @@ class ValidationResult:
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class PrecheckContext:
|
||||
"""Context for a pre-body-parse precheck on an operation.
|
||||
|
||||
Unlike :class:`RetainContext` / :class:`RecallContext` / etc., this
|
||||
context is constructed *before* the request body is deserialised. It
|
||||
therefore intentionally carries only the cheap, already-resolved
|
||||
pieces of request state:
|
||||
|
||||
- ``operation``: a short string identifying the route, e.g. ``"retain"``,
|
||||
``"recall"``, ``"reflect"``, ``"files_retain"``, ``"mental_model_create"``,
|
||||
``"mental_model_refresh"``.
|
||||
- ``bank_id``: parsed from the URL path.
|
||||
- ``request_context``: the authenticated :class:`RequestContext` (tenant
|
||||
already resolved by the tenant extension).
|
||||
|
||||
Implementations should keep precheck cheap and side-effect-free. The
|
||||
full per-request validators (``validate_retain`` / ``validate_recall``
|
||||
/ ``validate_reflect``) still run after the body is parsed and remain
|
||||
the source of truth for the precise per-call cost / quota arithmetic.
|
||||
"""
|
||||
|
||||
operation: str
|
||||
bank_id: str
|
||||
request_context: "RequestContext"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainContext:
|
||||
"""Context for a retain operation validation (pre-operation).
|
||||
@@ -407,6 +434,42 @@ class OperationValidatorExtension(Extension, ABC):
|
||||
- consolidate (mental models consolidation)
|
||||
"""
|
||||
|
||||
# =========================================================================
|
||||
# Pre-body-parse hook (optional - default no-op)
|
||||
# =========================================================================
|
||||
|
||||
async def precheck(self, ctx: PrecheckContext) -> ValidationResult:
|
||||
"""
|
||||
Cheap pre-body-parse check, called before the request body is read.
|
||||
|
||||
FastAPI resolves ``Depends`` callables before deserialising the route
|
||||
body; routes that wire ``precheck`` as a dependency therefore short
|
||||
-circuit here without ever materialising the JSON payload in memory.
|
||||
That makes this the right hook for "should this caller be allowed to
|
||||
spend resources on this request at all" decisions — e.g. a balance
|
||||
is exhausted, a key is revoked, or a tenant is rate-limited.
|
||||
|
||||
Implementations should:
|
||||
- Be cheap: prefer cached lookups, avoid heavy DB queries.
|
||||
- Use only data on ``ctx`` (operation name + bank_id + request_context);
|
||||
the body is not yet available.
|
||||
- Be conservative on errors: prefer ``ValidationResult.accept()`` so
|
||||
a transient lookup failure doesn't turn into a request rejection.
|
||||
The post-body ``validate_*`` hooks still run and remain the source
|
||||
of truth for the precise per-call cost check.
|
||||
|
||||
Default implementation accepts everything. Override to opt in.
|
||||
|
||||
Args:
|
||||
ctx: Pre-body context with operation name, bank_id, and
|
||||
request_context (tenant already resolved).
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the request may proceed to
|
||||
body parsing and the post-parse validators.
|
||||
"""
|
||||
return ValidationResult.accept()
|
||||
|
||||
# =========================================================================
|
||||
# Pre-operation validation hooks (abstract - must be implemented)
|
||||
# =========================================================================
|
||||
|
||||
@@ -24,7 +24,15 @@ import uvicorn
|
||||
from . import MemoryEngine, __version__
|
||||
from .api import create_app
|
||||
from .banner import print_banner
|
||||
from .config import DEFAULT_WORKERS, ENV_HOST, ENV_WORKERS, HindsightConfig, _get_raw_config
|
||||
from .config import (
|
||||
DEFAULT_ACCESS_LOG,
|
||||
DEFAULT_WORKERS,
|
||||
ENV_ACCESS_LOG,
|
||||
ENV_HOST,
|
||||
ENV_WORKERS,
|
||||
HindsightConfig,
|
||||
_get_raw_config,
|
||||
)
|
||||
from .daemon import (
|
||||
DEFAULT_DAEMON_PORT,
|
||||
DEFAULT_IDLE_TIMEOUT,
|
||||
@@ -65,29 +73,42 @@ def _signal_handler(signum, frame):
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
def resolve_daemon_host_port(*, args_host: str, args_port: int, config_host: str, config_port: int) -> tuple[str, int]:
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class ResolvedDaemonHostPort:
|
||||
host: str
|
||||
port: int
|
||||
|
||||
|
||||
def resolve_daemon_host_port(
|
||||
*,
|
||||
args_host: str,
|
||||
args_port: int,
|
||||
explicit_host: bool,
|
||||
explicit_port: bool,
|
||||
) -> ResolvedDaemonHostPort:
|
||||
"""Resolve host/port for daemon mode.
|
||||
|
||||
Defaults to 127.0.0.1 for security, but honors explicit user overrides
|
||||
via --host flag or HINDSIGHT_API_HOST env var. Uses DEFAULT_DAEMON_PORT
|
||||
unless the user specified a custom port.
|
||||
"""
|
||||
port = args_port if args_port != config_port else DEFAULT_DAEMON_PORT
|
||||
port = args_port if explicit_port else DEFAULT_DAEMON_PORT
|
||||
# Only force localhost if the user didn't explicitly set a host
|
||||
if args_host != config_host or os.environ.get(ENV_HOST):
|
||||
if explicit_host or os.environ.get(ENV_HOST):
|
||||
host = args_host
|
||||
else:
|
||||
host = "127.0.0.1"
|
||||
return host, port
|
||||
return ResolvedDaemonHostPort(host=host, port=port)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point for the CLI."""
|
||||
global _memory
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class ParsedCliArgs:
|
||||
args: argparse.Namespace
|
||||
explicit_host: bool
|
||||
explicit_port: bool
|
||||
|
||||
# Load configuration from environment (for CLI args defaults)
|
||||
config = _get_raw_config()
|
||||
|
||||
def _parse_cli_args(argv: list[str], config: HindsightConfig) -> ParsedCliArgs:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="hindsight-api",
|
||||
description="Hindsight API Server",
|
||||
@@ -95,12 +116,14 @@ def main():
|
||||
|
||||
# Server options
|
||||
parser.add_argument(
|
||||
"--host", default=config.host, help=f"Host to bind to (default: {config.host}, env: HINDSIGHT_API_HOST)"
|
||||
"--host",
|
||||
default=argparse.SUPPRESS,
|
||||
help=f"Host to bind to (default: {config.host}, env: HINDSIGHT_API_HOST)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=config.port,
|
||||
default=argparse.SUPPRESS,
|
||||
help=f"Port to bind to (default: {config.port}, env: HINDSIGHT_API_PORT)",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -120,9 +143,18 @@ def main():
|
||||
)
|
||||
|
||||
# Access log options
|
||||
parser.add_argument("--access-log", action="store_true", help="Enable access log")
|
||||
parser.add_argument("--no-access-log", dest="access_log", action="store_false", help="Disable access log (default)")
|
||||
parser.set_defaults(access_log=False)
|
||||
parser.add_argument(
|
||||
"--access-log",
|
||||
action="store_true",
|
||||
default=os.getenv(ENV_ACCESS_LOG, "").lower() in ("1", "true", "yes", "on") or DEFAULT_ACCESS_LOG,
|
||||
help=f"Enable access log (env: {ENV_ACCESS_LOG}, default: {DEFAULT_ACCESS_LOG})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-access-log",
|
||||
dest="access_log",
|
||||
action="store_false",
|
||||
help="Disable access log (overrides env and default)",
|
||||
)
|
||||
|
||||
# Proxy options
|
||||
parser.add_argument(
|
||||
@@ -149,7 +181,27 @@ def main():
|
||||
help=f"Idle timeout in seconds before auto-exit in daemon mode (default: {DEFAULT_IDLE_TIMEOUT})",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
explicit_host = hasattr(args, "host")
|
||||
explicit_port = hasattr(args, "port")
|
||||
if not explicit_host:
|
||||
args.host = config.host
|
||||
if not explicit_port:
|
||||
args.port = config.port
|
||||
|
||||
return ParsedCliArgs(args=args, explicit_host=explicit_host, explicit_port=explicit_port)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point for the CLI."""
|
||||
global _memory
|
||||
|
||||
# Load configuration from environment (for CLI args defaults)
|
||||
config = _get_raw_config()
|
||||
|
||||
parsed_cli_args = _parse_cli_args(sys.argv[1:], config)
|
||||
args = parsed_cli_args.args
|
||||
|
||||
# Daemon mode handling.
|
||||
# is_daemon_child is True when we are the re-exec'd child spawned by
|
||||
@@ -160,12 +212,14 @@ def main():
|
||||
is_daemon = args.daemon or is_daemon_child
|
||||
|
||||
if is_daemon:
|
||||
args.host, args.port = resolve_daemon_host_port(
|
||||
resolved_daemon_host_port = resolve_daemon_host_port(
|
||||
args_host=args.host,
|
||||
args_port=args.port,
|
||||
config_host=config.host,
|
||||
config_port=config.port,
|
||||
explicit_host=parsed_cli_args.explicit_host,
|
||||
explicit_port=parsed_cli_args.explicit_port,
|
||||
)
|
||||
args.host = resolved_daemon_host_port.host
|
||||
args.port = resolved_daemon_host_port.port
|
||||
|
||||
# Detach into background (parent re-execs and exits; child redirects
|
||||
# stdio to log file). No lockfile needed — port binding prevents
|
||||
@@ -239,8 +293,6 @@ def main():
|
||||
# When using workers or reload, we must use import string so each worker can import the app
|
||||
use_import_string = args.workers > 1 or args.reload
|
||||
# Check for uvloop/winloop availability
|
||||
import sys
|
||||
|
||||
loop_impl = "asyncio"
|
||||
if sys.platform == "win32":
|
||||
try:
|
||||
|
||||
@@ -44,6 +44,7 @@ _ALL_TOOLS: frozenset[str] = frozenset(
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
@@ -221,6 +222,7 @@ def register_mcp_tools(
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
@@ -277,6 +279,9 @@ def register_mcp_tools(
|
||||
if "refresh_mental_model" in tools_to_register:
|
||||
_register_refresh_mental_model(mcp, memory, config)
|
||||
|
||||
if "clear_mental_model" in tools_to_register:
|
||||
_register_clear_mental_model(mcp, memory, config)
|
||||
|
||||
# Directive tools
|
||||
if "list_directives" in tools_to_register:
|
||||
_register_list_directives(mcp, memory, config)
|
||||
@@ -438,6 +443,7 @@ _AUDITABLE_MCP_TOOLS: frozenset[str] = frozenset(
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
"delete_document",
|
||||
@@ -793,7 +799,8 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
{"tags": [...], "match": "any_strict"} or compound {"and": [...]}, {"or": [...]}, {"not": {...}}.
|
||||
Example: [{"not": {"tags": ["closeout"], "match": "any_strict"}}] excludes memories tagged closeout.
|
||||
Mutually exclusive with tags.
|
||||
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z'). Helps retrieve time-relevant memories.
|
||||
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z').
|
||||
Anchors relative temporal expressions and recency scoring.
|
||||
bank_id: Optional bank to search in (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
@@ -863,7 +870,8 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
{"tags": [...], "match": "any_strict"} or compound {"and": [...]}, {"or": [...]}, {"not": {...}}.
|
||||
Example: [{"not": {"tags": ["closeout"], "match": "any_strict"}}] excludes memories tagged closeout.
|
||||
Mutually exclusive with tags.
|
||||
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z'). Helps retrieve time-relevant memories.
|
||||
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z').
|
||||
Anchors relative temporal expressions and recency scoring.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
@@ -922,6 +930,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
response_schema: dict | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
include_based_on: bool = False,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -951,6 +960,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
response_schema: Optional JSON schema for structured output. When provided, the response includes a 'structured_output' field.
|
||||
tags: Optional tags to filter memories by (e.g., ['project:alpha'])
|
||||
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
|
||||
include_based_on: Include source facts used for synthesis. Defaults to false because broad reflections can exceed MCP client result limits.
|
||||
bank_id: Optional bank to reflect in (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
@@ -978,6 +988,8 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
reflect_result = await memory.reflect_async(**reflect_kwargs)
|
||||
|
||||
result_data = json.loads(reflect_result.model_dump_json(indent=2))
|
||||
if not include_based_on:
|
||||
result_data.pop("based_on", None)
|
||||
if response_schema is not None and hasattr(reflect_result, "structured_output"):
|
||||
result_data["structured_output"] = reflect_result.structured_output
|
||||
return json.dumps(result_data, indent=2)
|
||||
@@ -999,6 +1011,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
response_schema: dict | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
include_based_on: bool = False,
|
||||
) -> dict:
|
||||
"""
|
||||
Generate thoughtful analysis by synthesizing stored memories with the bank's personality.
|
||||
@@ -1027,6 +1040,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
response_schema: Optional JSON schema for structured output. When provided, the response includes a 'structured_output' field.
|
||||
tags: Optional tags to filter memories by (e.g., ['project:alpha'])
|
||||
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
|
||||
include_based_on: Include source facts used for synthesis. Defaults to false because broad reflections can exceed MCP client result limits.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
@@ -1053,6 +1067,8 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
reflect_result = await memory.reflect_async(**reflect_kwargs)
|
||||
|
||||
result_data = reflect_result.model_dump()
|
||||
if not include_based_on:
|
||||
result_data.pop("based_on", None)
|
||||
if response_schema is not None and hasattr(reflect_result, "structured_output"):
|
||||
result_data["structured_output"] = reflect_result.structured_output
|
||||
return result_data
|
||||
@@ -1765,6 +1781,98 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_clear_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the clear_mental_model tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def clear_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Clear a mental model's content so the next refresh performs a full re-synthesis.
|
||||
|
||||
This is useful for delta-mode models that have accumulated drift over many
|
||||
incremental refreshes. After clearing, call refresh_mental_model to trigger
|
||||
a clean full rebuild.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to clear
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
result = await memory.clear_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if result is None:
|
||||
return json.dumps({"error": f"Mental model '{mental_model_id}' not found"})
|
||||
return json.dumps(
|
||||
{
|
||||
"mental_model_id": result["id"],
|
||||
"status": "cleared",
|
||||
"message": f"Mental model '{mental_model_id}' content cleared. Call refresh_mental_model to rebuild.",
|
||||
}
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except ValueError as e:
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error clearing mental model: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def clear_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Clear a mental model's content so the next refresh performs a full re-synthesis.
|
||||
|
||||
This is useful for delta-mode models that have accumulated drift over many
|
||||
incremental refreshes. After clearing, call refresh_mental_model to trigger
|
||||
a clean full rebuild.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to clear
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
result = await memory.clear_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if result is None:
|
||||
return {"error": f"Mental model '{mental_model_id}' not found"}
|
||||
return {
|
||||
"mental_model_id": result["id"],
|
||||
"status": "cleared",
|
||||
"message": f"Mental model '{mental_model_id}' content cleared. Call refresh_mental_model to rebuild.",
|
||||
}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except ValueError as e:
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error clearing mental model: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# DIRECTIVE TOOLS
|
||||
# =========================================================================
|
||||
|
||||
@@ -27,6 +27,7 @@ from alembic.config import Config
|
||||
from alembic.script.revision import ResolutionError
|
||||
from sqlalchemy import Connection, create_engine, text
|
||||
|
||||
from ._pg_search import normalize_pg_search_tokenizer, pg_search_bm25_columns
|
||||
from ._vector_index import (
|
||||
bootstrap_extension,
|
||||
detect_vector_extension,
|
||||
@@ -803,6 +804,7 @@ def ensure_text_search_extension(
|
||||
database_url: str,
|
||||
text_search_extension: str = "native",
|
||||
schema: str | None = None,
|
||||
pg_search_tokenizer: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Ensure the text search columns and indexes match the configured extension.
|
||||
@@ -815,13 +817,18 @@ def ensure_text_search_extension(
|
||||
|
||||
Args:
|
||||
database_url: SQLAlchemy database URL
|
||||
text_search_extension: Configured text search extension ("native" or "vchord")
|
||||
text_search_extension: Configured text search extension — one of
|
||||
"native", "vchord", "pg_textsearch", "pgroonga", or "pg_search"
|
||||
schema: Target PostgreSQL schema name (None for public)
|
||||
pg_search_tokenizer: Optional ParadeDB tokenizer to apply to pg_search
|
||||
BM25 text fields when indexes are created. Empty keeps the
|
||||
ParadeDB default.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If extension mismatch with existing data
|
||||
"""
|
||||
schema_name = schema or "public"
|
||||
pg_search_tokenizer = normalize_pg_search_tokenizer(pg_search_tokenizer)
|
||||
|
||||
engine = create_engine(to_libpq_url(database_url))
|
||||
with engine.connect() as conn:
|
||||
@@ -838,6 +845,17 @@ def ensure_text_search_extension(
|
||||
elif text_search_extension == "pg_textsearch":
|
||||
target_column_type = "text"
|
||||
target_index_type = "bm25"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# pgroonga indexes the base text column directly. We keep a dummy
|
||||
# TEXT column named search_vector for symmetry with pg_textsearch
|
||||
# and so the column-type mismatch detection above keeps working.
|
||||
target_column_type = "text"
|
||||
target_index_type = "pgroonga"
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB: same column type / access method as pg_textsearch.
|
||||
# Disambiguated below by inspecting the index reloptions (key_field).
|
||||
target_column_type = "text"
|
||||
target_index_type = "bm25"
|
||||
else: # native
|
||||
target_column_type = "tsvector"
|
||||
target_index_type = "gin"
|
||||
@@ -875,16 +893,18 @@ def ensure_text_search_extension(
|
||||
|
||||
if not current_column_info:
|
||||
logger.warning(f"No search_vector column found for {table_name}, will create it")
|
||||
mismatched_tables.append((table_name, None, None))
|
||||
mismatched_tables.append((table_name, None, None, False))
|
||||
continue
|
||||
|
||||
# Check column type (udt_name contains the actual type: tsvector, bm25vector, etc.)
|
||||
current_column_type = current_column_info[1] # udt_name
|
||||
|
||||
# Get current index type
|
||||
# Get current index type and definition. The definition lets us
|
||||
# disambiguate pg_textsearch vs pg_search (both register a `bm25`
|
||||
# access method but only pg_search uses the `key_field` reloption).
|
||||
current_index_info = conn.execute(
|
||||
text("""
|
||||
SELECT am.amname
|
||||
SELECT am.amname, pi.indexdef
|
||||
FROM pg_indexes pi
|
||||
JOIN pg_class c ON c.relname = pi.indexname
|
||||
JOIN pg_am am ON am.oid = c.relam
|
||||
@@ -896,10 +916,21 @@ def ensure_text_search_extension(
|
||||
).fetchone()
|
||||
|
||||
current_index_type = current_index_info[0] if current_index_info else None
|
||||
current_index_def = current_index_info[1] if current_index_info else None
|
||||
|
||||
# Detect pg_search specifically (vs pg_textsearch) via the key_field reloption
|
||||
current_is_pg_search = bool(current_index_def and "key_field" in current_index_def)
|
||||
want_pg_search = text_search_extension == "pg_search"
|
||||
|
||||
# Check if column and index types match target
|
||||
column_matches = current_column_type == target_column_type
|
||||
index_matches = current_index_type == target_index_type if current_index_type else False
|
||||
# When both target and current sit at column=text/index=bm25, the
|
||||
# access-method check alone can't tell pg_textsearch from pg_search —
|
||||
# require the key_field reloption to agree with the configured backend.
|
||||
if column_matches and index_matches and target_index_type == "bm25" and target_column_type == "text":
|
||||
if current_is_pg_search != want_pg_search:
|
||||
index_matches = False
|
||||
|
||||
if not (column_matches and index_matches):
|
||||
logger.info(
|
||||
@@ -907,7 +938,7 @@ def ensure_text_search_extension(
|
||||
f"column={current_column_type} (want {target_column_type}), "
|
||||
f"index={current_index_type} (want {target_index_type})"
|
||||
)
|
||||
mismatched_tables.append((table_name, current_column_type, current_index_type))
|
||||
mismatched_tables.append((table_name, current_column_type, current_index_type, current_is_pg_search))
|
||||
|
||||
# Check if table has data
|
||||
row_count = conn.execute(text(f"SELECT COUNT(*) FROM {schema_name}.{table_name}")).scalar()
|
||||
@@ -925,14 +956,20 @@ def ensure_text_search_extension(
|
||||
# If there's data in any mismatched table, raise error
|
||||
if tables_with_data:
|
||||
table_list = ", ".join([f"{table}({count} rows)" for table, count in tables_with_data])
|
||||
# Detect current extension from column type
|
||||
# Detect current extension from column type, index type, and (for the
|
||||
# text/bm25 ambiguity) the key_field reloption. tsvector is
|
||||
# unambiguous; text could be pg_textsearch, pgroonga, or pg_search.
|
||||
current_col_type = mismatched_tables[0][1]
|
||||
current_idx_type = mismatched_tables[0][2]
|
||||
first_is_pg_search = mismatched_tables[0][3]
|
||||
if current_col_type == "tsvector":
|
||||
current_ext = "native"
|
||||
elif current_col_type == "bm25vector":
|
||||
current_ext = "vchord"
|
||||
elif current_col_type == "text" and current_idx_type == "pgroonga":
|
||||
current_ext = "pgroonga"
|
||||
elif current_col_type == "text":
|
||||
current_ext = "pg_textsearch"
|
||||
current_ext = "pg_search" if first_is_pg_search else "pg_textsearch"
|
||||
else:
|
||||
current_ext = "unknown"
|
||||
raise RuntimeError(
|
||||
@@ -947,7 +984,7 @@ def ensure_text_search_extension(
|
||||
# Tables are empty, safe to recreate columns/indexes
|
||||
logger.info(f"Recreating text search columns/indexes for {text_search_extension}")
|
||||
|
||||
for table_name, current_col_type, current_idx_type in mismatched_tables:
|
||||
for table_name, current_col_type, current_idx_type, _was_pg_search in mismatched_tables:
|
||||
# Drop existing index if it exists
|
||||
if current_idx_type:
|
||||
logger.info(f"Dropping {current_idx_type} index on {table_name}")
|
||||
@@ -1000,21 +1037,79 @@ def ensure_text_search_extension(
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
)
|
||||
else: # native
|
||||
logger.info(f"Creating tsvector column on {table_name}")
|
||||
# Different GENERATED expression for each table
|
||||
if table_name == "memory_units":
|
||||
generated_expr = "to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))"
|
||||
else: # reflections
|
||||
generated_expr = "to_tsvector('english', COALESCE(name, '') || ' ' || content)"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# Ensure pgroonga extension is available
|
||||
try:
|
||||
conn.execute(text("CREATE EXTENSION IF NOT EXISTS pgroonga CASCADE"))
|
||||
except Exception:
|
||||
# Extension might already exist or user lacks permissions — verify
|
||||
has_ext = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pgroonga'")).fetchone()
|
||||
if not has_ext:
|
||||
raise
|
||||
|
||||
logger.info(f"Creating dummy TEXT search_vector on {table_name} for pgroonga")
|
||||
# pgroonga indexes the base text column directly, but we keep a
|
||||
# dummy search_vector column for symmetry with pg_textsearch and
|
||||
# so the column-type mismatch detection above keeps working.
|
||||
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector TEXT"))
|
||||
|
||||
# pgroonga index expression mirrors pg_textsearch
|
||||
if table_name == "memory_units":
|
||||
index_expr = (
|
||||
"(COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, ''))"
|
||||
)
|
||||
else: # reflections
|
||||
index_expr = "(COALESCE(name, '') || ' ' || content)"
|
||||
|
||||
logger.info(f"Creating pgroonga index on {table_name}")
|
||||
# TokenBigram is the polyglot default — falls back to whitespace
|
||||
# tokenization for space-separated languages and bigram for CJK.
|
||||
# NormalizerNFKC150 handles Unicode normalization (full/half-width,
|
||||
# case folding, etc.) which materially improves Japanese recall.
|
||||
conn.execute(
|
||||
text(f"""
|
||||
ALTER TABLE {schema_name}.{table_name}
|
||||
ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS ({generated_expr}) STORED
|
||||
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
|
||||
ON {schema_name}.{table_name}
|
||||
USING pgroonga ({index_expr})
|
||||
WITH (tokenizer='TokenBigram', normalizer='NormalizerNFKC150')
|
||||
""")
|
||||
)
|
||||
elif text_search_extension == "pg_search":
|
||||
logger.info(f"Creating TEXT column on {table_name}")
|
||||
# Dummy TEXT column for schema symmetry; pg_search indexes operate on base columns.
|
||||
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector TEXT"))
|
||||
|
||||
# ParadeDB BM25 index over the table's primary key and text columns.
|
||||
# Column list mirrors what the initial / text_signals migrations create.
|
||||
if table_name == "memory_units":
|
||||
bm25_cols = pg_search_bm25_columns(
|
||||
"id",
|
||||
("text", "context", "text_signals"),
|
||||
pg_search_tokenizer,
|
||||
)
|
||||
else: # reflections
|
||||
bm25_cols = pg_search_bm25_columns(
|
||||
"id",
|
||||
("name", "content"),
|
||||
pg_search_tokenizer,
|
||||
)
|
||||
|
||||
logger.info(f"Creating ParadeDB BM25 index on {table_name}")
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
|
||||
ON {schema_name}.{table_name}
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
)
|
||||
else: # native
|
||||
logger.info(f"Creating tsvector column on {table_name}")
|
||||
# Plain tsvector column. The application populates search_vector
|
||||
# at INSERT time via to_tsvector($lang, ...) using the configured
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE — see
|
||||
# ops_postgresql.insert_facts_batch.
|
||||
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector tsvector"))
|
||||
|
||||
# Create GIN index
|
||||
logger.info(f"Creating GIN index on {table_name}")
|
||||
|
||||
@@ -129,6 +129,7 @@ PROVIDER_NAME_MAPPING = {
|
||||
"vertexai": "google",
|
||||
"groq": "groq",
|
||||
"ollama": "ollama",
|
||||
"ollama-cloud": "ollama",
|
||||
"lmstudio": "lmstudio",
|
||||
"openai-codex": "openai",
|
||||
"claude-code": "anthropic",
|
||||
|
||||
@@ -256,6 +256,7 @@ def main():
|
||||
tenant_extension=tenant_extension,
|
||||
max_slots=config.worker_max_slots,
|
||||
slot_reservations=config.worker_slot_reservations,
|
||||
consolidation_bank_priority=config.worker_consolidation_bank_priority or None,
|
||||
)
|
||||
|
||||
# Create the HTTP app for metrics/health
|
||||
|
||||
@@ -24,7 +24,7 @@ from .exceptions import DeferOperation, RetryTaskAt
|
||||
from .stage import StageHolder, bind_holder
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.db.base import DatabaseBackend
|
||||
from hindsight_api.engine.db.base import DatabaseBackend, DatabaseConnection
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -133,6 +133,7 @@ class WorkerPoller:
|
||||
tenant_extension: "TenantExtension | None" = None,
|
||||
max_slots: int = 10,
|
||||
slot_reservations: dict[str, int] | None = None,
|
||||
consolidation_bank_priority: dict[str, int] | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize the worker poller.
|
||||
@@ -150,6 +151,11 @@ class WorkerPoller:
|
||||
"retain": 3}). Reserved slots guarantee capacity for that operation type.
|
||||
Remaining slots (max_slots - sum of reservations) form a shared pool usable
|
||||
by any operation type. Defaults to {"consolidation": 2} if None.
|
||||
consolidation_bank_priority: Per-bank priority for consolidation scheduling.
|
||||
Maps bank name patterns to integer priorities (higher = claimed first).
|
||||
Patterns support ``*`` as wildcard. A bare ``*`` key is the catch-all default.
|
||||
When set, consolidation tasks are claimed in priority tiers rather than
|
||||
pure created_at order. None or empty dict preserves current behavior.
|
||||
"""
|
||||
self._backend = backend
|
||||
self._worker_id = worker_id
|
||||
@@ -168,6 +174,9 @@ class WorkerPoller:
|
||||
self._slot_reservations: dict[str, int] = (
|
||||
slot_reservations if slot_reservations is not None else {"consolidation": 2}
|
||||
)
|
||||
self._consolidation_bank_priority: dict[str, int] | None = (
|
||||
consolidation_bank_priority if consolidation_bank_priority else None
|
||||
)
|
||||
# Cache of which optional PG routines are installed on the server
|
||||
# (probed once, memoised for the life of the poller).
|
||||
from ..engine.db.optional_routines import OptionalRoutines
|
||||
@@ -187,13 +196,18 @@ class WorkerPoller:
|
||||
# schema we serviced so a busy tenant can't monopolize the poll order.
|
||||
self._next_schema_idx: int = 0
|
||||
|
||||
async def _get_schemas(self) -> list[str | None]:
|
||||
"""Get list of schemas to poll. Returns [None] for default schema (no prefix)."""
|
||||
@staticmethod
|
||||
def _normalize_poll_schema(schema: str | None) -> str | None:
|
||||
"""Use None internally for the default schema because SQL helpers omit that prefix."""
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA
|
||||
|
||||
return None if schema == DEFAULT_DATABASE_SCHEMA else schema
|
||||
|
||||
async def _get_schemas(self) -> list[str | None]:
|
||||
"""Get list of schemas to poll. Returns [None] for default schema (no prefix)."""
|
||||
tenants = await self._tenant_extension.list_tenants()
|
||||
# Convert default schema to None for SQL compatibility (no prefix), keep others as-is
|
||||
return [t.schema if t.schema != DEFAULT_DATABASE_SCHEMA else None for t in tenants]
|
||||
return [self._normalize_poll_schema(t.schema) for t in tenants]
|
||||
|
||||
async def _scan_active_schemas(self, schemas: list[str | None]) -> set[str | None]:
|
||||
"""Find which schemas have pending work.
|
||||
@@ -213,22 +227,57 @@ class WorkerPoller:
|
||||
async with self._backend.acquire() as conn:
|
||||
if await self._optional_routines.is_installed(conn, "schemas_with_pending_work"):
|
||||
rows = await conn.fetch("SELECT * FROM public.schemas_with_pending_work()")
|
||||
return {r[0] for r in rows}
|
||||
|
||||
# Fallback: per-schema EXISTS checks from Python
|
||||
active: set[str | None] = set()
|
||||
for schema in schemas:
|
||||
table = fq_table("async_operations", schema)
|
||||
try:
|
||||
has_work = await conn.fetchval(
|
||||
f"SELECT EXISTS(SELECT 1 FROM {table} "
|
||||
f"WHERE status = 'pending' AND task_payload IS NOT NULL LIMIT 1)"
|
||||
routine_active = {self._normalize_poll_schema(r[0]) for r in rows}
|
||||
known_schemas = set(schemas)
|
||||
active = routine_active & known_schemas
|
||||
unknown = routine_active - known_schemas
|
||||
if unknown:
|
||||
logger.warning(
|
||||
"Optional PG routine public.schemas_with_pending_work() returned schema(s) "
|
||||
"not present in tenant discovery: %s",
|
||||
sorted(str(s) for s in unknown),
|
||||
)
|
||||
if has_work:
|
||||
active.add(schema)
|
||||
except Exception:
|
||||
pass
|
||||
return active
|
||||
|
||||
# The optional routine returns PostgreSQL schema names, but the poller uses
|
||||
# None for the default schema. Older operator-supplied implementations also
|
||||
# commonly scan tenant_% only; when the default schema is in scope but absent
|
||||
# from the routine result, verify via the fully-correct per-schema fallback so
|
||||
# public single-tenant deployments cannot silently starve.
|
||||
should_verify_with_fallback = (None in known_schemas and None not in active) or (
|
||||
bool(routine_active) and not active
|
||||
)
|
||||
if not should_verify_with_fallback:
|
||||
return active
|
||||
|
||||
fallback_active = await self._scan_active_schemas_by_exists(conn, schemas)
|
||||
missed = fallback_active - active
|
||||
if missed:
|
||||
logger.warning(
|
||||
"Optional PG routine public.schemas_with_pending_work() missed claimable schema(s) %s; "
|
||||
"using per-schema fallback for this poll",
|
||||
sorted(str(s) for s in missed),
|
||||
)
|
||||
return fallback_active
|
||||
|
||||
return await self._scan_active_schemas_by_exists(conn, schemas)
|
||||
|
||||
async def _scan_active_schemas_by_exists(
|
||||
self, conn: "DatabaseConnection", schemas: list[str | None]
|
||||
) -> set[str | None]:
|
||||
"""Find active schemas using per-schema EXISTS checks."""
|
||||
active: set[str | None] = set()
|
||||
for schema in schemas:
|
||||
table = fq_table("async_operations", schema)
|
||||
try:
|
||||
has_work = await conn.fetchval(
|
||||
f"SELECT EXISTS(SELECT 1 FROM {table} "
|
||||
f"WHERE status = 'pending' AND task_payload IS NOT NULL LIMIT 1)"
|
||||
)
|
||||
if has_work:
|
||||
active.add(schema)
|
||||
except Exception:
|
||||
pass
|
||||
return active
|
||||
|
||||
async def _get_available_slots(self) -> SlotAvailability:
|
||||
"""
|
||||
@@ -428,6 +477,7 @@ class WorkerPoller:
|
||||
self._worker_id,
|
||||
reserved_limits,
|
||||
shared_limit,
|
||||
consolidation_bank_priority=self._consolidation_bank_priority,
|
||||
)
|
||||
|
||||
if not all_rows:
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api-slim"
|
||||
version = "0.6.2"
|
||||
version = "0.7.1"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -21,7 +21,7 @@ dependencies = [
|
||||
"sqlalchemy>=2.0.44",
|
||||
"alembic>=1.17.1",
|
||||
"pgvector>=0.4.1",
|
||||
"greenlet>=3.2.4,<3.4.0", # 3.4.0 lacks arm64 wheels for manylinux_2_41
|
||||
"greenlet>=3.2.4,<3.4.0", # 3.4.0 lacks arm64 wheels for manylinux_2_41
|
||||
"psycopg2-binary>=2.9.11",
|
||||
"tiktoken>=0.12.0",
|
||||
"httpx>=0.27.0",
|
||||
@@ -68,7 +68,7 @@ dependencies = [
|
||||
"tornado>=6.5.5", # DoS multipart/incomplete cookie validation fix
|
||||
"aiohttp>=3.13.3", # Multiple DoS vulnerabilities
|
||||
"pygments>=2.20.0", # ReDoS via inefficient GUID regex fix
|
||||
"claude-agent-sdk>=0.1.27",
|
||||
"claude-agent-sdk>=0.2.82",
|
||||
"boto3>=1.42.74",
|
||||
]
|
||||
|
||||
@@ -142,6 +142,7 @@ addopts = "--timeout 300 -n 8 --dist loadgroup --durations=10 -v"
|
||||
markers = [
|
||||
"oracle: Oracle 23ai integration tests (require ORACLE_TEST_DSN env var)",
|
||||
"hs_llm_mat: LLM minimum acceptance tests — run in CI matrix across multiple providers",
|
||||
"hs_llm_core: Core pipeline tests that need a real LLM but only one provider",
|
||||
]
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
|
||||
@@ -413,21 +413,48 @@ def query_analyzer():
|
||||
@pytest_asyncio.fixture(scope="function")
|
||||
async def memory(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
"""
|
||||
Provide a MemoryEngine instance for each test.
|
||||
Provide a MemoryEngine instance using a mock LLM for deterministic tests.
|
||||
|
||||
Must be function-scoped because:
|
||||
1. pytest-xdist runs tests in separate processes with different event loops
|
||||
2. asyncpg pools are bound to the event loop that created them
|
||||
3. Each test needs its own pool in its own event loop
|
||||
The mock LLM returns canned facts derived from input text, allowing the
|
||||
full retain → recall → reflect pipeline to work without real LLM calls.
|
||||
This makes core tests fast, deterministic, and free from LLM flakiness.
|
||||
|
||||
Uses small pool sizes since tests run in parallel.
|
||||
Uses pg0_db_url (a postgresql:// URL) directly, so MemoryEngine won't try to
|
||||
manage pg0 lifecycle - that's handled by the session-scoped pg0_db_url fixture.
|
||||
Migrations are disabled here since they're run once at session scope in pg0_db_url.
|
||||
Uses SyncTaskBackend so async tasks execute immediately (no worker needed).
|
||||
Tests that need real LLM output quality should use `memory_real_llm` instead.
|
||||
"""
|
||||
mem = MemoryEngine(
|
||||
db_url=pg0_db_url, # Direct postgresql:// URL, not pg0://
|
||||
db_url=pg0_db_url,
|
||||
memory_llm_provider="mock",
|
||||
memory_llm_api_key="",
|
||||
memory_llm_model="mock",
|
||||
embeddings=embeddings,
|
||||
cross_encoder=cross_encoder,
|
||||
query_analyzer=query_analyzer,
|
||||
pool_min_size=1,
|
||||
pool_max_size=5,
|
||||
run_migrations=False,
|
||||
task_backend=SyncTaskBackend(),
|
||||
)
|
||||
await mem.initialize()
|
||||
yield mem
|
||||
try:
|
||||
if mem._pool and not mem._pool._closing:
|
||||
await mem.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope="function")
|
||||
async def memory_real_llm(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
"""
|
||||
Provide a MemoryEngine instance using a real LLM provider.
|
||||
|
||||
Use this fixture ONLY for tests that assert on LLM output quality
|
||||
(fact extraction accuracy, language preservation, consolidation decisions, etc.).
|
||||
These tests are non-deterministic and should be marked with @pytest.mark.hs_llm_core
|
||||
(or @pytest.mark.hs_llm_mat for provider matrix acceptance tests).
|
||||
"""
|
||||
mem = MemoryEngine(
|
||||
db_url=pg0_db_url,
|
||||
memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"),
|
||||
memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"),
|
||||
memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"),
|
||||
@@ -437,8 +464,8 @@ async def memory(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
query_analyzer=query_analyzer,
|
||||
pool_min_size=1,
|
||||
pool_max_size=5,
|
||||
run_migrations=False, # Migrations already run at session scope
|
||||
task_backend=SyncTaskBackend(), # Execute tasks immediately in tests
|
||||
run_migrations=False,
|
||||
task_backend=SyncTaskBackend(),
|
||||
)
|
||||
await mem.initialize()
|
||||
yield mem
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
"""
|
||||
LLM-as-a-judge utility for hs_llm_core tests.
|
||||
|
||||
Replaces brittle string-matching assertions (e.g., `assert "alice" in answer`)
|
||||
with semantic evaluation via a frontier mini model. This makes tests resilient
|
||||
to phrasing variations while still verifying LLM output quality.
|
||||
|
||||
Usage in tests:
|
||||
result = await llm_judge.assert_response_meets_criteria(
|
||||
response="Alice is a researcher at Stanford...",
|
||||
criteria="The response mentions Alice and her role",
|
||||
)
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from hindsight_api.engine.llm_wrapper import create_llm_provider
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Judge model configuration — always uses Gemini by default since GEMINI_API_KEY
|
||||
# is available in all CI jobs. The judge must be independent of the test provider
|
||||
# (hs_llm_mat tests run across openai, groq, bedrock, etc.).
|
||||
# Override with HINDSIGHT_TEST_JUDGE_PROVIDER / MODEL / API_KEY env vars.
|
||||
_JUDGE_PROVIDER = os.getenv("HINDSIGHT_TEST_JUDGE_PROVIDER", "gemini")
|
||||
_raw_model = os.getenv("HINDSIGHT_TEST_JUDGE_MODEL", "gemini-2.5-flash-lite")
|
||||
# Strip "google/" prefix — gemini API key auth expects bare model names.
|
||||
_JUDGE_MODEL = _raw_model.removeprefix("google/") if _JUDGE_PROVIDER == "gemini" else _raw_model
|
||||
_JUDGE_API_KEY = os.getenv(
|
||||
"HINDSIGHT_TEST_JUDGE_API_KEY",
|
||||
os.getenv("GEMINI_API_KEY", os.getenv("HINDSIGHT_API_LLM_API_KEY", "")),
|
||||
)
|
||||
_JUDGE_BASE_URL = os.getenv("HINDSIGHT_TEST_JUDGE_BASE_URL", "")
|
||||
|
||||
|
||||
class JudgeVerdict(BaseModel):
|
||||
meets_criteria: bool
|
||||
reasoning: str
|
||||
|
||||
|
||||
_judge_instance = None
|
||||
|
||||
|
||||
def _get_judge():
|
||||
global _judge_instance
|
||||
if _judge_instance is None:
|
||||
_judge_instance = create_llm_provider(
|
||||
provider=_JUDGE_PROVIDER,
|
||||
api_key=_JUDGE_API_KEY,
|
||||
base_url=_JUDGE_BASE_URL or "",
|
||||
model=_JUDGE_MODEL,
|
||||
reasoning_effort="low",
|
||||
)
|
||||
return _judge_instance
|
||||
|
||||
|
||||
async def evaluate(
|
||||
response: str,
|
||||
criteria: str,
|
||||
context: str | None = None,
|
||||
) -> JudgeVerdict:
|
||||
"""Ask the judge LLM whether a response meets the given criteria.
|
||||
|
||||
Args:
|
||||
response: The LLM-generated text to evaluate.
|
||||
criteria: Plain-English description of what the response should contain/satisfy.
|
||||
context: Optional context (e.g., the stored memories or query) for the judge.
|
||||
|
||||
Returns:
|
||||
JudgeVerdict with meets_criteria bool and reasoning string.
|
||||
"""
|
||||
judge = _get_judge()
|
||||
|
||||
context_block = f"\n\nContext provided to the system:\n{context}" if context else ""
|
||||
|
||||
result = await judge.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You are a test evaluation judge. Given a response and evaluation criteria, "
|
||||
"determine whether the response meets the criteria. "
|
||||
"Respond with JSON: {\"meets_criteria\": true/false, \"reasoning\": \"brief explanation\"}"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
f"## Response to evaluate\n{response}\n"
|
||||
f"{context_block}\n"
|
||||
f"## Criteria\n{criteria}\n\n"
|
||||
"Does the response meet the criteria?"
|
||||
),
|
||||
},
|
||||
],
|
||||
response_format=JudgeVerdict,
|
||||
max_completion_tokens=256,
|
||||
temperature=0.0,
|
||||
scope="test_judge",
|
||||
)
|
||||
|
||||
if isinstance(result, JudgeVerdict):
|
||||
return result
|
||||
|
||||
# Fallback: parse raw dict/string
|
||||
if isinstance(result, dict):
|
||||
return JudgeVerdict(**result)
|
||||
return JudgeVerdict(**json.loads(str(result)))
|
||||
|
||||
|
||||
async def assert_meets_criteria(
|
||||
response: str,
|
||||
criteria: str,
|
||||
context: str | None = None,
|
||||
msg: str | None = None,
|
||||
) -> JudgeVerdict:
|
||||
"""Assert that a response meets criteria, with a clear failure message.
|
||||
|
||||
Raises AssertionError if the judge says criteria are not met.
|
||||
"""
|
||||
verdict = await evaluate(response=response, criteria=criteria, context=context)
|
||||
if not verdict.meets_criteria:
|
||||
fail_msg = msg or f"LLM judge: criteria not met"
|
||||
raise AssertionError(
|
||||
f"{fail_msg}\n"
|
||||
f" Criteria: {criteria}\n"
|
||||
f" Judge reasoning: {verdict.reasoning}\n"
|
||||
f" Response (first 300 chars): {response[:300]}"
|
||||
)
|
||||
return verdict
|
||||
@@ -328,8 +328,11 @@ async def test_run_migration_without_schema_discovers_and_deduplicates_schemas(m
|
||||
database_url: str,
|
||||
text_search_extension: str = "native",
|
||||
schema: str | None = None,
|
||||
pg_search_tokenizer: str | None = None,
|
||||
) -> None:
|
||||
calls["ensure_text_search_extension"].append((database_url, text_search_extension, schema))
|
||||
calls["ensure_text_search_extension"].append(
|
||||
(database_url, text_search_extension, pg_search_tokenizer, schema)
|
||||
)
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_DATABASE_URL", "postgresql://test")
|
||||
monkeypatch.setattr(admin_cli, "load_extension", lambda *args, **kwargs: MockTenantExtension())
|
||||
@@ -353,8 +356,8 @@ async def test_run_migration_without_schema_discovers_and_deduplicates_schemas(m
|
||||
("resolved::postgresql://test", "pgvector", "tenant_demo"),
|
||||
]
|
||||
assert calls["ensure_text_search_extension"] == [
|
||||
("resolved::postgresql://test", "native", "public"),
|
||||
("resolved::postgresql://test", "native", "tenant_demo"),
|
||||
("resolved::postgresql://test", "native", "", "public"),
|
||||
("resolved::postgresql://test", "native", "", "tenant_demo"),
|
||||
]
|
||||
|
||||
|
||||
@@ -398,8 +401,11 @@ async def test_run_migration_without_schema_runs_optional_post_migration_hooks(m
|
||||
database_url: str,
|
||||
text_search_extension: str = "native",
|
||||
schema: str | None = None,
|
||||
pg_search_tokenizer: str | None = None,
|
||||
) -> None:
|
||||
calls["ensure_text_search_extension"].append((database_url, text_search_extension, schema))
|
||||
calls["ensure_text_search_extension"].append(
|
||||
(database_url, text_search_extension, pg_search_tokenizer, schema)
|
||||
)
|
||||
|
||||
monkeypatch.setattr(admin_cli, "load_extension", lambda *args, **kwargs: MockTenantExtension())
|
||||
monkeypatch.setattr(admin_cli, "resolve_database_url", fake_resolve_database_url)
|
||||
@@ -431,8 +437,8 @@ async def test_run_migration_without_schema_runs_optional_post_migration_hooks(m
|
||||
("resolved::postgresql://test", "pgvector", "tenant_demo"),
|
||||
]
|
||||
assert calls["ensure_text_search_extension"] == [
|
||||
("resolved::postgresql://test", "native", "public"),
|
||||
("resolved::postgresql://test", "native", "tenant_demo"),
|
||||
("resolved::postgresql://test", "native", "", "public"),
|
||||
("resolved::postgresql://test", "native", "", "tenant_demo"),
|
||||
]
|
||||
|
||||
|
||||
@@ -467,8 +473,11 @@ async def test_run_migration_with_schema_only_runs_requested_schema(monkeypatch)
|
||||
database_url: str,
|
||||
text_search_extension: str = "native",
|
||||
schema: str | None = None,
|
||||
pg_search_tokenizer: str | None = None,
|
||||
) -> None:
|
||||
calls["ensure_text_search_extension"].append((database_url, text_search_extension, schema))
|
||||
calls["ensure_text_search_extension"].append(
|
||||
(database_url, text_search_extension, pg_search_tokenizer, schema)
|
||||
)
|
||||
|
||||
monkeypatch.setattr(admin_cli, "load_extension", lambda *args, **kwargs: MockTenantExtension())
|
||||
monkeypatch.setattr(admin_cli, "resolve_database_url", fake_resolve_database_url)
|
||||
@@ -484,4 +493,4 @@ async def test_run_migration_with_schema_only_runs_requested_schema(monkeypatch)
|
||||
assert schemas == ["tenant_demo"]
|
||||
assert calls["run_migrations"] == [("resolved::postgresql://test", "tenant_demo")]
|
||||
assert calls["ensure_vector_extension"] == [("resolved::postgresql://test", "pgvector", "tenant_demo")]
|
||||
assert calls["ensure_text_search_extension"] == [("resolved::postgresql://test", "native", "tenant_demo")]
|
||||
assert calls["ensure_text_search_extension"] == [("resolved::postgresql://test", "native", "", "tenant_demo")]
|
||||
|
||||
@@ -106,6 +106,7 @@ async def test_small_async_batch_no_splitting(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(600)
|
||||
async def test_large_async_batch_auto_splits(memory, request_context):
|
||||
"""Test that large async batches automatically split into sub-batches with parent operation."""
|
||||
from hindsight_api.engine.memory_engine import count_tokens
|
||||
@@ -831,6 +832,173 @@ async def test_submit_async_operation_leaves_claimable_row_when_submit_task_fail
|
||||
assert payload["contents"] == [{"content": "hello", "document_id": "d1"}]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Regression tests for issue #1795: submit_async_retain must NOT fragment a
|
||||
# single oversized item across multiple child async-operations. Workers have
|
||||
# no per-document serialization for retain (the busy-bank guard in
|
||||
# claim_tasks only covers consolidation), so concurrent siblings sharing one
|
||||
# document_id race on handle_document_tracking(is_first_batch=True), cascade-
|
||||
# delete each other's memory_units, and trip FK violations on memory_links in
|
||||
# the final ANN pass. The fix: keep the oversized item un-chunked in a single
|
||||
# child; the worker's in-process splitter handles intra-document chunking
|
||||
# sequentially with correct is_first_batch semantics.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oversized_single_item_creates_one_child_not_many(memory, request_context, monkeypatch):
|
||||
"""One oversized document → one child operation holding the un-chunked content."""
|
||||
from hindsight_api.config import get_config
|
||||
from hindsight_api.engine.memory_engine import count_tokens
|
||||
|
||||
bank_id = f"test_oversized_one_child_{uuid.uuid4().hex[:8]}"
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
|
||||
# No-op the worker dispatch: this is a structural assertion on the
|
||||
# async_operations rows that submit_async_retain inserts (those rows are
|
||||
# committed before submit_task is invoked), so we don't need the
|
||||
# SyncTaskBackend to drive the slow LLM pipeline to completion.
|
||||
async def noop_submit_task(_task_dict):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(memory._task_backend, "submit_task", noop_submit_task)
|
||||
|
||||
tokens_per_batch = get_config().retain_batch_tokens
|
||||
# Build content that comfortably exceeds the per-batch token budget.
|
||||
# The pre-fix code would split this into many siblings, all sharing
|
||||
# the same document_id and racing on handle_document_tracking.
|
||||
unit = "The quick brown fox jumps over the lazy dog. " * 50 # ~570 tokens
|
||||
repetitions = max(1, (tokens_per_batch * 4) // count_tokens(unit) + 1)
|
||||
big_content = unit * repetitions
|
||||
document_id = f"doc-oversize-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
total_tokens = count_tokens(big_content)
|
||||
assert total_tokens > tokens_per_batch, (
|
||||
f"Test setup error: content has {total_tokens} tokens, must exceed the batch budget of {tokens_per_batch}"
|
||||
)
|
||||
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": big_content, "document_id": document_id}],
|
||||
request_context=request_context,
|
||||
)
|
||||
parent_operation_id = result["operation_id"]
|
||||
|
||||
# Parent metadata must say one child — even though token count is huge,
|
||||
# the un-fragmenting splitter keeps the single item in one child.
|
||||
parent_row = await pool.fetchrow(
|
||||
"SELECT result_metadata FROM async_operations WHERE operation_id = $1",
|
||||
uuid.UUID(parent_operation_id),
|
||||
)
|
||||
parent_meta = (
|
||||
json.loads(parent_row["result_metadata"])
|
||||
if isinstance(parent_row["result_metadata"], str)
|
||||
else parent_row["result_metadata"]
|
||||
)
|
||||
assert parent_meta["num_sub_batches"] == 1, (
|
||||
f"Expected 1 child for an oversized single item, got {parent_meta['num_sub_batches']}. "
|
||||
"Issue #1795: per-chunk children race on the shared document_id."
|
||||
)
|
||||
|
||||
# Exactly one retain child row.
|
||||
children = await pool.fetch(
|
||||
"""
|
||||
SELECT operation_id, status, task_payload, result_metadata
|
||||
FROM async_operations
|
||||
WHERE bank_id = $1 AND operation_type = 'retain'
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
assert len(children) == 1, (
|
||||
f"Expected exactly 1 child retain operation for bank_id={bank_id}, "
|
||||
f"got {len(children)}. Issue #1795 regression: oversized item was "
|
||||
f"fragmented into per-chunk children."
|
||||
)
|
||||
|
||||
child = children[0]
|
||||
payload = child["task_payload"]
|
||||
payload = json.loads(payload) if isinstance(payload, str) else payload
|
||||
assert payload["type"] == "batch_retain"
|
||||
assert payload["bank_id"] == bank_id
|
||||
# Child holds the full un-chunked content — the worker's in-process
|
||||
# splitter will re-chunk it sequentially with is_first_batch=(i==1).
|
||||
assert len(payload["contents"]) == 1, (
|
||||
f"Child payload should hold the original single item, got {len(payload['contents'])} items"
|
||||
)
|
||||
assert payload["contents"][0]["document_id"] == document_id
|
||||
assert payload["contents"][0]["content"] == big_content, (
|
||||
"Child must carry the FULL un-chunked content — chunking happens inside the worker, not at submit time"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oversized_item_among_small_items_keeps_small_items_packed(memory, request_context, monkeypatch):
|
||||
"""Mixed batch: small items pack together; oversized item isolates to its own child."""
|
||||
from hindsight_api.config import get_config
|
||||
from hindsight_api.engine.memory_engine import count_tokens
|
||||
|
||||
bank_id = f"test_mixed_pack_{uuid.uuid4().hex[:8]}"
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
|
||||
# Structural assertion on async_operations rows only — skip worker dispatch.
|
||||
async def noop_submit_task(_task_dict):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(memory._task_backend, "submit_task", noop_submit_task)
|
||||
|
||||
tokens_per_batch = get_config().retain_batch_tokens
|
||||
big_unit = "The quick brown fox jumps over the lazy dog. " * 50
|
||||
big_repetitions = max(1, (tokens_per_batch * 3) // count_tokens(big_unit) + 1)
|
||||
big_content = big_unit * big_repetitions
|
||||
|
||||
big_doc = f"doc-big-{uuid.uuid4().hex[:8]}"
|
||||
small_docs = [f"doc-small-{i}-{uuid.uuid4().hex[:6]}" for i in range(3)]
|
||||
contents = [
|
||||
{"content": "Alice works at Google.", "document_id": small_docs[0]},
|
||||
{"content": big_content, "document_id": big_doc},
|
||||
{"content": "Bob loves Python.", "document_id": small_docs[1]},
|
||||
{"content": "Carol writes Rust.", "document_id": small_docs[2]},
|
||||
]
|
||||
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
children = await pool.fetch(
|
||||
"""
|
||||
SELECT task_payload
|
||||
FROM async_operations
|
||||
WHERE bank_id = $1 AND operation_type = 'retain'
|
||||
ORDER BY created_at, operation_id
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
# Find the child holding the big doc — it must hold ONLY the big doc.
|
||||
big_child = None
|
||||
small_doc_ids_seen: list[str] = []
|
||||
for row in children:
|
||||
payload = row["task_payload"]
|
||||
payload = json.loads(payload) if isinstance(payload, str) else payload
|
||||
items = payload["contents"]
|
||||
doc_ids = [item["document_id"] for item in items]
|
||||
if big_doc in doc_ids:
|
||||
assert big_child is None, f"Big doc {big_doc} should appear in exactly one child, found in 2"
|
||||
big_child = items
|
||||
assert doc_ids == [big_doc], f"Child holding the oversized item must hold ONLY that item, got {doc_ids}"
|
||||
assert items[0]["content"] == big_content
|
||||
else:
|
||||
small_doc_ids_seen.extend(doc_ids)
|
||||
|
||||
assert big_child is not None, "Big doc must appear in some child"
|
||||
# Every small input present, exactly once, across the other children.
|
||||
assert sorted(small_doc_ids_seen) == sorted(small_docs)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_async_batch_retain_rolls_back_parent_on_child_failure(
|
||||
memory_no_llm_verify, request_context, monkeypatch
|
||||
|
||||
@@ -4,7 +4,7 @@ Test OpenAI Batch API integration for retain fact extraction.
|
||||
Tests cover:
|
||||
- Normal batch API flow (submit, poll, complete)
|
||||
- Crash recovery (resume from existing batch_id)
|
||||
- Provider fallback (when batch API not supported)
|
||||
- Hard error when provider doesn't support the batch API (no silent fallback)
|
||||
- Worker recovery on restart
|
||||
"""
|
||||
import pytest
|
||||
@@ -332,20 +332,19 @@ async def test_batch_api_crash_recovery(mock_llm_config, test_contents, hindsigh
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_fallback_unsupported_provider(mock_llm_config, test_contents, hindsight_config):
|
||||
"""Test fallback to sync mode when provider doesn't support batch API."""
|
||||
async def test_batch_api_raises_for_unsupported_provider(mock_llm_config, test_contents, hindsight_config):
|
||||
"""Batch extraction must surface a hard error (not silently fall back) when
|
||||
the configured provider doesn't support the batch API.
|
||||
|
||||
# Mock provider that doesn't support batch API
|
||||
The silent-fallback behavior was removed in #1463 because it created a
|
||||
mutual-recursion path between sync and batch extraction. Misconfiguration
|
||||
should fail loudly and be caught at startup; this test guards that
|
||||
contract.
|
||||
"""
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=False)
|
||||
mock_llm_config.provider = "groq" # Example of provider
|
||||
mock_llm_config.provider = "groq"
|
||||
|
||||
# Patch the sync mode function to verify it's called
|
||||
with patch(
|
||||
"hindsight_api.engine.retain.fact_extraction.extract_facts_from_contents"
|
||||
) as mock_sync_extract:
|
||||
mock_sync_extract.return_value = ([], [], MagicMock())
|
||||
|
||||
# Call batch API extraction (should fallback to sync)
|
||||
with pytest.raises(RuntimeError, match="does not.*support the batch API"):
|
||||
await extract_facts_from_contents_batch_api(
|
||||
contents=test_contents,
|
||||
llm_config=mock_llm_config,
|
||||
@@ -356,13 +355,7 @@ async def test_batch_api_fallback_unsupported_provider(mock_llm_config, test_con
|
||||
schema=None,
|
||||
)
|
||||
|
||||
# Verify fallback occurred
|
||||
mock_sync_extract.assert_called_once()
|
||||
|
||||
# Verify batch API methods were NOT called
|
||||
mock_llm_config._provider_impl.submit_batch.assert_not_called()
|
||||
|
||||
logger.info("✅ Fallback to sync mode test passed")
|
||||
mock_llm_config._provider_impl.submit_batch.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -1,38 +1,103 @@
|
||||
"""
|
||||
Test validation for batch API + synchronous retain.
|
||||
Test validation for batch API configuration.
|
||||
|
||||
When HINDSIGHT_API_RETAIN_BATCH_ENABLED=true, synchronous retain operations
|
||||
should be rejected with a 400 error since they will timeout.
|
||||
When HINDSIGHT_API_RETAIN_BATCH_ENABLED=true but the LLM provider does not
|
||||
support the batch API, the server should fail at startup with a clear error
|
||||
message telling the user exactly what config is wrong.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
from hindsight_api.config import HindsightConfig
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_validation(memory, request_context):
|
||||
async def test_startup_rejects_batch_enabled_with_non_batch_provider():
|
||||
"""
|
||||
Test that attempting synchronous retain with batch API enabled
|
||||
raises an error at the HTTP layer.
|
||||
verify_llm() should raise RuntimeError at startup when
|
||||
retain_batch_enabled=True but the provider doesn't support batch API.
|
||||
"""
|
||||
mock_provider = AsyncMock()
|
||||
mock_provider.supports_batch_api = AsyncMock(return_value=False)
|
||||
|
||||
mock_llm_config = MagicMock()
|
||||
mock_llm_config.provider = "gemini"
|
||||
mock_llm_config._provider_impl = mock_provider
|
||||
mock_llm_config.verify_connection = AsyncMock()
|
||||
|
||||
This test verifies the validation logic exists - actual HTTP testing
|
||||
would require full FastAPI app setup.
|
||||
"""
|
||||
# Create config with batch API enabled
|
||||
config = HindsightConfig.from_env()
|
||||
config.retain_batch_enabled = True
|
||||
config.retain_batch_poll_interval_seconds = 1
|
||||
|
||||
# Verify the validation exists in memory engine
|
||||
# The actual HTTP validation happens in http.py api_retain()
|
||||
# This test documents the expected behavior
|
||||
with patch("hindsight_api.engine.memory_engine.get_config", return_value=config):
|
||||
supports_batch = await mock_provider.supports_batch_api()
|
||||
assert supports_batch is False
|
||||
|
||||
assert config.retain_batch_enabled is True
|
||||
assert config.retain_batch_poll_interval_seconds == 1
|
||||
with pytest.raises(RuntimeError, match="HINDSIGHT_API_RETAIN_BATCH_ENABLED=true"):
|
||||
if config.retain_batch_enabled and not supports_batch:
|
||||
raise RuntimeError(
|
||||
f"Configuration error: HINDSIGHT_API_RETAIN_BATCH_ENABLED=true "
|
||||
f"but the retain LLM provider '{mock_llm_config.provider}' "
|
||||
f"does not support the batch API. Either switch to a provider "
|
||||
f"that supports batch operations (e.g. 'openai', 'groq') or "
|
||||
f"set HINDSIGHT_API_RETAIN_BATCH_ENABLED=false."
|
||||
)
|
||||
|
||||
# When batch API is enabled and async=false, the HTTP endpoint
|
||||
# should return 400 with message:
|
||||
# "Batch API is enabled (HINDSIGHT_API_RETAIN_BATCH_ENABLED=true) but async=false"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_startup_allows_batch_enabled_with_batch_provider():
|
||||
"""
|
||||
verify_llm() should NOT raise when retain_batch_enabled=True and the
|
||||
provider supports batch API (e.g. OpenAI).
|
||||
"""
|
||||
mock_provider = AsyncMock()
|
||||
mock_provider.supports_batch_api = AsyncMock(return_value=True)
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
config.retain_batch_enabled = True
|
||||
|
||||
supports_batch = await mock_provider.supports_batch_api()
|
||||
assert supports_batch is True
|
||||
|
||||
# No error should be raised
|
||||
if config.retain_batch_enabled and not supports_batch:
|
||||
pytest.fail("Should not reach here -- provider supports batch API")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_startup_allows_batch_disabled_with_non_batch_provider():
|
||||
"""
|
||||
verify_llm() should NOT raise when retain_batch_enabled=False,
|
||||
regardless of provider batch support.
|
||||
"""
|
||||
mock_provider = AsyncMock()
|
||||
mock_provider.supports_batch_api = AsyncMock(return_value=False)
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
config.retain_batch_enabled = False
|
||||
|
||||
supports_batch = await mock_provider.supports_batch_api()
|
||||
assert supports_batch is False
|
||||
|
||||
# No error should be raised when batch is disabled
|
||||
if config.retain_batch_enabled and not supports_batch:
|
||||
pytest.fail("Should not reach here -- batch is disabled")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_raises_if_batch_unsupported():
|
||||
"""
|
||||
extract_facts_from_contents_batch_api() should raise RuntimeError
|
||||
if somehow called with a non-batch provider (startup check bypassed).
|
||||
"""
|
||||
mock_provider = AsyncMock()
|
||||
mock_provider.supports_batch_api = AsyncMock(return_value=False)
|
||||
|
||||
with pytest.raises(RuntimeError, match="does not support the batch API"):
|
||||
if not await mock_provider.supports_batch_api():
|
||||
raise RuntimeError(
|
||||
"retain_batch_enabled=True but provider 'gemini' does not "
|
||||
"support the batch API. This should have been caught at startup -- check "
|
||||
"HINDSIGHT_API_RETAIN_BATCH_ENABLED and your LLM provider configuration."
|
||||
)
|
||||
|
||||
@@ -1,9 +1,275 @@
|
||||
"""Test automatic batch chunking based on character count."""
|
||||
|
||||
import asyncio
|
||||
import pytest
|
||||
from hindsight_api import MemoryEngine
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.engine.memory_engine import (
|
||||
_split_contents_into_async_children,
|
||||
_split_contents_into_sub_batches,
|
||||
count_tokens,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Regression tests for issue #1571: the splitter must actually chunk an
|
||||
# oversized single item instead of passing it through as one giant
|
||||
# 1/1 sub-batch. The latter behavior contradicts the "splitting into
|
||||
# ~10K-token sub-batches" log message and OOMs the orchestrator under
|
||||
# realistic memory limits when one retain payload exceeds the budget.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_split_single_oversized_item_produces_multiple_sub_batches():
|
||||
"""A single item that exceeds tokens_per_batch must be chunked."""
|
||||
tokens_per_batch = 1_000
|
||||
# ~250 tokens per repetition × 100 ≈ 25k tokens — well over the budget.
|
||||
big_content = "The quick brown fox jumps over the lazy dog. " * 1_000
|
||||
assert count_tokens(big_content) > tokens_per_batch
|
||||
|
||||
split = _split_contents_into_sub_batches(
|
||||
[{"content": big_content, "document_id": "doc-oversize"}],
|
||||
tokens_per_batch,
|
||||
)
|
||||
|
||||
assert len(split.sub_batches) > 1, (
|
||||
f"Expected >1 sub-batches for a single oversize item, got {len(split.sub_batches)}. "
|
||||
"Splitter is regressing to the pre-#1571 'pass-through as 1/1' behavior."
|
||||
)
|
||||
# Every sub-batch is itself bounded by the token budget (modulo the
|
||||
# char-vs-token conversion headroom inside the helper).
|
||||
for batch in split.sub_batches:
|
||||
batch_tokens = sum(count_tokens(item.get("content", "")) for item in batch)
|
||||
assert batch_tokens <= tokens_per_batch, (
|
||||
f"Sub-batch with {batch_tokens} tokens exceeds budget {tokens_per_batch}"
|
||||
)
|
||||
# Every chunked sub-batch must trace back to the single source item.
|
||||
assert all(origins == [0] for origins in split.origin_indices)
|
||||
|
||||
|
||||
def test_split_oversized_item_preserves_document_id_and_metadata():
|
||||
"""Chunked sub-batches must inherit the original item's metadata."""
|
||||
tokens_per_batch = 500
|
||||
big_content = "Alice met Bob at the coffee shop. " * 500
|
||||
item = {
|
||||
"content": big_content,
|
||||
"document_id": "doc-42",
|
||||
"context": "shared-context",
|
||||
"tags": ["t1", "t2"],
|
||||
}
|
||||
|
||||
split = _split_contents_into_sub_batches([item], tokens_per_batch)
|
||||
|
||||
assert len(split.sub_batches) > 1
|
||||
for batch in split.sub_batches:
|
||||
assert len(batch) == 1
|
||||
chunk = batch[0]
|
||||
assert chunk["document_id"] == "doc-42"
|
||||
assert chunk["context"] == "shared-context"
|
||||
assert chunk["tags"] == ["t1", "t2"]
|
||||
# And the content is a non-empty substring (no chunk lost its text).
|
||||
assert chunk["content"]
|
||||
|
||||
|
||||
def test_split_mixed_batch_chunks_only_oversized_items():
|
||||
"""In a mixed batch, only the oversized item is chunked; others pack normally."""
|
||||
tokens_per_batch = 1_000
|
||||
small_a = "Alice works at Google. " * 5 # tiny
|
||||
small_b = "Bob loves Python. " * 5 # tiny
|
||||
big = "The quick brown fox jumps over the lazy dog. " * 1_000 # huge
|
||||
|
||||
contents = [
|
||||
{"content": small_a, "document_id": "doc-a"},
|
||||
{"content": big, "document_id": "doc-b"},
|
||||
{"content": small_b, "document_id": "doc-c"},
|
||||
]
|
||||
|
||||
split = _split_contents_into_sub_batches(contents, tokens_per_batch)
|
||||
|
||||
# We expect: [small_a packed] then N chunks of big, then [small_b packed].
|
||||
# At minimum: > 2 sub-batches (a + multiple big chunks + c).
|
||||
assert len(split.sub_batches) > 2
|
||||
|
||||
# Every original input must appear in origin_indices at least once.
|
||||
flat_origins = [idx for origins in split.origin_indices for idx in origins]
|
||||
assert 0 in flat_origins # small_a
|
||||
assert 1 in flat_origins # big (likely many times)
|
||||
assert 2 in flat_origins # small_b
|
||||
|
||||
# The oversized input (index 1) appears in more sub-batches than the
|
||||
# small ones — that's the chunked-fan-out signature.
|
||||
big_origin_count = sum(1 for origins in split.origin_indices if origins == [1])
|
||||
small_a_origin_count = sum(1 for origins in split.origin_indices if 0 in origins)
|
||||
assert big_origin_count > small_a_origin_count
|
||||
|
||||
|
||||
def test_split_small_batch_returns_single_sub_batch():
|
||||
"""A batch under the budget stays as a single sub-batch."""
|
||||
tokens_per_batch = 10_000
|
||||
contents = [
|
||||
{"content": "Alice works at Google", "document_id": "doc-1"},
|
||||
{"content": "Bob loves Python", "document_id": "doc-2"},
|
||||
]
|
||||
|
||||
split = _split_contents_into_sub_batches(contents, tokens_per_batch)
|
||||
|
||||
assert len(split.sub_batches) == 1
|
||||
assert split.sub_batches[0] == contents
|
||||
assert split.origin_indices == [[0, 1]]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Regression tests for issue #1795: the async-submit splitter must NEVER
|
||||
# fragment a single oversized item across multiple children. Each child
|
||||
# becomes an independent async_operations row claimed by workers in
|
||||
# parallel with no per-document gate, so siblings would race on the same
|
||||
# document_id, cascade-delete each other's memory_units, and trip FK
|
||||
# violations on memory_links in the final ANN pass.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_async_children_oversized_single_item_stays_in_one_child():
|
||||
"""An oversized single item must NOT be split across children."""
|
||||
tokens_per_batch = 1_000
|
||||
big_content = "The quick brown fox jumps over the lazy dog. " * 1_000
|
||||
item = {"content": big_content, "document_id": "doc-oversize"}
|
||||
assert count_tokens(big_content) > tokens_per_batch
|
||||
|
||||
children = _split_contents_into_async_children([item], tokens_per_batch)
|
||||
|
||||
# Exactly one child holding the full un-chunked item — the worker
|
||||
# will sequentially chunk it inside retain_batch_async.
|
||||
assert len(children) == 1, (
|
||||
f"Oversized single item must become exactly one child, got {len(children)}. "
|
||||
"Regressing to per-chunk children causes the issue #1795 race."
|
||||
)
|
||||
assert children[0] == [item]
|
||||
assert children[0][0]["content"] == big_content
|
||||
|
||||
|
||||
def test_async_children_oversized_item_preserves_metadata():
|
||||
"""The single-child payload must carry every original key untouched."""
|
||||
tokens_per_batch = 500
|
||||
item = {
|
||||
"content": "Alice met Bob at the coffee shop. " * 500,
|
||||
"document_id": "doc-42",
|
||||
"context": "shared-context",
|
||||
"tags": ["t1", "t2"],
|
||||
"metadata": {"source": "test"},
|
||||
}
|
||||
|
||||
children = _split_contents_into_async_children([item], tokens_per_batch)
|
||||
|
||||
assert len(children) == 1
|
||||
assert children[0] == [item]
|
||||
|
||||
|
||||
def test_async_children_packs_small_items_by_budget():
|
||||
"""Small items pack into shared children up to the token budget."""
|
||||
tokens_per_batch = 100
|
||||
item_text = "Alice works at Google and Bob loves Python programming language."
|
||||
item_tokens = count_tokens(item_text)
|
||||
# Pick enough items to force at least 2 children at this budget.
|
||||
num_items = max(4, (tokens_per_batch // max(item_tokens, 1)) * 3)
|
||||
contents = [{"content": item_text, "document_id": f"doc-{i}"} for i in range(num_items)]
|
||||
total = sum(count_tokens(c["content"]) for c in contents)
|
||||
assert total > tokens_per_batch, (
|
||||
f"Test setup error: {total} tokens does not exceed budget {tokens_per_batch}"
|
||||
)
|
||||
|
||||
children = _split_contents_into_async_children(contents, tokens_per_batch)
|
||||
|
||||
assert len(children) >= 2
|
||||
# Every child stays within budget.
|
||||
for child in children:
|
||||
child_tokens = sum(count_tokens(item["content"]) for item in child)
|
||||
assert child_tokens <= tokens_per_batch, (
|
||||
f"Child holding {child_tokens} tokens exceeds budget {tokens_per_batch}"
|
||||
)
|
||||
# Every input item appears exactly once across all children.
|
||||
flat = [item for child in children for item in child]
|
||||
assert flat == contents
|
||||
|
||||
|
||||
def test_async_children_mixed_small_and_oversized():
|
||||
"""Small items pack together; the oversized item gets its own child as-is."""
|
||||
tokens_per_batch = 1_000
|
||||
small_a = {"content": "Alice works at Google. " * 5, "document_id": "doc-a"}
|
||||
big = {"content": "The quick brown fox jumps over the lazy dog. " * 1_000, "document_id": "doc-big"}
|
||||
small_b = {"content": "Bob loves Python. " * 5, "document_id": "doc-b"}
|
||||
small_c = {"content": "Carol writes Rust. " * 5, "document_id": "doc-c"}
|
||||
|
||||
children = _split_contents_into_async_children([small_a, big, small_b, small_c], tokens_per_batch)
|
||||
|
||||
# The big item must be in its own child, un-chunked.
|
||||
big_children = [c for c in children if any(item is big or item.get("document_id") == "doc-big" for item in c)]
|
||||
assert len(big_children) == 1, (
|
||||
f"Oversized item must occupy exactly one child, found {len(big_children)} containing it"
|
||||
)
|
||||
assert big_children[0] == [big]
|
||||
assert big_children[0][0]["content"] == big["content"]
|
||||
|
||||
# Every small item is present somewhere.
|
||||
flat_doc_ids = [item["document_id"] for child in children for item in child]
|
||||
for small in (small_a, small_b, small_c):
|
||||
assert small["document_id"] in flat_doc_ids
|
||||
|
||||
# No child fragments the big item.
|
||||
for child in children:
|
||||
big_count = sum(1 for item in child if item.get("document_id") == "doc-big")
|
||||
assert big_count <= 1
|
||||
|
||||
|
||||
def test_async_children_each_oversized_item_gets_own_child():
|
||||
"""Multiple oversized items → one child per item, each un-chunked."""
|
||||
tokens_per_batch = 500
|
||||
big_a = {"content": "Alpha sentence. " * 500, "document_id": "doc-a"}
|
||||
big_b = {"content": "Bravo sentence. " * 500, "document_id": "doc-b"}
|
||||
big_c = {"content": "Charlie sentence. " * 500, "document_id": "doc-c"}
|
||||
|
||||
children = _split_contents_into_async_children([big_a, big_b, big_c], tokens_per_batch)
|
||||
|
||||
assert len(children) == 3
|
||||
assert children == [[big_a], [big_b], [big_c]]
|
||||
|
||||
|
||||
def test_async_children_single_small_item_returns_one_child():
|
||||
"""A single small item is one child of one item — no splitting."""
|
||||
tokens_per_batch = 10_000
|
||||
item = {"content": "Alice works at Google", "document_id": "doc-1"}
|
||||
|
||||
children = _split_contents_into_async_children([item], tokens_per_batch)
|
||||
|
||||
assert children == [[item]]
|
||||
|
||||
|
||||
def test_async_children_empty_returns_empty():
|
||||
"""Empty input → empty output (no spurious child)."""
|
||||
assert _split_contents_into_async_children([], 1_000) == []
|
||||
|
||||
|
||||
def test_async_children_oversized_at_boundaries():
|
||||
"""Oversized items at the start, middle, and end all isolate correctly."""
|
||||
tokens_per_batch = 1_000
|
||||
small_a = {"content": "Small A. " * 5, "document_id": "doc-sa"}
|
||||
small_b = {"content": "Small B. " * 5, "document_id": "doc-sb"}
|
||||
big_start = {"content": "First big. " * 1_000, "document_id": "doc-bs"}
|
||||
big_end = {"content": "Last big. " * 1_000, "document_id": "doc-be"}
|
||||
|
||||
children = _split_contents_into_async_children([big_start, small_a, small_b, big_end], tokens_per_batch)
|
||||
|
||||
# First child holds big_start alone; last child holds big_end alone.
|
||||
assert children[0] == [big_start]
|
||||
assert children[-1] == [big_end]
|
||||
# Middle children contain only small items.
|
||||
middle_items = [item for child in children[1:-1] for item in child]
|
||||
assert all(item["document_id"] in {"doc-sa", "doc-sb"} for item in middle_items)
|
||||
# Every small input is present.
|
||||
flat_ids = [item["document_id"] for child in children for item in child]
|
||||
assert "doc-sa" in flat_ids
|
||||
assert "doc-sb" in flat_ids
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_large_batch_auto_chunks(memory, request_context):
|
||||
@@ -11,10 +277,7 @@ async def test_large_batch_auto_chunks(memory, request_context):
|
||||
# Create a large batch that should trigger chunking
|
||||
# Each item is ~2000 chars, so 30 items = 60k chars (exceeds 50k threshold)
|
||||
large_content = "Alice met with Bob at the coffee shop. " * 50 # ~2000 chars
|
||||
contents = [
|
||||
{"content": large_content, "context": f"conversation_{i}"}
|
||||
for i in range(30)
|
||||
]
|
||||
contents = [{"content": large_content, "context": f"conversation_{i}"} for i in range(30)]
|
||||
|
||||
# Calculate total chars
|
||||
total_chars = sum(len(item["content"]) for item in contents)
|
||||
@@ -40,7 +303,7 @@ async def test_small_batch_no_chunking(memory, request_context):
|
||||
# Create a small batch that should NOT trigger chunking
|
||||
contents = [
|
||||
{"content": "Alice works at Google", "context": "conversation_1"},
|
||||
{"content": "Bob loves Python", "context": "conversation_2"}
|
||||
{"content": "Bob loves Python", "context": "conversation_2"},
|
||||
]
|
||||
|
||||
# Calculate total chars
|
||||
|
||||
@@ -15,6 +15,8 @@ from hindsight_api import LLMConfig
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
|
||||
|
||||
pytestmark = pytest.mark.hs_llm_core
|
||||
|
||||
|
||||
class TestCausalRelationsValidation:
|
||||
"""Tests for causal relations index validation."""
|
||||
|
||||
@@ -13,6 +13,8 @@ from hindsight_api import LLMConfig
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
|
||||
|
||||
pytestmark = pytest.mark.hs_llm_core
|
||||
|
||||
|
||||
class TestCausalRelationships:
|
||||
"""Tests for causal relationship extraction and validation."""
|
||||
|
||||
@@ -0,0 +1,176 @@
|
||||
"""Regression test for the Claude Code provider's subprocess isolation.
|
||||
|
||||
When an operator runs the API with HINDSIGHT_API_*_LLM_PROVIDER=claude-code,
|
||||
the Claude Agent SDK spawns the `claude` CLI as a subprocess for every LLM
|
||||
call. If that subprocess inherits the host's CLAUDE_CONFIG_DIR, it loads any
|
||||
user-installed plugins (e.g. hindsight-memory). Their Stop hooks then retain
|
||||
the subprocess's own transcript back into the same bank, causing a recursive
|
||||
feedback loop documented in issue #1751.
|
||||
|
||||
The fix redirects the subprocess to an isolated config dir via env. This test
|
||||
mocks the SDK to capture the ClaudeAgentOptions passed at call time and asserts
|
||||
the isolation env is present on both code paths (`call` and `call_with_tools`).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeOptions:
|
||||
"""Stand-in for ClaudeAgentOptions; captures kwargs without importing SDK."""
|
||||
|
||||
system_prompt: str | None = None
|
||||
max_turns: int | None = None
|
||||
allowed_tools: list[str] = field(default_factory=list)
|
||||
tools: list[str] = field(default_factory=list)
|
||||
env: dict[str, str] = field(default_factory=dict)
|
||||
mcp_servers: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class _FakeAssistantMessage:
|
||||
def __init__(self, content: list[Any]) -> None:
|
||||
self.content = content
|
||||
|
||||
|
||||
class _FakeTextBlock:
|
||||
def __init__(self, text: str) -> None:
|
||||
self.text = text
|
||||
|
||||
|
||||
def _instantiate_provider():
|
||||
from hindsight_api.engine.providers.claude_code_llm import ClaudeCodeLLM
|
||||
|
||||
return ClaudeCodeLLM(
|
||||
provider="claude-code",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="claude-haiku-4-5",
|
||||
reasoning_effort="low",
|
||||
)
|
||||
|
||||
|
||||
def _assert_isolation_env(env: dict[str, str]) -> None:
|
||||
assert "CLAUDE_CONFIG_DIR" in env, "Subprocess env missing CLAUDE_CONFIG_DIR redirect"
|
||||
assert env["CLAUDE_CONFIG_DIR"], "CLAUDE_CONFIG_DIR must be a non-empty path"
|
||||
assert env["CLAUDE_CONFIG_DIR"].startswith("/"), "CLAUDE_CONFIG_DIR must be an absolute path"
|
||||
assert env.get("CLAUDE_SECURESTORAGE_CONFIG_DIR") == "", (
|
||||
"CLAUDE_SECURESTORAGE_CONFIG_DIR must be set to '' to force the un-suffixed keychain entry"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_passes_isolation_env_to_sdk_options(monkeypatch):
|
||||
"""call() must pass CLAUDE_CONFIG_DIR + SECURESTORAGE_CONFIG_DIR='' to the spawned CLI."""
|
||||
import claude_agent_sdk
|
||||
|
||||
captured: dict[str, _FakeOptions] = {}
|
||||
|
||||
async def fake_query(prompt: str, options: _FakeOptions):
|
||||
captured["options"] = options
|
||||
yield _FakeAssistantMessage(content=[_FakeTextBlock(text="ok")])
|
||||
|
||||
monkeypatch.setattr(claude_agent_sdk, "ClaudeAgentOptions", _FakeOptions)
|
||||
monkeypatch.setattr(claude_agent_sdk, "AssistantMessage", _FakeAssistantMessage)
|
||||
monkeypatch.setattr(claude_agent_sdk, "TextBlock", _FakeTextBlock)
|
||||
monkeypatch.setattr(claude_agent_sdk, "query", fake_query)
|
||||
|
||||
provider = _instantiate_provider()
|
||||
result = await provider.call(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
max_retries=0,
|
||||
scope="test",
|
||||
)
|
||||
|
||||
assert result == "ok"
|
||||
assert "options" in captured, "fake query was not called"
|
||||
_assert_isolation_env(captured["options"].env)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_tools_passes_isolation_env_to_sdk_options(monkeypatch):
|
||||
"""call_with_tools() must apply the same isolation env."""
|
||||
import claude_agent_sdk
|
||||
|
||||
captured: dict[str, _FakeOptions] = {}
|
||||
|
||||
class _FakeClient:
|
||||
def __init__(self, options: _FakeOptions) -> None:
|
||||
captured["options"] = options
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def query(self, prompt: str) -> None:
|
||||
return None
|
||||
|
||||
async def receive_response(self):
|
||||
yield _FakeAssistantMessage(content=[_FakeTextBlock(text="ok")])
|
||||
|
||||
@dataclass
|
||||
class _FakeSdkMcpTool:
|
||||
name: str
|
||||
description: str
|
||||
input_schema: dict[str, Any]
|
||||
handler: Any
|
||||
|
||||
def fake_create_sdk_mcp_server(name: str, version: str, tools=None):
|
||||
return {"name": name, "version": version, "tools": tools}
|
||||
|
||||
monkeypatch.setattr(claude_agent_sdk, "ClaudeAgentOptions", _FakeOptions)
|
||||
monkeypatch.setattr(claude_agent_sdk, "AssistantMessage", _FakeAssistantMessage)
|
||||
monkeypatch.setattr(claude_agent_sdk, "TextBlock", _FakeTextBlock)
|
||||
monkeypatch.setattr(claude_agent_sdk, "ToolUseBlock", type("ToolUseBlock", (), {}))
|
||||
monkeypatch.setattr(claude_agent_sdk, "ClaudeSDKClient", _FakeClient)
|
||||
monkeypatch.setattr(claude_agent_sdk, "SdkMcpTool", _FakeSdkMcpTool)
|
||||
monkeypatch.setattr(claude_agent_sdk, "create_sdk_mcp_server", fake_create_sdk_mcp_server)
|
||||
|
||||
provider = _instantiate_provider()
|
||||
result = await provider.call_with_tools(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
tools=[
|
||||
{
|
||||
"function": {
|
||||
"name": "noop",
|
||||
"description": "no-op",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
}
|
||||
}
|
||||
],
|
||||
max_retries=0,
|
||||
scope="test",
|
||||
)
|
||||
|
||||
assert result.content == "ok"
|
||||
assert "options" in captured, "fake client was not constructed"
|
||||
_assert_isolation_env(captured["options"].env)
|
||||
|
||||
|
||||
def test_isolation_dir_is_reused_across_calls():
|
||||
"""The isolated dir is created once per process; concurrent calls share it."""
|
||||
from hindsight_api.engine.providers.claude_code_llm import _get_isolated_claude_env
|
||||
|
||||
env_a = _get_isolated_claude_env()
|
||||
env_b = _get_isolated_claude_env()
|
||||
assert env_a is env_b, "Isolated env dict must be cached for the process lifetime"
|
||||
assert env_a["CLAUDE_CONFIG_DIR"] == env_b["CLAUDE_CONFIG_DIR"]
|
||||
|
||||
|
||||
def test_isolation_dir_is_not_home():
|
||||
"""Guard against the isolated dir accidentally pointing at $HOME (would not isolate plugins)."""
|
||||
import os
|
||||
|
||||
from hindsight_api.engine.providers.claude_code_llm import _get_isolated_claude_env
|
||||
|
||||
env = _get_isolated_claude_env()
|
||||
assert env["CLAUDE_CONFIG_DIR"] != os.path.expanduser("~"), (
|
||||
"Isolated dir must not be $HOME, or it would still load installed_plugins.json"
|
||||
)
|
||||
@@ -0,0 +1,30 @@
|
||||
"""Tests for Codex provider base URL handling."""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from hindsight_api.engine.providers.codex_llm import CodexLLM
|
||||
|
||||
|
||||
def _make(base_url: str) -> CodexLLM:
|
||||
with (
|
||||
patch.object(CodexLLM, "_load_codex_auth", return_value=("at", "acct")),
|
||||
patch.object(CodexLLM, "_load_codex_refresh_token", return_value="rt"),
|
||||
):
|
||||
return CodexLLM(
|
||||
provider="openai-codex",
|
||||
api_key="ignored",
|
||||
base_url=base_url,
|
||||
model="gpt-5.4-mini",
|
||||
)
|
||||
|
||||
|
||||
def test_codex_uses_chatgpt_backend_when_base_url_empty():
|
||||
assert _make("").base_url == "https://chatgpt.com/backend-api"
|
||||
|
||||
|
||||
def test_codex_ignores_inherited_openai_compatible_v1_base_url():
|
||||
assert _make("https://newapi.example.com/v1").base_url == "https://chatgpt.com/backend-api"
|
||||
|
||||
|
||||
def test_codex_preserves_explicit_codex_backend_base_url_without_trailing_slash():
|
||||
assert _make("https://chatgpt.example.com/backend-api/").base_url == "https://chatgpt.example.com/backend-api"
|
||||
@@ -29,6 +29,7 @@ import stat
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
@@ -254,7 +255,7 @@ async def test_refresh_sends_canonical_request_shape(tmp_path: Path):
|
||||
fresh_access = _make_jwt(int(time.time()) + 3600)
|
||||
refresh_resp = _refresh_response(200, {"access_token": fresh_access, "refresh_token": "rt-rotated"})
|
||||
|
||||
with patch.object(llm._client, "post", new_callable=AsyncMock, return_value=refresh_resp) as mock_post:
|
||||
with patch.object(llm._auth_manager._http_client, "post", return_value=refresh_resp) as mock_post:
|
||||
await llm._refresh_oauth_tokens()
|
||||
|
||||
call_args = mock_post.call_args
|
||||
@@ -277,7 +278,7 @@ async def test_refresh_updates_in_memory_credentials(tmp_path: Path):
|
||||
new_access = _make_jwt(int(time.time()) + 3600)
|
||||
refresh_resp = _refresh_response(200, {"access_token": new_access, "refresh_token": "rt-new"})
|
||||
|
||||
with patch.object(llm._client, "post", new_callable=AsyncMock, return_value=refresh_resp):
|
||||
with patch.object(llm._auth_manager._http_client, "post", return_value=refresh_resp):
|
||||
await llm._refresh_oauth_tokens()
|
||||
|
||||
assert llm.access_token == new_access
|
||||
@@ -295,7 +296,7 @@ async def test_refresh_keeps_existing_refresh_token_when_server_omits_one(tmp_pa
|
||||
new_access = _make_jwt(int(time.time()) + 3600)
|
||||
refresh_resp = _refresh_response(200, {"access_token": new_access})
|
||||
|
||||
with patch.object(llm._client, "post", new_callable=AsyncMock, return_value=refresh_resp):
|
||||
with patch.object(llm._auth_manager._http_client, "post", return_value=refresh_resp):
|
||||
await llm._refresh_oauth_tokens()
|
||||
|
||||
assert llm.refresh_token == "rt-keep"
|
||||
@@ -310,7 +311,7 @@ async def test_refresh_raises_permanent_error_on_terminal_oauth_code(tmp_path: P
|
||||
|
||||
bad_resp = _refresh_response(401, {"error": {"code": "refresh_token_expired"}})
|
||||
|
||||
with patch.object(llm._client, "post", new_callable=AsyncMock, return_value=bad_resp):
|
||||
with patch.object(llm._auth_manager._http_client, "post", return_value=bad_resp):
|
||||
with pytest.raises(CodexRefreshExpiredError):
|
||||
await llm._refresh_oauth_tokens()
|
||||
|
||||
@@ -325,7 +326,7 @@ async def test_refresh_raises_permanent_error_on_unknown_401(tmp_path: Path):
|
||||
|
||||
bad_resp = _refresh_response(401, {"error": "something_else"})
|
||||
|
||||
with patch.object(llm._client, "post", new_callable=AsyncMock, return_value=bad_resp):
|
||||
with patch.object(llm._auth_manager._http_client, "post", return_value=bad_resp):
|
||||
with pytest.raises(CodexRefreshExpiredError):
|
||||
await llm._refresh_oauth_tokens()
|
||||
|
||||
@@ -340,7 +341,7 @@ async def test_refresh_raises_runtime_error_on_5xx(tmp_path: Path):
|
||||
|
||||
bad_resp = _refresh_response(503, "service unavailable")
|
||||
|
||||
with patch.object(llm._client, "post", new_callable=AsyncMock, return_value=bad_resp):
|
||||
with patch.object(llm._auth_manager._http_client, "post", return_value=bad_resp):
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
await llm._refresh_oauth_tokens()
|
||||
assert not isinstance(exc_info.value, CodexRefreshExpiredError)
|
||||
@@ -357,7 +358,7 @@ async def test_refresh_does_not_log_token_values(tmp_path: Path, caplog):
|
||||
new_access = _make_jwt(int(time.time()) + 3600)
|
||||
refresh_resp = _refresh_response(200, {"access_token": new_access, "refresh_token": "rt-also-secret"})
|
||||
|
||||
with patch.object(llm._client, "post", new_callable=AsyncMock, return_value=refresh_resp):
|
||||
with patch.object(llm._auth_manager._http_client, "post", return_value=refresh_resp):
|
||||
with caplog.at_level("DEBUG"):
|
||||
await llm._refresh_oauth_tokens()
|
||||
|
||||
@@ -382,14 +383,14 @@ async def test_concurrent_ensure_fresh_token_calls_produce_one_refresh(tmp_path:
|
||||
new_access = _make_jwt(int(time.time()) + 3600)
|
||||
call_count = 0
|
||||
|
||||
async def fake_post(*args, **kwargs):
|
||||
def fake_post(*args, **kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
# Simulate non-zero refresh latency so concurrent callers actually queue.
|
||||
await asyncio.sleep(0.01)
|
||||
time.sleep(0.01)
|
||||
return _refresh_response(200, {"access_token": new_access, "refresh_token": "rt-new"})
|
||||
|
||||
with patch.object(llm._client, "post", new=fake_post):
|
||||
with patch.object(llm._auth_manager._http_client, "post", new=fake_post):
|
||||
await asyncio.gather(*(llm._ensure_fresh_token() for _ in range(10)))
|
||||
|
||||
assert call_count == 1, f"expected 1 network refresh under contention, got {call_count}"
|
||||
@@ -410,42 +411,33 @@ async def test_call_reactively_refreshes_on_401_and_retries(tmp_path: Path):
|
||||
|
||||
new_access = _make_jwt(int(time.time()) + 3600)
|
||||
|
||||
# First post → 401 (backend rejects the token). After refresh, second post → 200.
|
||||
success_resp = MagicMock()
|
||||
success_resp.status_code = 200
|
||||
success_resp.raise_for_status.return_value = None
|
||||
success_resp.raise_for_status = MagicMock(return_value=None)
|
||||
|
||||
fail_response = MagicMock()
|
||||
fail_response.status_code = 401
|
||||
fail_response.text = "unauthorized"
|
||||
fail_exc = httpx.HTTPStatusError("401", request=MagicMock(), response=fail_response)
|
||||
success_resp.raise_for_status = MagicMock(return_value=None)
|
||||
|
||||
post_responses = [fail_exc, success_resp]
|
||||
|
||||
async def fake_post(*args, **kwargs):
|
||||
item = post_responses.pop(0)
|
||||
if isinstance(item, Exception):
|
||||
raise item
|
||||
return item
|
||||
|
||||
refresh_resp = _refresh_response(200, {"access_token": new_access, "refresh_token": "rt-new"})
|
||||
|
||||
call_count = {"refresh": 0, "post": 0}
|
||||
|
||||
async def counting_post(url, **kwargs):
|
||||
if url == _CODEX_REFRESH_TOKEN_URL:
|
||||
call_count["refresh"] += 1
|
||||
return refresh_resp
|
||||
# Sync mock for the auth manager's HTTP client (used for token refresh).
|
||||
def fake_refresh_post(*args, **kwargs):
|
||||
call_count["refresh"] += 1
|
||||
return refresh_resp
|
||||
|
||||
# Async mock for the LLM's HTTP client (used for backend calls).
|
||||
async def fake_backend_post(url, **kwargs):
|
||||
call_count["post"] += 1
|
||||
# First backend call fails with 401 wrapped in an HTTPStatusError-style response,
|
||||
# second succeeds.
|
||||
if call_count["post"] == 1:
|
||||
raise httpx.HTTPStatusError("401", request=MagicMock(), response=fail_response)
|
||||
return success_resp
|
||||
|
||||
with (
|
||||
patch.object(llm._client, "post", new=counting_post),
|
||||
patch.object(llm._auth_manager._http_client, "post", new=fake_refresh_post),
|
||||
patch.object(llm._client, "post", new=fake_backend_post),
|
||||
patch.object(llm, "_parse_sse_stream", new_callable=AsyncMock, return_value="ok"),
|
||||
):
|
||||
result = await llm.call(
|
||||
@@ -478,17 +470,19 @@ async def test_call_proactively_refreshes_when_token_is_stale(tmp_path: Path):
|
||||
|
||||
call_order: list[str] = []
|
||||
|
||||
async def fake_post(url, **kwargs):
|
||||
if url == _CODEX_REFRESH_TOKEN_URL:
|
||||
call_order.append("refresh")
|
||||
return refresh_resp
|
||||
def fake_refresh_post(*args, **kwargs):
|
||||
call_order.append("refresh")
|
||||
return refresh_resp
|
||||
|
||||
async def fake_backend_post(url, **kwargs):
|
||||
call_order.append("backend")
|
||||
# Assert that by the time the backend is called, the new token is in use.
|
||||
assert kwargs["headers"]["Authorization"] == f"Bearer {new_access}"
|
||||
return success_resp
|
||||
|
||||
with (
|
||||
patch.object(llm._client, "post", new=fake_post),
|
||||
patch.object(llm._auth_manager._http_client, "post", new=fake_refresh_post),
|
||||
patch.object(llm._client, "post", new=fake_backend_post),
|
||||
patch.object(llm, "_parse_sse_stream", new_callable=AsyncMock, return_value="ok"),
|
||||
):
|
||||
await llm.call(
|
||||
@@ -512,17 +506,14 @@ async def test_call_does_not_refresh_when_token_is_fresh(tmp_path: Path):
|
||||
success_resp.status_code = 200
|
||||
success_resp.raise_for_status.return_value = None
|
||||
|
||||
call_count = {"refresh": 0, "backend": 0}
|
||||
call_count = {"backend": 0}
|
||||
|
||||
async def fake_post(url, **kwargs):
|
||||
if url == _CODEX_REFRESH_TOKEN_URL:
|
||||
call_count["refresh"] += 1
|
||||
raise AssertionError("refresh endpoint should not be hit for a fresh token")
|
||||
async def fake_backend_post(url, **kwargs):
|
||||
call_count["backend"] += 1
|
||||
return success_resp
|
||||
|
||||
with (
|
||||
patch.object(llm._client, "post", new=fake_post),
|
||||
patch.object(llm._client, "post", new=fake_backend_post),
|
||||
patch.object(llm, "_parse_sse_stream", new_callable=AsyncMock, return_value="ok"),
|
||||
):
|
||||
await llm.call(
|
||||
@@ -532,4 +523,100 @@ async def test_call_does_not_refresh_when_token_is_fresh(tmp_path: Path):
|
||||
max_backoff=0.0,
|
||||
)
|
||||
|
||||
assert call_count == {"refresh": 0, "backend": 1}
|
||||
assert call_count == {"backend": 1}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CodexOAuthEmbeddings — proactive + reactive token refresh
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_codex_auth_file(tmp_path: Path, access_token: str, refresh_token: str = "rt-initial") -> Path:
|
||||
"""Write a minimal ~/.codex/auth.json in tmp_path and return its path."""
|
||||
codex_dir = tmp_path / ".codex"
|
||||
codex_dir.mkdir(parents=True, exist_ok=True)
|
||||
auth_file = codex_dir / "auth.json"
|
||||
auth_file.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"auth_mode": "chatgpt",
|
||||
"tokens": {
|
||||
"access_token": access_token,
|
||||
"refresh_token": refresh_token,
|
||||
"account_id": "acct-test",
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
return auth_file
|
||||
|
||||
|
||||
def test_codex_oauth_embeddings_picks_up_refreshed_token_on_encode(tmp_path: Path, monkeypatch):
|
||||
"""encode() calls ensure_fresh_token() and updates api_key when the token rotated."""
|
||||
from hindsight_api.engine.embeddings import CodexOAuthEmbeddings
|
||||
|
||||
expired = _make_jwt(int(time.time()) - 60)
|
||||
new_access = _make_jwt(int(time.time()) + 3600)
|
||||
|
||||
_make_codex_auth_file(tmp_path, expired, refresh_token="rt-embed")
|
||||
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
||||
|
||||
emb = CodexOAuthEmbeddings(model="text-embedding-3-small", batch_size=10)
|
||||
|
||||
refresh_resp = _refresh_response(200, {"access_token": new_access, "refresh_token": "rt-new"})
|
||||
|
||||
fake_embeddings = [SimpleNamespace(index=0, embedding=[0.1] * 1536)]
|
||||
fake_create_resp = SimpleNamespace(data=fake_embeddings)
|
||||
|
||||
with patch.object(emb._auth_manager._http_client, "post", return_value=refresh_resp):
|
||||
emb._client = SimpleNamespace(embeddings=SimpleNamespace(create=lambda **kw: fake_create_resp))
|
||||
emb._dimension = 1536
|
||||
result = emb.encode(["hello"])
|
||||
|
||||
assert result == [[0.1] * 1536]
|
||||
# After proactive refresh the manager's token should be the new one.
|
||||
assert emb._auth_manager.access_token == new_access
|
||||
# api_key on the embeddings object should also be updated.
|
||||
assert emb.api_key == new_access
|
||||
|
||||
|
||||
def test_codex_oauth_embeddings_reactive_refresh_on_401(tmp_path: Path, monkeypatch):
|
||||
"""On AuthenticationError from OpenAI, encode() refreshes and retries once."""
|
||||
from openai import AuthenticationError as OAIAuthError
|
||||
|
||||
from hindsight_api.engine.embeddings import CodexOAuthEmbeddings
|
||||
|
||||
fresh = _make_jwt(int(time.time()) + 3600)
|
||||
new_access = _make_jwt(int(time.time()) + 7200)
|
||||
|
||||
_make_codex_auth_file(tmp_path, fresh, refresh_token="rt-embed")
|
||||
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
||||
|
||||
emb = CodexOAuthEmbeddings(model="text-embedding-3-small", batch_size=10)
|
||||
emb._dimension = 1536
|
||||
|
||||
refresh_resp = _refresh_response(200, {"access_token": new_access, "refresh_token": "rt-rotated"})
|
||||
|
||||
call_count = {"create": 0}
|
||||
|
||||
def fake_create(**kwargs):
|
||||
call_count["create"] += 1
|
||||
if call_count["create"] == 1:
|
||||
# Simulate OpenAI returning 401.
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 401
|
||||
raise OAIAuthError(
|
||||
message="invalid api key",
|
||||
response=mock_response,
|
||||
body={"error": {"message": "invalid api key"}},
|
||||
)
|
||||
return SimpleNamespace(data=[SimpleNamespace(index=0, embedding=[0.2] * 1536)])
|
||||
|
||||
with patch.object(emb._auth_manager._http_client, "post", return_value=refresh_resp):
|
||||
emb._client = SimpleNamespace(embeddings=SimpleNamespace(create=fake_create))
|
||||
result = emb.encode(["world"])
|
||||
|
||||
assert result == [[0.2] * 1536]
|
||||
assert call_count["create"] == 2 # first failed with 401, second succeeded
|
||||
assert emb._auth_manager.access_token == new_access
|
||||
assert emb.api_key == new_access
|
||||
|
||||
@@ -21,6 +21,7 @@ def setup_test_env():
|
||||
"HINDSIGHT_API_RETAIN_CHUNK_SIZE",
|
||||
"HINDSIGHT_API_LLM_PROVIDER",
|
||||
"HINDSIGHT_API_LLM_MODEL",
|
||||
"HINDSIGHT_API_LLM_REASONING_EFFORT",
|
||||
"HINDSIGHT_API_DATABASE_URL",
|
||||
"HINDSIGHT_API_MIGRATION_DATABASE_URL",
|
||||
]
|
||||
@@ -204,3 +205,174 @@ def test_log_config_masks_read_database_url(monkeypatch, caplog):
|
||||
#
|
||||
# The config validation tests above ensure users get early feedback
|
||||
# about invalid configurations before runtime errors occur.
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Multilingual BM25 configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_native_language_defaults_to_english(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.delenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE", raising=False)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension_native_language == "english"
|
||||
|
||||
|
||||
def test_native_language_loaded_from_env(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE", "french")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension_native_language == "french"
|
||||
|
||||
|
||||
def test_native_language_lowercased(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE", "Spanish")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension_native_language == "spanish"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bad_value",
|
||||
["en glish", "english;DROP TABLE", "english'", "1english", "english-extra", ""],
|
||||
)
|
||||
def test_native_language_rejects_invalid_identifiers(monkeypatch, bad_value):
|
||||
"""text_search_extension_native_language is embedded into raw SQL — non-identifiers must be rejected."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE", bad_value)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid text_search_extension_native_language"):
|
||||
HindsightConfig.from_env()
|
||||
|
||||
|
||||
def test_text_search_extension_accepts_pgroonga(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "pgroonga")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension == "pgroonga"
|
||||
|
||||
|
||||
def test_text_search_extension_rejects_unknown(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "bogus")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid text_search_extension"):
|
||||
HindsightConfig.from_env()
|
||||
|
||||
|
||||
def test_pg_search_tokenizer_defaults_to_empty(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.delenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER", raising=False)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension_pg_search_tokenizer == ""
|
||||
|
||||
|
||||
def test_pg_search_tokenizer_loaded_from_env(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER", "Jieba")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension_pg_search_tokenizer == "jieba"
|
||||
|
||||
|
||||
def test_pg_search_tokenizer_accepts_lindera_alias(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER", "chinese_lindera")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension_pg_search_tokenizer == "lindera(chinese)"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_value", ["jieba;DROP TABLE", "ngram(3,2)", "unknown", "pdb.jieba"])
|
||||
def test_pg_search_tokenizer_rejects_invalid_values(monkeypatch, bad_value):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER", bad_value)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"):
|
||||
HindsightConfig.from_env()
|
||||
|
||||
|
||||
def test_pg_search_bm25_columns_apply_tokenizer():
|
||||
from hindsight_api._pg_search import pg_search_bm25_columns
|
||||
|
||||
assert pg_search_bm25_columns("id", ("text", "context"), "") == "id, text, context"
|
||||
assert pg_search_bm25_columns("id", ("text", "context"), "jieba") == "id, (text::pdb.jieba), (context::pdb.jieba)"
|
||||
assert pg_search_bm25_columns("id", ("text",), "ngram(2, 3)") == "id, (text::pdb.ngram(2,3))"
|
||||
assert pg_search_bm25_columns("id", ("text",), "edge_ngram(2, 5)") == "id, (text::pdb.edge_ngram(2,5))"
|
||||
|
||||
|
||||
def test_llm_output_language_defaults_to_none(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.delenv("HINDSIGHT_API_LLM_OUTPUT_LANGUAGE", raising=False)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_output_language is None
|
||||
|
||||
|
||||
def test_llm_output_language_loaded_from_env(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_OUTPUT_LANGUAGE", "Japanese")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_output_language == "Japanese"
|
||||
|
||||
|
||||
def test_llm_output_language_empty_string_is_unset(monkeypatch):
|
||||
"""Empty env var (e.g. from Helm) should be treated as unset, not literal ''."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_OUTPUT_LANGUAGE", "")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_output_language is None
|
||||
|
||||
|
||||
def test_llm_reasoning_effort_defaults_to_low(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.delenv("HINDSIGHT_API_LLM_REASONING_EFFORT", raising=False)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_reasoning_effort == "low"
|
||||
|
||||
|
||||
def test_llm_reasoning_effort_loaded_from_env(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_REASONING_EFFORT", "xhigh")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_reasoning_effort == "xhigh"
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,267 @@
|
||||
"""
|
||||
Regression tests for the consolidation retry path.
|
||||
|
||||
When a consolidation task hits a transient error, the worker's generic retry
|
||||
path raises ``RetryTaskAt`` to re-queue the operation (memory_engine.py, see
|
||||
``execute_task``'s exception branch). Two pieces of behaviour live there:
|
||||
|
||||
1. **Dedup guard.** During a long upstream outage, every retain on the same
|
||||
bank enqueues a fresh consolidation op via ``submit_async_consolidation``.
|
||||
Without dedup, each op independently consumes its own retry budget,
|
||||
multiplying load on the broken dependency. The guard: if another
|
||||
consolidation op is already ``pending`` for the same bank when the retry
|
||||
decision is made, skip the retry and let the other op cover the work.
|
||||
|
||||
2. **Indefinite retry with capped exponential backoff** (60s, 120, 240, 480,
|
||||
960, 1800-cap). Distinct from the generic 60s × 3 used by other task
|
||||
types. Transient consolidation failures (LLM/DB outage) must eventually
|
||||
recover; capping after a fixed number of attempts silently dead-letters
|
||||
the bank's backlog. The dedup-by-bank guard above prevents a retry
|
||||
storm when multiple ops exist for the same bank.
|
||||
"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.memory_engine import (
|
||||
_CONSOLIDATION_RETRY_BACKOFF_MAX_SECONDS,
|
||||
_consolidation_retry_backoff_seconds,
|
||||
)
|
||||
from hindsight_api.worker.exceptions import RetryTaskAt
|
||||
|
||||
|
||||
async def _ensure_bank(pool, bank_id: str) -> None:
|
||||
await pool.execute(
|
||||
"INSERT INTO banks (bank_id, name) VALUES ($1, $2) ON CONFLICT DO NOTHING",
|
||||
bank_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
async def _insert_consolidation_op(pool, bank_id: str, operation_id: uuid.UUID, status: str) -> None:
|
||||
"""Insert a consolidation row with the given ``status``."""
|
||||
payload = json.dumps(
|
||||
{
|
||||
"type": "consolidation",
|
||||
"operation_id": str(operation_id),
|
||||
"bank_id": bank_id,
|
||||
}
|
||||
)
|
||||
await pool.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
|
||||
VALUES ($1, $2, 'consolidation', $3, $4::jsonb)
|
||||
""",
|
||||
operation_id,
|
||||
bank_id,
|
||||
status,
|
||||
payload,
|
||||
)
|
||||
|
||||
|
||||
async def _cleanup(pool, bank_id: str, *operation_ids: uuid.UUID) -> None:
|
||||
if operation_ids:
|
||||
await pool.execute(
|
||||
"DELETE FROM async_operations WHERE operation_id = ANY($1::uuid[])",
|
||||
list(operation_ids),
|
||||
)
|
||||
await pool.execute("DELETE FROM banks WHERE bank_id = $1", bank_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_skipped_when_other_pending_consolidation_exists(memory):
|
||||
"""
|
||||
With another pending consolidation op already covering the same bank, a
|
||||
transient failure on the currently-executing op must NOT raise
|
||||
``RetryTaskAt`` — the pending op will process the same unconsolidated
|
||||
rows when the worker picks it up. Without this guard, a long outage
|
||||
causes a retry storm: every retain enqueues a new op, and each op burns
|
||||
three retry slots against the same broken dependency.
|
||||
"""
|
||||
bank_id = f"test-dedup-{uuid.uuid4().hex[:8]}"
|
||||
running_op_id = uuid.uuid4()
|
||||
peer_pending_op_id = uuid.uuid4()
|
||||
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
# The op currently executing (drives execute_task below).
|
||||
await _insert_consolidation_op(pool, bank_id, running_op_id, status="processing")
|
||||
# A peer consolidation already pending for the same bank — the dedup guard
|
||||
# should fire because of this row.
|
||||
await _insert_consolidation_op(pool, bank_id, peer_pending_op_id, status="pending")
|
||||
|
||||
task_dict = {
|
||||
"type": "consolidation",
|
||||
"operation_id": str(running_op_id),
|
||||
"bank_id": bank_id,
|
||||
}
|
||||
|
||||
transient = RuntimeError("upstream LLM 503")
|
||||
with patch.object(memory, "_handle_consolidation", side_effect=transient):
|
||||
# The original exception must propagate as-is (the poller catches it
|
||||
# and calls _mark_failed). The point of the test is that it is NOT
|
||||
# converted to RetryTaskAt — pytest.raises(RuntimeError) asserts that
|
||||
# implicitly, since RetryTaskAt is not a RuntimeError subclass.
|
||||
with pytest.raises(RuntimeError, match="upstream LLM 503"):
|
||||
await memory.execute_task(task_dict)
|
||||
|
||||
# Peer pending op must be untouched by the dedup guard.
|
||||
peer_row = await pool.fetchrow(
|
||||
"SELECT status FROM async_operations WHERE operation_id = $1",
|
||||
peer_pending_op_id,
|
||||
)
|
||||
assert peer_row["status"] == "pending"
|
||||
|
||||
await _cleanup(pool, bank_id, running_op_id, peer_pending_op_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_happens_when_no_other_pending_consolidation(memory):
|
||||
"""
|
||||
Sanity check: with no peer pending op, the normal retry path still fires
|
||||
and ``RetryTaskAt`` is raised — the dedup guard must not regress the
|
||||
common-case behaviour.
|
||||
"""
|
||||
bank_id = f"test-dedup-{uuid.uuid4().hex[:8]}"
|
||||
running_op_id = uuid.uuid4()
|
||||
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
await _insert_consolidation_op(pool, bank_id, running_op_id, status="processing")
|
||||
|
||||
task_dict = {
|
||||
"type": "consolidation",
|
||||
"operation_id": str(running_op_id),
|
||||
"bank_id": bank_id,
|
||||
}
|
||||
|
||||
transient = RuntimeError("transient blip")
|
||||
with patch.object(memory, "_handle_consolidation", side_effect=transient):
|
||||
with pytest.raises(RetryTaskAt):
|
||||
await memory.execute_task(task_dict)
|
||||
|
||||
await _cleanup(pool, bank_id, running_op_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_peer_only_dedups_when_same_bank(memory):
|
||||
"""
|
||||
A pending consolidation for a DIFFERENT bank must not suppress the retry
|
||||
— the dedup guard is per-bank, not global.
|
||||
"""
|
||||
bank_id_a = f"test-dedup-a-{uuid.uuid4().hex[:8]}"
|
||||
bank_id_b = f"test-dedup-b-{uuid.uuid4().hex[:8]}"
|
||||
running_op_id = uuid.uuid4()
|
||||
unrelated_pending_op_id = uuid.uuid4()
|
||||
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id_a)
|
||||
await _ensure_bank(pool, bank_id_b)
|
||||
await _insert_consolidation_op(pool, bank_id_a, running_op_id, status="processing")
|
||||
await _insert_consolidation_op(pool, bank_id_b, unrelated_pending_op_id, status="pending")
|
||||
|
||||
task_dict = {
|
||||
"type": "consolidation",
|
||||
"operation_id": str(running_op_id),
|
||||
"bank_id": bank_id_a,
|
||||
}
|
||||
|
||||
with patch.object(memory, "_handle_consolidation", side_effect=RuntimeError("blip")):
|
||||
with pytest.raises(RetryTaskAt):
|
||||
await memory.execute_task(task_dict)
|
||||
|
||||
await _cleanup(pool, bank_id_a, running_op_id)
|
||||
await _cleanup(pool, bank_id_b, unrelated_pending_op_id)
|
||||
|
||||
|
||||
def test_backoff_helper_schedule():
|
||||
"""
|
||||
The backoff helper produces capped exponential growth:
|
||||
60, 120, 240, 480, 960, then pinned at the 1800s cap.
|
||||
"""
|
||||
assert _consolidation_retry_backoff_seconds(0) == 60
|
||||
assert _consolidation_retry_backoff_seconds(1) == 120
|
||||
assert _consolidation_retry_backoff_seconds(2) == 240
|
||||
assert _consolidation_retry_backoff_seconds(3) == 480
|
||||
assert _consolidation_retry_backoff_seconds(4) == 960
|
||||
assert _consolidation_retry_backoff_seconds(5) == _CONSOLIDATION_RETRY_BACKOFF_MAX_SECONDS
|
||||
# Cap holds at arbitrarily high attempt counts (we never give up).
|
||||
assert _consolidation_retry_backoff_seconds(50) == _CONSOLIDATION_RETRY_BACKOFF_MAX_SECONDS
|
||||
assert _consolidation_retry_backoff_seconds(1000) == _CONSOLIDATION_RETRY_BACKOFF_MAX_SECONDS
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("retry_count", [0, 1, 2, 5])
|
||||
async def test_backoff_matches_schedule_by_retry_count(memory, retry_count):
|
||||
"""
|
||||
Each retry schedules `retry_at = now + backoff(retry_count)`. The tolerance
|
||||
covers the few seconds between ``now()`` capture in the test and inside
|
||||
the retry handler. Includes retry_count=5 to verify the cap branch.
|
||||
"""
|
||||
bank_id = f"test-backoff-{uuid.uuid4().hex[:8]}"
|
||||
op_id = uuid.uuid4()
|
||||
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
await _insert_consolidation_op(pool, bank_id, op_id, status="processing")
|
||||
|
||||
task_dict = {
|
||||
"type": "consolidation",
|
||||
"operation_id": str(op_id),
|
||||
"bank_id": bank_id,
|
||||
"_retry_count": retry_count,
|
||||
}
|
||||
|
||||
expected_backoff = _consolidation_retry_backoff_seconds(retry_count)
|
||||
before = datetime.now(UTC)
|
||||
with patch.object(memory, "_handle_consolidation", side_effect=RuntimeError("transient")):
|
||||
with pytest.raises(RetryTaskAt) as excinfo:
|
||||
await memory.execute_task(task_dict)
|
||||
|
||||
delta = (excinfo.value.retry_at - before).total_seconds()
|
||||
assert expected_backoff <= delta <= expected_backoff + 10, (
|
||||
f"retry_count={retry_count}: expected backoff ~{expected_backoff}s, "
|
||||
f"got delta={delta:.2f}s"
|
||||
)
|
||||
|
||||
await _cleanup(pool, bank_id, op_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_is_indefinite(memory):
|
||||
"""
|
||||
Consolidation has no attempt cap — even at very high retry_count, a
|
||||
transient failure still raises RetryTaskAt (capped at the max backoff).
|
||||
Without this property, a long outage silently dead-letters the bank's
|
||||
unconsolidated rows once the budget is exhausted.
|
||||
"""
|
||||
bank_id = f"test-indefinite-{uuid.uuid4().hex[:8]}"
|
||||
op_id = uuid.uuid4()
|
||||
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
await _insert_consolidation_op(pool, bank_id, op_id, status="processing")
|
||||
|
||||
task_dict = {
|
||||
"type": "consolidation",
|
||||
"operation_id": str(op_id),
|
||||
"bank_id": bank_id,
|
||||
"_retry_count": 100, # well past any plausible attempt cap
|
||||
}
|
||||
|
||||
before = datetime.now(UTC)
|
||||
with patch.object(memory, "_handle_consolidation", side_effect=RuntimeError("still failing")):
|
||||
with pytest.raises(RetryTaskAt) as excinfo:
|
||||
await memory.execute_task(task_dict)
|
||||
|
||||
delta = (excinfo.value.retry_at - before).total_seconds()
|
||||
cap = _CONSOLIDATION_RETRY_BACKOFF_MAX_SECONDS
|
||||
assert cap <= delta <= cap + 10, (
|
||||
f"At retry_count=100 expected backoff at cap (~{cap}s), got {delta:.2f}s"
|
||||
)
|
||||
|
||||
await _cleanup(pool, bank_id, op_id)
|
||||
@@ -77,7 +77,9 @@ async def test_round_limit_caps_processed_memories(memory: MemoryEngine, request
|
||||
assert result["memories_processed"] <= round_limit
|
||||
|
||||
# Must have re-queued consolidation for remaining work
|
||||
mock_requeue.assert_called_once_with(bank_id=bank_id, request_context=request_context)
|
||||
mock_requeue.assert_called_once_with(
|
||||
bank_id=bank_id, request_context=request_context, observation_scopes=None
|
||||
)
|
||||
|
||||
# Mental model refresh should be skipped on intermediate round
|
||||
assert result.get("mental_models_refreshed", 0) == 0
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user