Compare commits
183
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0022d427d3 | ||
|
|
1f9bad0858 | ||
|
|
138bf02f29 | ||
|
|
df178aae8a | ||
|
|
a43026b8f4 | ||
|
|
5c425e276e | ||
|
|
265192e509 | ||
|
|
8f2cee4568 | ||
|
|
92f433c904 | ||
|
|
f8ce15b9bf | ||
|
|
d68f618969 | ||
|
|
33e9db64a1 | ||
|
|
0c8699dc20 | ||
|
|
251c451fc3 | ||
|
|
8ed49e4387 | ||
|
|
82afa76182 | ||
|
|
4c307ce4e7 | ||
|
|
252b013243 | ||
|
|
a6b0f82124 | ||
|
|
a27754fb15 | ||
|
|
40fe7aac86 | ||
|
|
7393400f34 | ||
|
|
3b7d18d474 | ||
|
|
6e18858e32 | ||
|
|
84e67efbf4 | ||
|
|
ef2e8ab7ff | ||
|
|
cc45e16904 | ||
|
|
82b01ace5e | ||
|
|
962140eef6 | ||
|
|
5e73d5ff62 | ||
|
|
2c47b8b0d5 | ||
|
|
00968a1ce4 | ||
|
|
b7080a16cf | ||
|
|
c0aed313f4 | ||
|
|
2c53629420 | ||
|
|
760bfc7447 | ||
|
|
cce0a2cb39 | ||
|
|
1c1cf4ce56 | ||
|
|
a99a1ebf9b | ||
|
|
12a6739fc9 | ||
|
|
d8ee10a78d | ||
|
|
ab01144b26 | ||
|
|
072b3278ba | ||
|
|
74e82a3ea8 | ||
|
|
a5b752a983 | ||
|
|
7b878f89a7 | ||
|
|
1b92c8230f | ||
|
|
0178d91333 | ||
|
|
017b8d7271 | ||
|
|
21176f8ee8 | ||
|
|
74bdfc9475 | ||
|
|
b0038e9855 | ||
|
|
6eb85570af | ||
|
|
85599f3ef5 | ||
|
|
dd83bffeef | ||
|
|
911d27fc5f | ||
|
|
a0af096081 | ||
|
|
fcb2c958e7 | ||
|
|
758f346d30 | ||
|
|
78d32cd16c | ||
|
|
fb475cc5bc | ||
|
|
1621e5d261 | ||
|
|
91e095afa9 | ||
|
|
815d99f5ba | ||
|
|
2452f72e75 | ||
|
|
dae18b1faf | ||
|
|
6a10b6241d | ||
|
|
47992d843b | ||
|
|
93100ed314 | ||
|
|
01eda51880 | ||
|
|
e63d028a5a | ||
|
|
58b5677617 | ||
|
|
b6608076ff | ||
|
|
4fe477eaa3 | ||
|
|
a7d1f26f98 | ||
|
|
0673131a80 | ||
|
|
6e02a0829f | ||
|
|
bc813692c6 | ||
|
|
701de3293d | ||
|
|
9dafadc7eb | ||
|
|
d0b77f5bee | ||
|
|
7194f98b19 | ||
|
|
34ba3c676e | ||
|
|
0379b4c823 | ||
|
|
422e0fd809 | ||
|
|
91bf32842e | ||
|
|
e4afa5a61b | ||
|
|
680305aea4 | ||
|
|
9e06237e40 | ||
|
|
8e66c397a9 | ||
|
|
f7c7a62e5f | ||
|
|
c056edaa90 | ||
|
|
63a92bef5f | ||
|
|
1533c0915d | ||
|
|
8eb2937cdb | ||
|
|
d9a372a92e | ||
|
|
199ae146ab | ||
|
|
b4874672fa | ||
|
|
f8d277697d | ||
|
|
4db8a12362 | ||
|
|
cabcb3bb0b | ||
|
|
5f0b715517 | ||
|
|
0ba613c3ce | ||
|
|
04703d2153 | ||
|
|
20da6d7609 | ||
|
|
387c09e91e | ||
|
|
cbce937042 | ||
|
|
246803bcfe | ||
|
|
625c331e80 | ||
|
|
f21944d789 | ||
|
|
0672fba279 | ||
|
|
53a52afe8b | ||
|
|
a2166ee4ff | ||
|
|
26bfd2ece4 | ||
|
|
0d60f0c638 | ||
|
|
735172f806 | ||
|
|
2c2a20b290 | ||
|
|
1a09a9cccd | ||
|
|
f183b09b93 | ||
|
|
5543992d7d | ||
|
|
dcabd76911 | ||
|
|
1a51184a32 | ||
|
|
c7e5095a86 | ||
|
|
f187d32351 | ||
|
|
ee81c65e4b | ||
|
|
ccd3eb24c9 | ||
|
|
51cb32896f | ||
|
|
af42382983 | ||
|
|
955b0c523c | ||
|
|
80281e2543 | ||
|
|
bb4dd4f393 | ||
|
|
ba5ddd59af | ||
|
|
aab7032071 | ||
|
|
adb6dcd683 | ||
|
|
ae93c93182 | ||
|
|
1c0c53c062 | ||
|
|
731add1fdf | ||
|
|
a7f82453a3 | ||
|
|
2bd6be8d85 | ||
|
|
e1014cc790 | ||
|
|
da2125cf13 | ||
|
|
f4bac2d41d | ||
|
|
39abf0ad3f | ||
|
|
8426b0c359 | ||
|
|
f4a0a31f70 | ||
|
|
65862c4fef | ||
|
|
ef548833fd | ||
|
|
55f70e1d27 | ||
|
|
b8cfddd7b6 | ||
|
|
52cb9a2bae | ||
|
|
539101af38 | ||
|
|
efa37cb15f | ||
|
|
faaa97d4a0 | ||
|
|
70804fe2c6 | ||
|
|
ae2532b165 | ||
|
|
81865bf873 | ||
|
|
9e47759347 | ||
|
|
d8665d7ab0 | ||
|
|
9bde15331e | ||
|
|
2fb2de1aa8 | ||
|
|
d68bd07423 | ||
|
|
44972d3215 | ||
|
|
aa308ad201 | ||
|
|
cb73790c27 | ||
|
|
b1fe23fbe4 | ||
|
|
ce81217381 | ||
|
|
94619ce52b | ||
|
|
27aa6bbf46 | ||
|
|
acf4d5c860 | ||
|
|
551932991d | ||
|
|
5ee53c512f | ||
|
|
4efa204727 | ||
|
|
c3bb647640 | ||
|
|
12851bc7ee | ||
|
|
a5f4d30ea6 | ||
|
|
0c9bc765ce | ||
|
|
cd34efa596 | ||
|
|
2b521c3a09 | ||
|
|
9681d96195 | ||
|
|
a32ecfeb33 | ||
|
|
bf73a1dfbe | ||
|
|
ca2ce5c16d | ||
|
|
7b17da7a0c |
@@ -1,6 +1,7 @@
|
||||
{
|
||||
"$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
|
||||
"name": "hindsight",
|
||||
"version": "0.7.2",
|
||||
"description": "Official Hindsight integrations for Claude Code",
|
||||
"owner": {
|
||||
"name": "vectorize-io"
|
||||
|
||||
+46
-1
@@ -2,7 +2,7 @@
|
||||
# Copy this file to .env and fill in your values
|
||||
|
||||
# LLM Configuration (Required)
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, zai, volcano
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, zai, atlas, volcano
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
@@ -10,6 +10,17 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# Reasoning effort for providers/models that support it. Examples: low, medium, high, xhigh.
|
||||
# HINDSIGHT_API_LLM_REASONING_EFFORT=low
|
||||
|
||||
# Sampling temperature for internal LLM calls. Set a number in [0.0, 2.0], or `none`
|
||||
# to omit the temperature parameter entirely (required for models that reject explicit
|
||||
# temperatures, e.g. Azure gpt-5.5). The global override below applies to every operation;
|
||||
# per-operation overrides (defaults: verification=0.0, retain=0.1, reflect=0.9,
|
||||
# consolidation=0.0) take precedence.
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE=none
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_VERIFICATION=0.0
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_RETAIN=0.1
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_REFLECT=0.9
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_CONSOLIDATION=0.0
|
||||
|
||||
# Example: Anthropic Claude configuration
|
||||
# HINDSIGHT_API_LLM_PROVIDER=anthropic
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-anthropic-api-key
|
||||
@@ -37,12 +48,30 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-zai-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=glm-4.5-flash # or glm-4.5-air for the paid tier
|
||||
|
||||
# Example: Atlas Cloud configuration (OpenAI-compatible, https://www.atlascloud.ai)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=atlas
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-atlascloud-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=deepseek-ai/deepseek-v4-pro # reasoning model; also Qwen / GLM / Kimi / MiniMax, etc.
|
||||
|
||||
# Example: LM Studio local configuration (Qwen 2.5 32B recommended)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=lmstudio
|
||||
# HINDSIGHT_API_LLM_API_KEY=lmstudio
|
||||
# HINDSIGHT_API_LLM_BASE_URL=http://localhost:1234/v1
|
||||
# HINDSIGHT_API_LLM_MODEL=qwen2.5-32b-instruct
|
||||
|
||||
# Multi-LLM strategies: configure extra LLMs by index alongside the primary above,
|
||||
# then pick a routing strategy. Unset = single primary LLM (default). Members are
|
||||
# numbered from 1; indices must be contiguous. Each operation can override with a
|
||||
# RETAIN_/REFLECT_/CONSOLIDATION_ prefix (e.g. HINDSIGHT_API_RETAIN_LLM_1_PROVIDER).
|
||||
# HINDSIGHT_API_LLM_1_PROVIDER=groq
|
||||
# HINDSIGHT_API_LLM_1_API_KEY=your-groq-api-key
|
||||
# HINDSIGHT_API_LLM_1_MODEL=openai/gpt-oss-120b
|
||||
# HINDSIGHT_API_LLM_2_PROVIDER=anthropic
|
||||
# HINDSIGHT_API_LLM_2_API_KEY=your-anthropic-api-key
|
||||
# Strategy JSON: {"mode": "failover"} or {"mode": "round-robin"}.
|
||||
# Round-robin accepts optional positive-int "weights" (one per member, primary first).
|
||||
# HINDSIGHT_API_LLM_STRATEGY={"mode": "failover"}
|
||||
|
||||
# API Configuration (Optional)
|
||||
HINDSIGHT_API_HOST=0.0.0.0
|
||||
HINDSIGHT_API_PORT=8888
|
||||
@@ -87,6 +116,18 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# korean_lindera/lindera(korean), ngram(min,max), edge_ngram(min,max)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER=
|
||||
|
||||
# File Parser (Optional - uses markitdown by default)
|
||||
# HINDSIGHT_API_FILE_PARSER=markitdown
|
||||
# Enable image OCR for MarkItDown using an OpenAI-compatible OCR/vision endpoint.
|
||||
# These OCR settings are independent from HINDSIGHT_API_LLM_* because MarkItDown
|
||||
# uses the OpenAI SDK directly and requires Chat Completions image input support.
|
||||
# When OCR is enabled, API_KEY, BASE_URL, and MODEL are required.
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED=false
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT=
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default), "onnx", "tei", "openai", "cohere", "google", "openrouter", "zeroentropy", "litellm", or "litellm-sdk"
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
@@ -150,6 +191,10 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Custom service name and environment (optional, defaults: hindsight-api, development)
|
||||
# HINDSIGHT_API_OTEL_SERVICE_NAME=hindsight-production
|
||||
# HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT=production
|
||||
#
|
||||
# Expose async-operation queue + consolidation-backlog gauges on /metrics.
|
||||
# Runs periodic per-schema COUNT queries on a background task (disabled by default).
|
||||
# HINDSIGHT_API_METRICS_BACKLOG_ENABLED=true
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Control Plane (Optional)
|
||||
|
||||
@@ -1,6 +0,0 @@
|
||||
version: 2
|
||||
updates:
|
||||
- package-ecosystem: "github-actions"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
@@ -75,7 +75,7 @@ jobs:
|
||||
if: steps.type.outputs.type == 'plugin'
|
||||
run: |
|
||||
echo "Plugin integration ${{ steps.info.outputs.integration }} v${{ steps.info.outputs.version }} — no package to publish."
|
||||
echo "Users install via: claude plugin marketplace add vectorize-io/hindsight --sparse hindsight-integrations"
|
||||
echo "Users install via: claude plugin marketplace add vectorize-io/hindsight"
|
||||
|
||||
# ── TypeScript integrations (ai-sdk, chat, openclaw) ────────────────────
|
||||
|
||||
|
||||
@@ -266,7 +266,7 @@ jobs:
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- os: ubuntu-latest
|
||||
- os: ubuntu-22.04
|
||||
target: x86_64-unknown-linux-gnu
|
||||
artifact_name: hindsight
|
||||
asset_name: hindsight-linux-amd64
|
||||
@@ -278,7 +278,7 @@ jobs:
|
||||
target: aarch64-apple-darwin
|
||||
artifact_name: hindsight
|
||||
asset_name: hindsight-darwin-arm64
|
||||
- os: ubuntu-24.04-arm
|
||||
- os: ubuntu-22.04-arm
|
||||
target: aarch64-unknown-linux-gnu
|
||||
artifact_name: hindsight
|
||||
asset_name: hindsight-linux-arm64
|
||||
|
||||
@@ -38,6 +38,7 @@ jobs:
|
||||
integrations-claude-code: ${{ steps.filter.outputs.integrations-claude-code }}
|
||||
integrations-cline: ${{ steps.filter.outputs.integrations-cline }}
|
||||
integrations-codex: ${{ steps.filter.outputs.integrations-codex }}
|
||||
integrations-github-copilot: ${{ steps.filter.outputs.integrations-github-copilot }}
|
||||
integrations-continue: ${{ steps.filter.outputs.integrations-continue }}
|
||||
integrations-cursor-cli: ${{ steps.filter.outputs.integrations-cursor-cli }}
|
||||
integrations-crewai: ${{ steps.filter.outputs.integrations-crewai }}
|
||||
@@ -45,17 +46,22 @@ jobs:
|
||||
integrations-pydantic-ai: ${{ steps.filter.outputs.integrations-pydantic-ai }}
|
||||
integrations-ag2: ${{ steps.filter.outputs.integrations-ag2 }}
|
||||
integrations-autogen: ${{ steps.filter.outputs.integrations-autogen }}
|
||||
integrations-aider: ${{ steps.filter.outputs.integrations-aider }}
|
||||
integrations-langgraph: ${{ steps.filter.outputs.integrations-langgraph }}
|
||||
integrations-llamaindex: ${{ steps.filter.outputs.integrations-llamaindex }}
|
||||
integrations-paperclip: ${{ steps.filter.outputs.integrations-paperclip }}
|
||||
integrations-opencode: ${{ steps.filter.outputs.integrations-opencode }}
|
||||
integrations-eve: ${{ steps.filter.outputs.integrations-eve }}
|
||||
integrations-cursor: ${{ steps.filter.outputs.integrations-cursor }}
|
||||
integrations-zed: ${{ steps.filter.outputs.integrations-zed }}
|
||||
integrations-n8n: ${{ steps.filter.outputs.integrations-n8n }}
|
||||
integrations-zapier: ${{ steps.filter.outputs.integrations-zapier }}
|
||||
integrations-cloudflare-oauth-proxy: ${{ steps.filter.outputs.integrations-cloudflare-oauth-proxy }}
|
||||
integrations-superagent: ${{ steps.filter.outputs.integrations-superagent }}
|
||||
integrations-lockfiles: ${{ steps.filter.outputs.integrations-lockfiles }}
|
||||
integrations-openai-agents: ${{ steps.filter.outputs.integrations-openai-agents }}
|
||||
integrations-openhands: ${{ steps.filter.outputs.integrations-openhands }}
|
||||
integrations-devin-desktop: ${{ steps.filter.outputs.integrations-devin-desktop }}
|
||||
integrations-pipecat: ${{ steps.filter.outputs.integrations-pipecat }}
|
||||
integrations-agentcore: ${{ steps.filter.outputs.integrations-agentcore }}
|
||||
integrations-smolagents: ${{ steps.filter.outputs.integrations-smolagents }}
|
||||
@@ -140,6 +146,8 @@ jobs:
|
||||
- 'hindsight-integrations/cline/**'
|
||||
integrations-codex:
|
||||
- 'hindsight-integrations/codex/**'
|
||||
integrations-github-copilot:
|
||||
- 'hindsight-integrations/github-copilot/**'
|
||||
integrations-continue:
|
||||
- 'hindsight-integrations/continue/**'
|
||||
integrations-cursor-cli:
|
||||
@@ -154,6 +162,8 @@ jobs:
|
||||
- 'hindsight-integrations/ag2/**'
|
||||
integrations-autogen:
|
||||
- 'hindsight-integrations/autogen/**'
|
||||
integrations-aider:
|
||||
- 'hindsight-integrations/aider/**'
|
||||
integrations-langgraph:
|
||||
- 'hindsight-integrations/langgraph/**'
|
||||
integrations-llamaindex:
|
||||
@@ -164,8 +174,12 @@ jobs:
|
||||
- 'hindsight-integrations/paperclip/**'
|
||||
integrations-opencode:
|
||||
- 'hindsight-integrations/opencode/**'
|
||||
integrations-eve:
|
||||
- 'hindsight-integrations/eve/**'
|
||||
integrations-cursor:
|
||||
- 'hindsight-integrations/cursor/**'
|
||||
integrations-zed:
|
||||
- 'hindsight-integrations/zed/**'
|
||||
integrations-n8n:
|
||||
- 'hindsight-integrations/n8n/**'
|
||||
integrations-zapier:
|
||||
@@ -180,6 +194,10 @@ jobs:
|
||||
- 'scripts/check-integration-lockfiles.sh'
|
||||
integrations-openai-agents:
|
||||
- 'hindsight-integrations/openai-agents/**'
|
||||
integrations-openhands:
|
||||
- 'hindsight-integrations/openhands/**'
|
||||
integrations-devin-desktop:
|
||||
- 'hindsight-integrations/devin-desktop/**'
|
||||
integrations-pipecat:
|
||||
- 'hindsight-integrations/pipecat/**'
|
||||
integrations-agentcore:
|
||||
@@ -488,6 +506,37 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/cursor
|
||||
run: python -m pytest tests/ -v
|
||||
|
||||
test-zed-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-zed == '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 package and pytest
|
||||
working-directory: ./hindsight-integrations/zed
|
||||
# Installs the package (incl. the zstandard runtime dep) so the threads.db
|
||||
# reader tests can decompress Zed's zstd blobs.
|
||||
run: pip install -e . pytest
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/zed
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: python -m pytest tests/ -v -m "not requires_real_llm"
|
||||
|
||||
test-omo-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -551,6 +600,45 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/cline
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-github-copilot-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-github-copilot == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build github-copilot integration
|
||||
working-directory: ./hindsight-integrations/github-copilot
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/github-copilot
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/github-copilot
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-codex-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -708,6 +796,37 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/opencode
|
||||
run: npm run build
|
||||
|
||||
test-eve-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-eve == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '24'
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/eve
|
||||
run: npm ci
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/eve
|
||||
run: npm test
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/eve
|
||||
run: npm run build
|
||||
|
||||
test-n8n-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -3021,6 +3140,45 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/ag2
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-aider-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-aider == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build aider integration
|
||||
working-directory: ./hindsight-integrations/aider
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/aider
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/aider
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-autogen-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -3669,6 +3827,84 @@ jobs:
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-openhands-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-openhands == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build openhands integration
|
||||
working-directory: ./hindsight-integrations/openhands
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/openhands
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/openhands
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-devin-desktop-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-devin-desktop == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build devin-desktop integration
|
||||
working-directory: ./hindsight-integrations/devin-desktop
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/devin-desktop
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/devin-desktop
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-claude-agent-sdk-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -4670,11 +4906,13 @@ jobs:
|
||||
- test-claude-code-integration
|
||||
- test-cursor-integration
|
||||
- test-cline-integration
|
||||
- test-github-copilot-integration
|
||||
- test-codex-integration
|
||||
- test-cursor-cli-integration
|
||||
- build-ai-sdk-integration
|
||||
- test-ai-sdk-integration-deno
|
||||
- test-opencode-integration
|
||||
- test-eve-integration
|
||||
- test-omo-integration
|
||||
- test-cloudflare-oauth-proxy-integration
|
||||
- build-chat-integration
|
||||
@@ -4703,6 +4941,7 @@ jobs:
|
||||
- test-openclaw-integration
|
||||
- test-integration
|
||||
- test-ag2-integration
|
||||
- test-aider-integration
|
||||
- test-autogen-integration
|
||||
- test-continue-integration
|
||||
- test-smolagents-integration
|
||||
@@ -4717,6 +4956,8 @@ jobs:
|
||||
- test-pydantic-ai-integration
|
||||
- test-llamaindex-integration
|
||||
- test-openai-agents-integration
|
||||
- test-openhands-integration
|
||||
- test-devin-desktop-integration
|
||||
- test-agentcore-integration
|
||||
- test-haystack-integration
|
||||
- test-pip-slim
|
||||
|
||||
@@ -6,6 +6,7 @@ dist/
|
||||
wheels/
|
||||
*.egg-info
|
||||
.mcp.json
|
||||
.playwright-mcp/
|
||||
.osgrep
|
||||
# Virtual environments
|
||||
.venv
|
||||
|
||||
@@ -70,7 +70,7 @@ docker run -it --pull always --name hindsight --restart unless-stopped -p 8888:8
|
||||
>API: http://localhost:8888
|
||||
>UI: http://localhost:9999
|
||||
|
||||
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio`, and `minimax`. The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio`, `minimax`, and `atlas` ([Atlas Cloud](https://www.atlascloud.ai/?utm_source=github&utm_medium=link&utm_campaign=hindsight)). The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
|
||||
|
||||
|
||||
@@ -250,7 +250,7 @@ Recall performs 4 retrieval strategies in parallel:
|
||||
- Graph: Entity/temporal/causal links
|
||||
- Temporal: Time range filtering
|
||||
|
||||

|
||||

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

|
||||

|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.8.2
|
||||
appVersion: "0.8.2"
|
||||
version: 0.8.4
|
||||
appVersion: "0.8.4"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.8.2",
|
||||
"version": "0.8.4",
|
||||
"description": "Node.js programmatic lifecycle manager for Hindsight — embeds a local hindsight daemon in a Node application. Pair with @vectorize-io/hindsight-client for memory operations.",
|
||||
"main": "dist/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.8.2"
|
||||
version = "0.8.4"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim==0.8.2",
|
||||
"hindsight-api-slim==0.8.4",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.8.2"
|
||||
version = "0.8.4"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim[all]==0.8.2",
|
||||
"hindsight-api-slim[all]==0.8.4",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
@@ -21,7 +21,7 @@ hindsight-embed = { workspace = true }
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-llm = [
|
||||
"hindsight-api-slim[local-llm]==0.8.2",
|
||||
"hindsight-api-slim[local-llm]==0.8.4",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -121,7 +121,7 @@ This runs a stdio-based MCP server that can be used directly with MCP-compatible
|
||||
- **Entity Graph** — Automatic entity extraction and relationship tracking
|
||||
- **Temporal Reasoning** — Native support for time-based queries
|
||||
- **Disposition Traits** — Configurable skepticism, literalism, and empathy influence opinion formation
|
||||
- **Three Memory Types** — World facts, bank actions, and formed opinions with confidence scores
|
||||
- **Three Memory Types** — World facts, experience facts (the bank's own actions), and observations
|
||||
|
||||
## Documentation
|
||||
|
||||
|
||||
@@ -53,4 +53,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.8.2"
|
||||
__version__ = "0.8.4"
|
||||
|
||||
@@ -56,6 +56,7 @@ BACKUP_TABLES = [
|
||||
"observation_history",
|
||||
"mental_models",
|
||||
"mental_model_history",
|
||||
"knowledge_pages",
|
||||
"directives",
|
||||
"async_operations",
|
||||
"webhooks",
|
||||
@@ -256,6 +257,7 @@ async def _run_migration(
|
||||
schema: str | None = None,
|
||||
base_schema: str = DEFAULT_DATABASE_SCHEMA,
|
||||
embedding_dimension: int | None = None,
|
||||
ensure_extensions: bool = True,
|
||||
) -> list[str]:
|
||||
"""Resolve database URL and run migrations for one schema or all discovered schemas."""
|
||||
from ..migrations import run_migrations_for_schemas
|
||||
@@ -292,7 +294,7 @@ async def _run_migration(
|
||||
vector_extension=config.vector_extension,
|
||||
text_search_extension=config.text_search_extension,
|
||||
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
|
||||
ensure_extensions=True,
|
||||
ensure_extensions=ensure_extensions,
|
||||
)
|
||||
|
||||
return schemas
|
||||
@@ -311,6 +313,18 @@ def run_db_migration(
|
||||
"--embedding-dimension",
|
||||
help="Expected embedding dimension to enforce after migrations. Omit to skip dimension sync.",
|
||||
),
|
||||
skip_extension_reconcile: bool = typer.Option(
|
||||
False,
|
||||
"--skip-extension-reconcile",
|
||||
help=(
|
||||
"Skip the post-migration vector / text-search index reconcile. This step only does "
|
||||
"work when the configured backend (HINDSIGHT_API_VECTOR_EXTENSION / "
|
||||
"HINDSIGHT_API_TEXT_SEARCH_EXTENSION) differs from a schema's existing indexes — a "
|
||||
"rare, operator-driven change. Skipping it makes a no-change re-migration over many "
|
||||
"tenant schemas much faster. Only use when you have NOT changed the backend; a "
|
||||
"backend change still needs a normal run to reshape the indexes."
|
||||
),
|
||||
),
|
||||
):
|
||||
"""Run database migrations to the latest version."""
|
||||
config = HindsightConfig.from_env()
|
||||
@@ -324,6 +338,8 @@ def run_db_migration(
|
||||
typer.echo(f"Running database migrations for schema: {schema}...")
|
||||
else:
|
||||
typer.echo("Running database migrations for base schema and all discovered tenant schemas...")
|
||||
if skip_extension_reconcile:
|
||||
typer.echo("Skipping post-migration extension reconcile (--skip-extension-reconcile).")
|
||||
|
||||
schemas = asyncio.run(
|
||||
_run_migration(
|
||||
@@ -331,6 +347,7 @@ def run_db_migration(
|
||||
schema=schema,
|
||||
base_schema=config.database_schema,
|
||||
embedding_dimension=embedding_dimension,
|
||||
ensure_extensions=not skip_extension_reconcile,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
+52
@@ -0,0 +1,52 @@
|
||||
"""Add managed flag to knowledge_pages.
|
||||
|
||||
The knowledge base is managed by clients (CRUD over folders/pages). ``managed``
|
||||
lets a client tag a node as system-owned vs. hand-authored; it carries no
|
||||
server-side behaviour.
|
||||
|
||||
Revision ID: a5b6c7d8e9f0
|
||||
Revises: a9b8c7d6e5f4
|
||||
Create Date: 2026-06-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a5b6c7d8e9f0"
|
||||
down_revision: str | Sequence[str] | None = "a9b8c7d6e5f4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}knowledge_pages ADD COLUMN IF NOT EXISTS managed BOOLEAN NOT NULL DEFAULT false")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}knowledge_pages DROP COLUMN IF EXISTS managed")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
op.execute("ALTER TABLE knowledge_pages ADD (managed NUMBER(1) DEFAULT 0 NOT NULL)")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("ALTER TABLE knowledge_pages DROP COLUMN managed")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+110
@@ -0,0 +1,110 @@
|
||||
"""Add knowledge_pages table (knowledge-base hierarchy).
|
||||
|
||||
The knowledge base organizes synthesized mental models into a navigable tree of
|
||||
**folders** and **pages**. A page references the mental model that holds its
|
||||
content (``mental_model_id``); a folder is a pure container (``mental_model_id``
|
||||
NULL). Hierarchy is a single self-referential ``parent_id`` so folders can nest
|
||||
arbitrarily. Content stays in ``mental_models`` — this table is metadata + tree
|
||||
structure only.
|
||||
|
||||
Revision ID: a9b8c7d6e5f4
|
||||
Revises: b57a7c9e0d13
|
||||
Create Date: 2026-06-25
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a9b8c7d6e5f4"
|
||||
down_revision: str | Sequence[str] | None = "b57a7c9e0d13"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# parent_id self-FK cascades so deleting a folder row removes its whole
|
||||
# subtree of rows in one shot. The mental_model FK is composite (matches the
|
||||
# mental_models (id, bank_id) PK) and cascades too, so deleting a page's
|
||||
# mental model removes the page row — folders skip the FK because a NULL
|
||||
# column in a composite FK is not enforced (MATCH SIMPLE).
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}knowledge_pages (
|
||||
id VARCHAR(64) NOT NULL,
|
||||
bank_id TEXT NOT NULL,
|
||||
parent_id VARCHAR(64),
|
||||
kind VARCHAR(16) NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
mental_model_id VARCHAR(64),
|
||||
sort_order INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
CONSTRAINT pk_knowledge_pages PRIMARY KEY (id),
|
||||
CONSTRAINT ck_knowledge_pages_kind CHECK (kind IN ('folder', 'page')),
|
||||
CONSTRAINT fk_kp_bank FOREIGN KEY (bank_id)
|
||||
REFERENCES {schema}banks(bank_id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_parent FOREIGN KEY (parent_id)
|
||||
REFERENCES {schema}knowledge_pages(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_mm FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES {schema}mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_kp_bank_parent ON {schema}knowledge_pages (bank_id, parent_id, sort_order)"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_kp_bank_parent")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}knowledge_pages")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
op.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS knowledge_pages (
|
||||
id VARCHAR2(64) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
parent_id VARCHAR2(64),
|
||||
kind VARCHAR2(16) NOT NULL,
|
||||
name CLOB NOT NULL,
|
||||
mental_model_id VARCHAR2(64),
|
||||
sort_order NUMBER DEFAULT 0 NOT NULL,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_knowledge_pages PRIMARY KEY (id),
|
||||
CONSTRAINT ck_knowledge_pages_kind CHECK (kind IN ('folder', 'page')),
|
||||
CONSTRAINT fk_kp_bank FOREIGN KEY (bank_id)
|
||||
REFERENCES banks(bank_id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_parent FOREIGN KEY (parent_id)
|
||||
REFERENCES knowledge_pages(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_mm FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute("CREATE INDEX idx_kp_bank_parent ON knowledge_pages (bank_id, parent_id, sort_order)")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("DROP TABLE knowledge_pages CASCADE CONSTRAINTS")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+61
@@ -0,0 +1,61 @@
|
||||
"""Add bank_stats_cache table for distributed get_bank_stats caching
|
||||
|
||||
Revision ID: b57a7c9e0d13
|
||||
Revises: c3f7a1b9d2e4
|
||||
Create Date: 2026-07-01
|
||||
|
||||
get_bank_stats aggregates over memory_links / unit_entities — a multi-second scan
|
||||
on banks with millions of rows. The result was cached per-process (in-memory), so
|
||||
every API worker recomputed it once per TTL and the first caller after expiry
|
||||
stalled. This table backs a shared, cross-process TTL cache: one worker's compute
|
||||
is written here and served to all the others.
|
||||
|
||||
PostgreSQL only. Oracle keeps the in-process cache (the runtime picks the backing
|
||||
store by dialect), so the Oracle upgrade slot is intentionally absent.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b57a7c9e0d13"
|
||||
down_revision: str | Sequence[str] | None = "c3f7a1b9d2e4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# One row per bank: payload is the full get_bank_stats result, computed_at
|
||||
# drives logical TTL expiry. Rows are overwritten in place (ON CONFLICT), so
|
||||
# the table never grows beyond the number of banks and needs no purge job.
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}bank_stats_cache (
|
||||
bank_id TEXT PRIMARY KEY,
|
||||
payload JSONB NOT NULL,
|
||||
computed_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}bank_stats_cache")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade) # oracle slot intentionally absent → no-op
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+71
@@ -0,0 +1,71 @@
|
||||
"""Unique page name per folder in knowledge_pages.
|
||||
|
||||
The folder curator can fire concurrently (folder-create trigger + the
|
||||
post-consolidation sweep), and an in-process lock can't serialize runs that
|
||||
execute in different threads/loops. A partial unique index on
|
||||
(bank_id, parent, lower(name)) for pages makes duplicate-named pages in the same
|
||||
folder impossible at the DB level — the second concurrent insert fails and the
|
||||
curator treats it as "already exists".
|
||||
|
||||
PostgreSQL only: the Oracle ``name`` column is a CLOB and cannot back a
|
||||
functional unique index; Oracle relies on the in-process serialization instead.
|
||||
|
||||
Revision ID: c3d4e5f6a7b8
|
||||
Revises: a5b6c7d8e9f0
|
||||
Create Date: 2026-06-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c3d4e5f6a7b8"
|
||||
down_revision: str | Sequence[str] | None = "a5b6c7d8e9f0"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# First drop any pre-existing duplicate pages (created by the racy curator
|
||||
# before this guard existed), keeping the earliest row of each duplicate set,
|
||||
# so the unique index can be built. Their backing mental models are left in
|
||||
# place (harmless orphans).
|
||||
op.execute(
|
||||
f"""
|
||||
DELETE FROM {schema}knowledge_pages a
|
||||
USING {schema}knowledge_pages b
|
||||
WHERE a.kind = 'page' AND b.kind = 'page'
|
||||
AND a.bank_id = b.bank_id
|
||||
AND COALESCE(a.parent_id, '') = COALESCE(b.parent_id, '')
|
||||
AND lower(a.name) = lower(b.name)
|
||||
AND a.ctid > b.ctid
|
||||
"""
|
||||
)
|
||||
# COALESCE(parent_id, '') so root-level pages (NULL parent) are also unique by
|
||||
# name — NULLs would otherwise compare distinct and allow duplicates.
|
||||
op.execute(
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS uq_kp_folder_pagename "
|
||||
f"ON {schema}knowledge_pages (bank_id, COALESCE(parent_id, ''), lower(name)) "
|
||||
"WHERE kind = 'page'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}uq_kp_folder_pagename")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade) # oracle slot intentionally absent (CLOB name)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+144
@@ -0,0 +1,144 @@
|
||||
"""Backfill search_vector for native-backend observations.
|
||||
|
||||
Observations created or updated by the consolidator landed with a NULL
|
||||
``search_vector`` under the ``native`` text-search backend: the
|
||||
single-row INSERT/UPDATE paths in ``consolidator.py`` never populated the
|
||||
tsvector (only the batch raw-fact path in ``ops_postgresql.insert_facts_batch``
|
||||
did). Those observations were therefore invisible to the BM25 retrieval arm
|
||||
until they were re-written by a later consolidation pass. The writer is fixed
|
||||
in the same change set (all four consolidator sites now call
|
||||
``to_tsvector($lang, COALESCE(text, ''))``); this migration repairs the
|
||||
historical residue so existing observations become BM25-searchable without a
|
||||
re-ingest.
|
||||
|
||||
Scope mirrors the writer fix exactly:
|
||||
* Only the ``native`` backend is touched. The gate is the column *type*:
|
||||
under ``native`` ``search_vector`` is a regular (non-generated) tsvector
|
||||
column; under ``vchord`` it is a ``bm25vector`` and under
|
||||
``pg_textsearch`` / ``pgroonga`` / ``pg_search`` it is a dummy ``text``
|
||||
column. ``_is_regular_tsvector`` is true only for ``native``, so every
|
||||
other backend is a no-op.
|
||||
* The tsvector is built from the observation's own ``text`` only — matching
|
||||
the consolidator INSERT/UPDATE paths (entity / source / temporal signals
|
||||
are intentionally excluded; the other retrieval arms cover those).
|
||||
* Only ``fact_type = 'observation'`` rows with a NULL ``search_vector`` are
|
||||
rewritten. Raw facts already carry a populated tsvector, and the
|
||||
``IS NULL`` predicate makes the migration idempotent and re-runnable.
|
||||
|
||||
The configured ``HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE`` is used
|
||||
so backfilled rows are lexically identical to newly-created observations. The
|
||||
value is validated as a PG identifier (mirroring
|
||||
``HindsightConfig.validate``) before being embedded as a SQL literal.
|
||||
|
||||
This is a single UPDATE per schema: it locks the targeted observation rows for
|
||||
its duration. It is one-time and only touches unpopulated rows, so subsequent
|
||||
online writes (which now carry the tsvector via the writer fix) are unaffected.
|
||||
|
||||
Oracle slot is intentionally absent: the consolidator INSERT/UPDATE paths that
|
||||
this repairs are PostgreSQL-specific (``ops_postgresql``), and the native
|
||||
tsvector ``search_vector`` column only exists on PostgreSQL. There is no Oracle
|
||||
residue to repair.
|
||||
|
||||
Revision ID: c3f7a1b9d2e4
|
||||
Revises: f4d1c2b3a5e6
|
||||
Create Date: 2026-06-29
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import Connection, text
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
from hindsight_api.config import (
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
)
|
||||
|
||||
revision: str = "c3f7a1b9d2e4"
|
||||
down_revision: str | Sequence[str] | None = "f4d1c2b3a5e6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
# Matches HindsightConfig.validate(): a tsvector regconfig name embedded as a
|
||||
# SQL literal must be a bare PG identifier.
|
||||
_PG_IDENTIFIER = re.compile(r"[a-zA-Z_][a-zA-Z0-9_]*")
|
||||
|
||||
|
||||
def _schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _schema_name() -> str:
|
||||
return (context.config.get_main_option("target_schema") or "public").strip('"')
|
||||
|
||||
|
||||
def _native_language() -> str:
|
||||
"""Configured native tsvector language, validated as a PG identifier."""
|
||||
lang = os.getenv(
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
)
|
||||
if not _PG_IDENTIFIER.fullmatch(lang):
|
||||
return DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE
|
||||
return lang
|
||||
|
||||
|
||||
def _is_regular_tsvector(conn: Connection, schema: str, table: str) -> bool:
|
||||
"""True iff ``schema.table.search_vector`` is a non-generated tsvector column.
|
||||
|
||||
This is the ``native`` backend signature. ``vchord`` (bm25vector) and
|
||||
``pg_textsearch`` / ``pgroonga`` / ``pg_search`` (dummy text column) all
|
||||
fail this check, so the backfill is a no-op for them.
|
||||
"""
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT is_generated, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table
|
||||
AND column_name = 'search_vector'
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
is_generated, udt_name = row[0], row[1]
|
||||
return udt_name == "tsvector" and is_generated != "ALWAYS"
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
schema_name = _schema_name()
|
||||
if not _is_regular_tsvector(conn, schema_name, "memory_units"):
|
||||
# Non-native backend (or column absent) — nothing to backfill.
|
||||
return
|
||||
schema_prefix = _schema_prefix()
|
||||
lang = _native_language()
|
||||
op.execute(
|
||||
f"""
|
||||
UPDATE {schema_prefix}memory_units
|
||||
SET search_vector = to_tsvector('{lang}'::regconfig, COALESCE(text, ''))
|
||||
WHERE fact_type = 'observation' AND search_vector IS NULL
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# No-op: backfilled rows are indistinguishable from observations that were
|
||||
# populated by the post-fix writer, and reverting either to NULL would
|
||||
# re-break BM25 retrieval. The column simply stays populated.
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+158
@@ -0,0 +1,158 @@
|
||||
"""Make maintenance routines resilient to schemas that vanish mid-scan.
|
||||
|
||||
``public.banks_needing_consolidation()`` and
|
||||
``public.schemas_with_expired_rows(...)`` snapshot the set of schemas owning a
|
||||
target table from ``pg_class`` and then run a dynamic query against each schema
|
||||
in turn. That is a time-of-check/time-of-use race: a schema (or its tables) can
|
||||
be dropped — a tenant being deleted, or a tenant migration that recreates
|
||||
tables — between the snapshot and the per-schema query, which then aborts the
|
||||
whole routine with::
|
||||
|
||||
relation "<schema>.memory_units" does not exist
|
||||
relation "<schema>.audit_log" does not exist
|
||||
|
||||
In the test suite this surfaces as cross-worker contamination: the multi-tenant
|
||||
maintenance test creates and drops ~100 ``mt<hash>_NNN`` schemas while
|
||||
``test_maintenance_routines`` (on another xdist worker, same DB) calls the
|
||||
routines. In production the background maintenance loop hits the same race when
|
||||
a tenant is removed or mid-migration.
|
||||
|
||||
Wrap each per-schema query in its own ``BEGIN ... EXCEPTION`` block so a schema
|
||||
that disappears (``undefined_table`` / ``invalid_schema_name`` /
|
||||
``undefined_column``) is skipped instead of aborting the scan. The routines stay
|
||||
``CREATE OR REPLACE`` and PostgreSQL-only, and are (re)installed only on the run
|
||||
that targets the shared ``public`` schema — same gating as the original
|
||||
install (``e5f6a7b8c9d0``) and its repair (``b2d4f6a8c1e3``).
|
||||
|
||||
Revision ID: c7e9f1a3b5d2
|
||||
Revises: e1f2a3b4c5d6
|
||||
Create Date: 2026-06-19
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c7e9f1a3b5d2"
|
||||
down_revision: str | Sequence[str] | None = "e1f2a3b4c5d6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _should_install_public_routines(target_schema: str | None) -> bool:
|
||||
"""True for the run that must (re)create the shared ``public.*`` routines.
|
||||
|
||||
The routines physically live in ``public``, so they are installed exactly
|
||||
once — on the base run (no ``target_schema``) or the run that explicitly
|
||||
targets ``public``. Mirrors ``b2d4f6a8c1e3``.
|
||||
"""
|
||||
return not target_schema or target_schema == "public"
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if not _should_install_public_routines(context.config.get_main_option("target_schema")):
|
||||
return
|
||||
|
||||
# Same body as b2d4f6a8c1e3, but each per-schema query runs in its own
|
||||
# subtransaction so a schema dropped mid-scan is skipped, not fatal.
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.banks_needing_consolidation()
|
||||
RETURNS TABLE(schema_name text, bank_id text)
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
BEGIN
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = 'memory_units' AND c.relkind = 'r'
|
||||
LOOP
|
||||
BEGIN
|
||||
RETURN QUERY EXECUTE format($q$
|
||||
SELECT %1$L::text, m.bank_id
|
||||
FROM %1$I.memory_units m
|
||||
JOIN %1$I.banks b ON b.bank_id = m.bank_id
|
||||
WHERE m.consolidated_at IS NULL
|
||||
AND m.consolidation_failed_at IS NULL
|
||||
AND m.fact_type IN ('experience', 'world')
|
||||
AND COALESCE(b.config -> 'enable_auto_consolidation', 'true'::jsonb) <> 'false'::jsonb
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM %1$I.async_operations o
|
||||
WHERE o.bank_id = m.bank_id
|
||||
AND o.operation_type = 'consolidation'
|
||||
AND o.status IN ('pending', 'processing')
|
||||
)
|
||||
GROUP BY m.bank_id
|
||||
$q$, sch);
|
||||
EXCEPTION
|
||||
-- Schema or its tables vanished between the pg_class
|
||||
-- snapshot and this query (tenant dropped or migrating).
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.schemas_with_expired_rows(
|
||||
p_table text, p_ts_col text, p_days int
|
||||
)
|
||||
RETURNS SETOF text
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
has_expired boolean;
|
||||
BEGIN
|
||||
IF p_days IS NULL OR p_days <= 0 THEN
|
||||
RETURN;
|
||||
END IF;
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = p_table AND c.relkind = 'r'
|
||||
LOOP
|
||||
BEGIN
|
||||
EXECUTE format(
|
||||
'SELECT EXISTS (SELECT 1 FROM %I.%I WHERE %I < NOW() - make_interval(days => $1))',
|
||||
sch, p_table, p_ts_col
|
||||
) INTO has_expired USING p_days;
|
||||
EXCEPTION
|
||||
-- Schema or its table vanished mid-scan; skip it.
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
IF has_expired THEN
|
||||
RETURN NEXT sch;
|
||||
END IF;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# No-op: e5f6a7b8c9d0 owns these functions' lifecycle and drops them on its
|
||||
# own downgrade. This migration only re-installs them (the resilient body is
|
||||
# a strict superset of the previous behaviour), so there is nothing to undo
|
||||
# without racing that migration's DROP.
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+110
@@ -0,0 +1,110 @@
|
||||
"""Add server-side routine for cron-scheduled mental model refresh.
|
||||
|
||||
Installs ``public.mental_models_with_cron()`` — a discovery routine that returns
|
||||
every mental model carrying a non-empty ``trigger->>'refresh_cron'`` across all
|
||||
tenant schemas in one round-trip (the same per-schema scan as the other
|
||||
maintenance routines from ``e5f6a7b8c9d0``). The maintenance loop evaluates each
|
||||
candidate's cron expression in Python (``croniter``) against ``last_refreshed_at``
|
||||
to decide whether a scheduled refresh is due — cron arithmetic isn't expressible
|
||||
in plain SQL — and only the cron *candidate set* is discovered here.
|
||||
|
||||
Models that already have a ``refresh_mental_model`` operation pending/processing
|
||||
are excluded so a slow refresh isn't double-queued (mirrors the in-flight guard
|
||||
in ``banks_needing_consolidation``). Each per-schema query runs in its own
|
||||
``BEGIN ... EXCEPTION`` subtransaction so a schema dropped mid-scan (tenant
|
||||
deletion / migration) is skipped, not fatal — same resilience as
|
||||
``c7e9f1a3b5d2``.
|
||||
|
||||
Read-only (STABLE) discovery routine — the caller performs the refresh enqueue —
|
||||
so installing it never mutates data. PostgreSQL only: the worker poller and the
|
||||
maintenance loop are PG-only (Oracle slot intentionally absent, mirroring
|
||||
``e5f6a7b8c9d0``). The routine lives in ``public`` and is CREATE OR REPLACE, so
|
||||
it is installed exactly once (base / ``public`` run) to avoid the
|
||||
``tuple concurrently updated`` race on concurrent per-tenant runs.
|
||||
|
||||
Revision ID: f4d1c2b3a5e6
|
||||
Revises: c7e9f1a3b5d2
|
||||
Create Date: 2026-06-23
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "f4d1c2b3a5e6"
|
||||
down_revision: str | Sequence[str] | None = "c7e9f1a3b5d2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _should_install_public_routines(target_schema: str | None) -> bool:
|
||||
"""True for the run that must (re)create the shared ``public.*`` routine.
|
||||
|
||||
The routine physically lives in ``public``, so it is installed exactly once —
|
||||
on the base run (no ``target_schema``) or the run that explicitly targets
|
||||
``public``. Mirrors ``c7e9f1a3b5d2``.
|
||||
"""
|
||||
return not target_schema or target_schema == "public"
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if not _should_install_public_routines(context.config.get_main_option("target_schema")):
|
||||
return
|
||||
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.mental_models_with_cron()
|
||||
RETURNS TABLE(schema_name text, bank_id text, mental_model_id text,
|
||||
refresh_cron text, last_refreshed_at timestamptz)
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
BEGIN
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = 'mental_models' AND c.relkind = 'r'
|
||||
LOOP
|
||||
BEGIN
|
||||
RETURN QUERY EXECUTE format($q$
|
||||
SELECT %1$L::text, mm.bank_id::text, mm.id::text,
|
||||
mm.trigger->>'refresh_cron', mm.last_refreshed_at
|
||||
FROM %1$I.mental_models mm
|
||||
WHERE COALESCE(mm.trigger->>'refresh_cron', '') <> ''
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM %1$I.async_operations o
|
||||
WHERE o.bank_id = mm.bank_id
|
||||
AND o.operation_type = 'refresh_mental_model'
|
||||
AND o.status IN ('pending', 'processing')
|
||||
AND o.task_payload->>'mental_model_id' = mm.id::text
|
||||
)
|
||||
$q$, sch);
|
||||
EXCEPTION
|
||||
-- Schema or its tables vanished between the pg_class
|
||||
-- snapshot and this query (tenant dropped or migrating).
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
if not _should_install_public_routines(context.config.get_main_option("target_schema")):
|
||||
return
|
||||
op.execute("DROP FUNCTION IF EXISTS public.mental_models_with_cron()")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
@@ -18,6 +18,7 @@ from typing import Any, Literal, TypeVar
|
||||
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Query, Request, UploadFile
|
||||
from fastapi.middleware.gzip import GZipMiddleware
|
||||
|
||||
from hindsight_api.api import okf
|
||||
from hindsight_api.api.disconnect import ClientDisconnectCancellationMiddleware, get_scope_cancellation_token
|
||||
from hindsight_api.cancellation import OperationCancelledError
|
||||
from hindsight_api.engine.audit import (
|
||||
@@ -27,7 +28,7 @@ from hindsight_api.engine.audit import (
|
||||
AuditLogStatsResponse,
|
||||
)
|
||||
from hindsight_api.engine.llm_trace import LLMRequestListResponse, LLMRequestStatsResponse
|
||||
from hindsight_api.extensions import AuthenticationError
|
||||
from hindsight_api.extensions import AuthenticationError, PrecheckOperation
|
||||
|
||||
|
||||
def _parse_metadata(metadata: Any) -> dict[str, Any]:
|
||||
@@ -154,11 +155,18 @@ from hindsight_api.engine.response_models import (
|
||||
VALID_RECALL_FACT_TYPES,
|
||||
DryRunExtractionResult,
|
||||
MemoryFact,
|
||||
MinScores,
|
||||
RecallScores,
|
||||
TokenUsage,
|
||||
)
|
||||
from hindsight_api.engine.search.tags import TagGroup, TagsMatch
|
||||
from hindsight_api.extensions import HttpExtension, OperationValidationError, load_extension
|
||||
from hindsight_api.metrics import create_metrics_collector, get_metrics_collector, initialize_metrics
|
||||
from hindsight_api.metrics import (
|
||||
create_metrics_collector,
|
||||
get_metrics_collector,
|
||||
initialize_metrics,
|
||||
normalize_http_endpoint,
|
||||
)
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -265,6 +273,16 @@ class RecallRequest(BaseModel):
|
||||
default=None,
|
||||
description="List of fact types to recall: 'world', 'experience', 'observation'. Defaults to world and experience if not specified.",
|
||||
)
|
||||
prefer_observations: bool = Field(
|
||||
default=False,
|
||||
description=(
|
||||
"When recalling raw facts ('world'/'experience') together with 'observation', drop any raw "
|
||||
"fact that an observation in the results was consolidated from, so the observation supersedes "
|
||||
"it and you don't get duplicate content. The freed slots are backfilled with the next results, "
|
||||
"keeping the result count at the requested budget. Disabled by default; set to true to enable. "
|
||||
"No effect unless 'observation' and at least one raw type are both requested."
|
||||
),
|
||||
)
|
||||
budget: Budget = Budget.MID
|
||||
max_tokens: int = 4096
|
||||
trace: bool = False
|
||||
@@ -281,18 +299,31 @@ class RecallRequest(BaseModel):
|
||||
)
|
||||
tags: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Filter memories by tags. If not specified, all memories are returned.",
|
||||
description="Filter memories by tags. If not specified, all memories are returned. "
|
||||
"Omitting tags (or passing []) together with tags_match='exact' filters to "
|
||||
"untagged/global observations only (the scope written by observation_scopes='shared').",
|
||||
)
|
||||
tags_match: TagsMatch = Field(
|
||||
default="any",
|
||||
description="How to match tags: 'any' (OR, includes untagged), 'all' (AND, includes untagged), "
|
||||
"'any_strict' (OR, excludes untagged), 'all_strict' (AND, excludes untagged).",
|
||||
"'any_strict' (OR, excludes untagged), 'all_strict' (AND, excludes untagged), "
|
||||
"'exact' (set-equality on the full scope, excludes untagged). With 'exact' and no tags "
|
||||
"(or []), the empty global scope is selected and only untagged memories match.",
|
||||
)
|
||||
tag_groups: list[TagGroup] | None = Field(
|
||||
default=None,
|
||||
description="Compound tag filter using boolean groups. Groups in the list are AND-ed. "
|
||||
"Each group is a leaf {tags, match} or compound {and: [...]}, {or: [...]}, {not: ...}.",
|
||||
)
|
||||
min_scores: MinScores | None = Field(
|
||||
default=None,
|
||||
description="Optional per-stage score floors (all inclusive, AND-ed). `semantic` and `keyword` are "
|
||||
"retrieval-level cutoffs pushed into the SQL arms (overriding the global similarity/BM25 minimums for "
|
||||
"this request); `reranker` and `final` are post-ranking filters on the scored results. Any field left "
|
||||
"unset imposes no floor; omitting `min_scores` entirely (the default) applies no score filtering. Use "
|
||||
"with care — the reranker's absolute scores are not calibrated across queries (a clearly-relevant match "
|
||||
"may score ~0.001 even though it is ranked first).",
|
||||
)
|
||||
|
||||
@field_validator("query")
|
||||
@classmethod
|
||||
@@ -348,6 +379,7 @@ class RecallResult(BaseModel):
|
||||
source_fact_ids: list[str] | None = (
|
||||
None # IDs of source facts (observation type only, when source_facts is enabled)
|
||||
)
|
||||
scores: RecallScores | None = None # Per-stage recall scores (final/reranker/semantic/text)
|
||||
|
||||
|
||||
class EntityObservationResponse(BaseModel):
|
||||
@@ -1442,6 +1474,13 @@ class DryRunExtractRequest(BaseModel):
|
||||
entities_allow_free_form: bool | None = None
|
||||
llm_output_language: str | None = None
|
||||
|
||||
@field_validator("content")
|
||||
@classmethod
|
||||
def validate_content(cls, v: str) -> str:
|
||||
if not v.strip():
|
||||
raise ValueError("content cannot be empty")
|
||||
return v
|
||||
|
||||
|
||||
class ListDocumentsResponse(BaseModel):
|
||||
"""Response model for list documents endpoint."""
|
||||
@@ -1941,6 +1980,17 @@ class MentalModelTrigger(BaseModel):
|
||||
default=False,
|
||||
description="If true, refresh this mental model after observations consolidation (real-time mode)",
|
||||
)
|
||||
refresh_cron: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Cron expression (UTC, standard 5-field syntax, e.g. '0 3 * * *' for daily at 03:00 UTC) "
|
||||
"for refreshing this mental model on a fixed schedule. Mutually exclusive with "
|
||||
"refresh_after_consolidation — a model refreshes either after consolidation or on a cron "
|
||||
"schedule, not both. A scheduled refresh only runs when the model is stale (new memories in "
|
||||
"its scope since the last refresh); if nothing changed, the tick is skipped to avoid a "
|
||||
"wasted LLM call. null = no schedule."
|
||||
),
|
||||
)
|
||||
fact_types: list[Literal["world", "experience", "observation"]] | None = Field(
|
||||
default=None,
|
||||
description="Filter which fact types are retrieved during reflect. None means all types (world, experience, observation).",
|
||||
@@ -1999,6 +2049,31 @@ class MentalModelTrigger(BaseModel):
|
||||
raise ValueError("fact_types must not be empty. Use null to include all fact types.")
|
||||
return v
|
||||
|
||||
@field_validator("refresh_cron")
|
||||
@classmethod
|
||||
def validate_refresh_cron(cls, v: str | None) -> str | None:
|
||||
if v is None:
|
||||
return v
|
||||
v = v.strip()
|
||||
if not v:
|
||||
return None
|
||||
from croniter import croniter
|
||||
|
||||
if not croniter.is_valid(v):
|
||||
raise ValueError(f"refresh_cron is not a valid cron expression: {v!r}")
|
||||
return v
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_refresh_exclusivity(self) -> "MentalModelTrigger":
|
||||
# A mental model refreshes either after consolidation (real-time) or on a
|
||||
# cron schedule, never both — the two triggers would race and double-refresh.
|
||||
if self.refresh_after_consolidation and self.refresh_cron:
|
||||
raise ValueError(
|
||||
"refresh_after_consolidation and refresh_cron are mutually exclusive: "
|
||||
"a mental model refreshes either after consolidation or on a cron schedule, not both."
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class MentalModelResponse(BaseModel):
|
||||
"""Response model for a mental model (stored reflect response)."""
|
||||
@@ -2036,6 +2111,150 @@ class MentalModelListResponse(BaseModel):
|
||||
items: list[MentalModelResponse]
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# KNOWLEDGE BASE (folders + pages over mental models, projected to OKF)
|
||||
# =========================================================================
|
||||
|
||||
|
||||
class KnowledgeNode(BaseModel):
|
||||
"""A node in the knowledge-base tree — a folder or a page.
|
||||
|
||||
Pages carry ``description``/``tags`` from their backing mental model. The
|
||||
knowledge base is client-managed (CRUD); ``managed`` lets a client tag a node
|
||||
as system-owned vs. hand-authored.
|
||||
"""
|
||||
|
||||
id: str
|
||||
kind: Literal["folder", "page"]
|
||||
name: str
|
||||
parent_id: str | None = None
|
||||
mental_model_id: str | None = Field(default=None, description="Backing mental model id (pages only).")
|
||||
managed: bool = Field(default=False, description="Client-set flag: true = system-owned, false = hand-authored.")
|
||||
description: str | None = Field(default=None, description="Page source query (OKF `description`).")
|
||||
tags: list[str] = FieldWithDefault(list)
|
||||
timestamp: str | None = Field(default=None, description="Last refresh (page) or last update (folder).")
|
||||
children: list["KnowledgeNode"] = FieldWithDefault(list)
|
||||
|
||||
|
||||
class KnowledgeTreeResponse(BaseModel):
|
||||
"""The knowledge base as a nested folder/page tree."""
|
||||
|
||||
roots: list[KnowledgeNode]
|
||||
|
||||
|
||||
class CreateFolderRequest(BaseModel):
|
||||
"""Create a folder under an optional parent folder."""
|
||||
|
||||
name: str
|
||||
parent_id: str | None = None
|
||||
|
||||
|
||||
class CreatePageRequest(BaseModel):
|
||||
"""Create a page (a mental model + tree node) under an optional parent folder."""
|
||||
|
||||
name: str
|
||||
source_query: str
|
||||
parent_id: str | None = None
|
||||
tags: list[str] | None = None
|
||||
max_tokens: int | None = None
|
||||
trigger: MentalModelTrigger | None = None
|
||||
|
||||
|
||||
class UpdateNodeRequest(BaseModel):
|
||||
"""Rename and/or move a node. Each field applies only when present."""
|
||||
|
||||
name: str | None = None
|
||||
parent_id: str | None = None
|
||||
|
||||
|
||||
class CreateKnowledgePageResponse(BaseModel):
|
||||
"""Result of creating a page: the node id, its mental model, and the refresh op."""
|
||||
|
||||
page_id: str
|
||||
mental_model_id: str
|
||||
operation_id: str | None = None
|
||||
|
||||
|
||||
class KnowledgePageResponse(BaseModel):
|
||||
"""A knowledge page rendered as an OKF document."""
|
||||
|
||||
id: str
|
||||
name: str
|
||||
type: str = Field(description="OKF document type — from a `type:<x>` tag, else 'knowledge-page'.")
|
||||
description: str | None = Field(default=None, description="The source query that rebuilds the page.")
|
||||
tags: list[str] = FieldWithDefault(list)
|
||||
timestamp: str | None = Field(default=None, description="Last refresh time (falls back to creation).")
|
||||
body: str | None = Field(default=None, description="The page's synthesized markdown body.")
|
||||
markdown: str = Field(description="The full OKF document: YAML frontmatter + markdown body.")
|
||||
|
||||
|
||||
class KnowledgePageGraphResponse(BaseModel):
|
||||
"""Constellation graph of knowledge pages linked by shared tags."""
|
||||
|
||||
nodes: list[dict[str, Any]]
|
||||
edges: list[dict[str, Any]]
|
||||
total_pages: int
|
||||
total_edges: int
|
||||
|
||||
|
||||
class KnowledgePageBundleFile(BaseModel):
|
||||
"""One file in a portable OKF bundle."""
|
||||
|
||||
path: str
|
||||
content: str
|
||||
|
||||
|
||||
class KnowledgePageBundleResponse(BaseModel):
|
||||
"""A portable OKF bundle — a flat set of markdown files (index + pages + logs)."""
|
||||
|
||||
files: list[KnowledgePageBundleFile]
|
||||
|
||||
|
||||
def _knowledge_node_model(node: dict[str, Any]) -> KnowledgeNode:
|
||||
"""Project an engine node dict into a (childless) KnowledgeNode."""
|
||||
is_page = node.get("kind") == "page"
|
||||
return KnowledgeNode(
|
||||
id=node["id"],
|
||||
kind=node["kind"],
|
||||
name=node["name"],
|
||||
parent_id=node.get("parent_id"),
|
||||
mental_model_id=node.get("mental_model_id"),
|
||||
managed=bool(node.get("managed")),
|
||||
description=node.get("source_query") if is_page else None,
|
||||
tags=list(node.get("tags") or []) if is_page else [],
|
||||
timestamp=(node.get("last_refreshed_at") if is_page else node.get("updated_at")),
|
||||
)
|
||||
|
||||
|
||||
def _build_knowledge_tree(nodes: list[dict[str, Any]]) -> list[KnowledgeNode]:
|
||||
"""Assemble the flat node list into a nested tree of roots."""
|
||||
models = {n["id"]: _knowledge_node_model(n) for n in nodes}
|
||||
roots: list[KnowledgeNode] = []
|
||||
for node in nodes:
|
||||
model = models[node["id"]]
|
||||
parent_id = node.get("parent_id")
|
||||
if parent_id and parent_id in models:
|
||||
models[parent_id].children.append(model)
|
||||
else:
|
||||
roots.append(model)
|
||||
return roots
|
||||
|
||||
|
||||
def _knowledge_page_response(node: dict[str, Any]) -> KnowledgePageResponse:
|
||||
"""Project a page node (with merged mental-model content) into an OKF document."""
|
||||
page = okf.page_type(node.get("tags"))
|
||||
return KnowledgePageResponse(
|
||||
id=node["id"],
|
||||
name=node["name"],
|
||||
type=page.type,
|
||||
description=node.get("source_query"),
|
||||
tags=page.display_tags,
|
||||
timestamp=node.get("last_refreshed_at") or node.get("created_at"),
|
||||
body=node.get("content"),
|
||||
markdown=okf.render_document(node),
|
||||
)
|
||||
|
||||
|
||||
class CreateMentalModelRequest(BaseModel):
|
||||
"""Request model for creating a mental model."""
|
||||
|
||||
@@ -2534,6 +2753,10 @@ class OperationResponse(BaseModel):
|
||||
task_type: str
|
||||
items_count: int
|
||||
document_id: str | None = None
|
||||
filename: str | None = Field(
|
||||
default=None,
|
||||
description="Original filename for file-conversion operations (file_convert_retain); null for other task types.",
|
||||
)
|
||||
created_at: str
|
||||
updated_at: str | None = Field(
|
||||
default=None,
|
||||
@@ -3043,8 +3266,12 @@ def create_app(
|
||||
# All current backends (PostgreSQL, Oracle) support async worker/poller.
|
||||
if config.worker_enabled and memory._backend.supports_worker_poller:
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA
|
||||
from ..utils import warn_if_container_default_worker_id
|
||||
|
||||
warn_if_container_default_worker_id(config.worker_id)
|
||||
worker_id = config.worker_id or socket.gethostname()
|
||||
worker_id_source = "HINDSIGHT_API_WORKER_ID" if config.worker_id else "hostname (default)"
|
||||
logging.info(f"Worker id: {worker_id} (source: {worker_id_source})")
|
||||
# Convert default schema to None for SQL compatibility (no schema prefix)
|
||||
schema = None if config.database_schema == DEFAULT_DATABASE_SCHEMA else config.database_schema
|
||||
poller = WorkerPoller(
|
||||
@@ -3237,15 +3464,9 @@ def create_app(
|
||||
@app.middleware("http")
|
||||
async def http_metrics_middleware(request, call_next):
|
||||
"""Record HTTP request metrics."""
|
||||
# Normalize endpoint path to reduce cardinality
|
||||
# Replace UUIDs and numeric IDs with placeholders
|
||||
import re
|
||||
|
||||
path = request.url.path
|
||||
# Replace UUIDs
|
||||
path = re.sub(r"/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", "/{id}", path)
|
||||
# Replace numeric IDs
|
||||
path = re.sub(r"/\d+(?=/|$)", "/{id}", path)
|
||||
# Template id segments (bank ids, UUIDs, numeric ids) so the endpoint
|
||||
# metric label stays bounded-cardinality.
|
||||
path = normalize_http_endpoint(request.url.path)
|
||||
|
||||
status_code = [500] # Default to 500, will be updated
|
||||
metrics_collector = get_metrics_collector()
|
||||
@@ -3304,7 +3525,7 @@ def _register_routes(app: FastAPI):
|
||||
api_key = authorization.strip()
|
||||
return RequestContext(api_key=api_key)
|
||||
|
||||
def precheck_for(operation: str):
|
||||
def precheck_for(operation: PrecheckOperation):
|
||||
"""
|
||||
Build a FastAPI dependency that runs ``OperationValidator.precheck``.
|
||||
|
||||
@@ -3333,6 +3554,7 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
async def _precheck_dep(
|
||||
bank_id: str,
|
||||
request: Request,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
) -> None:
|
||||
validator = getattr(app.state.memory, "_operation_validator", None)
|
||||
@@ -3341,10 +3563,20 @@ def _register_routes(app: FastAPI):
|
||||
from hindsight_api.extensions import PrecheckContext
|
||||
|
||||
await app.state.memory._authenticate_tenant(request_context)
|
||||
cl_header = request.headers.get("content-length")
|
||||
content_length: int | None = None
|
||||
if cl_header is not None:
|
||||
try:
|
||||
parsed = int(cl_header)
|
||||
except ValueError:
|
||||
parsed = -1
|
||||
if parsed >= 0:
|
||||
content_length = parsed
|
||||
ctx = PrecheckContext(
|
||||
operation=operation,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
content_length=content_length,
|
||||
)
|
||||
result = await validator.precheck(ctx)
|
||||
if not result.allowed:
|
||||
@@ -3448,7 +3680,7 @@ def _register_routes(app: FastAPI):
|
||||
async def api_graph(
|
||||
bank_id: str,
|
||||
type: str | None = None,
|
||||
limit: int = 1000,
|
||||
limit: int = Query(default=1000, ge=0),
|
||||
q: str | None = None,
|
||||
tags: list[str] | None = Query(None),
|
||||
tags_match: str = "all_strict",
|
||||
@@ -3496,8 +3728,8 @@ def _register_routes(app: FastAPI):
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
document_id: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
limit: int = Query(default=100, ge=0),
|
||||
offset: int = Query(default=0, ge=0),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""
|
||||
@@ -3544,7 +3776,7 @@ def _register_routes(app: FastAPI):
|
||||
async def _require_dry_run_enabled() -> None:
|
||||
"""Feature-flag gate for dry-run extraction.
|
||||
|
||||
Declared as a dependency BEFORE ``precheck_for("dry_run_extract")`` so a
|
||||
Declared as a dependency BEFORE ``precheck_for(PrecheckOperation.DRY_RUN_EXTRACT)`` so a
|
||||
disabled route returns 404 regardless of tenant/billing state — FastAPI
|
||||
resolves path-operation dependencies in signature order, so this runs
|
||||
first and preserves the original "disabled → 404" contract.
|
||||
@@ -3574,7 +3806,7 @@ def _register_routes(app: FastAPI):
|
||||
body: DryRunExtractRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_enabled: None = Depends(_require_dry_run_enabled),
|
||||
_precheck: None = Depends(precheck_for("dry_run_extract")),
|
||||
_precheck: None = Depends(precheck_for(PrecheckOperation.DRY_RUN_EXTRACT)),
|
||||
):
|
||||
try:
|
||||
override_fields = (
|
||||
@@ -3742,7 +3974,7 @@ def _register_routes(app: FastAPI):
|
||||
request: RecallRequest,
|
||||
http_request: Request,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("recall")),
|
||||
_precheck: None = Depends(precheck_for(PrecheckOperation.RECALL)),
|
||||
):
|
||||
"""Run a recall and return results with trace."""
|
||||
import time
|
||||
@@ -3809,6 +4041,7 @@ def _register_routes(app: FastAPI):
|
||||
max_tokens=request.max_tokens,
|
||||
enable_trace=request.trace,
|
||||
fact_type=fact_types,
|
||||
prefer_observations=request.prefer_observations,
|
||||
question_date=question_date,
|
||||
include_entities=include_entities,
|
||||
max_entity_tokens=max_entity_tokens,
|
||||
@@ -3821,6 +4054,7 @@ def _register_routes(app: FastAPI):
|
||||
tags=request.tags,
|
||||
tags_match=request.tags_match,
|
||||
tag_groups=request.tag_groups,
|
||||
min_scores=request.min_scores,
|
||||
),
|
||||
operation="recall",
|
||||
bank_id=bank_id,
|
||||
@@ -3842,6 +4076,7 @@ def _register_routes(app: FastAPI):
|
||||
chunk_id=fact.chunk_id,
|
||||
tags=fact.tags,
|
||||
source_fact_ids=fact.source_fact_ids,
|
||||
scores=fact.scores,
|
||||
)
|
||||
|
||||
recall_results = [_fact_to_result(fact) for fact in core_result.results]
|
||||
@@ -3943,7 +4178,7 @@ def _register_routes(app: FastAPI):
|
||||
request: ReflectRequest,
|
||||
http_request: Request,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("reflect")),
|
||||
_precheck: None = Depends(precheck_for(PrecheckOperation.REFLECT)),
|
||||
):
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
@@ -4101,11 +4336,17 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
async def api_stats(
|
||||
bank_id: str,
|
||||
refresh: bool = Query(
|
||||
default=False,
|
||||
description="Force a fresh recompute, bypassing the cached value (and refreshing the cache).",
|
||||
),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Get statistics about memory nodes and links for a memory bank."""
|
||||
try:
|
||||
stats = await app.state.memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
stats = await app.state.memory.get_bank_stats(
|
||||
bank_id, request_context=request_context, force_refresh=refresh
|
||||
)
|
||||
nodes_by_type = stats["node_counts"]
|
||||
links_by_type = stats["link_counts"]
|
||||
links_by_fact_type = stats["link_counts_by_fact_type"]
|
||||
@@ -4231,8 +4472,8 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
async def api_list_entities(
|
||||
bank_id: str,
|
||||
limit: int = Query(default=100, description="Maximum number of entities to return"),
|
||||
offset: int = Query(default=0, description="Offset for pagination"),
|
||||
limit: int = Query(default=100, ge=0, description="Maximum number of entities to return"),
|
||||
offset: int = Query(default=0, ge=0, description="Offset for pagination"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""List entities for a memory bank with pagination."""
|
||||
@@ -4267,7 +4508,7 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
async def api_entity_graph(
|
||||
bank_id: str,
|
||||
limit: int = Query(default=1000, description="Maximum number of co-occurrence edges to return"),
|
||||
limit: int = Query(default=1000, ge=0, description="Maximum number of co-occurrence edges to return"),
|
||||
min_count: int = Query(default=1, description="Minimum cooccurrence_count to include an edge"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
@@ -4490,7 +4731,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")),
|
||||
_precheck: None = Depends(precheck_for(PrecheckOperation.MENTAL_MODEL_CREATE)),
|
||||
):
|
||||
"""Create a mental model (async - returns operation_id)."""
|
||||
try:
|
||||
@@ -4539,7 +4780,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")),
|
||||
_precheck: None = Depends(precheck_for(PrecheckOperation.MENTAL_MODEL_REFRESH)),
|
||||
):
|
||||
"""Refresh a mental model by re-running its source query (async)."""
|
||||
try:
|
||||
@@ -4685,6 +4926,333 @@ def _register_routes(app: FastAPI):
|
||||
logger.error(f"Error in DELETE /v1/default/banks/{bank_id}/mental-models/{mental_model_id}: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
# =========================================================================
|
||||
# KNOWLEDGE BASE ENDPOINTS (folders + pages, Open Knowledge Format)
|
||||
# =========================================================================
|
||||
# A hierarchy of folders and pages over mental models. Pages project to OKF
|
||||
# documents (markdown body + YAML frontmatter); see api/okf.py. The static
|
||||
# sub-paths (/tree, /folders, /pages, /graph, /export) are declared before
|
||||
# the /pages/{id} and /nodes/{id} path-parameter routes so they win.
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/tree",
|
||||
response_model=KnowledgeTreeResponse,
|
||||
summary="Get the knowledge-base tree",
|
||||
description="Return the knowledge base as a nested tree of folders and pages.",
|
||||
operation_id="get_knowledge_base_tree",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_knowledge_base_tree(
|
||||
bank_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Return the folder/page tree for a bank."""
|
||||
try:
|
||||
nodes = await app.state.memory.list_knowledge_nodes(bank_id=bank_id, request_context=request_context)
|
||||
return KnowledgeTreeResponse(roots=_build_knowledge_tree(nodes))
|
||||
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 GET /v1/default/banks/{bank_id}/knowledge-base/tree: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/folders",
|
||||
response_model=KnowledgeNode,
|
||||
status_code=201,
|
||||
summary="Create a knowledge-base folder",
|
||||
description="Create a folder, optionally nested under a parent folder.",
|
||||
operation_id="create_knowledge_folder",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_create_knowledge_folder(
|
||||
bank_id: str,
|
||||
body: CreateFolderRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Create a folder node."""
|
||||
try:
|
||||
node = await app.state.memory.create_knowledge_folder(
|
||||
bank_id=bank_id,
|
||||
name=body.name,
|
||||
parent_id=body.parent_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
return _knowledge_node_model(node)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
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}/knowledge-base/folders: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/pages",
|
||||
response_model=CreateKnowledgePageResponse,
|
||||
status_code=201,
|
||||
summary="Create a knowledge-base page",
|
||||
description="Create a page (a mental model + tree node). Content is generated asynchronously; "
|
||||
"use the returned operation_id to track completion.",
|
||||
operation_id="create_knowledge_page",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_create_knowledge_page(
|
||||
bank_id: str,
|
||||
body: CreatePageRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Create a page node (async content generation)."""
|
||||
try:
|
||||
node = await app.state.memory.create_knowledge_page(
|
||||
bank_id=bank_id,
|
||||
name=body.name,
|
||||
source_query=body.source_query,
|
||||
content="Generating content...",
|
||||
parent_id=body.parent_id,
|
||||
tags=body.tags if body.tags else None,
|
||||
max_tokens=body.max_tokens,
|
||||
trigger=body.trigger.model_dump() if body.trigger else None,
|
||||
request_context=request_context,
|
||||
)
|
||||
if node is None:
|
||||
raise HTTPException(status_code=409, detail=f"A page named '{body.name}' already exists in this folder")
|
||||
result = await app.state.memory.submit_async_refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=node["mental_model_id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
return CreateKnowledgePageResponse(
|
||||
page_id=node["id"],
|
||||
mental_model_id=node["mental_model_id"],
|
||||
operation_id=result["operation_id"],
|
||||
)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
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}/knowledge-base/pages: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/graph",
|
||||
response_model=KnowledgePageGraphResponse,
|
||||
summary="Knowledge-base constellation graph",
|
||||
description="Return pages as nodes linked by shared tags, for the constellation view.",
|
||||
operation_id="get_knowledge_base_graph",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_knowledge_base_graph(
|
||||
bank_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Return the shared-tag constellation graph for a bank's pages."""
|
||||
try:
|
||||
nodes = await app.state.memory.list_knowledge_nodes(bank_id=bank_id, request_context=request_context)
|
||||
pages = [n for n in nodes if n.get("kind") == "page"]
|
||||
# Cluster the constellation by parent folder (the knowledge base's own
|
||||
# structure) rather than by the retired type: tag.
|
||||
folder_names = {n["id"]: n["name"] for n in nodes if n.get("kind") == "folder"}
|
||||
graph = okf.knowledge_graph(pages, cluster_for=lambda p: folder_names.get(p.get("parent_id"), "Ungrouped"))
|
||||
return KnowledgePageGraphResponse(
|
||||
nodes=graph.nodes,
|
||||
edges=graph.edges,
|
||||
total_pages=len(graph.nodes),
|
||||
total_edges=len(graph.edges),
|
||||
)
|
||||
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 GET /v1/default/banks/{bank_id}/knowledge-base/graph: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/export",
|
||||
response_model=KnowledgePageBundleResponse,
|
||||
summary="Export the knowledge base as an OKF bundle",
|
||||
description="Return a portable OKF bundle: a nested index.md, one <id>.md per page, and history logs.",
|
||||
operation_id="export_knowledge_base",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_export_knowledge_base(
|
||||
bank_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Export a bank's knowledge base as a flat OKF markdown bundle."""
|
||||
try:
|
||||
nodes = await app.state.memory.list_knowledge_nodes(bank_id=bank_id, request_context=request_context)
|
||||
files = [KnowledgePageBundleFile(path=okf.INDEX_FILENAME, content=okf.render_index(nodes))]
|
||||
for node in nodes:
|
||||
if node.get("kind") != "page":
|
||||
continue
|
||||
page = await app.state.memory.get_knowledge_page(
|
||||
bank_id=bank_id, page_id=node["id"], request_context=request_context
|
||||
)
|
||||
if page is None:
|
||||
continue
|
||||
files.append(
|
||||
KnowledgePageBundleFile(path=okf.page_filename(node["id"]), content=okf.render_document(page))
|
||||
)
|
||||
if node.get("mental_model_id"):
|
||||
history = (
|
||||
await app.state.memory.get_mental_model_history(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=node["mental_model_id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
or []
|
||||
)
|
||||
if history:
|
||||
files.append(
|
||||
KnowledgePageBundleFile(
|
||||
path=okf.log_filename(node["id"]), content=okf.render_log(page, history)
|
||||
)
|
||||
)
|
||||
return KnowledgePageBundleResponse(files=files)
|
||||
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 GET /v1/default/banks/{bank_id}/knowledge-base/export: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/pages/{page_id}",
|
||||
response_model=KnowledgePageResponse,
|
||||
summary="Get a knowledge-base page",
|
||||
description="Return a single page as an OKF document (frontmatter + markdown body).",
|
||||
operation_id="get_knowledge_page",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_get_knowledge_page(
|
||||
bank_id: str,
|
||||
page_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Get a single page as an OKF document."""
|
||||
try:
|
||||
node = await app.state.memory.get_knowledge_page(
|
||||
bank_id=bank_id, page_id=page_id, request_context=request_context
|
||||
)
|
||||
if node is None:
|
||||
raise HTTPException(status_code=404, detail=f"Knowledge page '{page_id}' not found")
|
||||
return _knowledge_page_response(node)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in GET /v1/default/banks/{bank_id}/knowledge-base/pages/{page_id}: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.patch(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}",
|
||||
response_model=KnowledgeNode,
|
||||
summary="Rename or move a knowledge-base node",
|
||||
description="Rename a node (set `name`) and/or move it under another folder (set `parent_id`, "
|
||||
"null for the root).",
|
||||
operation_id="update_knowledge_node",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_update_knowledge_node(
|
||||
bank_id: str,
|
||||
node_id: str,
|
||||
body: UpdateNodeRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Rename and/or move a node."""
|
||||
try:
|
||||
updated: dict[str, Any] | None = None
|
||||
did_change = False
|
||||
if body.name is not None:
|
||||
did_change = True
|
||||
updated = await app.state.memory.rename_knowledge_node(
|
||||
bank_id=bank_id, node_id=node_id, name=body.name, request_context=request_context
|
||||
)
|
||||
# parent_id is applied only when present in the body, so passing null
|
||||
# moves the node to the root (distinct from "not provided").
|
||||
if "parent_id" in body.model_fields_set:
|
||||
did_change = True
|
||||
updated = await app.state.memory.move_knowledge_node(
|
||||
bank_id=bank_id, node_id=node_id, new_parent_id=body.parent_id, request_context=request_context
|
||||
)
|
||||
if not did_change:
|
||||
raise HTTPException(status_code=400, detail="Provide name and/or parent_id to update")
|
||||
if updated is None:
|
||||
raise HTTPException(status_code=404, detail=f"Knowledge node '{node_id}' not found")
|
||||
return _knowledge_node_model(updated)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in PATCH /v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.delete(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}",
|
||||
summary="Delete a knowledge-base node",
|
||||
description="Delete a folder or page and its whole subtree (pages' mental models are removed too).",
|
||||
operation_id="delete_knowledge_node",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_delete_knowledge_node(
|
||||
bank_id: str,
|
||||
node_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Delete a node and its subtree."""
|
||||
try:
|
||||
deleted = await app.state.memory.delete_knowledge_node(
|
||||
bank_id=bank_id, node_id=node_id, request_context=request_context
|
||||
)
|
||||
if not deleted:
|
||||
raise HTTPException(status_code=404, detail=f"Knowledge node '{node_id}' not found")
|
||||
return {"status": "deleted"}
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in DELETE /v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
# =========================================================================
|
||||
# DIRECTIVES ENDPOINTS
|
||||
# =========================================================================
|
||||
@@ -4894,8 +5462,8 @@ def _register_routes(app: FastAPI):
|
||||
tags_match: str = Query(
|
||||
"any_strict", description="How to match tags: 'any', 'all', 'any_strict', 'all_strict'"
|
||||
),
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
limit: int = Query(default=100, ge=0),
|
||||
offset: int = Query(default=0, ge=0),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""
|
||||
@@ -5079,8 +5647,8 @@ def _register_routes(app: FastAPI):
|
||||
default="memories",
|
||||
description="Where to read tags from: 'memories' (memory_units, default) or 'mental_models'.",
|
||||
),
|
||||
limit: int = Query(default=100, description="Maximum number of tags to return"),
|
||||
offset: int = Query(default=0, description="Offset for pagination"),
|
||||
limit: int = Query(default=100, ge=0, description="Maximum number of tags to return"),
|
||||
offset: int = Query(default=0, ge=0, description="Offset for pagination"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""
|
||||
@@ -5608,8 +6176,13 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""Partially update an agent's profile (name, mission, disposition)."""
|
||||
try:
|
||||
# Ensure bank exists
|
||||
await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
# PATCH is update-only; missing banks must not be created as a
|
||||
# side effect of reading the profile.
|
||||
existing_profile = await app.state.memory.get_bank_profile(
|
||||
bank_id, request_context=request_context, create_if_missing=False
|
||||
)
|
||||
if existing_profile is None:
|
||||
raise HTTPException(status_code=404, detail=f"Bank '{bank_id}' not found")
|
||||
|
||||
# Update name if provided (stored in DB for display only, deprecated)
|
||||
if request.name is not None:
|
||||
@@ -5625,7 +6198,11 @@ def _register_routes(app: FastAPI):
|
||||
await app.state.memory._config_resolver.update_bank_config(bank_id, config_updates, request_context)
|
||||
|
||||
# Get final profile
|
||||
final_profile = await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
final_profile = await app.state.memory.get_bank_profile(
|
||||
bank_id, request_context=request_context, create_if_missing=False
|
||||
)
|
||||
if final_profile is None:
|
||||
raise HTTPException(status_code=404, detail=f"Bank '{bank_id}' not found")
|
||||
disposition_dict = (
|
||||
final_profile["disposition"].model_dump()
|
||||
if hasattr(final_profile["disposition"], "model_dump")
|
||||
@@ -6116,9 +6693,11 @@ def _register_routes(app: FastAPI):
|
||||
# Authenticate and set schema context for multi-tenant DB queries
|
||||
await app.state.memory._authenticate_tenant(request_context)
|
||||
if app.state.memory._operation_validator:
|
||||
from hindsight_api.extensions import BankReadContext
|
||||
from hindsight_api.extensions import BankReadContext, BankReadOperation
|
||||
|
||||
ctx = BankReadContext(bank_id=bank_id, operation="get_bank_config", request_context=request_context)
|
||||
ctx = BankReadContext(
|
||||
bank_id=bank_id, operation=BankReadOperation.GET_BANK_CONFIG, request_context=request_context
|
||||
)
|
||||
await app.state.memory._validate_operation(
|
||||
app.state.memory._operation_validator.validate_bank_read(ctx)
|
||||
)
|
||||
@@ -6164,9 +6743,11 @@ def _register_routes(app: FastAPI):
|
||||
# Authenticate and set schema context for multi-tenant DB queries
|
||||
await app.state.memory._authenticate_tenant(request_context)
|
||||
if app.state.memory._operation_validator:
|
||||
from hindsight_api.extensions import BankWriteContext
|
||||
from hindsight_api.extensions import BankWriteContext, BankWriteOperation
|
||||
|
||||
ctx = BankWriteContext(bank_id=bank_id, operation="update_bank_config", request_context=request_context)
|
||||
ctx = BankWriteContext(
|
||||
bank_id=bank_id, operation=BankWriteOperation.UPDATE_BANK_CONFIG, request_context=request_context
|
||||
)
|
||||
await app.state.memory._validate_operation(
|
||||
app.state.memory._operation_validator.validate_bank_write(ctx)
|
||||
)
|
||||
@@ -6223,9 +6804,11 @@ def _register_routes(app: FastAPI):
|
||||
# Authenticate and set schema context for multi-tenant DB queries
|
||||
await app.state.memory._authenticate_tenant(request_context)
|
||||
if app.state.memory._operation_validator:
|
||||
from hindsight_api.extensions import BankWriteContext
|
||||
from hindsight_api.extensions import BankWriteContext, BankWriteOperation
|
||||
|
||||
ctx = BankWriteContext(bank_id=bank_id, operation="reset_bank_config", request_context=request_context)
|
||||
ctx = BankWriteContext(
|
||||
bank_id=bank_id, operation=BankWriteOperation.RESET_BANK_CONFIG, request_context=request_context
|
||||
)
|
||||
await app.state.memory._validate_operation(
|
||||
app.state.memory._operation_validator.validate_bank_write(ctx)
|
||||
)
|
||||
@@ -6596,7 +7179,7 @@ def _register_routes(app: FastAPI):
|
||||
bank_id: str,
|
||||
request: RetainRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("retain")),
|
||||
_precheck: None = Depends(precheck_for(PrecheckOperation.RETAIN)),
|
||||
):
|
||||
"""Retain memories with optional async processing."""
|
||||
metrics = get_metrics_collector()
|
||||
@@ -6749,7 +7332,7 @@ def _register_routes(app: FastAPI):
|
||||
description="Upload files (PDF, DOCX, etc.), convert them to markdown, and retain as memories.\n\n"
|
||||
"This endpoint handles file upload, conversion, and memory creation in a single operation.\n\n"
|
||||
"**Features:**\n"
|
||||
"- Supports PDF, DOCX, PPTX, XLSX, images (with OCR), audio (with transcription)\n"
|
||||
"- Supports PDF, DOCX, PPTX, XLSX, images (parser-dependent OCR), audio (with transcription)\n"
|
||||
"- Automatic file-to-markdown conversion using pluggable parsers\n"
|
||||
"- Files stored in object storage (PostgreSQL by default, S3 for production)\n"
|
||||
"- Each file becomes a separate document with optional metadata/tags\n"
|
||||
@@ -6779,7 +7362,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")),
|
||||
_precheck: None = Depends(precheck_for(PrecheckOperation.FILES_RETAIN)),
|
||||
):
|
||||
"""Upload and convert files to memories."""
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
@@ -9,7 +9,7 @@ from fastmcp import FastMCP
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api import __version__ as HINDSIGHT_VERSION
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.config import DEFAULT_MCP_RECALL_DESCRIPTION, DEFAULT_MCP_RETAIN_DESCRIPTION, _get_raw_config
|
||||
from hindsight_api.engine.memory_engine import _current_schema
|
||||
from hindsight_api.extensions import MCPExtension, load_extension
|
||||
from hindsight_api.extensions.tenant import AuthenticationError
|
||||
@@ -78,6 +78,19 @@ def get_current_mcp_authenticated() -> bool:
|
||||
return _current_mcp_authenticated.get()
|
||||
|
||||
|
||||
def _build_mcp_tool_descriptions(extra_instructions: str | None) -> tuple[str | None, str | None]:
|
||||
"""Return custom retain/recall descriptions when server-level MCP instructions are set."""
|
||||
if not isinstance(extra_instructions, str):
|
||||
return None, None
|
||||
|
||||
extra_instructions = extra_instructions.strip()
|
||||
if not extra_instructions:
|
||||
return None, None
|
||||
|
||||
suffix = f"\n\nAdditional instructions: {extra_instructions}"
|
||||
return DEFAULT_MCP_RETAIN_DESCRIPTION + suffix, DEFAULT_MCP_RECALL_DESCRIPTION + suffix
|
||||
|
||||
|
||||
def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"""
|
||||
Create and configure the Hindsight MCP server.
|
||||
@@ -135,6 +148,10 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
allowed = frozenset(global_config.mcp_enabled_tools)
|
||||
base_tools = (base_tools if base_tools is not None else _ALL_TOOLS) & allowed
|
||||
|
||||
retain_description, recall_description = _build_mcp_tool_descriptions(
|
||||
getattr(global_config, "mcp_instructions", None)
|
||||
)
|
||||
|
||||
# Configure and register tools using shared module
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=get_current_bank_id,
|
||||
@@ -144,6 +161,8 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
mcp_authenticated_resolver=get_current_mcp_authenticated, # Propagate MCP pre-auth flag
|
||||
include_bank_id_param=multi_bank,
|
||||
tools=base_tools,
|
||||
retain_description=retain_description,
|
||||
recall_description=recall_description,
|
||||
)
|
||||
|
||||
register_mcp_tools(mcp, memory, config)
|
||||
|
||||
@@ -0,0 +1,263 @@
|
||||
"""Open Knowledge Format (OKF) projection for knowledge pages.
|
||||
|
||||
Knowledge pages are a *read-only* OKF view over the existing mental models: each
|
||||
mental model is projected into an OKF document — a markdown body with YAML
|
||||
frontmatter (``type`` required; ``title``/``description``/``tags``/``timestamp``
|
||||
optional) — and pages are linked into a constellation graph via shared tags.
|
||||
|
||||
See the Open Knowledge Format spec:
|
||||
https://github.com/GoogleCloudPlatform/knowledge-catalog/tree/main/okf
|
||||
|
||||
This module is intentionally pure: every function transforms the mental-model
|
||||
dicts returned by ``MemoryEngine.list_mental_models`` / ``get_mental_model`` and
|
||||
never touches the database. That keeps the OKF contract unit-testable without a
|
||||
DB or LLM and lets the HTTP layer stay a thin wrapper.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
# OKF requires exactly one frontmatter field — ``type``. We default to this when
|
||||
# a page does not declare one via a ``type:<x>`` tag.
|
||||
DEFAULT_PAGE_TYPE = "knowledge-page"
|
||||
|
||||
# A page declares its OKF ``type`` through a tag of the form ``type:runbook``.
|
||||
# This keeps the projection schema-free (no new mental_models column): the type
|
||||
# is lifted from the existing tags array.
|
||||
TYPE_TAG_PREFIX = "type:"
|
||||
|
||||
INDEX_FILENAME = "index.md"
|
||||
|
||||
# Deterministic, colour-blind-friendly palette. Type → colour is stable across
|
||||
# requests so the constellation keeps the same colours between reloads.
|
||||
_PALETTE = (
|
||||
"#0074d9", # blue
|
||||
"#2ecc40", # green
|
||||
"#b10dc9", # purple
|
||||
"#ff851b", # orange
|
||||
"#39cccc", # teal
|
||||
"#f012be", # magenta
|
||||
"#3d9970", # olive
|
||||
"#ff4136", # red
|
||||
)
|
||||
|
||||
_EDGE_COLOR = "#9aa5b1"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PageType:
|
||||
"""A page's OKF ``type`` and the tags that remain after the type tag is split off."""
|
||||
|
||||
type: str
|
||||
display_tags: list[str]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class KnowledgeGraph:
|
||||
"""Cytoscape-style node/edge graph of knowledge pages linked by shared tags."""
|
||||
|
||||
nodes: list[dict[str, Any]] = field(default_factory=list)
|
||||
edges: list[dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
|
||||
def _color_for(key: str) -> str:
|
||||
"""Stable colour for a string key (FNV-ish hash into the fixed palette)."""
|
||||
h = 0
|
||||
for ch in key:
|
||||
h = (h * 31 + ord(ch)) & 0xFFFFFFFF
|
||||
return _PALETTE[h % len(_PALETTE)]
|
||||
|
||||
|
||||
def _scalar(value: Any) -> str:
|
||||
"""Emit a YAML-safe double-quoted scalar.
|
||||
|
||||
We always double-quote so arbitrary page names / source queries can't be
|
||||
misread as YAML special forms (``true``, ``2026-01-01``, ``- x``, etc.).
|
||||
"""
|
||||
text = str(value)
|
||||
escaped = text.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n").replace("\r", "")
|
||||
return f'"{escaped}"'
|
||||
|
||||
|
||||
def page_type(tags: list[str] | None) -> PageType:
|
||||
"""Split an OKF ``type`` out of the tag list.
|
||||
|
||||
The first ``type:<x>`` tag wins; all ``type:`` tags are removed from the
|
||||
returned ``display_tags`` so they don't pollute the constellation's
|
||||
shared-tag edges. Falls back to :data:`DEFAULT_PAGE_TYPE`.
|
||||
"""
|
||||
resolved = DEFAULT_PAGE_TYPE
|
||||
display: list[str] = []
|
||||
for tag in tags or []:
|
||||
if tag.startswith(TYPE_TAG_PREFIX):
|
||||
suffix = tag[len(TYPE_TAG_PREFIX) :].strip()
|
||||
if suffix and resolved == DEFAULT_PAGE_TYPE:
|
||||
resolved = suffix
|
||||
continue
|
||||
display.append(tag)
|
||||
return PageType(type=resolved, display_tags=display)
|
||||
|
||||
|
||||
def _timestamp(mm: dict[str, Any]) -> str | None:
|
||||
return mm.get("last_refreshed_at") or mm.get("created_at")
|
||||
|
||||
|
||||
def frontmatter(mm: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Build the ordered OKF frontmatter mapping for a mental model.
|
||||
|
||||
``None``/empty values are dropped by :func:`render_frontmatter`.
|
||||
"""
|
||||
pt = page_type(mm.get("tags"))
|
||||
return {
|
||||
"id": mm.get("id"),
|
||||
"type": pt.type,
|
||||
"title": mm.get("name"),
|
||||
"description": mm.get("source_query"),
|
||||
"tags": pt.display_tags,
|
||||
"timestamp": _timestamp(mm),
|
||||
}
|
||||
|
||||
|
||||
def render_frontmatter(fm: dict[str, Any]) -> str:
|
||||
"""Render a frontmatter mapping into a ``---`` fenced YAML block."""
|
||||
lines = ["---"]
|
||||
for key, value in fm.items():
|
||||
if value is None:
|
||||
continue
|
||||
if isinstance(value, list):
|
||||
if not value:
|
||||
continue
|
||||
lines.append(f"{key}:")
|
||||
lines.extend(f" - {_scalar(item)}" for item in value)
|
||||
else:
|
||||
lines.append(f"{key}: {_scalar(value)}")
|
||||
lines.append("---")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def render_document(mm: dict[str, Any]) -> str:
|
||||
"""Render a full OKF document: frontmatter block + markdown body."""
|
||||
body = (mm.get("content") or "").strip()
|
||||
return f"{render_frontmatter(frontmatter(mm))}\n\n{body}\n" if body else f"{render_frontmatter(frontmatter(mm))}\n"
|
||||
|
||||
|
||||
def page_filename(page_id: str) -> str:
|
||||
"""OKF bundle filename for a page id."""
|
||||
return f"{page_id}.md"
|
||||
|
||||
|
||||
def log_filename(page_id: str) -> str:
|
||||
"""OKF reserved per-page history filename."""
|
||||
return f"{page_id}.log.md"
|
||||
|
||||
|
||||
def render_index(nodes: list[dict[str, Any]]) -> str:
|
||||
"""Render the reserved ``index.md`` — nested OKF navigation over the tree.
|
||||
|
||||
``nodes`` is the flat folder/page list (each with ``id``, ``kind``, ``name``,
|
||||
``parent_id``); folders nest their children, pages link to their ``.md``.
|
||||
"""
|
||||
fm = render_frontmatter({"type": "index", "title": "Knowledge base"})
|
||||
lines = [fm, "", "# Knowledge base", ""]
|
||||
|
||||
children: dict[Any, list[dict[str, Any]]] = {}
|
||||
for node in nodes:
|
||||
children.setdefault(node.get("parent_id"), []).append(node)
|
||||
|
||||
def walk(parent: Any, depth: int) -> None:
|
||||
ordered = sorted(children.get(parent, []), key=lambda n: (n.get("sort_order", 0), n.get("name") or ""))
|
||||
for node in ordered:
|
||||
indent = " " * depth
|
||||
if node.get("kind") == "folder":
|
||||
lines.append(f"{indent}- **{node['name']}/**")
|
||||
walk(node["id"], depth + 1)
|
||||
else:
|
||||
description = node.get("source_query") or node.get("description")
|
||||
link = f"{indent}- [{node['name']}](./{page_filename(node['id'])})"
|
||||
lines.append(f"{link} — {description}" if description else link)
|
||||
|
||||
walk(None, 0)
|
||||
if len(lines) == 4:
|
||||
lines.append("_No knowledge pages yet._")
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def render_log(mm: dict[str, Any], history: list[dict[str, Any]]) -> str:
|
||||
"""Render the reserved per-page ``log.md`` from refresh history.
|
||||
|
||||
Each history entry is ``{previous_content, previous_reflect_response,
|
||||
changed_at}`` (newest first), capturing the content *before* a refresh.
|
||||
"""
|
||||
name = mm.get("name") or mm.get("id")
|
||||
fm = render_frontmatter({"type": "log", "title": f"{name} — history"})
|
||||
lines = [fm, "", f"# {name} — history", ""]
|
||||
if not history:
|
||||
lines.append("_No refresh history._")
|
||||
return "\n".join(lines) + "\n"
|
||||
for entry in history:
|
||||
changed_at = entry.get("changed_at") or "unknown"
|
||||
previous = (entry.get("previous_content") or "").strip()
|
||||
lines.append(f"## {changed_at}")
|
||||
lines.append("")
|
||||
lines.append(previous if previous else "_(empty)_")
|
||||
lines.append("")
|
||||
return "\n".join(lines).rstrip() + "\n"
|
||||
|
||||
|
||||
def knowledge_graph(
|
||||
pages: list[dict[str, Any]],
|
||||
cluster_for: "Callable[[dict[str, Any]], str] | None" = None,
|
||||
) -> KnowledgeGraph:
|
||||
"""Derive the constellation graph: pages as nodes, shared tags as edges.
|
||||
|
||||
Two pages are linked when they share at least one (non-``type:``) tag; the
|
||||
edge weight is the number of shared tags. Each node's cluster (``type`` field
|
||||
+ colour) comes from ``cluster_for(page)`` — the knowledge base groups by
|
||||
parent folder; the default groups by OKF ``type``.
|
||||
"""
|
||||
nodes: list[dict[str, Any]] = []
|
||||
tag_sets: list[tuple[str, frozenset[str]]] = []
|
||||
for mm in pages:
|
||||
page_id = mm["id"]
|
||||
pt = page_type(mm.get("tags"))
|
||||
cluster = cluster_for(mm) if cluster_for else pt.type
|
||||
tag_sets.append((page_id, frozenset(pt.display_tags)))
|
||||
nodes.append(
|
||||
{
|
||||
"data": {
|
||||
"id": page_id,
|
||||
"label": mm.get("name") or page_id,
|
||||
"type": cluster,
|
||||
"tagCount": len(pt.display_tags),
|
||||
"color": _color_for(cluster),
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
edges: list[dict[str, Any]] = []
|
||||
for i in range(len(tag_sets)):
|
||||
source_id, source_tags = tag_sets[i]
|
||||
if not source_tags:
|
||||
continue
|
||||
for j in range(i + 1, len(tag_sets)):
|
||||
target_id, target_tags = tag_sets[j]
|
||||
shared = source_tags & target_tags
|
||||
if not shared:
|
||||
continue
|
||||
edges.append(
|
||||
{
|
||||
"data": {
|
||||
"id": f"{source_id}--{target_id}",
|
||||
"source": source_id,
|
||||
"target": target_id,
|
||||
"sharedTags": sorted(shared),
|
||||
"weight": len(shared),
|
||||
"color": _EDGE_COLOR,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
return KnowledgeGraph(nodes=nodes, edges=edges)
|
||||
@@ -142,11 +142,35 @@ ENV_LLM_REASONING_EFFORT = "HINDSIGHT_API_LLM_REASONING_EFFORT"
|
||||
ENV_LLM_GROQ_SERVICE_TIER = "HINDSIGHT_API_LLM_GROQ_SERVICE_TIER"
|
||||
ENV_LLM_OPENAI_SERVICE_TIER = "HINDSIGHT_API_LLM_OPENAI_SERVICE_TIER"
|
||||
ENV_LLM_BEDROCK_SERVICE_TIER = "HINDSIGHT_API_LLM_BEDROCK_SERVICE_TIER"
|
||||
ENV_LLM_GEMINI_SERVICE_TIER = "HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER"
|
||||
ENV_LLM_EXTRA_BODY = "HINDSIGHT_API_LLM_EXTRA_BODY"
|
||||
ENV_LLM_DEFAULT_HEADERS = "HINDSIGHT_API_LLM_DEFAULT_HEADERS"
|
||||
ENV_LLM_STRICT_SCHEMA = "HINDSIGHT_API_LLM_STRICT_SCHEMA"
|
||||
ENV_LLM_SEND_BANK_AS_USER = "HINDSIGHT_API_LLM_SEND_BANK_AS_USER"
|
||||
|
||||
# Per-operation sampling temperature. Each internal LLM call uses a temperature
|
||||
# tuned for its task (deterministic extraction vs. creative reflection). These
|
||||
# expose those as overridable knobs. Resolution per operation:
|
||||
# per-operation env -> global env (ENV_LLM_TEMPERATURE) -> built-in default.
|
||||
# A value of "none"/"default"/"" (or "off") omits the temperature parameter
|
||||
# entirely, for models that reject explicit temperatures (e.g. Azure GPT-5.5,
|
||||
# which only accepts the default value) -- see issue #2459.
|
||||
ENV_LLM_TEMPERATURE = "HINDSIGHT_API_LLM_TEMPERATURE"
|
||||
ENV_LLM_TEMPERATURE_VERIFICATION = "HINDSIGHT_API_LLM_TEMPERATURE_VERIFICATION"
|
||||
ENV_LLM_TEMPERATURE_RETAIN = "HINDSIGHT_API_LLM_TEMPERATURE_RETAIN"
|
||||
ENV_LLM_TEMPERATURE_REFLECT = "HINDSIGHT_API_LLM_TEMPERATURE_REFLECT"
|
||||
ENV_LLM_TEMPERATURE_CONSOLIDATION = "HINDSIGHT_API_LLM_TEMPERATURE_CONSOLIDATION"
|
||||
|
||||
# Multi-LLM strategy. Extra LLMs are configured by index alongside the unindexed
|
||||
# primary (e.g. HINDSIGHT_API_LLM_1_PROVIDER, HINDSIGHT_API_LLM_2_PROVIDER, ...),
|
||||
# and HINDSIGHT_API_LLM_STRATEGY (JSON) selects how to route across them — see
|
||||
# _parse_llm_members / _parse_llm_strategy below. Each operation can override the
|
||||
# global chain with its own HINDSIGHT_API_<OP>_LLM_<n>_* members + _STRATEGY.
|
||||
ENV_LLM_STRATEGY = "HINDSIGHT_API_LLM_STRATEGY"
|
||||
ENV_RETAIN_LLM_STRATEGY = "HINDSIGHT_API_RETAIN_LLM_STRATEGY"
|
||||
ENV_REFLECT_LLM_STRATEGY = "HINDSIGHT_API_REFLECT_LLM_STRATEGY"
|
||||
ENV_CONSOLIDATION_LLM_STRATEGY = "HINDSIGHT_API_CONSOLIDATION_LLM_STRATEGY"
|
||||
|
||||
# LiteLLM Router chain — provider-specific config consumed by the "litellmrouter"
|
||||
# provider. Each entry is a deployment; the Router tries them in declared order and
|
||||
# falls back to the next on transient errors (5xx, rate-limit, timeout).
|
||||
@@ -155,15 +179,75 @@ ENV_LLM_SEND_BANK_AS_USER = "HINDSIGHT_API_LLM_SEND_BANK_AS_USER"
|
||||
# disambiguates from the embeddings/reranker LITELLM_* settings.
|
||||
ENV_LLM_LITELLMROUTER_CONFIG = "HINDSIGHT_API_LLM_LITELLMROUTER_CONFIG"
|
||||
|
||||
# Per-operation temperature defaults (preserve historical hardcoded values).
|
||||
DEFAULT_LLM_TEMPERATURE_VERIFICATION = 0.0 # connection check
|
||||
DEFAULT_LLM_TEMPERATURE_RETAIN = 0.1 # fact extraction
|
||||
DEFAULT_LLM_TEMPERATURE_REFLECT = 0.9 # reflect "thinking"
|
||||
DEFAULT_LLM_TEMPERATURE_CONSOLIDATION = 0.0 # mental-model delta / dedup
|
||||
|
||||
# Defaults for service tiers
|
||||
DEFAULT_LLM_GROQ_SERVICE_TIER = "auto" # "on_demand", "flex", or "auto"
|
||||
DEFAULT_LLM_OPENAI_SERVICE_TIER = None # None (default) or "flex" (50% cheaper)
|
||||
DEFAULT_LLM_BEDROCK_SERVICE_TIER = None # None (default), "flex", "priority", or "reserved"
|
||||
DEFAULT_LLM_GEMINI_SERVICE_TIER = None # None (default) or "flex" (50% cheaper best-effort tier)
|
||||
DEFAULT_LLM_EXTRA_BODY = None # None = no extra body params; JSON dict merged into OpenAI extra_body
|
||||
DEFAULT_LLM_DEFAULT_HEADERS = (
|
||||
None # None = no extra headers; JSON dict passed as default_headers to provider SDK clients
|
||||
)
|
||||
|
||||
|
||||
def parse_gemini_service_tier(value: str | None) -> str | None:
|
||||
"""Normalize and validate the Gemini service tier."""
|
||||
tier = value or None
|
||||
valid_tiers = (None, "flex")
|
||||
if tier not in valid_tiers:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER: "
|
||||
f"{tier!r}. Must be one of: {', '.join(t for t in valid_tiers if t is not None)}."
|
||||
)
|
||||
return tier
|
||||
|
||||
|
||||
# Sentinel strings that, as a temperature value, mean "omit the temperature
|
||||
# parameter entirely" rather than a numeric setting.
|
||||
_TEMPERATURE_OMIT_VALUES = frozenset({"", "none", "default", "off", "unset"})
|
||||
|
||||
|
||||
def _parse_temperature(raw: str) -> float | None:
|
||||
"""Parse a raw temperature env value into a float, or None to omit it.
|
||||
|
||||
Returns None for the omit sentinels (so the temperature parameter is dropped
|
||||
from the LLM call); otherwise parses a float and validates the 0.0-2.0 range.
|
||||
"""
|
||||
if raw.strip().lower() in _TEMPERATURE_OMIT_VALUES:
|
||||
return None
|
||||
try:
|
||||
value = float(raw)
|
||||
except ValueError as e:
|
||||
raise ValueError(
|
||||
f"Invalid LLM temperature {raw!r}: must be a number in [0.0, 2.0] "
|
||||
f"or one of {sorted(_TEMPERATURE_OMIT_VALUES)} to omit it."
|
||||
) from e
|
||||
if not 0.0 <= value <= 2.0:
|
||||
raise ValueError(f"Invalid LLM temperature {value}: must be in [0.0, 2.0].")
|
||||
return value
|
||||
|
||||
|
||||
def _resolve_operation_temperature(operation_env: str, default: float) -> float | None:
|
||||
"""Resolve a per-operation temperature: per-op env -> global env -> default.
|
||||
|
||||
The omit sentinels resolve to None at any layer, so a single
|
||||
``HINDSIGHT_API_LLM_TEMPERATURE=none`` drops temperature from every operation
|
||||
that has no explicit per-operation override.
|
||||
"""
|
||||
raw = os.getenv(operation_env)
|
||||
if raw is None:
|
||||
raw = os.getenv(ENV_LLM_TEMPERATURE)
|
||||
if raw is None:
|
||||
return default
|
||||
return _parse_temperature(raw)
|
||||
|
||||
|
||||
# Per-operation LLM configuration (optional, falls back to global LLM config)
|
||||
ENV_RETAIN_LLM_PROVIDER = "HINDSIGHT_API_RETAIN_LLM_PROVIDER"
|
||||
ENV_RETAIN_LLM_API_KEY = "HINDSIGHT_API_RETAIN_LLM_API_KEY"
|
||||
@@ -257,6 +341,11 @@ ENV_RERANKER_OPENROUTER_API_KEY = "HINDSIGHT_API_RERANKER_OPENROUTER_API_KEY"
|
||||
ENV_RERANKER_OPENROUTER_MODEL = "HINDSIGHT_API_RERANKER_OPENROUTER_MODEL"
|
||||
ENV_RERANKER_OPENROUTER_BASE_URL = "HINDSIGHT_API_RERANKER_OPENROUTER_BASE_URL"
|
||||
|
||||
# Requesty configuration (OpenAI-compatible gateway; embeddings)
|
||||
ENV_REQUESTY_API_KEY = "HINDSIGHT_API_REQUESTY_API_KEY"
|
||||
ENV_EMBEDDINGS_REQUESTY_API_KEY = "HINDSIGHT_API_EMBEDDINGS_REQUESTY_API_KEY"
|
||||
ENV_EMBEDDINGS_REQUESTY_MODEL = "HINDSIGHT_API_EMBEDDINGS_REQUESTY_MODEL"
|
||||
|
||||
# ZeroEntropy configuration (embeddings)
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_API_KEY = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_API_KEY"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_MODEL = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_MODEL"
|
||||
@@ -354,6 +443,7 @@ 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"
|
||||
ENV_MCP_INSTRUCTIONS = "HINDSIGHT_API_MCP_INSTRUCTIONS"
|
||||
ENV_ENABLE_BANK_CONFIG_API = "HINDSIGHT_API_ENABLE_BANK_CONFIG_API"
|
||||
ENV_ENABLE_BANK_LLM_HEALTH = "HINDSIGHT_API_ENABLE_BANK_LLM_HEALTH"
|
||||
ENV_ENABLE_DRY_RUN_EXTRACT = "HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT"
|
||||
@@ -375,6 +465,7 @@ ENV_OTEL_EXPORTER_OTLP_HEADERS = "HINDSIGHT_API_OTEL_EXPORTER_OTLP_HEADERS"
|
||||
ENV_OTEL_SERVICE_NAME = "HINDSIGHT_API_OTEL_SERVICE_NAME"
|
||||
ENV_OTEL_DEPLOYMENT_ENVIRONMENT = "HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT"
|
||||
ENV_METRICS_INCLUDE_BANK_ID = "HINDSIGHT_API_METRICS_INCLUDE_BANK_ID"
|
||||
ENV_METRICS_BACKLOG_ENABLED = "HINDSIGHT_API_METRICS_BACKLOG_ENABLED"
|
||||
|
||||
# Vertex AI configuration
|
||||
ENV_LLM_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID"
|
||||
@@ -424,6 +515,11 @@ ENV_FILE_STORAGE_AZURE_ACCOUNT_NAME = "HINDSIGHT_API_FILE_STORAGE_AZURE_ACCOUNT_
|
||||
ENV_FILE_STORAGE_AZURE_ACCOUNT_KEY = "HINDSIGHT_API_FILE_STORAGE_AZURE_ACCOUNT_KEY"
|
||||
ENV_FILE_PARSER = "HINDSIGHT_API_FILE_PARSER"
|
||||
ENV_FILE_PARSER_ALLOWLIST = "HINDSIGHT_API_FILE_PARSER_ALLOWLIST"
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_ENABLED = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED"
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_API_KEY = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY"
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_BASE_URL = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL"
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_MODEL = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL"
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_PROMPT = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT"
|
||||
ENV_FILE_PARSER_IRIS_TOKEN = "HINDSIGHT_API_FILE_PARSER_IRIS_TOKEN"
|
||||
ENV_FILE_PARSER_IRIS_ORG_ID = "HINDSIGHT_API_FILE_PARSER_IRIS_ORG_ID"
|
||||
ENV_FILE_PARSER_LLAMA_PARSE_API_KEY = "HINDSIGHT_API_FILE_PARSER_LLAMA_PARSE_API_KEY"
|
||||
@@ -477,7 +573,6 @@ ENV_LLAMACPP_EXTRA_ARGS = "HINDSIGHT_API_LLAMACPP_EXTRA_ARGS"
|
||||
|
||||
# Optimization flags
|
||||
ENV_SKIP_LLM_VERIFICATION = "HINDSIGHT_API_SKIP_LLM_VERIFICATION"
|
||||
ENV_LAZY_RERANKER = "HINDSIGHT_API_LAZY_RERANKER"
|
||||
|
||||
# Database migrations
|
||||
ENV_RUN_MIGRATIONS_ON_STARTUP = "HINDSIGHT_API_RUN_MIGRATIONS_ON_STARTUP"
|
||||
@@ -553,6 +648,14 @@ ENV_RECALL_MAX_CANDIDATES_PER_SOURCE = "HINDSIGHT_API_RECALL_MAX_CANDIDATES_PER_
|
||||
# Empty disables the feature.
|
||||
ENV_RECALL_STRATEGY_BOOSTS = "HINDSIGHT_API_RECALL_STRATEGY_BOOSTS"
|
||||
|
||||
# Recency decay used by recall reranking (engine/search/reranking.py). The decay
|
||||
# function maps a memory's age onto a freshness signal that nudges its final
|
||||
# ranking via a small multiplicative boost. "linear" (default) preserves the
|
||||
# historical behaviour; "exponential" decays by half-life; "none" disables it.
|
||||
ENV_RECENCY_DECAY_FUNCTION = "HINDSIGHT_API_RECENCY_DECAY_FUNCTION"
|
||||
ENV_RECENCY_DECAY_LINEAR_WINDOW_DAYS = "HINDSIGHT_API_RECENCY_DECAY_LINEAR_WINDOW_DAYS"
|
||||
ENV_RECENCY_DECAY_HALFLIFE_DAYS = "HINDSIGHT_API_RECENCY_DECAY_HALFLIFE_DAYS"
|
||||
|
||||
# Audit log settings
|
||||
ENV_AUDIT_LOG_ENABLED = "HINDSIGHT_API_AUDIT_LOG_ENABLED"
|
||||
ENV_AUDIT_LOG_ACTIONS = "HINDSIGHT_API_AUDIT_LOG_ACTIONS"
|
||||
@@ -566,6 +669,7 @@ ENV_LLM_TRACE_MAX_CHARS = "HINDSIGHT_API_LLM_TRACE_MAX_CHARS"
|
||||
|
||||
# Background maintenance settings
|
||||
ENV_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS = "HINDSIGHT_API_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS"
|
||||
ENV_MENTAL_MODEL_REFRESH_TICK_SECONDS = "HINDSIGHT_API_MENTAL_MODEL_REFRESH_TICK_SECONDS"
|
||||
|
||||
# Disposition settings
|
||||
ENV_DISPOSITION_SKEPTICISM = "HINDSIGHT_API_DISPOSITION_SKEPTICISM"
|
||||
@@ -588,6 +692,7 @@ PROVIDER_DEFAULT_MODELS = {
|
||||
"deepseek": "deepseek-v4-flash",
|
||||
"zai": "glm-4.5-flash",
|
||||
"opencode-go": "deepseek-v4-flash",
|
||||
"atlas": "deepseek-ai/deepseek-v4-pro",
|
||||
"ollama": "gemma3:12b",
|
||||
"ollama-cloud": "gemma3:12b",
|
||||
"llamacpp": "gemma-4-e2b-it",
|
||||
@@ -601,6 +706,7 @@ PROVIDER_DEFAULT_MODELS = {
|
||||
"bedrock": "us.amazon.nova-2-lite-v1:0",
|
||||
"volcano": "doubao-pro-32k",
|
||||
"openrouter": "qwen/qwen3.5-9b",
|
||||
"requesty": "openai/gpt-4o-mini",
|
||||
"fireworks": "accounts/fireworks/models/llama-v3p1-8b-instruct",
|
||||
"nous": "deepseek/deepseek-v4-flash",
|
||||
}
|
||||
@@ -691,6 +797,14 @@ DEFAULT_RECALL_MAX_CANDIDATES_PER_SOURCE = 0
|
||||
# "graph:high,semantic:low"). Empty disables the feature. See
|
||||
# ENV_RECALL_STRATEGY_BOOSTS for the full rationale.
|
||||
DEFAULT_RECALL_STRATEGY_BOOSTS = ""
|
||||
# Recency decay shape used by recall reranking. "linear" reproduces the
|
||||
# historical straight-line decay; defaults below keep behaviour unchanged.
|
||||
RECENCY_DECAY_FUNCTIONS = ("linear", "exponential", "none")
|
||||
DEFAULT_RECENCY_DECAY_FUNCTION = "linear"
|
||||
# Linear: days over which freshness decays from 1.0 to its 0.1 floor.
|
||||
DEFAULT_RECENCY_DECAY_LINEAR_WINDOW_DAYS = 365.0
|
||||
# Exponential: age (days) at which the recency signal is neutral (0.5).
|
||||
DEFAULT_RECENCY_DECAY_HALFLIFE_DAYS = 90.0
|
||||
# Retrieval arms that can be boosted; mirrors fusion.py source_names.
|
||||
RECALL_STRATEGY_NAMES = ("semantic", "bm25", "graph", "temporal")
|
||||
# User-facing priority levels. Kept in sync with recall_boost.BOOST_LEVELS by a
|
||||
@@ -749,6 +863,9 @@ DEFAULT_EMBEDDINGS_OPENROUTER_MODEL = "perplexity/pplx-embed-v1-0.6b"
|
||||
DEFAULT_RERANKER_OPENROUTER_MODEL = "cohere/rerank-v3.5"
|
||||
DEFAULT_RERANKER_OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1/rerank"
|
||||
|
||||
# Requesty defaults
|
||||
DEFAULT_EMBEDDINGS_REQUESTY_MODEL = "openai/text-embedding-3-small"
|
||||
|
||||
# ZeroEntropy defaults
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL = "zembed-1"
|
||||
# Shared between embeddings (zembed-1) and reranker (zerank-*) — the host is the same.
|
||||
@@ -804,6 +921,7 @@ 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)
|
||||
DEFAULT_MCP_INSTRUCTIONS = None
|
||||
DEFAULT_ENABLE_BANK_CONFIG_API = True
|
||||
# Dry-run extraction is a preview tool that makes a real LLM call but stores nothing. Enabled by
|
||||
# default; set HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT=false to remove the endpoint (e.g. to cap
|
||||
@@ -847,6 +965,10 @@ DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60 # Batch API polling interval in
|
||||
DEFAULT_FILE_STORAGE_TYPE = "native" # PostgreSQL BYTEA storage
|
||||
DEFAULT_FILE_PARSER = "markitdown" # Default parser fallback chain (comma-separated, e.g. "iris,markitdown")
|
||||
DEFAULT_FILE_PARSER_ALLOWLIST = None # Allowlist of parsers clients may request (None = all registered parsers)
|
||||
DEFAULT_FILE_PARSER_MARKITDOWN_OCR_ENABLED = False
|
||||
DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT = """You are a precise OCR transcription engine.
|
||||
|
||||
Transcribe only the visible text in the image. Do not describe the image, summarize it, translate it, infer missing content, or add commentary. Preserve the original language, wording, numbers, punctuation, capitalization, and reading order. Reconstruct headings, lists, key-value fields, stamps, and tables as clean Markdown when the layout is clear. If text is unreadable or uncertain, write [unclear] for that span. Return only the extracted Markdown."""
|
||||
DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE_MB = 100 # Max total batch size in MB (all files combined)
|
||||
DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE = 10 # Max files per batch upload
|
||||
DEFAULT_ENABLE_FILE_UPLOAD_API = True # Enable file upload endpoint
|
||||
@@ -965,6 +1087,7 @@ DEFAULT_OTEL_TRACES_ENABLED = False # Disabled by default for backward compatib
|
||||
DEFAULT_OTEL_SERVICE_NAME = "hindsight-api"
|
||||
DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT = "development"
|
||||
DEFAULT_METRICS_INCLUDE_BANK_ID = False # Disabled by default to avoid high-cardinality OTel metric growth
|
||||
DEFAULT_METRICS_BACKLOG_ENABLED = False # Disabled by default: runs periodic per-schema COUNT queries
|
||||
|
||||
# Audit log defaults
|
||||
DEFAULT_AUDIT_LOG_ENABLED = False # Disabled by default
|
||||
@@ -983,6 +1106,11 @@ DEFAULT_LLM_TRACE_MAX_CHARS = 50000 # Truncate stored input/output beyond this
|
||||
# 0 disables the reconcile sweep.
|
||||
DEFAULT_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS = 300
|
||||
|
||||
# How often the maintenance loop checks for cron-scheduled mental models that are
|
||||
# due for a refresh. This is the *check* cadence; the actual schedule is the
|
||||
# per-model cron expression in the mental model's trigger. 0 disables the sweep.
|
||||
DEFAULT_MENTAL_MODEL_REFRESH_TICK_SECONDS = 60
|
||||
|
||||
# Default MCP tool descriptions (can be customized via env vars)
|
||||
DEFAULT_MCP_RETAIN_DESCRIPTION = """Store important information to long-term memory.
|
||||
|
||||
@@ -1180,6 +1308,18 @@ def _validate_recall_budget_function(function: str) -> str:
|
||||
return function_lower
|
||||
|
||||
|
||||
def _validate_recency_decay_function(function: str) -> str:
|
||||
"""Validate and normalize the recency decay function."""
|
||||
function_lower = function.lower()
|
||||
if function_lower not in RECENCY_DECAY_FUNCTIONS:
|
||||
logger.warning(
|
||||
f"Invalid recency decay function '{function}', must be one of {RECENCY_DECAY_FUNCTIONS}. "
|
||||
f"Defaulting to '{DEFAULT_RECENCY_DECAY_FUNCTION}'."
|
||||
)
|
||||
return DEFAULT_RECENCY_DECAY_FUNCTION
|
||||
return function_lower
|
||||
|
||||
|
||||
def _parse_bank_priority(raw: str) -> dict[str, int]:
|
||||
"""Parse ``bank-pattern:priority,...`` into ``{pattern: priority}``.
|
||||
|
||||
@@ -1235,6 +1375,132 @@ def _parse_llm_router_config(env_var: str) -> dict | None:
|
||||
raise ValueError(f"Invalid {env_var}: invalid JSON: {e}") from e
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMMemberConfig:
|
||||
"""One extra LLM in a multi-LLM chain, configured via indexed env vars.
|
||||
|
||||
Mirrors the subset of LLM settings an indexed member supports
|
||||
(``HINDSIGHT_API_<OP>LLM_<n>_*``). The unindexed config remains the primary
|
||||
member (index 0); these describe members 1..N.
|
||||
"""
|
||||
|
||||
provider: str
|
||||
api_key: str | None
|
||||
model: str
|
||||
base_url: str | None
|
||||
reasoning_effort: str | None
|
||||
extra_body: dict | None
|
||||
default_headers: dict | None
|
||||
bedrock_service_tier: str | None
|
||||
gemini_service_tier: str | None
|
||||
vertexai_project_id: str | None = None
|
||||
vertexai_region: str | None = None
|
||||
vertexai_service_account_key: str | None = None
|
||||
litellmrouter_config: dict | None = None
|
||||
|
||||
|
||||
# Valid multi-LLM strategy modes.
|
||||
LLM_STRATEGY_FAILOVER = "failover"
|
||||
LLM_STRATEGY_ROUND_ROBIN = "round-robin"
|
||||
_VALID_LLM_STRATEGY_MODES = (LLM_STRATEGY_FAILOVER, LLM_STRATEGY_ROUND_ROBIN)
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMStrategyConfig:
|
||||
"""How to route a request across the members of a multi-LLM chain.
|
||||
|
||||
``mode`` is "failover" (try members in order) or "round-robin" (rotate the
|
||||
starting member per request, then fall through the rest on error). ``weights``
|
||||
is round-robin only: positive integers, one per member (primary first), giving
|
||||
an unbalanced rotation; ``None`` means uniform.
|
||||
"""
|
||||
|
||||
mode: str
|
||||
weights: list[int] | None = None
|
||||
|
||||
|
||||
def _parse_llm_strategy(raw: str | None) -> LLMStrategyConfig | None:
|
||||
"""Parse a multi-LLM strategy from a JSON env var.
|
||||
|
||||
Returns ``None`` when unset. The value must be a JSON object with a ``mode``
|
||||
of "failover" or "round-robin"; ``weights`` (round-robin only) must be a list
|
||||
of positive ints. Raises ``ValueError`` on any malformed input so
|
||||
misconfiguration fails fast at startup rather than silently degrading.
|
||||
"""
|
||||
text = (raw or "").strip()
|
||||
if not text:
|
||||
return None
|
||||
try:
|
||||
parsed = json.loads(text)
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(f"Invalid {ENV_LLM_STRATEGY}: invalid JSON: {e}") from e
|
||||
if not isinstance(parsed, dict):
|
||||
raise ValueError(f"Invalid LLM strategy: expected a JSON object, got {type(parsed).__name__}")
|
||||
|
||||
mode = parsed.get("mode")
|
||||
if mode not in _VALID_LLM_STRATEGY_MODES:
|
||||
raise ValueError(f"Invalid LLM strategy mode {mode!r}. Must be one of: {', '.join(_VALID_LLM_STRATEGY_MODES)}.")
|
||||
|
||||
weights = parsed.get("weights")
|
||||
if weights is not None:
|
||||
if mode != LLM_STRATEGY_ROUND_ROBIN:
|
||||
raise ValueError(f"LLM strategy 'weights' is only valid with mode '{LLM_STRATEGY_ROUND_ROBIN}'.")
|
||||
if not isinstance(weights, list) or not weights or not all(isinstance(w, int) and w > 0 for w in weights):
|
||||
raise ValueError("LLM strategy 'weights' must be a non-empty list of positive integers.")
|
||||
|
||||
return LLMStrategyConfig(mode=mode, weights=weights)
|
||||
|
||||
|
||||
def _parse_llm_members(prefix: str) -> list[LLMMemberConfig]:
|
||||
"""Parse indexed extra-LLM members for an operation env prefix.
|
||||
|
||||
``prefix`` is the operation segment in the env name: ``""`` (global),
|
||||
``"RETAIN_"``, ``"REFLECT_"`` or ``"CONSOLIDATION_"``. Members are read from
|
||||
``HINDSIGHT_API_{prefix}LLM_{n}_PROVIDER`` for n = 1, 2, ... and scanning
|
||||
stops at the first index whose ``_PROVIDER`` is unset (so indices must be
|
||||
contiguous from 1). ``MODEL`` defaults to the provider's default model.
|
||||
"""
|
||||
from .engine.llm_wrapper import requires_api_key
|
||||
|
||||
members: list[LLMMemberConfig] = []
|
||||
index = 1
|
||||
while True:
|
||||
base = f"HINDSIGHT_API_{prefix}LLM_{index}_"
|
||||
provider = os.getenv(base + "PROVIDER")
|
||||
if not provider:
|
||||
break
|
||||
|
||||
api_key = os.getenv(base + "API_KEY") or None
|
||||
if not api_key and requires_api_key(provider):
|
||||
raise ValueError(
|
||||
f"{base}API_KEY is required for provider '{provider}' (member {index} of the multi-LLM chain)."
|
||||
)
|
||||
|
||||
gemini_service_tier = os.getenv(base + "GEMINI_SERVICE_TIER")
|
||||
members.append(
|
||||
LLMMemberConfig(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
model=os.getenv(base + "MODEL") or _get_default_model_for_provider(provider),
|
||||
base_url=os.getenv(base + "BASE_URL") or None,
|
||||
reasoning_effort=os.getenv(base + "REASONING_EFFORT") or None,
|
||||
extra_body=json.loads(os.getenv(base + "EXTRA_BODY", "null")),
|
||||
default_headers=json.loads(os.getenv(base + "DEFAULT_HEADERS", "null")),
|
||||
bedrock_service_tier=os.getenv(base + "BEDROCK_SERVICE_TIER") or None,
|
||||
gemini_service_tier=(
|
||||
parse_gemini_service_tier(gemini_service_tier) if provider.lower() == "gemini" else None
|
||||
),
|
||||
vertexai_project_id=os.getenv(base + "VERTEXAI_PROJECT_ID") or None,
|
||||
vertexai_region=os.getenv(base + "VERTEXAI_REGION") or None,
|
||||
vertexai_service_account_key=os.getenv(base + "VERTEXAI_SERVICE_ACCOUNT_KEY") or None,
|
||||
litellmrouter_config=_parse_llm_router_config(base + "LITELLMROUTER_CONFIG"),
|
||||
)
|
||||
)
|
||||
index += 1
|
||||
|
||||
return members
|
||||
|
||||
|
||||
def _parse_default_bank_template(raw: str | None) -> dict | None:
|
||||
"""
|
||||
Parse HINDSIGHT_API_DEFAULT_BANK_TEMPLATE as JSON.
|
||||
@@ -1298,6 +1564,7 @@ class HindsightConfig:
|
||||
llm_groq_service_tier: str # Groq: "on_demand", "flex", or "auto"
|
||||
llm_openai_service_tier: str | None # OpenAI: None (default) or "flex" (50% cheaper)
|
||||
llm_bedrock_service_tier: str | None # Bedrock: None (default), "flex", "priority", or "reserved"
|
||||
llm_gemini_service_tier: str | None # Gemini: None (default) or "flex" (50% cheaper)
|
||||
llm_extra_body: (
|
||||
dict | None
|
||||
) # Extra body params merged into OpenAI-compatible API calls (e.g. {"chat_template_kwargs": {"enable_thinking": true}})
|
||||
@@ -1311,6 +1578,14 @@ class HindsightConfig:
|
||||
# overrides a `user` the caller already set.
|
||||
llm_send_bank_as_user: bool
|
||||
|
||||
# Per-operation sampling temperature. None means the temperature parameter is
|
||||
# omitted from the call (for models that reject explicit temperatures). See
|
||||
# ENV_LLM_TEMPERATURE and _resolve_operation_temperature.
|
||||
llm_temperature_verification: float | None
|
||||
llm_temperature_retain: float | None
|
||||
llm_temperature_reflect: float | None
|
||||
llm_temperature_consolidation: float | None
|
||||
|
||||
# LiteLLM Router chain (provider-specific; consumed by the "litellmrouter" provider).
|
||||
# List of deployment dicts evaluated in order with fallback on transient errors.
|
||||
# Each entry: {"provider": str, "model": str, "api_key": str | None, "base_url": str | None}.
|
||||
@@ -1400,6 +1675,8 @@ class HindsightConfig:
|
||||
embeddings_cohere_output_dimensions: int | None
|
||||
embeddings_openrouter_api_key: str | None
|
||||
embeddings_openrouter_model: str
|
||||
embeddings_requesty_api_key: str | None
|
||||
embeddings_requesty_model: str
|
||||
embeddings_litellm_api_base: str
|
||||
embeddings_litellm_api_key: str | None
|
||||
embeddings_litellm_model: str
|
||||
@@ -1435,6 +1712,9 @@ class HindsightConfig:
|
||||
bm25_min_score: float
|
||||
recall_max_candidates_per_source: int
|
||||
recall_strategy_boosts: dict[str, str]
|
||||
recency_decay_function: str
|
||||
recency_decay_linear_window_days: float
|
||||
recency_decay_halflife_days: float
|
||||
reranker_cohere_api_key: str | None
|
||||
reranker_cohere_model: str
|
||||
reranker_cohere_base_url: str | None
|
||||
@@ -1478,6 +1758,7 @@ class HindsightConfig:
|
||||
mcp_enabled: bool
|
||||
mcp_enabled_tools: list[str] | None # None = all tools; explicit list = allowlist
|
||||
mcp_stateless: bool # True = stateless HTTP (POST-only); False = stateful (supports GET/SSE)
|
||||
mcp_instructions: str | None # Additional instructions appended to retain/recall MCP tool descriptions
|
||||
enable_bank_config_api: bool
|
||||
enable_bank_llm_health: bool
|
||||
enable_dry_run_extract: bool
|
||||
@@ -1604,7 +1885,6 @@ class HindsightConfig:
|
||||
|
||||
# Optimization flags
|
||||
skip_llm_verification: bool
|
||||
lazy_reranker: bool
|
||||
|
||||
# Database migrations
|
||||
run_migrations_on_startup: bool
|
||||
@@ -1642,6 +1922,7 @@ class HindsightConfig:
|
||||
otel_service_name: str
|
||||
otel_deployment_environment: str
|
||||
metrics_include_bank_id: bool
|
||||
metrics_backlog_enabled: bool
|
||||
|
||||
# Audit log configuration (static - server-level only)
|
||||
audit_log_enabled: bool # Master switch for audit logging
|
||||
@@ -1658,6 +1939,9 @@ class HindsightConfig:
|
||||
# Interval for the periodic sweep that re-schedules consolidation for banks with
|
||||
# eligible-but-unscheduled facts. 0 = disabled.
|
||||
consolidation_reconcile_interval_seconds: int
|
||||
# How often the maintenance loop checks for cron-scheduled mental models due for
|
||||
# refresh (the per-model schedule lives in the mental model trigger). 0 = disabled.
|
||||
mental_model_refresh_tick_seconds: int
|
||||
|
||||
# Webhook configuration (static - server-level only, not per-bank)
|
||||
webhook_url: str | None # Global webhook URL (None = disabled)
|
||||
@@ -1676,6 +1960,25 @@ class HindsightConfig:
|
||||
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
|
||||
file_parser_markitdown_ocr_enabled: bool = DEFAULT_FILE_PARSER_MARKITDOWN_OCR_ENABLED
|
||||
file_parser_markitdown_ocr_api_key: str | None = None
|
||||
file_parser_markitdown_ocr_base_url: str | None = None
|
||||
file_parser_markitdown_ocr_model: str | None = None
|
||||
file_parser_markitdown_ocr_prompt: str = DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
|
||||
|
||||
# Multi-LLM chains (static, server-level). Index 0 of each chain is the
|
||||
# corresponding unindexed/base LLM config above; these hold the extra indexed
|
||||
# members and the routing strategy. Per-op members fall back to the global
|
||||
# members when unset (see MemoryEngine._build_llm). Credential fields (members
|
||||
# embed api_keys/base_urls).
|
||||
llm_members: list[LLMMemberConfig] = field(default_factory=list)
|
||||
llm_strategy: LLMStrategyConfig | None = None
|
||||
retain_llm_members: list[LLMMemberConfig] = field(default_factory=list)
|
||||
retain_llm_strategy: LLMStrategyConfig | None = None
|
||||
reflect_llm_members: list[LLMMemberConfig] = field(default_factory=list)
|
||||
reflect_llm_strategy: LLMStrategyConfig | None = None
|
||||
consolidation_llm_members: list[LLMMemberConfig] = field(default_factory=list)
|
||||
consolidation_llm_strategy: LLMStrategyConfig | None = None
|
||||
|
||||
# Class-level sets for configuration categorization
|
||||
|
||||
@@ -1691,6 +1994,11 @@ class HindsightConfig:
|
||||
"retain_llm_litellmrouter_config",
|
||||
"reflect_llm_litellmrouter_config",
|
||||
"consolidation_llm_litellmrouter_config",
|
||||
# Multi-LLM chains — members embed api_keys and base_urls
|
||||
"llm_members",
|
||||
"retain_llm_members",
|
||||
"reflect_llm_members",
|
||||
"consolidation_llm_members",
|
||||
# Base URLs (could expose infrastructure)
|
||||
"llm_base_url",
|
||||
"retain_llm_base_url",
|
||||
@@ -1716,6 +2024,8 @@ class HindsightConfig:
|
||||
"file_storage_gcs_service_account_key",
|
||||
"file_storage_azure_account_key",
|
||||
# File parser credentials
|
||||
"file_parser_markitdown_ocr_api_key",
|
||||
"file_parser_markitdown_ocr_base_url",
|
||||
"file_parser_iris_token",
|
||||
"file_parser_llama_parse_api_key",
|
||||
}
|
||||
@@ -1879,6 +2189,9 @@ class HindsightConfig:
|
||||
f"Note: 'standard' is not a valid Bedrock service tier -- use unset for default tier."
|
||||
)
|
||||
|
||||
# Validate gemini_service_tier
|
||||
self.llm_gemini_service_tier = parse_gemini_service_tier(self.llm_gemini_service_tier)
|
||||
|
||||
# When LLM provider is "none", force chunks-only mode and disable LLM-dependent features
|
||||
if self.llm_provider == "none":
|
||||
self.retain_extraction_mode = "chunks"
|
||||
@@ -1996,11 +2309,28 @@ class HindsightConfig:
|
||||
llm_groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
|
||||
llm_openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
|
||||
llm_bedrock_service_tier=os.getenv(ENV_LLM_BEDROCK_SERVICE_TIER) or None,
|
||||
llm_gemini_service_tier=(
|
||||
parse_gemini_service_tier(os.getenv(ENV_LLM_GEMINI_SERVICE_TIER) or DEFAULT_LLM_GEMINI_SERVICE_TIER)
|
||||
if llm_provider.lower() == "gemini"
|
||||
else None
|
||||
),
|
||||
llm_extra_body=json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null")),
|
||||
llm_default_headers=json.loads(os.getenv(ENV_LLM_DEFAULT_HEADERS, "null")),
|
||||
llm_strict_schema=os.getenv(ENV_LLM_STRICT_SCHEMA, str(DEFAULT_LLM_STRICT_SCHEMA)).lower() in ("true", "1"),
|
||||
llm_send_bank_as_user=os.getenv(ENV_LLM_SEND_BANK_AS_USER, str(DEFAULT_LLM_SEND_BANK_AS_USER)).lower()
|
||||
in ("true", "1"),
|
||||
llm_temperature_verification=_resolve_operation_temperature(
|
||||
ENV_LLM_TEMPERATURE_VERIFICATION, DEFAULT_LLM_TEMPERATURE_VERIFICATION
|
||||
),
|
||||
llm_temperature_retain=_resolve_operation_temperature(
|
||||
ENV_LLM_TEMPERATURE_RETAIN, DEFAULT_LLM_TEMPERATURE_RETAIN
|
||||
),
|
||||
llm_temperature_reflect=_resolve_operation_temperature(
|
||||
ENV_LLM_TEMPERATURE_REFLECT, DEFAULT_LLM_TEMPERATURE_REFLECT
|
||||
),
|
||||
llm_temperature_consolidation=_resolve_operation_temperature(
|
||||
ENV_LLM_TEMPERATURE_CONSOLIDATION, DEFAULT_LLM_TEMPERATURE_CONSOLIDATION
|
||||
),
|
||||
llm_litellmrouter_config=_parse_llm_router_config(ENV_LLM_LITELLMROUTER_CONFIG),
|
||||
# Vertex AI
|
||||
llm_vertexai_project_id=os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID) or DEFAULT_LLM_VERTEXAI_PROJECT_ID,
|
||||
@@ -2100,6 +2430,15 @@ class HindsightConfig:
|
||||
if os.getenv(ENV_CONSOLIDATION_LLM_TIMEOUT)
|
||||
else None,
|
||||
consolidation_llm_litellmrouter_config=_parse_llm_router_config(ENV_CONSOLIDATION_LLM_LITELLMROUTER_CONFIG),
|
||||
# Multi-LLM chains (indexed members + routing strategy)
|
||||
llm_members=_parse_llm_members(""),
|
||||
llm_strategy=_parse_llm_strategy(os.getenv(ENV_LLM_STRATEGY)),
|
||||
retain_llm_members=_parse_llm_members("RETAIN_"),
|
||||
retain_llm_strategy=_parse_llm_strategy(os.getenv(ENV_RETAIN_LLM_STRATEGY)),
|
||||
reflect_llm_members=_parse_llm_members("REFLECT_"),
|
||||
reflect_llm_strategy=_parse_llm_strategy(os.getenv(ENV_REFLECT_LLM_STRATEGY)),
|
||||
consolidation_llm_members=_parse_llm_members("CONSOLIDATION_"),
|
||||
consolidation_llm_strategy=_parse_llm_strategy(os.getenv(ENV_CONSOLIDATION_LLM_STRATEGY)),
|
||||
# Embeddings
|
||||
embeddings_provider=os.getenv(ENV_EMBEDDINGS_PROVIDER, DEFAULT_EMBEDDINGS_PROVIDER),
|
||||
embeddings_local_model=os.getenv(ENV_EMBEDDINGS_LOCAL_MODEL, DEFAULT_EMBEDDINGS_LOCAL_MODEL),
|
||||
@@ -2164,6 +2503,11 @@ 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),
|
||||
# Requesty embeddings (with fallback to shared Requesty key, then LLM key)
|
||||
embeddings_requesty_api_key=os.getenv(ENV_EMBEDDINGS_REQUESTY_API_KEY)
|
||||
or os.getenv(ENV_REQUESTY_API_KEY)
|
||||
or os.getenv(ENV_LLM_API_KEY),
|
||||
embeddings_requesty_model=os.getenv(ENV_EMBEDDINGS_REQUESTY_MODEL, DEFAULT_EMBEDDINGS_REQUESTY_MODEL),
|
||||
# ZeroEntropy embeddings
|
||||
embeddings_zeroentropy_api_key=os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_API_KEY)
|
||||
or os.getenv("ZEROENTROPY_API_KEY"),
|
||||
@@ -2270,6 +2614,15 @@ class HindsightConfig:
|
||||
recall_strategy_boosts=_parse_strategy_boosts(
|
||||
os.getenv(ENV_RECALL_STRATEGY_BOOSTS, DEFAULT_RECALL_STRATEGY_BOOSTS)
|
||||
),
|
||||
recency_decay_function=_validate_recency_decay_function(
|
||||
os.getenv(ENV_RECENCY_DECAY_FUNCTION, DEFAULT_RECENCY_DECAY_FUNCTION)
|
||||
),
|
||||
recency_decay_linear_window_days=float(
|
||||
os.getenv(ENV_RECENCY_DECAY_LINEAR_WINDOW_DAYS, str(DEFAULT_RECENCY_DECAY_LINEAR_WINDOW_DAYS))
|
||||
),
|
||||
recency_decay_halflife_days=float(
|
||||
os.getenv(ENV_RECENCY_DECAY_HALFLIFE_DAYS, str(DEFAULT_RECENCY_DECAY_HALFLIFE_DAYS))
|
||||
),
|
||||
# Cohere reranker (with backward-compatible fallback to shared API key)
|
||||
reranker_cohere_api_key=os.getenv(ENV_RERANKER_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
|
||||
reranker_cohere_model=os.getenv(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL),
|
||||
@@ -2345,6 +2698,7 @@ class HindsightConfig:
|
||||
if os.getenv(ENV_MCP_ENABLED_TOOLS)
|
||||
else DEFAULT_MCP_ENABLED_TOOLS,
|
||||
mcp_stateless=os.getenv(ENV_MCP_STATELESS, str(DEFAULT_MCP_STATELESS)).lower() == "true",
|
||||
mcp_instructions=os.getenv(ENV_MCP_INSTRUCTIONS) or DEFAULT_MCP_INSTRUCTIONS,
|
||||
enable_bank_llm_health=os.getenv(ENV_ENABLE_BANK_LLM_HEALTH, str(DEFAULT_ENABLE_BANK_LLM_HEALTH)).lower()
|
||||
== "true",
|
||||
enable_bank_config_api=os.getenv(ENV_ENABLE_BANK_CONFIG_API, str(DEFAULT_ENABLE_BANK_CONFIG_API)).lower()
|
||||
@@ -2374,7 +2728,6 @@ class HindsightConfig:
|
||||
),
|
||||
# Optimization flags
|
||||
skip_llm_verification=os.getenv(ENV_SKIP_LLM_VERIFICATION, "false").lower() == "true",
|
||||
lazy_reranker=os.getenv(ENV_LAZY_RERANKER, "false").lower() == "true",
|
||||
# Retain settings
|
||||
retain_max_completion_tokens=int(
|
||||
os.getenv(ENV_RETAIN_MAX_COMPLETION_TOKENS, str(DEFAULT_RETAIN_MAX_COMPLETION_TOKENS))
|
||||
@@ -2424,6 +2777,18 @@ class HindsightConfig:
|
||||
file_parser_allowlist=_parse_str_list(os.getenv(ENV_FILE_PARSER_ALLOWLIST))
|
||||
if os.getenv(ENV_FILE_PARSER_ALLOWLIST)
|
||||
else None,
|
||||
file_parser_markitdown_ocr_enabled=os.getenv(
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_ENABLED,
|
||||
str(DEFAULT_FILE_PARSER_MARKITDOWN_OCR_ENABLED),
|
||||
).lower()
|
||||
in ("1", "true", "yes", "on"),
|
||||
file_parser_markitdown_ocr_api_key=os.getenv(ENV_FILE_PARSER_MARKITDOWN_OCR_API_KEY) or None,
|
||||
file_parser_markitdown_ocr_base_url=os.getenv(ENV_FILE_PARSER_MARKITDOWN_OCR_BASE_URL) or None,
|
||||
file_parser_markitdown_ocr_model=os.getenv(ENV_FILE_PARSER_MARKITDOWN_OCR_MODEL) or None,
|
||||
file_parser_markitdown_ocr_prompt=os.getenv(
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
|
||||
DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
|
||||
),
|
||||
file_parser_iris_token=os.getenv(ENV_FILE_PARSER_IRIS_TOKEN) or None,
|
||||
file_parser_iris_org_id=os.getenv(ENV_FILE_PARSER_IRIS_ORG_ID) or None,
|
||||
file_parser_llama_parse_api_key=os.getenv(ENV_FILE_PARSER_LLAMA_PARSE_API_KEY) or None,
|
||||
@@ -2614,6 +2979,8 @@ class HindsightConfig:
|
||||
otel_deployment_environment=os.getenv(ENV_OTEL_DEPLOYMENT_ENVIRONMENT, DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT),
|
||||
metrics_include_bank_id=os.getenv(ENV_METRICS_INCLUDE_BANK_ID, str(DEFAULT_METRICS_INCLUDE_BANK_ID)).lower()
|
||||
in ("true", "1", "yes"),
|
||||
metrics_backlog_enabled=os.getenv(ENV_METRICS_BACKLOG_ENABLED, str(DEFAULT_METRICS_BACKLOG_ENABLED)).lower()
|
||||
in ("true", "1", "yes"),
|
||||
# Audit log configuration (static, server-level only)
|
||||
audit_log_enabled=os.getenv(ENV_AUDIT_LOG_ENABLED, str(DEFAULT_AUDIT_LOG_ENABLED)).lower() == "true",
|
||||
audit_log_actions=[
|
||||
@@ -2638,6 +3005,12 @@ class HindsightConfig:
|
||||
str(DEFAULT_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS),
|
||||
)
|
||||
),
|
||||
mental_model_refresh_tick_seconds=int(
|
||||
os.getenv(
|
||||
ENV_MENTAL_MODEL_REFRESH_TICK_SECONDS,
|
||||
str(DEFAULT_MENTAL_MODEL_REFRESH_TICK_SECONDS),
|
||||
)
|
||||
),
|
||||
# Webhook configuration (static, server-level only)
|
||||
webhook_url=os.getenv(ENV_WEBHOOK_URL) or DEFAULT_WEBHOOK_URL,
|
||||
webhook_secret=os.getenv(ENV_WEBHOOK_SECRET) or DEFAULT_WEBHOOK_SECRET,
|
||||
|
||||
@@ -8,6 +8,7 @@ Config values are resolved on every request to ensure consistency across
|
||||
multiple API servers.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import asdict, replace
|
||||
@@ -127,6 +128,21 @@ class ConfigResolver:
|
||||
# Return full config object (dataclass doesn't have __init__ that accepts kwargs, so we update the object)
|
||||
# Create a new config instance by copying the global config and updating fields
|
||||
resolved_config = HindsightConfig(**config_dict)
|
||||
# Multi-LLM chains are static credential fields (never tenant/bank-overridable),
|
||||
# but asdict() above flattened their member dataclasses into plain dicts. Restore
|
||||
# the original typed objects from the global config so the resolved object stays
|
||||
# well-typed for any consumer that reads them.
|
||||
resolved_config = replace(
|
||||
resolved_config,
|
||||
llm_members=self._global_config.llm_members,
|
||||
llm_strategy=self._global_config.llm_strategy,
|
||||
retain_llm_members=self._global_config.retain_llm_members,
|
||||
retain_llm_strategy=self._global_config.retain_llm_strategy,
|
||||
reflect_llm_members=self._global_config.reflect_llm_members,
|
||||
reflect_llm_strategy=self._global_config.reflect_llm_strategy,
|
||||
consolidation_llm_members=self._global_config.consolidation_llm_members,
|
||||
consolidation_llm_strategy=self._global_config.consolidation_llm_strategy,
|
||||
)
|
||||
validate_retain_chunking_config(
|
||||
resolved_config.retain_chunk_size,
|
||||
resolved_config.retain_structured_chunk_size,
|
||||
@@ -161,26 +177,83 @@ class ConfigResolver:
|
||||
resolved_config = await self.resolve_full_config(bank_id, context)
|
||||
config_dict = asdict(resolved_config)
|
||||
|
||||
# SECURITY: Filter to only configurable fields (exclude static/infrastructure)
|
||||
filtered = {k: v for k, v in config_dict.items() if k in self._configurable_fields}
|
||||
# SECURITY: drop static/infrastructure + credential fields, then permission-filter.
|
||||
filtered = self._strip_static_and_credential_fields(config_dict)
|
||||
return await self._apply_permission_filter(filtered, bank_id, context)
|
||||
|
||||
# SECURITY: Remove ALL credential fields (API keys, base URLs, etc.)
|
||||
filtered = {k: v for k, v in filtered.items() if k not in self._credential_fields}
|
||||
def _strip_static_and_credential_fields(self, config_dict: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Keep only configurable, non-credential fields.
|
||||
|
||||
# PERMISSIONS: Further filter based on tenant/bank permissions
|
||||
SECURITY: excludes static/infrastructure fields and ALL credential fields
|
||||
(API keys, base URLs, etc.) so a resolved config is safe to return over the API.
|
||||
"""
|
||||
return {
|
||||
k: v for k, v in config_dict.items() if k in self._configurable_fields and k not in self._credential_fields
|
||||
}
|
||||
|
||||
async def _apply_permission_filter(
|
||||
self, filtered: dict[str, Any], bank_id: str, context: RequestContext | None
|
||||
) -> dict[str, Any]:
|
||||
"""Further restrict already-stripped config to the tenant/bank permission allow-list.
|
||||
|
||||
On extension error, leaves ``filtered`` unchanged (parity with the historical
|
||||
single-bank path: a permissions lookup failure must not leak or drop fields).
|
||||
"""
|
||||
if not (self.tenant_extension and context):
|
||||
return filtered
|
||||
try:
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
filtered = {k: v for k, v in filtered.items() if k in allowed_fields}
|
||||
logger.debug(
|
||||
f"Applied permission filter for bank {bank_id}: allowed={len(allowed_fields)} fields, "
|
||||
f"returned={len(filtered)} fields"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load permissions for bank {bank_id}: {e}")
|
||||
return filtered
|
||||
|
||||
async def get_bank_configs(
|
||||
self, bank_ids: list[str], context: RequestContext | None = None
|
||||
) -> dict[str, dict[str, Any]]:
|
||||
"""Batch variant of :meth:`get_bank_config` for many banks.
|
||||
|
||||
Equivalent to calling ``get_bank_config`` per bank, but resolves the
|
||||
global + tenant base once and loads every bank's ``banks.config`` JSONB
|
||||
in a single query, instead of one config round-trip per bank. Used by
|
||||
``list_banks`` to overlay disposition + mission without an N+1.
|
||||
|
||||
Returns a mapping of bank_id -> filtered configurable-field dict. A bank
|
||||
with no config row still appears, mapped to the global+tenant base.
|
||||
"""
|
||||
if not bank_ids:
|
||||
return {}
|
||||
|
||||
# Global + tenant base, resolved once (tenant override is per-request, not per-bank).
|
||||
base_dict = asdict(self._global_config)
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
filtered = {k: v for k, v in filtered.items() if k in allowed_fields}
|
||||
logger.debug(
|
||||
f"Applied permission filter for bank {bank_id}: allowed={len(allowed_fields)} fields, "
|
||||
f"returned={len(filtered)} fields"
|
||||
)
|
||||
tenant_overrides = await self.tenant_extension.get_tenant_config(context)
|
||||
if tenant_overrides:
|
||||
normalized_tenant = normalize_config_dict(tenant_overrides)
|
||||
base_dict.update({k: v for k, v in normalized_tenant.items() if k in self._configurable_fields})
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load permissions for bank {bank_id}: {e}")
|
||||
logger.warning(f"Failed to load tenant config for bulk resolve: {e}")
|
||||
|
||||
return filtered
|
||||
# All bank overrides in one query, then merge + strip per bank.
|
||||
bank_overrides = await self._load_bank_configs(bank_ids)
|
||||
stripped = {
|
||||
bank_id: self._strip_static_and_credential_fields({**base_dict, **bank_overrides.get(bank_id, {})})
|
||||
for bank_id in bank_ids
|
||||
}
|
||||
|
||||
# Permission filter is per-bank; resolve concurrently when an extension is present.
|
||||
if not (self.tenant_extension and context):
|
||||
return stripped
|
||||
permission_filtered = await asyncio.gather(
|
||||
*(self._apply_permission_filter(stripped[bank_id], bank_id, context) for bank_id in bank_ids)
|
||||
)
|
||||
return dict(zip(bank_ids, permission_filtered, strict=True))
|
||||
|
||||
async def _load_bank_config(self, bank_id: str) -> dict[str, Any]:
|
||||
"""
|
||||
@@ -219,6 +292,45 @@ class ConfigResolver:
|
||||
|
||||
return {}
|
||||
|
||||
async def _load_bank_configs(self, bank_ids: list[str]) -> dict[str, dict[str, Any]]:
|
||||
"""Bulk variant of :meth:`_load_bank_config`: load many banks' overrides in one query.
|
||||
|
||||
Returns a mapping of bank_id -> normalized active overrides. Banks with no row
|
||||
(or an empty/all-tombstone config) are simply absent from the mapping.
|
||||
"""
|
||||
result: dict[str, dict[str, Any]] = {}
|
||||
if not bank_ids:
|
||||
return result
|
||||
try:
|
||||
async with self._backend.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT bank_id, config FROM {fq_table("banks")} WHERE bank_id = ANY($1)
|
||||
""",
|
||||
bank_ids,
|
||||
)
|
||||
for row in rows:
|
||||
config_data = row["config"]
|
||||
if not config_data:
|
||||
continue
|
||||
# Handle case where JSONB is returned as JSON string
|
||||
if isinstance(config_data, str):
|
||||
config_data = json.loads(config_data)
|
||||
|
||||
# Normalize keys (handle both env var format and Python field format)
|
||||
normalized = normalize_config_dict(config_data)
|
||||
|
||||
# Only active overrides for configurable fields. JSON null is a tombstone
|
||||
# for "Server Default" in the bank-config UI and must not override defaults.
|
||||
overrides = {
|
||||
k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None
|
||||
}
|
||||
if overrides:
|
||||
result[row["bank_id"]] = overrides
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to bulk-load bank configs: {e}")
|
||||
return result
|
||||
|
||||
async def update_bank_config(
|
||||
self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None
|
||||
) -> None:
|
||||
@@ -305,6 +417,9 @@ class ConfigResolver:
|
||||
# Validate recall budget fields
|
||||
_validate_recall_budget_updates(normalized_updates)
|
||||
|
||||
# Validate disposition trait fields (1-5 integer scale)
|
||||
_validate_disposition_updates(normalized_updates)
|
||||
|
||||
chunking_fields_updated = (
|
||||
"retain_chunk_size" in normalized_updates
|
||||
or "retain_structured_chunk_size" in normalized_updates
|
||||
@@ -419,6 +534,31 @@ def _validate_recall_budget_updates(updates: dict[str, Any]) -> None:
|
||||
)
|
||||
|
||||
|
||||
_DISPOSITION_KEYS = (
|
||||
"disposition_skepticism",
|
||||
"disposition_literalism",
|
||||
"disposition_empathy",
|
||||
)
|
||||
|
||||
|
||||
def _validate_disposition_updates(updates: dict[str, Any]) -> None:
|
||||
"""Validate disposition trait config updates. Raises ValueError on invalid input.
|
||||
|
||||
Each trait is an integer on a 1-5 scale (or None to clear the per-bank
|
||||
override). The read overlay injects the stored value verbatim into a strict
|
||||
``DispositionTraits(int, ge=1, le=5)``; an out-of-contract value (a float, a
|
||||
0-1 scale, or an int outside 1-5) accepted here would later 500 the whole
|
||||
bank list when any bank profile is serialized (issue #2348).
|
||||
"""
|
||||
for key in _DISPOSITION_KEYS:
|
||||
if key in updates:
|
||||
value = updates[key]
|
||||
if value is None:
|
||||
continue
|
||||
if not isinstance(value, int) or isinstance(value, bool) or not (1 <= value <= 5):
|
||||
raise ValueError(f"{key} must be an integer between 1 and 5, got {value!r}")
|
||||
|
||||
|
||||
def apply_strategy(config: HindsightConfig, strategy_name: str) -> HindsightConfig:
|
||||
"""
|
||||
Apply a named retain strategy's overrides on top of a resolved config.
|
||||
|
||||
@@ -13,9 +13,18 @@ in-flight task so that N concurrent callers produce one query rather than N.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Awaitable, Callable
|
||||
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .db.base import DatabaseBackend
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BankStatsCache:
|
||||
@@ -66,17 +75,28 @@ class BankStatsCache:
|
||||
schema: str,
|
||||
bank_id: str,
|
||||
loader: Callable[[], Awaitable[dict[str, Any]]],
|
||||
*,
|
||||
force_refresh: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""Return cached stats for `(schema, bank_id)` or call `loader()`.
|
||||
|
||||
Concurrent misses on the same key are coalesced onto a single
|
||||
in-flight loader.
|
||||
in-flight loader. When ``force_refresh`` is set the cached value is
|
||||
ignored: the loader runs and its result replaces the cached entry.
|
||||
"""
|
||||
if not self.enabled:
|
||||
return await loader()
|
||||
|
||||
key = (schema, bank_id)
|
||||
|
||||
if force_refresh:
|
||||
value = await loader()
|
||||
async with self._lock:
|
||||
self._store_unlocked(key, value)
|
||||
# Supersede any loader that was in flight for this key.
|
||||
self._in_flight.pop(key, None)
|
||||
return value
|
||||
|
||||
async with self._lock:
|
||||
cached = self._get_fresh_unlocked(key)
|
||||
if cached is not None:
|
||||
@@ -96,7 +116,10 @@ class BankStatsCache:
|
||||
value = await loader()
|
||||
except BaseException as exc:
|
||||
async with self._lock:
|
||||
self._in_flight.pop(key, None)
|
||||
# Invalidation may have detached this loader and allowed a new
|
||||
# one to claim the key. Never remove that newer loader's slot.
|
||||
if self._in_flight.get(key) is in_flight:
|
||||
self._in_flight.pop(key, None)
|
||||
if not in_flight.done():
|
||||
in_flight.set_exception(exc)
|
||||
# Suppress "Future exception was never retrieved" when no other
|
||||
@@ -106,8 +129,12 @@ class BankStatsCache:
|
||||
raise
|
||||
|
||||
async with self._lock:
|
||||
self._store_unlocked(key, value)
|
||||
self._in_flight.pop(key, None)
|
||||
# Only the loader that still owns the key may populate the cache.
|
||||
# An invalidated loader can finish for its original callers, but its
|
||||
# pre-invalidation result must not overwrite a newer load.
|
||||
if self._in_flight.get(key) is in_flight:
|
||||
self._store_unlocked(key, value)
|
||||
self._in_flight.pop(key, None)
|
||||
if not in_flight.done():
|
||||
in_flight.set_result(value)
|
||||
return value
|
||||
@@ -115,8 +142,113 @@ class BankStatsCache:
|
||||
async def invalidate(self, schema: str, bank_id: str) -> None:
|
||||
"""Drop any cached stats for `(schema, bank_id)`."""
|
||||
async with self._lock:
|
||||
self._entries.pop((schema, bank_id), None)
|
||||
key = (schema, bank_id)
|
||||
self._entries.pop(key, None)
|
||||
# Detach rather than cancel: existing callers may finish with the
|
||||
# snapshot they requested, while post-invalidation callers reload.
|
||||
self._in_flight.pop(key, None)
|
||||
|
||||
async def clear(self) -> None:
|
||||
async with self._lock:
|
||||
self._entries.clear()
|
||||
self._in_flight.clear()
|
||||
|
||||
|
||||
class DistributedBankStatsCache:
|
||||
"""Table-backed (cross-process) TTL cache for `get_bank_stats`.
|
||||
|
||||
Same ``get_or_load`` / ``invalidate`` / ``clear`` contract as
|
||||
:class:`BankStatsCache`, but the store is the per-schema ``bank_stats_cache``
|
||||
table instead of a per-process dict — so one worker's computation is shared
|
||||
with every other worker, and no caller recomputes while a fresh row exists.
|
||||
|
||||
On a hit, a call is a single primary-key ``SELECT`` (sub-millisecond); only a
|
||||
miss runs the (expensive) ``loader`` and writes the row back. Concurrent
|
||||
misses are *not* coalesced across processes (that would need a lock): they
|
||||
each compute and ``UPSERT``, last write wins — all results are correct, at the
|
||||
cost of a brief redundant compute at expiry.
|
||||
|
||||
Every DB touch is best-effort: if the cache table is unreachable or missing
|
||||
(e.g. a schema mid-migration), the call degrades to computing without caching
|
||||
rather than failing ``get_bank_stats``. PostgreSQL only — the engine keeps the
|
||||
in-process :class:`BankStatsCache` for Oracle.
|
||||
"""
|
||||
|
||||
def __init__(self, *, backend: "DatabaseBackend", ttl_seconds: float) -> None:
|
||||
self._backend = backend
|
||||
self._ttl = float(ttl_seconds)
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self._ttl > 0
|
||||
|
||||
@staticmethod
|
||||
def _qualified(schema: str) -> str:
|
||||
return f'"{schema}".bank_stats_cache' if schema else "bank_stats_cache"
|
||||
|
||||
async def get_or_load(
|
||||
self,
|
||||
schema: str,
|
||||
bank_id: str,
|
||||
loader: Callable[[], Awaitable[dict[str, Any]]],
|
||||
*,
|
||||
force_refresh: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
if not self.enabled:
|
||||
return await loader()
|
||||
|
||||
table = self._qualified(schema)
|
||||
|
||||
# 1. Fresh row? Single PK lookup; ``payload::text`` sidesteps any
|
||||
# jsonb->object codec so we always decode the same way. Skipped when
|
||||
# the caller forces a refresh — then we recompute and overwrite below.
|
||||
if not force_refresh:
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT payload::text AS payload FROM {table} "
|
||||
f"WHERE bank_id = $1 AND computed_at > now() - make_interval(secs => $2::double precision)",
|
||||
bank_id,
|
||||
self._ttl,
|
||||
)
|
||||
if row is not None:
|
||||
return json.loads(row["payload"])
|
||||
except Exception as exc: # noqa: BLE001 — cache read must never break the endpoint
|
||||
logger.debug("bank_stats_cache read failed for %s.%s (%s); computing uncached", schema, bank_id, exc)
|
||||
return await loader()
|
||||
|
||||
# 2. Miss — compute, then write the row back (best-effort).
|
||||
value = await loader()
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
await conn.execute(
|
||||
f"INSERT INTO {table} (bank_id, payload, computed_at) VALUES ($1, $2::jsonb, now()) "
|
||||
f"ON CONFLICT (bank_id) DO UPDATE SET payload = EXCLUDED.payload, computed_at = now()",
|
||||
bank_id,
|
||||
json.dumps(value),
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 — a failed write just means no caching this round
|
||||
logger.warning("bank_stats_cache write failed for %s.%s (%s)", schema, bank_id, exc)
|
||||
return value
|
||||
|
||||
async def invalidate(self, schema: str, bank_id: str) -> None:
|
||||
"""Drop the cached row so the next read recomputes."""
|
||||
if not self.enabled:
|
||||
return
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
await conn.execute(f"DELETE FROM {self._qualified(schema)} WHERE bank_id = $1", bank_id)
|
||||
except Exception as exc: # noqa: BLE001 — invalidation must never break the write path
|
||||
logger.debug("bank_stats_cache invalidate failed for %s.%s (%s)", schema, bank_id, exc)
|
||||
|
||||
async def clear(self) -> None:
|
||||
"""Drop all cached rows in the current schema (best-effort)."""
|
||||
if not self.enabled:
|
||||
return
|
||||
from .memory_engine import get_current_schema
|
||||
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
await conn.execute(f"DELETE FROM {self._qualified(get_current_schema())}")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("bank_stats_cache clear failed (%s)", exc)
|
||||
|
||||
@@ -98,7 +98,7 @@ _DEDUP_TOP_K = 5
|
||||
class _DedupDecision(BaseModel):
|
||||
"""Focused 1-by-1 verdict for whether a new observation duplicates an existing one."""
|
||||
|
||||
action: Literal["merge", "keep"]
|
||||
action: Literal["merge", "keep"] = "keep"
|
||||
text: str = "" # the synthesized merged observation (when action == "merge")
|
||||
reason: str = ""
|
||||
|
||||
@@ -224,13 +224,18 @@ async def _dedup_reconcile_create(
|
||||
# Fold the new source facts into the twin and persist the merged text. We keep the twin's
|
||||
# existing embedding: the merged text is >= threshold similar, so the stored vector stays
|
||||
# representative and we avoid a re-embed + a dialect-specific vector UPDATE.
|
||||
search_vector_clause = (
|
||||
f",\n search_vector = to_tsvector('{config.text_search_extension_native_language}'::regconfig, COALESCE($1, ''))"
|
||||
if config.text_search_extension == "native"
|
||||
else ""
|
||||
)
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
SET text = $1,
|
||||
source_memory_ids = (SELECT array_agg(DISTINCT e) FROM unnest(source_memory_ids || $2::uuid[]) e),
|
||||
proof_count = (SELECT count(DISTINCT e) FROM unnest(source_memory_ids || $2::uuid[]) e),
|
||||
updated_at = now()
|
||||
updated_at = now(){search_vector_clause}
|
||||
WHERE id = $3::uuid
|
||||
""",
|
||||
outcome.merged_text,
|
||||
@@ -279,6 +284,11 @@ async def _dedup_reconcile_update(
|
||||
# the create path) then delete the now-redundant updated row. The all_strict/any tag match
|
||||
# guarantees twin and updated share scope, so dropping the updated row's tags loses no
|
||||
# visibility. Temporal fields follow the surviving twin (minimal scope; matches create).
|
||||
search_vector_clause = (
|
||||
f",\n search_vector = to_tsvector('{config.text_search_extension_native_language}'::regconfig, COALESCE($1, ''))"
|
||||
if config.text_search_extension == "native"
|
||||
else ""
|
||||
)
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")} t
|
||||
@@ -289,7 +299,7 @@ async def _dedup_reconcile_update(
|
||||
proof_count = (
|
||||
SELECT count(DISTINCT e) FROM unnest(t.source_memory_ids || u.source_memory_ids) e
|
||||
),
|
||||
updated_at = now()
|
||||
updated_at = now(){search_vector_clause}
|
||||
FROM {fq_table("memory_units")} u
|
||||
WHERE t.id = $2::uuid AND u.id = $3::uuid
|
||||
""",
|
||||
@@ -449,6 +459,13 @@ class _CreateAction(BaseModel):
|
||||
def sanitize_text(cls, v: str) -> str:
|
||||
return sanitize_llm_output(v) or ""
|
||||
|
||||
@field_validator("source_fact_ids", mode="before")
|
||||
@classmethod
|
||||
def ensure_list(cls, v: str | list[str]) -> list[str]:
|
||||
if isinstance(v, str):
|
||||
return [v]
|
||||
return v
|
||||
|
||||
|
||||
class _UpdateAction(BaseModel):
|
||||
text: str
|
||||
@@ -461,6 +478,13 @@ class _UpdateAction(BaseModel):
|
||||
def sanitize_text(cls, v: str) -> str:
|
||||
return sanitize_llm_output(v) or ""
|
||||
|
||||
@field_validator("source_fact_ids", mode="before")
|
||||
@classmethod
|
||||
def ensure_list(cls, v: str | list[str]) -> list[str]:
|
||||
if isinstance(v, str):
|
||||
return [v]
|
||||
return v
|
||||
|
||||
|
||||
class _DeleteAction(BaseModel):
|
||||
observation_id: str # UUID of the observation to remove
|
||||
@@ -640,6 +664,7 @@ class ConsolidationPerfLog:
|
||||
self.start_time = time.time()
|
||||
self.lines: list[str] = []
|
||||
self.timings: dict[str, float] = {}
|
||||
self.timing_counts: dict[str, int] = {}
|
||||
self.llm_calls: int = 0
|
||||
self.total_obs_in_context: int = 0
|
||||
self.total_prompt_chars: int = 0
|
||||
@@ -649,11 +674,13 @@ class ConsolidationPerfLog:
|
||||
self.lines.append(message)
|
||||
|
||||
def record_timing(self, key: str, duration: float) -> None:
|
||||
"""Record a timing measurement."""
|
||||
if key in self.timings:
|
||||
self.timings[key] += duration
|
||||
else:
|
||||
self.timings[key] = duration
|
||||
"""Record a timing measurement.
|
||||
|
||||
Tracks both total seconds and call count so the summary can
|
||||
distinguish one slow call from many fast calls in aggregate.
|
||||
"""
|
||||
self.timings[key] = self.timings.get(key, 0.0) + duration
|
||||
self.timing_counts[key] = self.timing_counts.get(key, 0) + 1
|
||||
|
||||
def record_llm_call(self, obs_count: int, prompt_chars: int) -> None:
|
||||
"""Record stats for a single LLM call."""
|
||||
@@ -676,6 +703,8 @@ class ConsolidationPerfLog:
|
||||
"""
|
||||
for key, value in other.timings.items():
|
||||
self.timings[key] = self.timings.get(key, 0.0) + value
|
||||
for key, count in other.timing_counts.items():
|
||||
self.timing_counts[key] = self.timing_counts.get(key, 0) + count
|
||||
self.llm_calls += other.llm_calls
|
||||
self.total_obs_in_context += other.total_obs_in_context
|
||||
self.total_prompt_chars += other.total_prompt_chars
|
||||
@@ -1276,16 +1305,22 @@ async def _run_consolidation_job(
|
||||
f"{stats['skipped']} skipped)"
|
||||
)
|
||||
|
||||
# Add timing breakdown
|
||||
# Add timing breakdown. Each phase is recorded once per call, so the count
|
||||
# disambiguates a single slow call from many fast calls — important for
|
||||
# operators triaging "the recall phase took 15s" log lines, where the
|
||||
# total is the sum of many serial sub-calls rather than one slow query.
|
||||
def _fmt(key: str) -> str:
|
||||
total = perf.timings[key]
|
||||
count = perf.timing_counts.get(key, 0)
|
||||
if count > 1:
|
||||
avg_ms = total * 1000.0 / count
|
||||
return f"{key}={total:.3f}s ({count} calls, avg={avg_ms:.0f}ms)"
|
||||
return f"{key}={total:.3f}s"
|
||||
|
||||
timing_parts = []
|
||||
if "recall" in perf.timings:
|
||||
timing_parts.append(f"recall={perf.timings['recall']:.3f}s")
|
||||
if "llm" in perf.timings:
|
||||
timing_parts.append(f"llm={perf.timings['llm']:.3f}s")
|
||||
if "embedding" in perf.timings:
|
||||
timing_parts.append(f"embedding={perf.timings['embedding']:.3f}s")
|
||||
if "db_write" in perf.timings:
|
||||
timing_parts.append(f"db_write={perf.timings['db_write']:.3f}s")
|
||||
for key in ("recall", "llm", "embedding", "db_write"):
|
||||
if key in perf.timings:
|
||||
timing_parts.append(_fmt(key))
|
||||
|
||||
if perf.llm_calls > 0:
|
||||
timing_parts.append(f"avg_obs={perf.total_obs_in_context / perf.llm_calls:.1f}")
|
||||
@@ -1820,6 +1855,12 @@ async def _execute_update_action(
|
||||
|
||||
config = get_config()
|
||||
|
||||
search_vector_clause = (
|
||||
f",\n search_vector = to_tsvector('{config.text_search_extension_native_language}'::regconfig, COALESCE($1, ''))"
|
||||
if config.text_search_extension == "native"
|
||||
else ""
|
||||
)
|
||||
|
||||
t0 = time.time()
|
||||
await conn.execute(
|
||||
f"""
|
||||
@@ -1832,7 +1873,7 @@ async def _execute_update_action(
|
||||
updated_at = now(),
|
||||
occurred_start = LEAST(occurred_start, COALESCE($6, occurred_start)),
|
||||
occurred_end = GREATEST(occurred_end, COALESCE($7, occurred_end)),
|
||||
mentioned_at = GREATEST(mentioned_at, COALESCE($8, mentioned_at))
|
||||
mentioned_at = GREATEST(mentioned_at, COALESCE($8, mentioned_at)){search_vector_clause}
|
||||
WHERE id = $5
|
||||
""",
|
||||
new_text,
|
||||
@@ -2308,16 +2349,20 @@ async def _create_observation_directly(
|
||||
tokenize($3, 'llmlingua2')::bm25_catalog.bm25vector)
|
||||
RETURNING id
|
||||
"""
|
||||
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.
|
||||
elif config.text_search_extension == "native":
|
||||
# Native: search_vector is populated with to_tsvector() using the
|
||||
# configured native language dictionary, matching the batch insert
|
||||
# path in ops_postgresql.insert_facts_batch.
|
||||
query = f"""
|
||||
INSERT INTO {fq_table("memory_units")} (
|
||||
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids,
|
||||
tags, event_date, occurred_start, occurred_end, mentioned_at, search_vector
|
||||
)
|
||||
VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, $6, $7, $8, $9, $10,
|
||||
to_tsvector('{config.text_search_extension_native_language}'::regconfig, COALESCE($3, '')))
|
||||
RETURNING id
|
||||
"""
|
||||
else: # pg_textsearch, pgroonga, pg_search: indexes operate on base text columns directly
|
||||
query = f"""
|
||||
INSERT INTO {fq_table("memory_units")} (
|
||||
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids,
|
||||
|
||||
@@ -212,7 +212,7 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
device = "cpu"
|
||||
logger.info("Reranker: forcing CPU mode (HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU=1)")
|
||||
else:
|
||||
# Check for GPU (CUDA) or Apple Silicon (MPS)
|
||||
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
|
||||
# Wrap in try-except to gracefully handle any device detection issues
|
||||
# (e.g., in CI environments or when PyTorch is built without GPU support)
|
||||
device = "cpu" # Default to CPU
|
||||
@@ -220,10 +220,13 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
has_gpu = torch.cuda.is_available() or (
|
||||
hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
)
|
||||
# Intel Arc XPU support — torch.xpu is available when the XPU build is loaded
|
||||
if not has_gpu and hasattr(torch, "xpu"):
|
||||
has_gpu = torch.xpu.is_available()
|
||||
if has_gpu:
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS/XPU
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}")
|
||||
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
|
||||
|
||||
# Patch transformers 5.x compatibility for models using XLM-RoBERTa
|
||||
# (e.g., jina-reranker-v2-base-multilingual). transformers 5.x removed
|
||||
|
||||
@@ -256,13 +256,21 @@ class OracleOps(DataAccessOps):
|
||||
# Oracle doesn't support ON CONFLICT; rely on the PK and the
|
||||
# IGNORE_ROW_ON_DUPKEY_INDEX hint to skip duplicates server-side.
|
||||
# The hint name must match the PK constraint exactly.
|
||||
#
|
||||
# Sort to enforce a global lock-acquisition order on the
|
||||
# (bank_id, unit_id) PK. Without this, two concurrent
|
||||
# transactions inserting overlapping unit_id sets in different
|
||||
# orders can deadlock on the unique-check row locks. Sorting
|
||||
# gives every concurrent caller the same lock order, so
|
||||
# conflicting inserts queue cleanly instead of cycling.
|
||||
sorted_unit_ids = sorted(unit_ids)
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT /*+ IGNORE_ROW_ON_DUPKEY_INDEX({table}, pk_graph_maintenance_queue) */
|
||||
INTO {table} (bank_id, unit_id)
|
||||
VALUES ($1, $2)
|
||||
""",
|
||||
[(bank_id, uid) for uid in unit_ids],
|
||||
[(bank_id, uid) for uid in sorted_unit_ids],
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
|
||||
@@ -348,6 +348,15 @@ class PostgreSQLOps(DataAccessOps):
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
# Sort to enforce a global lock-acquisition order on the
|
||||
# (bank_id, unit_id) unique-key. Without this, two concurrent
|
||||
# transactions inserting overlapping unit_id sets in different
|
||||
# orders can deadlock on the ON CONFLICT row locks — Postgres
|
||||
# acquires a short-lived lock per row being checked, and cycle
|
||||
# detection then aborts one transaction. Sorting gives every
|
||||
# concurrent caller the same lock order, so conflicting inserts
|
||||
# queue cleanly instead of cycling.
|
||||
sorted_unit_ids = sorted(unit_ids)
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (bank_id, unit_id)
|
||||
@@ -355,7 +364,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
ON CONFLICT (bank_id, unit_id) DO NOTHING
|
||||
""",
|
||||
bank_id,
|
||||
unit_ids,
|
||||
sorted_unit_ids,
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
|
||||
@@ -190,7 +190,7 @@ class LocalSTEmbeddings(Embeddings):
|
||||
device = "cpu"
|
||||
logger.info("Embeddings: forcing CPU mode")
|
||||
else:
|
||||
# Check for GPU (CUDA) or Apple Silicon (MPS)
|
||||
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
|
||||
# Wrap in try-except to gracefully handle any device detection issues
|
||||
# (e.g., in CI environments or when PyTorch is built without GPU support)
|
||||
device = "cpu" # Default to CPU
|
||||
@@ -198,10 +198,13 @@ class LocalSTEmbeddings(Embeddings):
|
||||
has_gpu = torch.cuda.is_available() or (
|
||||
hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
)
|
||||
# Intel Arc XPU support — torch.xpu is available when the XPU build is loaded
|
||||
if not has_gpu and hasattr(torch, "xpu"):
|
||||
has_gpu = torch.xpu.is_available()
|
||||
if has_gpu:
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS/XPU
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}")
|
||||
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
|
||||
|
||||
# Suppress verbose transformers warnings during model loading
|
||||
# This suppresses the "UNEXPECTED" warnings from BertModel which are harmless
|
||||
@@ -709,7 +712,8 @@ class OpenAIEmbeddings(Embeddings):
|
||||
|
||||
class CodexOAuthEmbeddings(OpenAIEmbeddings):
|
||||
"""
|
||||
OpenAI embeddings using the Codex/ChatGPT OAuth token from ``~/.codex/auth.json``.
|
||||
OpenAI embeddings using the Codex/ChatGPT OAuth token from the Codex
|
||||
``auth.json`` (``$CODEX_HOME/auth.json``, or ``~/.codex/auth.json`` when unset).
|
||||
|
||||
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
|
||||
@@ -1634,6 +1638,20 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "requesty":
|
||||
api_key = config.embeddings_requesty_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_EMBEDDINGS_REQUESTY_API_KEY, HINDSIGHT_API_REQUESTY_API_KEY, "
|
||||
f"or {ENV_LLM_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'requesty'"
|
||||
)
|
||||
return OpenAIEmbeddings(
|
||||
api_key=api_key,
|
||||
model=config.embeddings_requesty_model,
|
||||
base_url="https://router.requesty.ai/v1",
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "zeroentropy":
|
||||
api_key = config.embeddings_zeroentropy_api_key
|
||||
if not api_key:
|
||||
@@ -1697,6 +1715,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown embeddings provider: {provider}. "
|
||||
f"Supported: 'local', 'onnx', 'tei', 'openai', 'openai-codex', 'openrouter', 'cohere', 'google', "
|
||||
f"Supported: 'local', 'onnx', 'tei', 'openai', 'openai-codex', 'openrouter', 'requesty', 'cohere', 'google', "
|
||||
f"'zeroentropy', 'litellm', 'litellm-sdk'"
|
||||
)
|
||||
|
||||
@@ -782,236 +782,6 @@ class EntityResolver:
|
||||
|
||||
return entity_ids
|
||||
|
||||
async def resolve_entity(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_text: str,
|
||||
context: str,
|
||||
nearby_entities: list[dict],
|
||||
unit_event_date,
|
||||
) -> str:
|
||||
"""
|
||||
Resolve an entity to a canonical entity ID.
|
||||
|
||||
Args:
|
||||
bank_id: bank ID (entities are scoped to agents)
|
||||
entity_text: Entity text ("Alice", "Google", etc.)
|
||||
context: Context where entity appears
|
||||
nearby_entities: Other entities in the same unit
|
||||
unit_event_date: When this unit was created
|
||||
|
||||
Returns:
|
||||
Entity ID (creates new entity if needed)
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
# Find candidate entities with similar name
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, canonical_name, metadata, last_seen
|
||||
FROM {fq_table("entities")}
|
||||
WHERE bank_id = $1
|
||||
AND (
|
||||
canonical_name ILIKE $2
|
||||
OR canonical_name ILIKE $3
|
||||
OR $2 ILIKE canonical_name || '%%'
|
||||
)
|
||||
ORDER BY mention_count DESC
|
||||
""",
|
||||
bank_id,
|
||||
entity_text,
|
||||
f"%{entity_text}%",
|
||||
)
|
||||
|
||||
if not candidates:
|
||||
# New entity - create it
|
||||
return await self._create_entity(conn, bank_id, entity_text, unit_event_date)
|
||||
|
||||
# Score candidates based on:
|
||||
# 1. Name similarity
|
||||
# 2. Context overlap (TODO: could use embeddings)
|
||||
# 3. Co-occurring entities
|
||||
# 4. Temporal proximity
|
||||
|
||||
best_candidate = None
|
||||
best_score = 0.0
|
||||
|
||||
nearby_entity_set = {e["text"].lower() for e in nearby_entities if e["text"] != entity_text}
|
||||
|
||||
for row in candidates:
|
||||
candidate_id = row["id"]
|
||||
canonical_name = row["canonical_name"]
|
||||
last_seen = row["last_seen"]
|
||||
score = 0.0
|
||||
|
||||
# 1. Name similarity (0-1)
|
||||
name_similarity = SequenceMatcher(None, entity_text.lower(), canonical_name.lower()).ratio()
|
||||
score += name_similarity * 0.5
|
||||
|
||||
# 2. Co-occurring entities (0-0.5)
|
||||
# Get entities that co-occurred with this candidate before
|
||||
# Use the materialized co-occurrence cache for fast lookup
|
||||
co_entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.canonical_name, ec.cooccurrence_count
|
||||
FROM {fq_table("entity_cooccurrences")} ec
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
CASE
|
||||
WHEN ec.entity_id_1 = $1 THEN ec.entity_id_2
|
||||
WHEN ec.entity_id_2 = $1 THEN ec.entity_id_1
|
||||
END = e.id
|
||||
)
|
||||
WHERE ec.entity_id_1 = $1 OR ec.entity_id_2 = $1
|
||||
""",
|
||||
candidate_id,
|
||||
)
|
||||
co_entities = {r["canonical_name"].lower() for r in co_entity_rows}
|
||||
|
||||
# Check overlap with nearby entities
|
||||
overlap = len(nearby_entity_set & co_entities)
|
||||
if nearby_entity_set:
|
||||
co_entity_score = overlap / len(nearby_entity_set)
|
||||
score += co_entity_score * 0.3
|
||||
|
||||
# 3. Temporal proximity (0-0.2)
|
||||
if last_seen:
|
||||
# Normalize both to UTC-aware to avoid naive/aware mismatch
|
||||
# (Oracle returns naive datetimes from fromisoformat)
|
||||
_evt = unit_event_date if unit_event_date.tzinfo else unit_event_date.replace(tzinfo=UTC)
|
||||
_seen = last_seen if last_seen.tzinfo else last_seen.replace(tzinfo=UTC)
|
||||
days_diff = abs((_evt - _seen).total_seconds() / 86400)
|
||||
if days_diff < 7: # Within a week
|
||||
temporal_score = max(0, 1.0 - (days_diff / 7))
|
||||
score += temporal_score * 0.2
|
||||
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_candidate = candidate_id
|
||||
|
||||
# Threshold for considering it the same entity
|
||||
threshold = 0.6
|
||||
|
||||
if best_score > threshold:
|
||||
# Update entity
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("entities")}
|
||||
SET mention_count = mention_count + 1,
|
||||
last_seen = $1
|
||||
WHERE id = $2
|
||||
""",
|
||||
unit_event_date,
|
||||
best_candidate,
|
||||
)
|
||||
return best_candidate
|
||||
else:
|
||||
# Not confident - create new entity
|
||||
return await self._create_entity(conn, bank_id, entity_text, unit_event_date)
|
||||
|
||||
async def _create_entity(
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entity_text: str,
|
||||
event_date,
|
||||
) -> str:
|
||||
"""
|
||||
Create a new entity or get existing one if it already exists.
|
||||
|
||||
Uses INSERT ... ON CONFLICT to handle race conditions where
|
||||
two concurrent transactions try to create the same entity.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
bank_id: bank ID
|
||||
entity_text: Entity text
|
||||
event_date: When first seen
|
||||
|
||||
Returns:
|
||||
Entity ID
|
||||
"""
|
||||
entity_id = await conn.fetchval(
|
||||
f"""
|
||||
INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
VALUES ($1, $2, COALESCE($3, now()), COALESCE($4, now()), 1)
|
||||
ON CONFLICT (bank_id, LOWER(canonical_name))
|
||||
DO UPDATE SET
|
||||
mention_count = {fq_table("entities")}.mention_count + 1,
|
||||
last_seen = EXCLUDED.last_seen
|
||||
RETURNING id
|
||||
""",
|
||||
bank_id,
|
||||
entity_text,
|
||||
event_date,
|
||||
event_date,
|
||||
)
|
||||
return entity_id
|
||||
|
||||
async def link_unit_to_entity(self, unit_id: str, entity_id: str):
|
||||
"""
|
||||
Link a memory unit to an entity.
|
||||
Also updates co-occurrence cache with other entities in the same unit.
|
||||
|
||||
Args:
|
||||
unit_id: Memory unit ID
|
||||
entity_id: Entity ID
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
# Insert unit-entity link
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
unit_id,
|
||||
entity_id,
|
||||
)
|
||||
|
||||
# Update co-occurrence cache: find other entities in this unit
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT entity_id
|
||||
FROM {fq_table("unit_entities")}
|
||||
WHERE unit_id = $1 AND entity_id != $2
|
||||
""",
|
||||
unit_id,
|
||||
entity_id,
|
||||
)
|
||||
|
||||
other_entities = [row["entity_id"] for row in rows]
|
||||
|
||||
# Update co-occurrences for each pair
|
||||
for other_entity_id in other_entities:
|
||||
await self._update_cooccurrence(conn, entity_id, other_entity_id)
|
||||
|
||||
async def _update_cooccurrence(self, conn, entity_id_1: str, entity_id_2: str):
|
||||
"""
|
||||
Update the co-occurrence cache for two entities.
|
||||
|
||||
Uses CHECK constraint ordering (entity_id_1 < entity_id_2) to avoid duplicates.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
entity_id_1: First entity ID
|
||||
entity_id_2: Second entity ID
|
||||
"""
|
||||
# Ensure consistent ordering (smaller UUID first)
|
||||
if entity_id_1 > entity_id_2:
|
||||
entity_id_1, entity_id_2 = entity_id_2, entity_id_1
|
||||
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("entity_cooccurrences")} (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
VALUES ($1, $2, 1, NOW())
|
||||
ON CONFLICT (entity_id_1, entity_id_2)
|
||||
DO UPDATE SET
|
||||
cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + 1,
|
||||
last_cooccurred = NOW()
|
||||
""",
|
||||
entity_id_1,
|
||||
entity_id_2,
|
||||
)
|
||||
|
||||
async def link_units_to_entities_batch(
|
||||
self,
|
||||
unit_entity_pairs: list[tuple[str, str]] | list[tuple[str, str, datetime | None]],
|
||||
|
||||
@@ -449,6 +449,7 @@ class MemoryEngineInterface(ABC):
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
force_refresh: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Get statistics about memory nodes and links for a bank.
|
||||
@@ -456,6 +457,8 @@ class MemoryEngineInterface(ABC):
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
force_refresh: Bypass the cached value and recompute (also refreshes
|
||||
the cache for subsequent callers).
|
||||
|
||||
Returns:
|
||||
Dict with node_counts, link_counts, link_counts_by_fact_type
|
||||
|
||||
@@ -6,6 +6,7 @@ enabling support for multiple LLM backends (OpenAI, Anthropic, Gemini, Codex, et
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from .response_models import LLMToolCallResult
|
||||
@@ -252,3 +253,11 @@ class OutputTooLongError(Exception):
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ProviderRateLimitResetError(Exception):
|
||||
"""Raised when an upstream provider says quota will reopen at a known time."""
|
||||
|
||||
def __init__(self, retry_at: datetime, message: str = "") -> None:
|
||||
self.retry_at = retry_at
|
||||
super().__init__(message)
|
||||
|
||||
@@ -76,6 +76,51 @@ _request_ctx: ContextVar[dict[str, Any] | None] = ContextVar("hindsight_llm_requ
|
||||
_call_metadata_ctx: ContextVar[dict[str, Any] | None] = ContextVar("hindsight_llm_call_metadata_ctx", default=None)
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMResponseUsage:
|
||||
"""Provider-reported token usage for the in-flight LLM call.
|
||||
|
||||
Stashed by provider implementations as soon as a response is received —
|
||||
*before* local JSON parsing / schema validation, which may still fail. The
|
||||
wrapper reads it to attach real token counts to an error trace when the
|
||||
provider call itself succeeded but the structured output couldn't be parsed
|
||||
or validated (providers charge for those tokens regardless). See #2387.
|
||||
"""
|
||||
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cached_tokens: int = 0
|
||||
|
||||
|
||||
# Per-call provider usage, set by providers right after a response is received.
|
||||
_response_usage_ctx: ContextVar[LLMResponseUsage | None] = ContextVar("hindsight_llm_response_usage_ctx", default=None)
|
||||
|
||||
|
||||
def set_response_usage(usage: LLMResponseUsage | None) -> Token:
|
||||
"""Bind provider-reported usage for the current call. Returns a reset token."""
|
||||
return _response_usage_ctx.set(usage)
|
||||
|
||||
|
||||
def stash_response_usage(usage: LLMResponseUsage | None) -> None:
|
||||
"""Record provider-reported usage so an error trace can attach it later.
|
||||
|
||||
Called by provider implementations once a response (with usage) is in hand,
|
||||
before parsing/validation that may raise. Overwrites any prior value from an
|
||||
earlier retry attempt so the last attempt's usage wins.
|
||||
"""
|
||||
_response_usage_ctx.set(usage)
|
||||
|
||||
|
||||
def reset_response_usage(token: Token) -> None:
|
||||
"""Unwind a binding made by :func:`set_response_usage`."""
|
||||
_response_usage_ctx.reset(token)
|
||||
|
||||
|
||||
def current_response_usage() -> LLMResponseUsage | None:
|
||||
"""Return the active call's provider-reported usage, or None."""
|
||||
return _response_usage_ctx.get()
|
||||
|
||||
|
||||
def set_trace_context(ctx: LLMTraceContext | None) -> Token:
|
||||
"""Bind trace attribution to the current context. Returns a reset token."""
|
||||
return _trace_ctx.set(ctx)
|
||||
|
||||
@@ -10,7 +10,6 @@ import re
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import AsyncExitStack
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
# Vertex AI imports (conditional - for LLMProvider to pass credentials to GeminiLLM)
|
||||
@@ -253,6 +252,8 @@ def create_llm_provider(
|
||||
gemini_safety_settings: list | None = None,
|
||||
prompt_cache_enabled: bool = False,
|
||||
litellmrouter_config: dict[str, Any] | None = None,
|
||||
gemini_service_tier: str | None = None,
|
||||
timeout: float | None = None,
|
||||
) -> Any: # Returns LLMInterface
|
||||
"""
|
||||
Factory function to create the appropriate LLM provider implementation.
|
||||
@@ -266,17 +267,26 @@ def create_llm_provider(
|
||||
groq_service_tier: Groq service tier (for Groq provider) - "on_demand", "flex", or "auto".
|
||||
openai_service_tier: OpenAI service tier (for OpenAI provider) - None (default) or "flex" (50% cheaper).
|
||||
bedrock_service_tier: Bedrock service tier (for Bedrock provider) - None (default), "flex", "priority", or "reserved".
|
||||
gemini_service_tier: Gemini service tier (for Gemini provider) - None (default) or "flex" (50% cheaper).
|
||||
extra_body: Extra request-body params merged into the provider's native
|
||||
call. Threaded into OpenAI-compatible, Fireworks, Anthropic, Gemini/
|
||||
VertexAI and LiteLLM providers (each merges them in its own parameter
|
||||
space). Keys must use each provider's native names (e.g. ``max_tokens``
|
||||
for OpenAI/Anthropic vs ``max_output_tokens`` for Gemini).
|
||||
default_headers: Custom headers passed as ``default_headers`` to provider SDK clients
|
||||
(used by operators routing through proxies / request-tracing middleware). Currently
|
||||
wired into the Anthropic provider; other providers may opt in as needed.
|
||||
default_headers: Custom headers passed to provider SDK clients (used by operators
|
||||
routing through proxies / request-tracing middleware). Wired into the Anthropic
|
||||
provider (SDK ``default_headers``) and the LiteLLM-backed providers — ``litellm``,
|
||||
``litellmrouter`` and ``bedrock`` — as the LiteLLM ``extra_headers`` completion
|
||||
kwarg; other providers may opt in as needed.
|
||||
vertexai_project_id: Vertex AI project ID (for VertexAI provider).
|
||||
vertexai_region: Vertex AI region (for VertexAI provider).
|
||||
vertexai_credentials: Vertex AI credentials object (for VertexAI provider).
|
||||
timeout: Per-request LLM timeout in seconds (resolved by the caller from the
|
||||
per-operation/global config). Threaded into the providers that honour a
|
||||
configurable request timeout (LiteLLM, LiteLLM Router, OpenAI-compatible,
|
||||
Nous). ``None`` lets each provider fall back to its own default
|
||||
(``HINDSIGHT_API_LLM_TIMEOUT`` / ``DEFAULT_LLM_TIMEOUT`` for those four;
|
||||
Anthropic and Gemini keep their provider-specific defaults).
|
||||
|
||||
Returns:
|
||||
LLMInterface implementation for the specified provider.
|
||||
@@ -296,6 +306,12 @@ def create_llm_provider(
|
||||
)
|
||||
|
||||
provider_lower = provider.lower()
|
||||
if provider_lower == "gemini":
|
||||
from ..config import parse_gemini_service_tier
|
||||
|
||||
gemini_service_tier = parse_gemini_service_tier(gemini_service_tier)
|
||||
else:
|
||||
gemini_service_tier = None
|
||||
|
||||
if provider_lower == "openai-codex":
|
||||
return CodexLLM(
|
||||
@@ -344,6 +360,7 @@ def create_llm_provider(
|
||||
vertexai_region=vertexai_region,
|
||||
vertexai_credentials=vertexai_credentials,
|
||||
gemini_safety_settings=gemini_safety_settings,
|
||||
gemini_service_tier=gemini_service_tier,
|
||||
prompt_cache_enabled=prompt_cache_enabled,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
@@ -367,6 +384,8 @@ def create_llm_provider(
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
elif provider_lower == "litellmrouter":
|
||||
@@ -385,6 +404,8 @@ def create_llm_provider(
|
||||
config=litellmrouter_config,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
elif provider_lower == "bedrock":
|
||||
@@ -397,7 +418,9 @@ def create_llm_provider(
|
||||
model=bedrock_model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
bedrock_service_tier=bedrock_service_tier,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
elif provider_lower == "llamacpp":
|
||||
@@ -444,6 +467,7 @@ def create_llm_provider(
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
elif provider_lower in (
|
||||
@@ -456,8 +480,10 @@ def create_llm_provider(
|
||||
"deepseek",
|
||||
"volcano",
|
||||
"openrouter",
|
||||
"requesty",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
):
|
||||
return OpenAICompatibleLLM(
|
||||
provider=provider,
|
||||
@@ -468,6 +494,7 @@ def create_llm_provider(
|
||||
groq_service_tier=groq_service_tier,
|
||||
openai_service_tier=openai_service_tier,
|
||||
extra_body=extra_body,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
else:
|
||||
@@ -496,6 +523,14 @@ class LLMProvider:
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
default_headers: dict[str, str] | None = None,
|
||||
litellmrouter_config: dict[str, Any] | None = None,
|
||||
gemini_service_tier: str | None = None,
|
||||
vertexai_project_id: str | None = None,
|
||||
vertexai_region: str | None = None,
|
||||
vertexai_service_account_key: str | None = None,
|
||||
timeout: float | None = None,
|
||||
max_retries: int | None = None,
|
||||
initial_backoff: float | None = None,
|
||||
max_backoff: float | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize LLM provider.
|
||||
@@ -509,29 +544,60 @@ class LLMProvider:
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto") - from config.
|
||||
openai_service_tier: OpenAI service tier (None or "flex") - from config.
|
||||
bedrock_service_tier: Bedrock service tier (None, "flex", "priority", "reserved") - from config.
|
||||
gemini_service_tier: Gemini service tier (None or "flex") - from config.
|
||||
gemini_safety_settings: Safety settings for Gemini/VertexAI providers.
|
||||
extra_body: Extra request-body params merged into the provider's native call
|
||||
(OpenAI-compatible, Fireworks, Anthropic, Gemini/VertexAI, LiteLLM).
|
||||
default_headers: Custom headers passed as ``default_headers`` to provider SDK clients.
|
||||
Used by operators routing through proxies / request-tracing middleware. Falls
|
||||
back to ``HindsightConfig.llm_default_headers`` (env: ``HINDSIGHT_API_LLM_DEFAULT_HEADERS``)
|
||||
when ``None``.
|
||||
Used by operators routing through proxies / request-tracing middleware.
|
||||
litellmrouter_config: Provider-specific config for ``provider="litellmrouter"``.
|
||||
JSON object passed verbatim to ``litellm.Router(**config)`` — see
|
||||
https://docs.litellm.ai/docs/routing. Ignored unless ``provider == "litellmrouter"``.
|
||||
When None and the provider is ``litellmrouter``, falls back to
|
||||
``HindsightConfig.llm_litellmrouter_config``.
|
||||
vertexai_project_id: Vertex AI project ID for ``provider="vertexai"`` (required for
|
||||
that provider).
|
||||
vertexai_region: Vertex AI region for ``provider="vertexai"`` (defaults to
|
||||
``"us-central1"`` when ``None``).
|
||||
vertexai_service_account_key: Path to a Vertex AI service-account key file for
|
||||
``provider="vertexai"`` (uses ADC when ``None``).
|
||||
timeout: Per-request LLM timeout in seconds. Resolved by the caller from the
|
||||
per-operation/global config (``retain_llm_timeout`` falling back to
|
||||
``llm_timeout``, etc.). ``None`` lets each provider apply its own default.
|
||||
max_retries: Default retry-attempt budget for ``call`` / ``call_with_tools``
|
||||
when the per-call argument is omitted. Resolved by the caller from the
|
||||
per-operation/global config (``reflect_llm_max_retries`` falling back to
|
||||
``llm_max_retries``, etc.). ``None`` keeps each method's own fallback.
|
||||
initial_backoff: Default initial retry backoff (seconds), same resolution as
|
||||
``max_retries``. ``None`` keeps each method's own fallback.
|
||||
max_backoff: Default maximum retry backoff (seconds), same resolution as
|
||||
``max_retries``. ``None`` keeps each method's own fallback.
|
||||
|
||||
This constructor uses every argument as passed and does not read global
|
||||
``HindsightConfig``: resolving the server-level default for a ``None`` argument is the
|
||||
caller's responsibility (see ``MemoryEngine``'s per-op builds, ``_member_to_llm``, and
|
||||
``LLMProvider.from_env``). Keeping it config-free makes a provider's effective settings a
|
||||
pure function of its arguments — which is what lets each member of a multi-LLM chain be
|
||||
configured independently.
|
||||
"""
|
||||
self.provider = provider.lower()
|
||||
self.api_key = api_key
|
||||
self.base_url = base_url
|
||||
self.model = model
|
||||
self.reasoning_effort = reasoning_effort
|
||||
# Per-request timeout (seconds). Used verbatim — the caller resolves the
|
||||
# per-operation/global fallback. ``None`` defers to the provider default.
|
||||
self.timeout = timeout
|
||||
# Default retry policy for call()/call_with_tools(). The caller resolves the
|
||||
# per-operation/global fallback; ``None`` keeps each method's own fallback so
|
||||
# providers built without a resolved config (from_env, tests) are unchanged.
|
||||
self.max_retries = max_retries
|
||||
self.initial_backoff = initial_backoff
|
||||
self.max_backoff = max_backoff
|
||||
self.litellmrouter_config = litellmrouter_config
|
||||
# Service tiers from hierarchical config (not env vars)
|
||||
self.groq_service_tier = groq_service_tier
|
||||
self.openai_service_tier = openai_service_tier
|
||||
self.bedrock_service_tier = bedrock_service_tier
|
||||
self.gemini_service_tier = gemini_service_tier
|
||||
# Gemini safety settings (instance default; can be overridden per-request via context var)
|
||||
self.gemini_safety_settings = gemini_safety_settings
|
||||
# Gemini prompt caching: when True, retain extraction (and any future
|
||||
@@ -542,16 +608,9 @@ class LLMProvider:
|
||||
# Extra body params for OpenAI-compatible providers (e.g. chat_template_kwargs)
|
||||
self.extra_body = extra_body
|
||||
# Default headers passed to provider SDK clients (e.g. proxy auth, request tracing).
|
||||
# Same pattern as ``gemini_safety_settings``: explicit override wins; otherwise read
|
||||
# the static server-level default from ``HindsightConfig`` via ``_get_raw_config()``.
|
||||
# Used verbatim — callers resolve the global fallback (see _member_to_llm /
|
||||
# the per-op builds in MemoryEngine, and LLMProvider.from_env).
|
||||
self.default_headers = default_headers
|
||||
if self.default_headers is None:
|
||||
from ..config import _get_raw_config
|
||||
|
||||
try:
|
||||
self.default_headers = _get_raw_config().llm_default_headers
|
||||
except Exception:
|
||||
pass # Config may not be initialized in test environments
|
||||
|
||||
# Validate provider
|
||||
valid_providers = [
|
||||
@@ -575,8 +634,10 @@ class LLMProvider:
|
||||
"bedrock",
|
||||
"volcano",
|
||||
"openrouter",
|
||||
"requesty",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
"fireworks",
|
||||
"nous",
|
||||
]
|
||||
@@ -599,32 +660,31 @@ class LLMProvider:
|
||||
self.base_url = "https://api.deepseek.com"
|
||||
elif self.provider == "openrouter":
|
||||
self.base_url = "https://openrouter.ai/api/v1"
|
||||
elif self.provider == "requesty":
|
||||
self.base_url = "https://router.requesty.ai/v1"
|
||||
elif self.provider == "zai":
|
||||
self.base_url = "https://api.z.ai/api/coding/paas/v4"
|
||||
elif self.provider == "opencode-go":
|
||||
self.base_url = "https://opencode.ai/zen/go/v1"
|
||||
elif self.provider == "atlas":
|
||||
self.base_url = "https://api.atlascloud.ai/v1"
|
||||
elif self.provider == "nous":
|
||||
self.base_url = "https://inference-api.nousresearch.com/v1"
|
||||
|
||||
# Prepare Vertex AI config (if applicable)
|
||||
vertexai_project_id = None
|
||||
vertexai_region = None
|
||||
# Prepare Vertex AI config (if applicable). Values are used as passed; the
|
||||
# caller resolves the global-config fallback (MemoryEngine builds /
|
||||
# _member_to_llm / from_env). The region keeps a constant default here.
|
||||
vertexai_credentials = None
|
||||
|
||||
if self.provider == "vertexai":
|
||||
from ..config import get_config
|
||||
|
||||
config = get_config()
|
||||
|
||||
vertexai_project_id = config.llm_vertexai_project_id
|
||||
if not vertexai_project_id:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID is required for Vertex AI provider. "
|
||||
"Set it to your GCP project ID."
|
||||
)
|
||||
|
||||
vertexai_region = config.llm_vertexai_region or "us-central1"
|
||||
service_account_key = config.llm_vertexai_service_account_key
|
||||
vertexai_region = vertexai_region or "us-central1"
|
||||
service_account_key = vertexai_service_account_key
|
||||
|
||||
# Load explicit service account credentials if provided
|
||||
if service_account_key:
|
||||
@@ -648,45 +708,20 @@ class LLMProvider:
|
||||
f"model={self.model}, auth={'service_account' if service_account_key else 'ADC'}"
|
||||
)
|
||||
|
||||
# For Gemini/VertexAI providers: read safety settings from global config if not explicitly provided
|
||||
# Use _get_raw_config() to bypass StaticConfigProxy (which blocks configurable fields),
|
||||
# since LLMProvider initialization legitimately needs the server-level default.
|
||||
if self.provider in ("gemini", "vertexai") and self.gemini_safety_settings is None:
|
||||
from ..config import _get_raw_config
|
||||
# Normalize the Gemini service tier (pure: maps/validates the passed value,
|
||||
# no global config read). Non-Gemini providers never carry a tier. The
|
||||
# server-level default is resolved by the caller, like the other fields.
|
||||
if self.provider == "gemini":
|
||||
from ..config import parse_gemini_service_tier
|
||||
|
||||
try:
|
||||
raw_config = _get_raw_config()
|
||||
self.gemini_safety_settings = raw_config.llm_gemini_safety_settings
|
||||
except Exception:
|
||||
pass # Config may not be initialized in test environments
|
||||
self.gemini_service_tier = parse_gemini_service_tier(self.gemini_service_tier)
|
||||
else:
|
||||
self.gemini_service_tier = None
|
||||
|
||||
# Prompt-prefix caching is a provider-agnostic toggle (default on): resolve
|
||||
# it from the static server config for every provider when the caller didn't
|
||||
# pass an explicit override. Providers that don't support caching ignore the
|
||||
# value; only those that implement get_or_create_cached_prefix act on it.
|
||||
if not self.prompt_cache_enabled:
|
||||
from ..config import DEFAULT_LLM_PROMPT_CACHE_ENABLED, _get_raw_config
|
||||
|
||||
try:
|
||||
raw_config = _get_raw_config()
|
||||
self.prompt_cache_enabled = bool(
|
||||
getattr(raw_config, "llm_prompt_cache_enabled", DEFAULT_LLM_PROMPT_CACHE_ENABLED)
|
||||
)
|
||||
except Exception:
|
||||
pass # Config may not be initialized in test environments
|
||||
|
||||
# For litellmrouter: prefer an explicit chain from the caller (per-op
|
||||
# construction in MemoryEngine threads the right chain through). If the caller
|
||||
# didn't supply one, fall back to the global ``llm_litellmrouter_config`` so
|
||||
# ad-hoc constructions (e.g. ``LLMProvider.from_env()``) keep working.
|
||||
# gemini_safety_settings / prompt_cache_enabled / litellmrouter_config are
|
||||
# used as passed — the caller resolves the global-config fallback. Providers
|
||||
# that don't support prompt caching ignore the flag.
|
||||
router_config: dict[str, Any] | None = self.litellmrouter_config
|
||||
if self.provider == "litellmrouter" and router_config is None:
|
||||
from ..config import _get_raw_config
|
||||
|
||||
try:
|
||||
router_config = _get_raw_config().llm_litellmrouter_config
|
||||
except Exception:
|
||||
router_config = None
|
||||
|
||||
# Create provider implementation using factory
|
||||
self._provider_impl = create_llm_provider(
|
||||
@@ -698,6 +733,7 @@ class LLMProvider:
|
||||
groq_service_tier=self.groq_service_tier,
|
||||
openai_service_tier=self.openai_service_tier,
|
||||
bedrock_service_tier=self.bedrock_service_tier,
|
||||
gemini_service_tier=self.gemini_service_tier,
|
||||
extra_body=self.extra_body,
|
||||
default_headers=self.default_headers,
|
||||
vertexai_project_id=vertexai_project_id,
|
||||
@@ -706,6 +742,7 @@ class LLMProvider:
|
||||
gemini_safety_settings=self.gemini_safety_settings,
|
||||
prompt_cache_enabled=self.prompt_cache_enabled,
|
||||
litellmrouter_config=router_config,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
# Backward compatibility: Keep mock provider properties
|
||||
@@ -762,9 +799,9 @@ class LLMProvider:
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "memory",
|
||||
max_retries: int = 10,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 60.0,
|
||||
max_retries: int | None = None,
|
||||
initial_backoff: float | None = None,
|
||||
max_backoff: float | None = None,
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
@@ -779,9 +816,12 @@ class LLMProvider:
|
||||
max_completion_tokens: Maximum tokens in response.
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
max_retries: Maximum retry attempts. ``None`` uses the provider's configured
|
||||
default (per-operation/global ``llm_max_retries``), else 10.
|
||||
initial_backoff: Initial backoff time in seconds. ``None`` uses the provider's
|
||||
configured default (``llm_initial_backoff``), else 1.0.
|
||||
max_backoff: Maximum backoff time in seconds. ``None`` uses the provider's
|
||||
configured default (``llm_max_backoff``), else 60.0.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Per-call override requesting grammar-enforced (json_schema strict)
|
||||
structured output instead of the soft json_object path. The server-level
|
||||
@@ -806,6 +846,20 @@ class LLMProvider:
|
||||
structured = "+structured" if response_format is not None else ""
|
||||
set_stage(f"llm.{self.provider}.{scope}{structured}")
|
||||
|
||||
# Resolve the retry policy: explicit per-call arg wins, else the provider's
|
||||
# configured per-operation/global default, else this method's own fallback.
|
||||
max_retries = (
|
||||
max_retries if max_retries is not None else (self.max_retries if self.max_retries is not None else 10)
|
||||
)
|
||||
initial_backoff = (
|
||||
initial_backoff
|
||||
if initial_backoff is not None
|
||||
else (self.initial_backoff if self.initial_backoff is not None else 1.0)
|
||||
)
|
||||
max_backoff = (
|
||||
max_backoff if max_backoff is not None else (self.max_backoff if self.max_backoff is not None else 60.0)
|
||||
)
|
||||
|
||||
# Resolve strict-schema once, here, rather than in each provider: the
|
||||
# per-call argument OR the server-level HINDSIGHT_API_LLM_STRICT_SCHEMA
|
||||
# flag. Providers with a json_schema response_format (OpenAI-compatible,
|
||||
@@ -822,7 +876,13 @@ class LLMProvider:
|
||||
# The requested params are stashed in a contextvar (only what the caller
|
||||
# actually set) so the recorder can attach them to either path.
|
||||
from ..tracing import get_span_recorder
|
||||
from .llm_trace import reset_request_context, set_request_context
|
||||
from .llm_trace import (
|
||||
current_response_usage,
|
||||
reset_request_context,
|
||||
reset_response_usage,
|
||||
set_request_context,
|
||||
set_response_usage,
|
||||
)
|
||||
|
||||
call_start = time.monotonic()
|
||||
request_token = set_request_context(
|
||||
@@ -833,6 +893,9 @@ class LLMProvider:
|
||||
response_format=response_format,
|
||||
)
|
||||
)
|
||||
# Cleared per call; the provider stashes real usage once a response is in
|
||||
# hand so the error path below can attach it if parsing/validation fails.
|
||||
usage_token = set_response_usage(None)
|
||||
try:
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
@@ -860,14 +923,19 @@ class LLMProvider:
|
||||
**cache_kwarg,
|
||||
)
|
||||
except Exception as e:
|
||||
# The provider call may have succeeded (and incurred token
|
||||
# cost) before local parsing/validation raised; attach the
|
||||
# provider-reported usage to the error trace when available.
|
||||
usage = current_response_usage()
|
||||
get_span_recorder().record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=None,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
input_tokens=usage.input_tokens if usage else 0,
|
||||
output_tokens=usage.output_tokens if usage else 0,
|
||||
cached_tokens=usage.cached_tokens if usage else 0,
|
||||
duration=time.monotonic() - call_start,
|
||||
error=e,
|
||||
)
|
||||
@@ -883,6 +951,7 @@ class LLMProvider:
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
finally:
|
||||
reset_request_context(request_token)
|
||||
reset_response_usage(usage_token)
|
||||
|
||||
return result
|
||||
|
||||
@@ -893,9 +962,9 @@ class LLMProvider:
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "tools",
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
max_retries: int | None = None,
|
||||
initial_backoff: float | None = None,
|
||||
max_backoff: float | None = None,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
cached_prefix: str | None = None,
|
||||
) -> "LLMToolCallResult":
|
||||
@@ -908,9 +977,12 @@ class LLMProvider:
|
||||
max_completion_tokens: Maximum tokens in response.
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
max_retries: Maximum retry attempts. ``None`` uses the provider's configured
|
||||
default (per-operation/global ``llm_max_retries``), else 5.
|
||||
initial_backoff: Initial backoff time in seconds. ``None`` uses the provider's
|
||||
configured default (``llm_initial_backoff``), else 1.0.
|
||||
max_backoff: Maximum backoff time in seconds. ``None`` uses the provider's
|
||||
configured default (``llm_max_backoff``), else 30.0.
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or {"type": "function", "function": {"name": "..."}}
|
||||
|
||||
Returns:
|
||||
@@ -920,9 +992,29 @@ class LLMProvider:
|
||||
|
||||
set_stage(f"llm.{self.provider}.{scope}+tools")
|
||||
|
||||
# Resolve the retry policy: explicit per-call arg wins, else the provider's
|
||||
# configured per-operation/global default, else this method's own fallback.
|
||||
max_retries = (
|
||||
max_retries if max_retries is not None else (self.max_retries if self.max_retries is not None else 5)
|
||||
)
|
||||
initial_backoff = (
|
||||
initial_backoff
|
||||
if initial_backoff is not None
|
||||
else (self.initial_backoff if self.initial_backoff is not None else 1.0)
|
||||
)
|
||||
max_backoff = (
|
||||
max_backoff if max_backoff is not None else (self.max_backoff if self.max_backoff is not None else 30.0)
|
||||
)
|
||||
|
||||
# Failures forwarded to the GenAI recorder; successes recorded by providers.
|
||||
from ..tracing import get_span_recorder
|
||||
from .llm_trace import reset_request_context, set_request_context
|
||||
from .llm_trace import (
|
||||
current_response_usage,
|
||||
reset_request_context,
|
||||
reset_response_usage,
|
||||
set_request_context,
|
||||
set_response_usage,
|
||||
)
|
||||
|
||||
call_start = time.monotonic()
|
||||
request_token = set_request_context(
|
||||
@@ -933,6 +1025,9 @@ class LLMProvider:
|
||||
tool_choice=tool_choice,
|
||||
)
|
||||
)
|
||||
# Cleared per call; the provider stashes real usage once a response is in
|
||||
# hand so the error path below can attach it if parsing/validation fails.
|
||||
usage_token = set_response_usage(None)
|
||||
try:
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
@@ -957,14 +1052,19 @@ class LLMProvider:
|
||||
**cache_kwarg,
|
||||
)
|
||||
except Exception as e:
|
||||
# The provider call may have succeeded (and incurred token
|
||||
# cost) before local parsing/validation raised; attach the
|
||||
# provider-reported usage to the error trace when available.
|
||||
usage = current_response_usage()
|
||||
get_span_recorder().record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=None,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
input_tokens=usage.input_tokens if usage else 0,
|
||||
output_tokens=usage.output_tokens if usage else 0,
|
||||
cached_tokens=usage.cached_tokens if usage else 0,
|
||||
duration=time.monotonic() - call_start,
|
||||
error=e,
|
||||
)
|
||||
@@ -980,6 +1080,7 @@ class LLMProvider:
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
finally:
|
||||
reset_request_context(request_token)
|
||||
reset_response_usage(usage_token)
|
||||
|
||||
return result
|
||||
|
||||
@@ -1023,7 +1124,9 @@ class LLMProvider:
|
||||
|
||||
def _load_codex_auth(self) -> tuple[str, str]:
|
||||
"""
|
||||
Load OAuth credentials from ~/.codex/auth.json.
|
||||
Load OAuth credentials from the Codex ``auth.json``.
|
||||
|
||||
Honors ``CODEX_HOME`` (falling back to ``~/.codex``).
|
||||
|
||||
Returns:
|
||||
Tuple of (access_token, account_id).
|
||||
@@ -1032,7 +1135,9 @@ class LLMProvider:
|
||||
FileNotFoundError: If auth file doesn't exist.
|
||||
ValueError: If auth file is invalid.
|
||||
"""
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
from .providers.codex_auth import default_codex_auth_file
|
||||
|
||||
auth_file = default_codex_auth_file()
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(
|
||||
@@ -1134,18 +1239,38 @@ class LLMProvider:
|
||||
@classmethod
|
||||
def from_env(cls) -> "LLMProvider":
|
||||
"""Create provider from environment variables using config.py constants."""
|
||||
# Read every field straight from the environment. The constructor no longer
|
||||
# resolves global-config fallbacks, so this factory must supply them — and it
|
||||
# does so without building the full HindsightConfig, keeping from_env() a
|
||||
# lightweight env-only loader (see test_llm_provider_from_env_keeps_lightweight_loader).
|
||||
from ..config import (
|
||||
DEFAULT_LLM_GROQ_SERVICE_TIER,
|
||||
DEFAULT_LLM_OPENAI_SERVICE_TIER,
|
||||
DEFAULT_LLM_PROMPT_CACHE_ENABLED,
|
||||
DEFAULT_LLM_PROVIDER,
|
||||
DEFAULT_LLM_REASONING_EFFORT,
|
||||
DEFAULT_LLM_TIMEOUT,
|
||||
ENV_LLM_API_KEY,
|
||||
ENV_LLM_BASE_URL,
|
||||
ENV_LLM_BEDROCK_SERVICE_TIER,
|
||||
ENV_LLM_DEFAULT_HEADERS,
|
||||
ENV_LLM_EXTRA_BODY,
|
||||
ENV_LLM_GEMINI_SAFETY_SETTINGS,
|
||||
ENV_LLM_GEMINI_SERVICE_TIER,
|
||||
ENV_LLM_GROQ_SERVICE_TIER,
|
||||
ENV_LLM_LITELLMROUTER_CONFIG,
|
||||
ENV_LLM_MODEL,
|
||||
ENV_LLM_OPENAI_SERVICE_TIER,
|
||||
ENV_LLM_PROMPT_CACHE_ENABLED,
|
||||
ENV_LLM_PROVIDER,
|
||||
ENV_LLM_REASONING_EFFORT,
|
||||
ENV_LLM_TIMEOUT,
|
||||
ENV_LLM_VERTEXAI_PROJECT_ID,
|
||||
ENV_LLM_VERTEXAI_REGION,
|
||||
ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY,
|
||||
_get_default_model_for_provider,
|
||||
_parse_llm_router_config,
|
||||
parse_gemini_service_tier,
|
||||
)
|
||||
|
||||
provider = os.getenv(ENV_LLM_PROVIDER, DEFAULT_LLM_PROVIDER)
|
||||
@@ -1162,6 +1287,14 @@ class LLMProvider:
|
||||
model = os.getenv(ENV_LLM_MODEL) or _get_default_model_for_provider(provider)
|
||||
extra_body = json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null"))
|
||||
default_headers = json.loads(os.getenv(ENV_LLM_DEFAULT_HEADERS, "null"))
|
||||
prompt_cache_enabled = os.getenv(
|
||||
ENV_LLM_PROMPT_CACHE_ENABLED, str(DEFAULT_LLM_PROMPT_CACHE_ENABLED)
|
||||
).lower() in (
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
"on",
|
||||
)
|
||||
|
||||
return cls(
|
||||
provider=provider,
|
||||
@@ -1171,7 +1304,21 @@ class LLMProvider:
|
||||
reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
|
||||
openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
|
||||
bedrock_service_tier=os.getenv(ENV_LLM_BEDROCK_SERVICE_TIER) or None,
|
||||
gemini_service_tier=(
|
||||
parse_gemini_service_tier(os.getenv(ENV_LLM_GEMINI_SERVICE_TIER))
|
||||
if provider.lower() == "gemini"
|
||||
else None
|
||||
),
|
||||
gemini_safety_settings=json.loads(os.getenv(ENV_LLM_GEMINI_SAFETY_SETTINGS, "null")),
|
||||
prompt_cache_enabled=prompt_cache_enabled,
|
||||
litellmrouter_config=_parse_llm_router_config(ENV_LLM_LITELLMROUTER_CONFIG),
|
||||
vertexai_project_id=os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID) or None,
|
||||
vertexai_region=os.getenv(ENV_LLM_VERTEXAI_REGION) or None,
|
||||
vertexai_service_account_key=os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY) or None,
|
||||
timeout=float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT))),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -11,6 +11,12 @@ from one place, so we don't spawn a separate ``asyncio`` task per concern:
|
||||
consolidation operation failed terminally and left them with
|
||||
``consolidated_at IS NULL AND consolidation_failed_at IS NULL`` and nothing to
|
||||
re-trigger them.
|
||||
- **Scheduled mental model refresh** (configurable check cadence, default 60s):
|
||||
refresh mental models whose ``trigger.refresh_cron`` schedule is due, but only
|
||||
when the model is stale (new memories in its scope since its last refresh), so
|
||||
a scheduled tick never burns an LLM call to regenerate identical content. The
|
||||
per-model schedule lives in the cron expression; this loop only decides when to
|
||||
*check*.
|
||||
|
||||
The loop wakes on a short fixed tick and runs each job when its own
|
||||
``last_run + interval`` is due (run-at-start, then on interval), so adding jobs
|
||||
@@ -25,12 +31,14 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
from collections.abc import Coroutine
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..config import HindsightConfig, get_config
|
||||
from ..models import RequestContext
|
||||
from .db_utils import acquire_with_retry
|
||||
from .schema import _is_oracle
|
||||
from .schema import _is_oracle, fq_table
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .memory_engine import MemoryEngine
|
||||
@@ -91,7 +99,8 @@ class MaintenanceLoop:
|
||||
reconcile_on = cfg.consolidation_reconcile_interval_seconds > 0
|
||||
audit_on = cfg.audit_log_enabled and cfg.audit_log_retention_days > 0
|
||||
llm_on = cfg.llm_trace_enabled and cfg.llm_trace_retention_days > 0
|
||||
return reconcile_on or audit_on or llm_on
|
||||
mm_refresh_on = cfg.mental_model_refresh_tick_seconds > 0
|
||||
return reconcile_on or audit_on or llm_on or mm_refresh_on
|
||||
|
||||
# ── loop ───────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -118,10 +127,25 @@ class MaintenanceLoop:
|
||||
async def _tick(self) -> None:
|
||||
cfg = get_config()
|
||||
if self._is_due("retention", _RETENTION_INTERVAL_SECONDS):
|
||||
await self._run_retention(cfg)
|
||||
await self._run_timed("retention", self._run_retention(cfg))
|
||||
interval = cfg.consolidation_reconcile_interval_seconds
|
||||
if interval > 0 and self._is_due("reconcile", interval):
|
||||
await self._run_reconcile()
|
||||
await self._run_timed("consolidation reconcile", self._run_reconcile())
|
||||
mm_interval = cfg.mental_model_refresh_tick_seconds
|
||||
if mm_interval > 0 and self._is_due("mm_refresh", mm_interval):
|
||||
await self._run_timed("scheduled mental model refresh", self._run_scheduled_mm_refresh())
|
||||
|
||||
async def _run_timed(self, name: str, coro: Coroutine[Any, Any, None]) -> None:
|
||||
"""Run a maintenance job and emit one timing line for it.
|
||||
|
||||
Each job keeps its own summary log (counts of work done); this adds a
|
||||
single, uniform line per run so the cost of every sweep is observable.
|
||||
"""
|
||||
start = time.monotonic()
|
||||
try:
|
||||
await coro
|
||||
finally:
|
||||
logger.info(f"Maintenance: {name} took {time.monotonic() - start:.3f}s")
|
||||
|
||||
# ── retention ──────────────────────────────────────────────────────────
|
||||
|
||||
@@ -212,3 +236,112 @@ class MaintenanceLoop:
|
||||
f"Consolidation reconcile: scheduled {submitted} bank(s)"
|
||||
+ (f", skipped {skipped_unknown} in unrecognized schema(s)" if skipped_unknown else "")
|
||||
)
|
||||
|
||||
# ── scheduled mental model refresh ───────────────────────────────────────
|
||||
|
||||
async def _run_scheduled_mm_refresh(self) -> None:
|
||||
"""Refresh mental models whose ``trigger.refresh_cron`` is due.
|
||||
|
||||
Discovery (the set of cron-scheduled models, minus any with an in-flight
|
||||
refresh) is one cross-tenant round-trip via
|
||||
``public.mental_models_with_cron()``. Cron *due-ness* is evaluated here in
|
||||
Python — a scheduled fire has elapsed when the most recent cron boundary at
|
||||
or before now is later than ``last_refreshed_at`` — because cron arithmetic
|
||||
isn't expressible in plain SQL. Each due model is refreshed only when it is
|
||||
actually stale, so a schedule that fires while nothing changed costs a
|
||||
cheap staleness query, not an LLM call.
|
||||
"""
|
||||
engine = self._engine
|
||||
try:
|
||||
async with acquire_with_retry(engine._backend, max_retries=1) as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT schema_name, bank_id, mental_model_id, refresh_cron, last_refreshed_at "
|
||||
"FROM public.mental_models_with_cron()"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Scheduled mental model refresh discovery failed: {e}")
|
||||
return
|
||||
if not rows:
|
||||
return
|
||||
|
||||
from croniter import croniter
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
due = []
|
||||
for row in rows:
|
||||
cron = row["refresh_cron"]
|
||||
last = row["last_refreshed_at"]
|
||||
try:
|
||||
prev_fire = croniter(cron, now).get_prev(datetime)
|
||||
except (ValueError, KeyError) as e:
|
||||
logger.warning(
|
||||
f"Scheduled mental model refresh: skipping invalid cron {cron!r} for "
|
||||
f"{row['schema_name']}/{row['mental_model_id']}: {e}"
|
||||
)
|
||||
continue
|
||||
if last is None or prev_fire > last:
|
||||
due.append(row)
|
||||
if not due:
|
||||
return
|
||||
|
||||
# Only enqueue into schemas the worker actually polls (tenant discovery),
|
||||
# otherwise the op would never be claimed. The tenant_id (when provided)
|
||||
# lets config resolution honor tenant-level overrides.
|
||||
try:
|
||||
tenants = await engine._tenant_extension.list_tenants()
|
||||
except Exception as e:
|
||||
logger.warning(f"Scheduled mental model refresh tenant discovery failed: {e}")
|
||||
return
|
||||
tenant_by_schema = {t.schema: t for t in tenants}
|
||||
default_schema = get_config().database_schema
|
||||
|
||||
from .memory_engine import _current_schema
|
||||
|
||||
submitted = 0
|
||||
skipped_unknown = 0
|
||||
skipped_fresh = 0
|
||||
for row in due:
|
||||
schema = row["schema_name"]
|
||||
bank_id = row["bank_id"]
|
||||
mm_id = row["mental_model_id"]
|
||||
tenant = tenant_by_schema.get(schema)
|
||||
if tenant is None and schema != default_schema:
|
||||
skipped_unknown += 1
|
||||
continue
|
||||
tenant_id = tenant.tenant_id if tenant else None
|
||||
token = _current_schema.set(schema)
|
||||
try:
|
||||
context = RequestContext(internal=True, tenant_id=tenant_id)
|
||||
# Skip if nothing in the model's scope changed since its last
|
||||
# refresh — a scheduled refresh must not regenerate identical
|
||||
# content. compute_mental_model_is_stale needs the model's tags +
|
||||
# trigger, which the discovery routine doesn't return, so re-read
|
||||
# the row under the bank's schema context.
|
||||
async with acquire_with_retry(engine._backend, max_retries=1) as conn:
|
||||
mm_row = await conn.fetchrow(
|
||||
f"SELECT id, tags, trigger, last_refreshed_at FROM {fq_table('mental_models')} "
|
||||
"WHERE bank_id = $1 AND id = $2",
|
||||
bank_id,
|
||||
mm_id,
|
||||
)
|
||||
if mm_row is None:
|
||||
continue
|
||||
is_stale = await engine.compute_mental_model_is_stale(conn, bank_id, mm_row)
|
||||
if not is_stale:
|
||||
skipped_fresh += 1
|
||||
continue
|
||||
await engine.submit_async_refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm_id, request_context=context
|
||||
)
|
||||
submitted += 1
|
||||
except Exception as e:
|
||||
logger.warning(f"Scheduled mental model refresh failed for {mm_id} in {schema}: {e}")
|
||||
finally:
|
||||
_current_schema.reset(token)
|
||||
|
||||
if submitted or skipped_unknown or skipped_fresh:
|
||||
logger.info(
|
||||
f"Scheduled mental model refresh: scheduled {submitted} model(s)"
|
||||
+ (f", {skipped_fresh} up-to-date" if skipped_fresh else "")
|
||||
+ (f", skipped {skipped_unknown} in unrecognized schema(s)" if skipped_unknown else "")
|
||||
)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,204 @@
|
||||
"""Multi-LLM routing: failover and (weighted) round-robin across N providers.
|
||||
|
||||
``MultiLLMProvider`` wraps an ordered list of :class:`LLMProvider` members and a
|
||||
:class:`~hindsight_api.config.LLMStrategyConfig`, exposing the same public surface
|
||||
as a single ``LLMProvider`` so it drops into every existing call path (including
|
||||
``with_config()`` / ``ConfiguredLLMProvider``).
|
||||
|
||||
Member 0 is the **primary** (the operation's unindexed/base LLM); members 1..N are
|
||||
the indexed extras (``HINDSIGHT_API_<OP>LLM_<n>_*``). Each member keeps its own
|
||||
internal retry budget, so we only advance to the next member after a member has
|
||||
exhausted its retries and raised.
|
||||
|
||||
Strategies:
|
||||
- ``failover``: try members in declared order ``[0..N]``.
|
||||
- ``round-robin``: rotate the starting member per request (optionally weighted),
|
||||
then fall through the remaining members on error.
|
||||
|
||||
Batch retain and any direct ``_provider_impl`` access operate on the **primary
|
||||
member only** (via attribute passthrough) — failover/round-robin apply to the
|
||||
interactive ``call`` / ``call_with_tools`` paths.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..config import LLM_STRATEGY_FAILOVER, LLMStrategyConfig
|
||||
from .llm_wrapper import LLMProvider, OutputTooLongError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .llm_wrapper import ConfiguredLLMProvider, LLMToolCallResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _should_failover(exc: BaseException) -> bool:
|
||||
"""Whether ``exc`` from one member should trigger a try on the next member.
|
||||
|
||||
Generic ``Exception`` instances (network errors, provider 5xx, timeouts after
|
||||
a member's own retries) fail over. ``OutputTooLongError`` is propagated — a
|
||||
different provider won't fit an over-length output either. ``CancelledError``,
|
||||
``KeyboardInterrupt`` and ``SystemExit`` are ``BaseException`` (not
|
||||
``Exception``) and therefore propagate unchanged.
|
||||
"""
|
||||
if isinstance(exc, OutputTooLongError):
|
||||
return False
|
||||
return isinstance(exc, Exception)
|
||||
|
||||
|
||||
class _WeightedRoundRobin:
|
||||
"""Smooth weighted round-robin scheduler (nginx SWRR).
|
||||
|
||||
Produces a starting member index per request such that, over time, member
|
||||
``i`` is chosen in proportion to ``weights[i]`` while keeping selections
|
||||
interleaved rather than bursty. Uniform weights degrade to plain round-robin.
|
||||
The tiny selection critical section is mutex-guarded so concurrent callers
|
||||
don't corrupt the running totals (they may still interleave, which only
|
||||
affects distribution, never correctness).
|
||||
"""
|
||||
|
||||
def __init__(self, weights: list[int]) -> None:
|
||||
self._weights = list(weights)
|
||||
self._current = [0] * len(weights)
|
||||
self._total = sum(weights)
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def next(self) -> int:
|
||||
with self._lock:
|
||||
best = 0
|
||||
for i, w in enumerate(self._weights):
|
||||
self._current[i] += w
|
||||
if self._current[i] > self._current[best]:
|
||||
best = i
|
||||
self._current[best] -= self._total
|
||||
return best
|
||||
|
||||
|
||||
class MultiLLMProvider:
|
||||
"""Route LLM calls across multiple members per a failover / round-robin strategy."""
|
||||
|
||||
def __init__(self, members: list[LLMProvider], strategy: LLMStrategyConfig) -> None:
|
||||
if not members:
|
||||
raise ValueError("MultiLLMProvider requires at least one member")
|
||||
self._members = members
|
||||
self._strategy = strategy
|
||||
|
||||
weights = strategy.weights or [1] * len(members)
|
||||
if len(weights) != len(members):
|
||||
raise ValueError(
|
||||
f"LLM strategy 'weights' has {len(weights)} entries but the chain has "
|
||||
f"{len(members)} members (primary + indexed); they must match."
|
||||
)
|
||||
self._scheduler = _WeightedRoundRobin(weights)
|
||||
|
||||
# ── routing ────────────────────────────────────────────────────────────────
|
||||
|
||||
def _member_order(self) -> list[int]:
|
||||
"""Indices to try, in order, for one request."""
|
||||
n = len(self._members)
|
||||
if self._strategy.mode == LLM_STRATEGY_FAILOVER:
|
||||
return list(range(n))
|
||||
start = self._scheduler.next()
|
||||
return [(start + i) % n for i in range(n)]
|
||||
|
||||
async def _dispatch(self, method_name: str, **kwargs: Any) -> Any:
|
||||
last_exc: BaseException | None = None
|
||||
order = self._member_order()
|
||||
for position, idx in enumerate(order):
|
||||
member = self._members[idx]
|
||||
try:
|
||||
return await getattr(member, method_name)(**kwargs)
|
||||
except BaseException as e: # noqa: BLE001 - re-raised unless it should fail over
|
||||
if not _should_failover(e):
|
||||
raise
|
||||
last_exc = e
|
||||
remaining = len(order) - position - 1
|
||||
logger.warning(
|
||||
"LLM member %d (%s/%s) failed on %s: %s%s",
|
||||
idx,
|
||||
member.provider,
|
||||
member.model,
|
||||
method_name,
|
||||
e,
|
||||
f"; trying next member ({remaining} left)" if remaining else "; no members left",
|
||||
)
|
||||
# All members failed; surface the last error (loop ran at least once).
|
||||
assert last_exc is not None
|
||||
raise last_exc
|
||||
|
||||
async def call(self, messages: list[dict[str, Any]], **kwargs: Any) -> Any:
|
||||
return await self._dispatch("call", messages=messages, **kwargs)
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]],
|
||||
**kwargs: Any,
|
||||
) -> "LLMToolCallResult":
|
||||
return await self._dispatch("call_with_tools", messages=messages, tools=tools, **kwargs)
|
||||
|
||||
# ── lifecycle ────────────────────────────────────────────────────────────────
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""Strictly verify the primary; soft-verify the rest (warn, don't fail).
|
||||
|
||||
A failover member being unreachable at startup must not block the server —
|
||||
it may come back before it's needed. The primary is the steady-state path,
|
||||
so its failure is still surfaced (the caller already wraps this in a
|
||||
warn-only try/except at startup).
|
||||
"""
|
||||
await self._members[0].verify_connection()
|
||||
for member in self._members[1:]:
|
||||
try:
|
||||
await member.verify_connection()
|
||||
except Exception as e: # noqa: BLE001 - soft verification
|
||||
logger.warning(
|
||||
"Failover LLM member %s/%s failed connection verification: %s. "
|
||||
"It will be tried at request time if the primary fails.",
|
||||
member.provider,
|
||||
member.model,
|
||||
e,
|
||||
)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
for member in self._members:
|
||||
await member.cleanup()
|
||||
|
||||
def with_config(
|
||||
self,
|
||||
config: Any,
|
||||
*,
|
||||
bank_id: str | None = None,
|
||||
operation: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> "ConfiguredLLMProvider":
|
||||
"""Mirror ``LLMProvider.with_config`` so the strategy runs inside the
|
||||
per-operation configured wrapper (gemini-safety + trace contextvars wrap
|
||||
every member call)."""
|
||||
from .llm_trace import LLMTraceContext
|
||||
from .llm_wrapper import ConfiguredLLMProvider
|
||||
|
||||
trace_ctx = None
|
||||
if bank_id is not None or operation is not None or metadata:
|
||||
trace_ctx = LLMTraceContext(
|
||||
bank_id=bank_id,
|
||||
operation=operation,
|
||||
metadata=dict(metadata or {}),
|
||||
trace_id=str(uuid.uuid4()),
|
||||
operation_span_id=str(uuid.uuid4()),
|
||||
)
|
||||
return ConfiguredLLMProvider(self, config.llm_gemini_safety_settings, trace_ctx)
|
||||
|
||||
# ── attribute passthrough ────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def members(self) -> list[LLMProvider]:
|
||||
return self._members
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
# Anything not defined here (provider, model, api_key, base_url,
|
||||
# _provider_impl, mock helpers, batch helpers, ...) delegates to the
|
||||
# primary member so existing call sites keep working unchanged.
|
||||
return getattr(object.__getattribute__(self, "_members")[0], name)
|
||||
@@ -3,43 +3,138 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from hindsight_api.config import DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
|
||||
|
||||
from .base import FileParser
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from markitdown import StreamInfo
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Extensions whose markitdown converters decode the raw bytes as text. markitdown
|
||||
# samples only the first chunk for charset detection, so a UTF-8 file with a long
|
||||
# ASCII-only prefix is mis-detected as ASCII; the JSON/ipynb converter then crashes
|
||||
# decoding the first multibyte byte. Passing an explicit UTF-8 hint when the bytes
|
||||
# are valid UTF-8 sidesteps the faulty detection without affecting other encodings.
|
||||
_TEXT_EXTENSIONS = {
|
||||
".json",
|
||||
".jsonl",
|
||||
".ipynb",
|
||||
".txt",
|
||||
".text",
|
||||
".md",
|
||||
".markdown",
|
||||
".csv",
|
||||
".html",
|
||||
".htm",
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MarkitdownOcrOptions:
|
||||
"""OpenAI-compatible OCR options passed through to MarkItDown."""
|
||||
|
||||
# Keep this typed as object so the OpenAI SDK import stays lazy for non-OCR users.
|
||||
llm_client: object
|
||||
llm_model: str
|
||||
llm_prompt: str
|
||||
|
||||
|
||||
class MarkitdownParser(FileParser):
|
||||
"""
|
||||
Markitdown file parser.
|
||||
|
||||
Uses Microsoft's markitdown library to convert various file formats
|
||||
to markdown including PDF, Office docs, images (via OCR), audio, HTML.
|
||||
to markdown including PDF, Office docs, images with optional OCR,
|
||||
audio, HTML.
|
||||
|
||||
Supported formats:
|
||||
- PDF (.pdf)
|
||||
- Word (.docx, .doc)
|
||||
- PowerPoint (.pptx, .ppt)
|
||||
- Excel (.xlsx, .xls)
|
||||
- Images (.jpg, .jpeg, .png) - with OCR
|
||||
- Images (.jpg, .jpeg, .png) - optional OCR
|
||||
- HTML (.html, .htm)
|
||||
- Text (.txt, .md)
|
||||
- Audio (.mp3, .wav) - with transcription
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
ocr_enabled: bool = False,
|
||||
ocr_api_key: str | None = None,
|
||||
ocr_base_url: str | None = None,
|
||||
ocr_model: str | None = None,
|
||||
ocr_prompt: str | None = None,
|
||||
):
|
||||
"""Initialize markitdown parser."""
|
||||
# Lazy import to avoid requiring markitdown for all users
|
||||
try:
|
||||
from markitdown import MarkItDown
|
||||
|
||||
self._markitdown = MarkItDown()
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"markitdown package is required for file parsing. Install with: pip install markitdown"
|
||||
) from e
|
||||
|
||||
self._ocr_enabled = ocr_enabled
|
||||
if ocr_enabled:
|
||||
ocr_options = self._build_ocr_options(
|
||||
api_key=ocr_api_key,
|
||||
base_url=ocr_base_url,
|
||||
model=ocr_model,
|
||||
prompt=ocr_prompt,
|
||||
)
|
||||
self._markitdown = MarkItDown(
|
||||
llm_client=ocr_options.llm_client,
|
||||
llm_model=ocr_options.llm_model,
|
||||
llm_prompt=ocr_options.llm_prompt,
|
||||
)
|
||||
else:
|
||||
self._markitdown = MarkItDown()
|
||||
|
||||
def _build_ocr_options(
|
||||
self,
|
||||
*,
|
||||
api_key: str | None,
|
||||
base_url: str | None,
|
||||
model: str | None,
|
||||
prompt: str | None,
|
||||
) -> MarkitdownOcrOptions:
|
||||
"""Build MarkItDown options for OpenAI-compatible image OCR."""
|
||||
if not model or not model.strip():
|
||||
raise ValueError(
|
||||
"Markitdown OCR is enabled but no model is configured. "
|
||||
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL to an OpenAI-compatible OCR/vision model "
|
||||
"with image-input support."
|
||||
)
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"Markitdown OCR is enabled but no API key is configured. "
|
||||
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY."
|
||||
)
|
||||
if not base_url or not base_url.strip():
|
||||
raise ValueError(
|
||||
"Markitdown OCR is enabled but no base URL is configured. "
|
||||
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL to an OpenAI-compatible OCR/vision endpoint."
|
||||
)
|
||||
|
||||
try:
|
||||
from openai import OpenAI
|
||||
except ImportError as e:
|
||||
raise RuntimeError("openai package is required when Markitdown OCR is enabled.") from e
|
||||
|
||||
return MarkitdownOcrOptions(
|
||||
llm_client=OpenAI(api_key=api_key, base_url=base_url.strip()),
|
||||
llm_model=model.strip(),
|
||||
llm_prompt=prompt or DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
|
||||
)
|
||||
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
"""Parse file to markdown using markitdown."""
|
||||
# markitdown is synchronous, so we run it in executor to avoid blocking
|
||||
@@ -48,14 +143,22 @@ class MarkitdownParser(FileParser):
|
||||
|
||||
def _convert_sync(self, file_data: bytes, filename: str) -> str:
|
||||
"""Synchronous parsing (runs in thread pool)."""
|
||||
if self._is_image_file(filename) and not self._ocr_enabled:
|
||||
raise RuntimeError(
|
||||
"Image OCR is not enabled for the markitdown parser. "
|
||||
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED=true and configure an OpenAI-compatible "
|
||||
"OCR/vision endpoint with image-input support, or choose an OCR-capable parser."
|
||||
)
|
||||
|
||||
# Write to temp file (markitdown requires file path)
|
||||
with tempfile.NamedTemporaryFile(suffix=Path(filename).suffix, delete=False) as tmp:
|
||||
tmp.write(file_data)
|
||||
tmp_path = tmp.name
|
||||
|
||||
try:
|
||||
# Parse using markitdown
|
||||
result = self._markitdown.convert(tmp_path)
|
||||
# Parse using markitdown, passing an explicit charset hint for text
|
||||
# files to avoid markitdown's sample-based (and crash-prone) detection.
|
||||
result = self._markitdown.convert(tmp_path, stream_info=self._utf8_stream_info(file_data, filename))
|
||||
|
||||
if not result or not result.text_content:
|
||||
raise RuntimeError(f"No content extracted from '{filename}'")
|
||||
@@ -73,6 +176,28 @@ class MarkitdownParser(FileParser):
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _utf8_stream_info(file_data: bytes, filename: str) -> "StreamInfo | None":
|
||||
"""Return a UTF-8 charset hint for text files that decode cleanly as UTF-8.
|
||||
|
||||
Returns None for binary files or non-UTF-8 text so markitdown falls back
|
||||
to its own detection.
|
||||
"""
|
||||
if Path(filename).suffix.lower() not in _TEXT_EXTENSIONS:
|
||||
return None
|
||||
try:
|
||||
file_data.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return None
|
||||
from markitdown import StreamInfo
|
||||
|
||||
return StreamInfo(charset="utf-8")
|
||||
|
||||
@staticmethod
|
||||
def _is_image_file(filename: str) -> bool:
|
||||
"""Return whether the file type needs OCR to extract useful text."""
|
||||
return Path(filename).suffix.lower() in {".jpg", ".jpeg", ".png"}
|
||||
|
||||
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||
"""Check if markitdown supports this file type."""
|
||||
# Supported extensions (from markitdown docs)
|
||||
@@ -85,7 +210,7 @@ class MarkitdownParser(FileParser):
|
||||
".ppt",
|
||||
".xlsx",
|
||||
".xls",
|
||||
# Images (with OCR)
|
||||
# Images (optional OCR)
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".png",
|
||||
|
||||
@@ -15,12 +15,25 @@ import time
|
||||
from typing import Any
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _usage_from_anthropic_response(response: Any) -> LLMResponseUsage:
|
||||
"""Extract input/output/cached token counts from an Anthropic usage block."""
|
||||
usage = getattr(response, "usage", None)
|
||||
if not usage:
|
||||
return LLMResponseUsage()
|
||||
return LLMResponseUsage(
|
||||
input_tokens=usage.input_tokens or 0,
|
||||
output_tokens=usage.output_tokens or 0,
|
||||
cached_tokens=getattr(usage, "cache_read_input_tokens", 0) or 0,
|
||||
)
|
||||
|
||||
|
||||
class AnthropicLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using Anthropic's Claude models.
|
||||
@@ -136,7 +149,9 @@ class AnthropicLLM(LLMInterface):
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Use strict JSON schema enforcement (not supported by Anthropic).
|
||||
strict_schema: Route structured output through a forced tool_use tool for
|
||||
native constrained decoding (issue #1002). When False, falls back to
|
||||
schema-in-prompt + JSON parse.
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
|
||||
Returns:
|
||||
@@ -167,14 +182,21 @@ class AnthropicLLM(LLMInterface):
|
||||
else:
|
||||
anthropic_messages.append({"role": role, "content": content})
|
||||
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
# Structured output: prefer Anthropic-native constrained decoding via a single
|
||||
# forced tool_use tool (strict_schema) over text-injecting the schema and
|
||||
# parsing the reply. Native constrained decoding guarantees schema-valid JSON,
|
||||
# eliminating the invalid-JSON retry storm (issue #1002). When strict_schema is
|
||||
# off we keep the text-inject + json.loads fallback for backward compatibility.
|
||||
schema = None
|
||||
use_forced_tool = False
|
||||
_tool_name = "structured_response"
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
if system_prompt:
|
||||
system_prompt += schema_msg
|
||||
if strict_schema:
|
||||
use_forced_tool = True
|
||||
else:
|
||||
system_prompt = schema_msg
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
system_prompt = (system_prompt + schema_msg) if system_prompt else schema_msg
|
||||
|
||||
# Prepare parameters
|
||||
call_params: dict[str, Any] = {
|
||||
@@ -186,6 +208,14 @@ class AnthropicLLM(LLMInterface):
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if use_forced_tool:
|
||||
# Single tool whose input_schema IS the response schema; force the model to
|
||||
# emit it via tool_choice so the SDK does constrained decoding for us.
|
||||
call_params["tools"] = [
|
||||
{"name": _tool_name, "description": "Return the structured response.", "input_schema": schema}
|
||||
]
|
||||
call_params["tool_choice"] = {"type": "tool", "name": _tool_name}
|
||||
|
||||
if self._extra_body:
|
||||
call_params["extra_body"] = self._extra_body
|
||||
|
||||
@@ -194,40 +224,61 @@ class AnthropicLLM(LLMInterface):
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await self._client.messages.create(**call_params)
|
||||
# Stash usage before parse/validate, which may raise locally
|
||||
# even though the provider charged for these tokens (#2387).
|
||||
stash_response_usage(_usage_from_anthropic_response(response))
|
||||
|
||||
# Anthropic response content is a list of blocks
|
||||
content = ""
|
||||
for block in response.content:
|
||||
if block.type == "text":
|
||||
content += block.text
|
||||
|
||||
if response_format is not None:
|
||||
# Models may wrap JSON in markdown code blocks
|
||||
clean_content = content
|
||||
if "```json" in content:
|
||||
clean_content = content.split("```json")[1].split("```")[0].strip()
|
||||
elif "```" in content:
|
||||
clean_content = content.split("```")[1].split("```")[0].strip()
|
||||
|
||||
try:
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError:
|
||||
# Fallback to parsing raw content if markdown stripping failed
|
||||
json_data = json.loads(content)
|
||||
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
else:
|
||||
result = response_format.model_validate(json_data)
|
||||
if use_forced_tool:
|
||||
# Forced tool_use → the validated args are already a dict; no parsing,
|
||||
# no markdown-strip, no JSON-decode retry possible.
|
||||
tool_input = None
|
||||
for block in response.content:
|
||||
if block.type == "tool_use" and block.name == _tool_name:
|
||||
tool_input = block.input or {}
|
||||
break
|
||||
if tool_input is None:
|
||||
# Model ignored the forced tool (rare, e.g. a gateway that drops
|
||||
# tool_choice). Fall back to text parse so we don't hard-fail; the
|
||||
# existing retry loop still covers genuine errors.
|
||||
content = "".join(b.text for b in response.content if b.type == "text")
|
||||
tool_input = json.loads(content)
|
||||
content = json.dumps(tool_input)
|
||||
result = tool_input if skip_validation else response_format.model_validate(tool_input)
|
||||
else:
|
||||
result = content
|
||||
# Anthropic response content is a list of blocks
|
||||
content = ""
|
||||
for block in response.content:
|
||||
if block.type == "text":
|
||||
content += block.text
|
||||
|
||||
if response_format is not None:
|
||||
# Models may wrap JSON in markdown code blocks
|
||||
clean_content = content
|
||||
if "```json" in content:
|
||||
clean_content = content.split("```json")[1].split("```")[0].strip()
|
||||
elif "```" in content:
|
||||
clean_content = content.split("```")[1].split("```")[0].strip()
|
||||
|
||||
try:
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError:
|
||||
# Fallback to parsing raw content if markdown stripping failed
|
||||
json_data = json.loads(content)
|
||||
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
else:
|
||||
result = response_format.model_validate(json_data)
|
||||
else:
|
||||
result = content
|
||||
|
||||
# Record metrics and log slow calls
|
||||
duration = time.time() - start_time
|
||||
input_tokens = response.usage.input_tokens or 0 if response.usage else 0
|
||||
output_tokens = response.usage.output_tokens or 0 if response.usage else 0
|
||||
response_usage = _usage_from_anthropic_response(response)
|
||||
input_tokens = response_usage.input_tokens
|
||||
output_tokens = response_usage.output_tokens
|
||||
total_tokens = input_tokens + output_tokens
|
||||
cached_tokens = getattr(response.usage, "cache_read_input_tokens", 0) or 0 if response.usage else 0
|
||||
cached_tokens = response_usage.cached_tokens
|
||||
|
||||
# Record LLM metrics
|
||||
metrics = get_metrics_collector()
|
||||
@@ -415,6 +466,7 @@ class AnthropicLLM(LLMInterface):
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await self._client.messages.create(**call_params)
|
||||
stash_response_usage(_usage_from_anthropic_response(response))
|
||||
|
||||
# Extract content and tool calls
|
||||
content_parts = []
|
||||
|
||||
@@ -16,6 +16,7 @@ from typing import Any
|
||||
from pydantic import ValidationError
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
|
||||
@@ -118,12 +119,14 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
Raises:
|
||||
RuntimeError: If the connection test fails.
|
||||
"""
|
||||
from ...config import get_config
|
||||
|
||||
try:
|
||||
test_messages = [{"role": "user", "content": "test"}]
|
||||
await self.call(
|
||||
messages=test_messages,
|
||||
max_completion_tokens=10,
|
||||
temperature=0.0,
|
||||
temperature=get_config().llm_temperature_verification,
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
@@ -226,6 +229,16 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
if isinstance(block, TextBlock):
|
||||
full_text += block.text
|
||||
|
||||
# The Claude Agent SDK doesn't report exact counts; stash the same
|
||||
# char/4 estimate the success path traces so a later parse/validate
|
||||
# failure records consistent (estimated) tokens, not zero (#2387).
|
||||
stash_response_usage(
|
||||
LLMResponseUsage(
|
||||
input_tokens=sum(len(m.get("content", "")) for m in messages) // 4,
|
||||
output_tokens=len(full_text) // 4,
|
||||
)
|
||||
)
|
||||
|
||||
# Handle structured output
|
||||
if response_format is not None:
|
||||
# Models may wrap JSON in markdown
|
||||
|
||||
@@ -60,6 +60,22 @@ _CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
|
||||
)
|
||||
|
||||
|
||||
def default_codex_auth_file() -> Path:
|
||||
"""Return the path to Codex's ``auth.json``.
|
||||
|
||||
Honors the ``CODEX_HOME`` environment variable — the same variable the
|
||||
canonical ``@openai/codex`` CLI uses to relocate its config/credentials
|
||||
directory — and falls back to ``~/.codex`` when it is unset or empty.
|
||||
|
||||
Resolved lazily on each call (rather than cached at import time) so that
|
||||
the environment is read at the point of use.
|
||||
"""
|
||||
codex_home = os.environ.get("CODEX_HOME")
|
||||
if codex_home:
|
||||
return Path(codex_home) / "auth.json"
|
||||
return Path.home() / ".codex" / "auth.json"
|
||||
|
||||
|
||||
class CodexRefreshExpiredError(RuntimeError):
|
||||
"""Raised when the Codex refresh_token itself is no longer valid.
|
||||
|
||||
@@ -86,7 +102,7 @@ class CodexAuthManager:
|
||||
The OAuth refresh token. May be ``None`` when the auth file omits it;
|
||||
the provider still works as a one-shot loader in that case.
|
||||
auth_file:
|
||||
Path to ``~/.codex/auth.json``. Used for re-reading the refresh token
|
||||
Path to the Codex ``auth.json``. Used for re-reading the refresh token
|
||||
on demand and for atomic persistence of rotated credentials.
|
||||
"""
|
||||
|
||||
@@ -115,7 +131,8 @@ class CodexAuthManager:
|
||||
Parameters
|
||||
----------
|
||||
auth_file:
|
||||
Defaults to ``~/.codex/auth.json``.
|
||||
Defaults to ``$CODEX_HOME/auth.json`` (or ``~/.codex/auth.json``
|
||||
when ``CODEX_HOME`` is unset).
|
||||
|
||||
Raises
|
||||
------
|
||||
@@ -126,7 +143,7 @@ class CodexAuthManager:
|
||||
``auth_mode``.
|
||||
"""
|
||||
if auth_file is None:
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
auth_file = default_codex_auth_file()
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(f"Codex auth file not found: {auth_file}. Run 'codex auth login' to authenticate.")
|
||||
|
||||
@@ -2,8 +2,9 @@
|
||||
OpenAI Codex LLM provider using ChatGPT Plus/Pro OAuth authentication.
|
||||
|
||||
This provider enables using ChatGPT Plus/Pro subscriptions for API calls
|
||||
without separate OpenAI Platform API credits. It uses OAuth tokens from
|
||||
~/.codex/auth.json and communicates with the ChatGPT backend API.
|
||||
without separate OpenAI Platform API credits. It uses OAuth tokens from the
|
||||
Codex ``auth.json`` (``$CODEX_HOME/auth.json``, or ``~/.codex/auth.json`` when
|
||||
``CODEX_HOME`` is unset) and communicates with the ChatGPT backend API.
|
||||
|
||||
Tokens are refreshed automatically: the provider decodes the access_token
|
||||
JWT's ``exp`` claim and proactively refreshes via
|
||||
@@ -25,6 +26,7 @@ from typing import Any
|
||||
import httpx
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
|
||||
@@ -35,6 +37,7 @@ from .codex_auth import (
|
||||
_CODEX_TOKEN_REFRESH_SKEW_SECONDS,
|
||||
CodexAuthManager,
|
||||
CodexRefreshExpiredError,
|
||||
default_codex_auth_file,
|
||||
)
|
||||
|
||||
# Re-export for backward compatibility (tests import from this module).
|
||||
@@ -55,14 +58,15 @@ class CodexLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using OpenAI Codex OAuth authentication.
|
||||
|
||||
Authenticates using ChatGPT Plus/Pro credentials stored in ~/.codex/auth.json
|
||||
and makes API calls to chatgpt.com/backend-api/codex/responses.
|
||||
Authenticates using ChatGPT Plus/Pro credentials stored in the Codex
|
||||
``auth.json`` (honoring ``CODEX_HOME``, default ``~/.codex``) and makes API
|
||||
calls to chatgpt.com/backend-api/codex/responses.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str, # Will be ignored, reads from ~/.codex/auth.json
|
||||
api_key: str, # Will be ignored, reads from the Codex auth.json (CODEX_HOME or ~/.codex)
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
@@ -81,12 +85,14 @@ class CodexLLM(LLMInterface):
|
||||
refresh_token = self._load_codex_refresh_token()
|
||||
logger.info(f"Loaded Codex OAuth credentials for account: {account_id}")
|
||||
except Exception as e:
|
||||
auth_file = default_codex_auth_file()
|
||||
raise RuntimeError(
|
||||
f"Failed to load Codex OAuth credentials from ~/.codex/auth.json: {e}\n\n"
|
||||
f"Failed to load Codex OAuth credentials from {auth_file}: {e}\n\n"
|
||||
"To set up Codex authentication:\n"
|
||||
"1. Install Codex CLI: npm install -g @openai/codex\n"
|
||||
"2. Login: codex auth login\n"
|
||||
"3. Verify: ls ~/.codex/auth.json\n\n"
|
||||
f"3. Verify: ls {auth_file}\n\n"
|
||||
"(Set CODEX_HOME to use a credentials directory other than ~/.codex.)\n\n"
|
||||
"Or use a different provider (openai, anthropic, gemini) with API keys."
|
||||
) from e
|
||||
|
||||
@@ -94,7 +100,7 @@ class CodexLLM(LLMInterface):
|
||||
access_token=access_token,
|
||||
account_id=account_id,
|
||||
refresh_token=refresh_token,
|
||||
auth_file=Path.home() / ".codex" / "auth.json",
|
||||
auth_file=default_codex_auth_file(),
|
||||
)
|
||||
|
||||
# Use ChatGPT backend API endpoint. Codex auth is tied to
|
||||
@@ -156,7 +162,7 @@ class CodexLLM(LLMInterface):
|
||||
|
||||
def _load_codex_auth(self) -> tuple[str, str]:
|
||||
"""
|
||||
Load OAuth credentials from ~/.codex/auth.json.
|
||||
Load OAuth credentials from the Codex ``auth.json`` (CODEX_HOME or ~/.codex).
|
||||
|
||||
Returns:
|
||||
Tuple of (access_token, account_id).
|
||||
@@ -165,7 +171,7 @@ class CodexLLM(LLMInterface):
|
||||
FileNotFoundError: If auth file doesn't exist.
|
||||
ValueError: If auth file is invalid.
|
||||
"""
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
auth_file = default_codex_auth_file()
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(
|
||||
@@ -197,9 +203,7 @@ class CodexLLM(LLMInterface):
|
||||
pre- and post-``__init__`` because it does not depend on
|
||||
``_auth_manager`` being constructed yet.
|
||||
"""
|
||||
auth_file = (
|
||||
self._auth_manager._auth_file if hasattr(self, "_auth_manager") else Path.home() / ".codex" / "auth.json"
|
||||
)
|
||||
auth_file = self._auth_manager._auth_file if hasattr(self, "_auth_manager") else default_codex_auth_file()
|
||||
return CodexAuthManager.load_refresh_token_from_file(auth_file)
|
||||
|
||||
@staticmethod
|
||||
@@ -411,6 +415,16 @@ class CodexLLM(LLMInterface):
|
||||
# Parse SSE stream
|
||||
content = await self._parse_sse_stream(response)
|
||||
|
||||
# Codex SSE carries no usage block; stash the same char/4 estimate
|
||||
# the success path traces so a later parse/validate failure records
|
||||
# consistent (estimated) token counts rather than zero (#2387).
|
||||
stash_response_usage(
|
||||
LLMResponseUsage(
|
||||
input_tokens=sum(len(m.get("content", "")) for m in messages) // 4,
|
||||
output_tokens=len(content) // 4,
|
||||
)
|
||||
)
|
||||
|
||||
# Handle structured output
|
||||
if response_format is not None:
|
||||
# Models may wrap JSON in markdown
|
||||
|
||||
@@ -20,6 +20,7 @@ from google.genai import errors as genai_errors
|
||||
from google.genai import types as genai_types
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.llm_wrapper import parse_llm_json
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
@@ -50,6 +51,18 @@ def _to_int(value: Any) -> int:
|
||||
return 0
|
||||
|
||||
|
||||
def _usage_from_gemini_response(response: Any) -> LLMResponseUsage:
|
||||
"""Extract prompt/candidate/cached token counts from a Gemini usage_metadata block."""
|
||||
usage = getattr(response, "usage_metadata", None)
|
||||
if not usage:
|
||||
return LLMResponseUsage()
|
||||
return LLMResponseUsage(
|
||||
input_tokens=usage.prompt_token_count or 0,
|
||||
output_tokens=usage.candidates_token_count or 0,
|
||||
cached_tokens=getattr(usage, "cached_content_token_count", 0) or 0,
|
||||
)
|
||||
|
||||
|
||||
class GeminiLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider for Google Gemini and Vertex AI.
|
||||
@@ -76,6 +89,7 @@ class GeminiLLM(LLMInterface):
|
||||
|
||||
# Safety settings: None means use Gemini's defaults
|
||||
self._safety_settings: list | None = kwargs.get("gemini_safety_settings")
|
||||
self._service_tier: str | None = kwargs.get("gemini_service_tier")
|
||||
|
||||
# User-configured extra params merged into the GenerateContentConfig of
|
||||
# every call. Gemini's request body nests generation params, so we expose
|
||||
@@ -106,6 +120,16 @@ class GeminiLLM(LLMInterface):
|
||||
self._client = genai.Client(api_key=self.api_key)
|
||||
logger.info(f"Gemini API: model={self.model}")
|
||||
|
||||
def _apply_service_tier(self, config_kwargs: dict[str, Any]) -> None:
|
||||
if not self._service_tier:
|
||||
return
|
||||
|
||||
http_options = dict(config_kwargs.get("http_options") or {})
|
||||
extra_body = dict(http_options.get("extra_body") or {})
|
||||
extra_body.setdefault("service_tier", self._service_tier)
|
||||
http_options["extra_body"] = extra_body
|
||||
config_kwargs["http_options"] = http_options
|
||||
|
||||
def _init_vertexai(self, **kwargs: Any) -> None:
|
||||
"""Initialize Vertex AI client with project, region, and credentials."""
|
||||
# Extract Vertex AI config from kwargs
|
||||
@@ -247,16 +271,13 @@ class GeminiLLM(LLMInterface):
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
|
||||
|
||||
# Add the JSON schema as a textual hint in the system_instruction (matching
|
||||
# the normal uncached path). Structured output is still enforced via
|
||||
# response_schema regardless; this is just guidance text.
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
def _system_instruction_with_schema() -> str:
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
if system_instruction:
|
||||
system_instruction += schema_msg
|
||||
else:
|
||||
system_instruction = schema_msg
|
||||
schema_msg = (
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n"
|
||||
f"{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
)
|
||||
return (system_instruction + schema_msg) if system_instruction else schema_msg
|
||||
|
||||
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
|
||||
effective_safety_settings = _safety_settings_ctx.get()
|
||||
@@ -273,11 +294,18 @@ class GeminiLLM(LLMInterface):
|
||||
def _build_generation_config(use_cache: bool) -> "genai_types.GenerateContentConfig | None":
|
||||
# Seed with user-configured extra params; explicit settings below win.
|
||||
config_kwargs: dict[str, Any] = dict(self._extra_body)
|
||||
self._apply_service_tier(config_kwargs)
|
||||
if use_cache:
|
||||
config_kwargs["cached_content"] = cached_prefix
|
||||
elif (
|
||||
use_schema_prompt_fallback
|
||||
and response_format is not None
|
||||
and hasattr(response_format, "model_json_schema")
|
||||
):
|
||||
config_kwargs["system_instruction"] = _system_instruction_with_schema()
|
||||
elif system_instruction:
|
||||
config_kwargs["system_instruction"] = system_instruction
|
||||
if response_format is not None:
|
||||
if response_format is not None and not use_schema_prompt_fallback:
|
||||
config_kwargs["response_mime_type"] = "application/json"
|
||||
config_kwargs["response_schema"] = response_format
|
||||
if temperature is not None:
|
||||
@@ -295,6 +323,7 @@ class GeminiLLM(LLMInterface):
|
||||
return genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
|
||||
|
||||
cache_active = using_cache
|
||||
use_schema_prompt_fallback = False
|
||||
generation_config = _build_generation_config(cache_active)
|
||||
|
||||
last_exception = None
|
||||
@@ -311,6 +340,9 @@ class GeminiLLM(LLMInterface):
|
||||
),
|
||||
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
|
||||
)
|
||||
# Stash usage before parse/validate, which may raise locally
|
||||
# even though the provider charged for these tokens (#2387).
|
||||
stash_response_usage(_usage_from_gemini_response(response))
|
||||
|
||||
content = response.text
|
||||
|
||||
@@ -412,12 +444,26 @@ class GeminiLLM(LLMInterface):
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=input_tokens + output_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
last_exception = e
|
||||
if (
|
||||
attempt < max_retries
|
||||
and response_format is not None
|
||||
and hasattr(response_format, "model_json_schema")
|
||||
and not cache_active
|
||||
and not use_schema_prompt_fallback
|
||||
):
|
||||
logger.warning("Gemini returned invalid JSON, retrying with prompt-side schema guidance...")
|
||||
cache_active = False
|
||||
use_schema_prompt_fallback = True
|
||||
generation_config = _build_generation_config(cache_active)
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
if attempt < max_retries:
|
||||
logger.warning("Gemini returned invalid JSON, retrying...")
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
@@ -604,6 +650,7 @@ class GeminiLLM(LLMInterface):
|
||||
def _build_tools_config(use_cache: bool) -> "genai_types.GenerateContentConfig":
|
||||
# Seed with user-configured extra params; explicit settings below win.
|
||||
config_kwargs: dict[str, Any] = dict(self._extra_body)
|
||||
self._apply_service_tier(config_kwargs)
|
||||
if use_cache:
|
||||
config_kwargs["cached_content"] = cached_prefix
|
||||
else:
|
||||
@@ -662,6 +709,7 @@ class GeminiLLM(LLMInterface):
|
||||
),
|
||||
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
|
||||
)
|
||||
stash_response_usage(_usage_from_gemini_response(response))
|
||||
|
||||
# Extract content and tool calls
|
||||
content = None
|
||||
@@ -749,6 +797,8 @@ class GeminiLLM(LLMInterface):
|
||||
finish_reason=finish_reason,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cached_tokens=cached_input_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
|
||||
except genai_errors.APIError as e:
|
||||
|
||||
@@ -15,10 +15,15 @@ is handled automatically by LiteLLM.
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from litellm.exceptions import Timeout as LiteLLMTimeout
|
||||
|
||||
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
@@ -26,6 +31,22 @@ from hindsight_api.worker.stage import set_stage
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _usage_from_litellm_response(response: Any) -> LLMResponseUsage:
|
||||
"""Extract prompt/completion/cached token counts from a LiteLLM (OpenAI-shaped) usage block."""
|
||||
usage = getattr(response, "usage", None)
|
||||
if not usage:
|
||||
return LLMResponseUsage()
|
||||
cached_tokens = 0
|
||||
details = getattr(usage, "prompt_tokens_details", None)
|
||||
if details:
|
||||
cached_tokens = getattr(details, "cached_tokens", 0) or 0
|
||||
return LLMResponseUsage(
|
||||
input_tokens=getattr(usage, "prompt_tokens", 0) or 0,
|
||||
output_tokens=getattr(usage, "completion_tokens", 0) or 0,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
|
||||
|
||||
class LiteLLMLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using the LiteLLM SDK for universal model support.
|
||||
@@ -47,13 +68,16 @@ class LiteLLMLLM(LLMInterface):
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float = 300.0,
|
||||
timeout: float | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
bedrock_service_tier: str | None = None,
|
||||
default_headers: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
self.timeout = timeout
|
||||
# ``None`` falls back to HINDSIGHT_API_LLM_TIMEOUT, then DEFAULT_LLM_TIMEOUT — never None,
|
||||
# so the hard ``asyncio.wait_for`` backstop in ``call`` is always bounded.
|
||||
self.timeout = timeout if timeout is not None else float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT)))
|
||||
self._litellm: Any = None
|
||||
# User-configured extra params merged as top-level kwargs into every
|
||||
# completion call so LiteLLM normalizes them per-provider (e.g. maps
|
||||
@@ -61,6 +85,13 @@ class LiteLLMLLM(LLMInterface):
|
||||
# drops any the target model rejects (litellm.drop_params=True below).
|
||||
# Sourced from llm_extra_body (env: HINDSIGHT_API_LLM_EXTRA_BODY).
|
||||
self._extra_body: dict[str, Any] = extra_body or {}
|
||||
# Operator-configured default headers forwarded to litellm.acompletion as
|
||||
# ``extra_headers`` (used by deployments routing through proxies / request-
|
||||
# tracing middleware). Mirrors the Anthropic provider's default_headers
|
||||
# wiring. Sourced from llm_default_headers (env: HINDSIGHT_API_LLM_DEFAULT_HEADERS).
|
||||
# Copied so a caller-owned dict can't be mutated through us, and a fresh
|
||||
# copy is handed to each call below to avoid cross-request contamination.
|
||||
self._default_headers: dict[str, Any] = dict(default_headers or {})
|
||||
self.bedrock_service_tier = bedrock_service_tier
|
||||
|
||||
try:
|
||||
@@ -77,12 +108,14 @@ class LiteLLMLLM(LLMInterface):
|
||||
raise RuntimeError("LiteLLM SDK not installed. Run: uv add litellm or pip install litellm") from e
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
from ...config import get_config
|
||||
|
||||
try:
|
||||
test_messages = [{"role": "user", "content": "test"}]
|
||||
await self.call(
|
||||
messages=test_messages,
|
||||
max_completion_tokens=50,
|
||||
temperature=0.0,
|
||||
temperature=get_config().llm_temperature_verification,
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
@@ -121,6 +154,13 @@ class LiteLLMLLM(LLMInterface):
|
||||
for key, value in self._extra_body.items():
|
||||
kwargs.setdefault(key, value)
|
||||
|
||||
# Forward operator-configured default headers as ``extra_headers`` so they
|
||||
# reach the provider behind LiteLLM (proxies / request-tracing middleware).
|
||||
# ``setdefault`` keeps any explicit per-call ``extra_headers`` authoritative;
|
||||
# a per-call copy prevents LiteLLM/downstream from mutating the stored dict.
|
||||
if self._default_headers:
|
||||
kwargs.setdefault("extra_headers", dict(self._default_headers))
|
||||
|
||||
# Bedrock service tier: flex (50% cheaper), priority, or reserved
|
||||
if self.model.startswith("bedrock/") and self.bedrock_service_tier is not None:
|
||||
kwargs["service_tier"] = self.bedrock_service_tier
|
||||
@@ -209,7 +249,14 @@ class LiteLLMLLM(LLMInterface):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self._stage_label}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._acompletion(**call_kwargs)
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
timeout=self.timeout,
|
||||
)
|
||||
# Stash usage before the length check and parse/validate below,
|
||||
# which may raise locally even though the provider charged for
|
||||
# these tokens (#2387).
|
||||
stash_response_usage(_usage_from_litellm_response(response))
|
||||
|
||||
content = response.choices[0].message.content or ""
|
||||
finish_reason = response.choices[0].finish_reason
|
||||
@@ -240,8 +287,9 @@ class LiteLLMLLM(LLMInterface):
|
||||
result = content
|
||||
|
||||
# Extract usage
|
||||
input_tokens = getattr(response.usage, "prompt_tokens", 0) or 0
|
||||
output_tokens = getattr(response.usage, "completion_tokens", 0) or 0
|
||||
response_usage = _usage_from_litellm_response(response)
|
||||
input_tokens = response_usage.input_tokens
|
||||
output_tokens = response_usage.output_tokens
|
||||
total_tokens = input_tokens + output_tokens
|
||||
|
||||
# Record metrics
|
||||
@@ -304,6 +352,25 @@ class LiteLLMLLM(LLMInterface):
|
||||
logger.error(f"LiteLLM returned invalid JSON after {max_retries + 1} attempts")
|
||||
raise
|
||||
|
||||
except (TimeoutError, asyncio.TimeoutError, LiteLLMTimeout) as e:
|
||||
# litellm/httpx don't always honor their own ``timeout=`` (e.g. a connection held
|
||||
# open with no token progress), so ``wait_for`` is the hard cap that cancels a hung
|
||||
# call regardless — otherwise one straggler pins a worker slot and stalls its gather.
|
||||
last_exception = e
|
||||
exc_name = type(e).__name__
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"LiteLLM call exceeded timeout={self.timeout}s ({exc_name}, scope={scope}), retrying..."
|
||||
)
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
logger.error(
|
||||
f"LiteLLM call timed out after {self.timeout}s on {attempt + 1} attempts "
|
||||
f"({exc_name}, scope={scope})"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
error_str = str(e).lower()
|
||||
# Fast fail on auth errors
|
||||
@@ -354,7 +421,18 @@ class LiteLLMLLM(LLMInterface):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self._stage_label}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._acompletion(**call_kwargs)
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
timeout=self.timeout,
|
||||
)
|
||||
# Stash usage before the tool-call argument parse below, which
|
||||
# can raise json.JSONDecodeError locally even though the provider
|
||||
# already billed for these tokens; without this the error trace
|
||||
# records 0/0 tokens (#2387). Mirrors call() and the anthropic/
|
||||
# gemini call_with_tools paths so the litellm tool path (and the
|
||||
# LiteLLMRouterLLM subclass that inherits this method) completes
|
||||
# the #2396 usage-on-error coverage.
|
||||
stash_response_usage(_usage_from_litellm_response(response))
|
||||
|
||||
message = response.choices[0].message
|
||||
content = message.content
|
||||
@@ -424,6 +502,23 @@ class LiteLLMLLM(LLMInterface):
|
||||
output_tokens=output_tokens,
|
||||
)
|
||||
|
||||
except (TimeoutError, asyncio.TimeoutError, LiteLLMTimeout) as e:
|
||||
# See ``call`` — hard cap so a hung completion cannot block
|
||||
# forever and pin a worker slot / concurrency permit.
|
||||
last_exception = e
|
||||
exc_name = type(e).__name__
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"LiteLLM tool call exceeded timeout={self.timeout}s ({exc_name}, scope={scope}), retrying..."
|
||||
)
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
logger.error(
|
||||
f"LiteLLM tool call timed out after {self.timeout}s on {attempt + 1} attempts "
|
||||
f"({exc_name}, scope={scope})"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
error_str = str(e).lower()
|
||||
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
|
||||
|
||||
@@ -67,7 +67,7 @@ class LiteLLMRouterLLM(LiteLLMLLM):
|
||||
model: str,
|
||||
config: dict[str, Any],
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float = 300.0,
|
||||
timeout: float | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(
|
||||
@@ -146,16 +146,28 @@ class LiteLLMRouterLLM(LiteLLMLLM):
|
||||
kwargs["max_completion_tokens"] = self._cap_max_completion_tokens(max_completion_tokens)
|
||||
if temperature is not None:
|
||||
kwargs["temperature"] = temperature
|
||||
|
||||
# Forward operator-configured default headers as ``extra_headers`` so they
|
||||
# reach the provider behind the Router (proxies / request-tracing middleware).
|
||||
# This override deliberately omits api_key/base_url/extra_body (those live in
|
||||
# the per-deployment Router config), but headers are a cross-cutting operator
|
||||
# concern, so we inject them here too — mirroring the base provider.
|
||||
# ``setdefault`` keeps any explicit per-call ``extra_headers`` authoritative;
|
||||
# a per-call copy prevents LiteLLM/downstream from mutating the stored dict.
|
||||
if self._default_headers:
|
||||
kwargs.setdefault("extra_headers", dict(self._default_headers))
|
||||
return kwargs
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
from hindsight_api.engine.llm_interface import OutputTooLongError
|
||||
|
||||
from ...config import get_config
|
||||
|
||||
try:
|
||||
await self.call(
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
max_completion_tokens=50,
|
||||
temperature=0.0,
|
||||
temperature=get_config().llm_temperature_verification,
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
@@ -101,7 +101,7 @@ class MockLLM(LLMInterface):
|
||||
messages: List of message dicts with 'role' and 'content'.
|
||||
response_format: Optional Pydantic model for structured output.
|
||||
max_completion_tokens: Not used in mock.
|
||||
temperature: Not used in mock.
|
||||
temperature: Recorded on the call record for test assertions.
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Not used in mock.
|
||||
initial_backoff: Not used in mock.
|
||||
@@ -123,6 +123,9 @@ class MockLLM(LLMInterface):
|
||||
if response_format and hasattr(response_format, "__name__")
|
||||
else str(response_format),
|
||||
"scope": scope,
|
||||
# Record the temperature so tests can assert per-operation temperature
|
||||
# wiring (None means the parameter was omitted from the call).
|
||||
"temperature": temperature,
|
||||
}
|
||||
self._mock_calls.append(call_record)
|
||||
logger.debug(f"Mock LLM call recorded: scope={scope}, model={self.model}")
|
||||
@@ -208,7 +211,7 @@ class MockLLM(LLMInterface):
|
||||
messages: List of message dicts. Can include tool results with role='tool'.
|
||||
tools: List of tool definitions in OpenAI format.
|
||||
max_completion_tokens: Not used in mock.
|
||||
temperature: Not used in mock.
|
||||
temperature: Recorded on the call record for test assertions.
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Not used in mock.
|
||||
initial_backoff: Not used in mock.
|
||||
@@ -225,6 +228,9 @@ class MockLLM(LLMInterface):
|
||||
"messages": messages,
|
||||
"tools": [t.get("function", {}).get("name") for t in tools],
|
||||
"scope": scope,
|
||||
# Record the temperature so tests can assert per-operation temperature
|
||||
# wiring (None means the parameter was omitted from the call).
|
||||
"temperature": temperature,
|
||||
}
|
||||
self._mock_calls.append(call_record)
|
||||
|
||||
|
||||
@@ -26,6 +26,8 @@ import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from email.utils import parsedate_to_datetime
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qs, urlparse, urlunparse
|
||||
|
||||
@@ -34,7 +36,8 @@ from openai import APIConnectionError, APIStatusError, AsyncOpenAI, LengthFinish
|
||||
|
||||
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
|
||||
from hindsight_api.engine.bank_attribution import apply_bank_attribution
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError, ProviderRateLimitResetError
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
@@ -83,6 +86,49 @@ def _strip_code_fences(content: str) -> str:
|
||||
return content
|
||||
|
||||
|
||||
# Reasoning/thinking tags emitted by extended-thinking models. Some providers
|
||||
# (e.g. MiniMax-M3) leak the chain-of-thought wrapped in these tags into the
|
||||
# response body instead of a separate reasoning_content field. Each entry is
|
||||
# (open_tag, close_tag); the open tag also matches when the close tag is missing
|
||||
# (truncated output) so a dangling block is removed to end-of-string.
|
||||
_REASONING_TAG_PAIRS: tuple[tuple[str, str], ...] = (
|
||||
("<think>", "</think>"),
|
||||
("<thinking>", "</thinking>"),
|
||||
("<thought>", "</thought>"),
|
||||
("<reasoning>", "</reasoning>"),
|
||||
("|startthink|", "|endthink|"),
|
||||
)
|
||||
|
||||
|
||||
def _strip_reasoning_tags(text: str) -> str:
|
||||
"""Strip extended-thinking/reasoning blocks from an LLM response.
|
||||
|
||||
Removes the full set of tag styles emitted by reasoning models:
|
||||
``<think>``, ``<thinking>``, ``<thought>``, ``<reasoning>`` and the
|
||||
``|startthink|...|endthink|`` markers. Both the structured (JSON) path and
|
||||
the free-form path must call this — otherwise a non-structured response
|
||||
(e.g. a mental-model markdown blob from MiniMax-M3) leaks the raw
|
||||
``<think>...</think>`` verbatim into stored memories.
|
||||
|
||||
Handles two cases:
|
||||
1. Closed blocks: ``<think>...</think>`` removed wherever they appear.
|
||||
2. Unclosed blocks: a dangling ``<think>`` with no closing tag (model output
|
||||
truncated mid-thought) is removed from the open tag to end-of-string.
|
||||
|
||||
Returns the input unchanged (modulo surrounding whitespace) when no tags are
|
||||
present.
|
||||
"""
|
||||
if not text:
|
||||
return text
|
||||
for open_tag, close_tag in _REASONING_TAG_PAIRS:
|
||||
open_re = re.escape(open_tag)
|
||||
close_re = re.escape(close_tag)
|
||||
# Closed blocks first, then any remaining unclosed (truncated) block.
|
||||
text = re.sub(rf"{open_re}.*?{close_re}", "", text, flags=re.DOTALL)
|
||||
text = re.sub(rf"{open_re}.*", "", text, flags=re.DOTALL)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def _response_get(response: Any, key: str, default: Any = None) -> Any:
|
||||
if isinstance(response, dict):
|
||||
return response.get(key, default)
|
||||
@@ -187,6 +233,21 @@ def _content_or_error(response: Any, *, provider: str, model: str, scope: str) -
|
||||
return content, choice
|
||||
|
||||
|
||||
def _usage_from_openai_response(response: Any) -> LLMResponseUsage:
|
||||
"""Extract prompt/completion/cached token counts from an OpenAI-shaped usage block."""
|
||||
usage = getattr(response, "usage", None)
|
||||
input_tokens = (usage.prompt_tokens or 0) if usage else 0
|
||||
output_tokens = (usage.completion_tokens or 0) if usage else 0
|
||||
cached_tokens = 0
|
||||
if usage and getattr(usage, "prompt_tokens_details", None):
|
||||
cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
|
||||
return LLMResponseUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
)
|
||||
|
||||
|
||||
def _ensure_json_word_in_user_message(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Some OpenAI-compatible gateways require 'json' in a user message for json_object mode."""
|
||||
|
||||
@@ -234,6 +295,122 @@ def _summarize_status_error(e: APIStatusError, body_max: int = 400) -> str:
|
||||
return f"HTTP {e.status_code}: {body_str or '<no body>'}"
|
||||
|
||||
|
||||
_RATE_LIMIT_RESET_AT_RE = re.compile(
|
||||
r"\breset at\s+"
|
||||
r"(?P<reset_at>\d{4}-\d{2}-\d{2}[ T]\d{2}:\d{2}:\d{2}(?:\s*(?:Z|[+-]\d{2}:?\d{2}))?)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_RATE_LIMIT_WINDOW_RE = re.compile(
|
||||
r"\b(?:for|in)\s+(?P<amount>\d+)\s*(?P<unit>second|minute|hour|day)s?\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def _status_error_body_text(e: APIStatusError) -> str:
|
||||
body: Any = getattr(e, "body", None)
|
||||
if body is None:
|
||||
try:
|
||||
body = e.response.text
|
||||
except Exception:
|
||||
body = None
|
||||
if isinstance(body, (dict, list)):
|
||||
try:
|
||||
return json.dumps(body, default=str, ensure_ascii=False)
|
||||
except Exception:
|
||||
return str(body)
|
||||
return str(body or "").strip()
|
||||
|
||||
|
||||
def _parse_retry_after_header(value: str | None, now: datetime) -> datetime | None:
|
||||
if not value:
|
||||
return None
|
||||
raw = value.strip()
|
||||
try:
|
||||
seconds = float(raw)
|
||||
except ValueError:
|
||||
seconds = -1.0
|
||||
if seconds >= 0:
|
||||
return now + timedelta(seconds=seconds)
|
||||
|
||||
try:
|
||||
parsed = parsedate_to_datetime(raw)
|
||||
except (TypeError, ValueError, IndexError, OverflowError):
|
||||
return None
|
||||
if parsed.tzinfo is None:
|
||||
parsed = parsed.replace(tzinfo=UTC)
|
||||
return parsed.astimezone(UTC)
|
||||
|
||||
|
||||
def _parse_reset_at_datetime(value: str) -> datetime | None:
|
||||
raw = value.strip().replace(" ", "T")
|
||||
if raw.endswith("Z"):
|
||||
raw = f"{raw[:-1]}+00:00"
|
||||
elif re.search(r"[+-]\d{4}$", raw):
|
||||
raw = f"{raw[:-2]}:{raw[-2:]}"
|
||||
try:
|
||||
parsed = datetime.fromisoformat(raw)
|
||||
except ValueError:
|
||||
return None
|
||||
if parsed.tzinfo is None:
|
||||
# Some providers (z.ai included) return a wall-clock reset timestamp
|
||||
# without a zone. Interpret it in the host's local zone so logs, status
|
||||
# pages, and the queued next_retry_at describe the same operator-facing
|
||||
# clock instead of silently shifting by UTC offset.
|
||||
parsed = parsed.astimezone()
|
||||
return parsed.astimezone(UTC)
|
||||
|
||||
|
||||
def _rate_limit_retry_at(e: APIStatusError) -> datetime | None:
|
||||
now = datetime.now(UTC)
|
||||
response = getattr(e, "response", None)
|
||||
headers = getattr(response, "headers", None)
|
||||
if headers is not None:
|
||||
retry_at = _parse_retry_after_header(headers.get("retry-after") or headers.get("Retry-After"), now)
|
||||
if retry_at is not None and retry_at > now:
|
||||
return retry_at
|
||||
|
||||
body_text = _status_error_body_text(e)
|
||||
reset_match = _RATE_LIMIT_RESET_AT_RE.search(body_text)
|
||||
if reset_match:
|
||||
retry_at = _parse_reset_at_datetime(reset_match.group("reset_at"))
|
||||
if retry_at is not None and retry_at > now:
|
||||
return retry_at
|
||||
|
||||
window_match = _RATE_LIMIT_WINDOW_RE.search(body_text)
|
||||
if not window_match:
|
||||
return None
|
||||
amount = int(window_match.group("amount"))
|
||||
unit = window_match.group("unit").lower()
|
||||
if unit == "second":
|
||||
seconds = amount
|
||||
elif unit == "minute":
|
||||
seconds = amount * 60
|
||||
elif unit == "hour":
|
||||
seconds = amount * 3600
|
||||
else:
|
||||
seconds = amount * 86400
|
||||
return now + timedelta(seconds=seconds)
|
||||
|
||||
|
||||
def _raise_provider_quota_defer(
|
||||
e: APIStatusError, *, provider: str, model: str, scope: str, max_backoff: float
|
||||
) -> None:
|
||||
if e.status_code != 429:
|
||||
return
|
||||
retry_at = _rate_limit_retry_at(e)
|
||||
if retry_at is None:
|
||||
return
|
||||
if (retry_at - datetime.now(UTC)).total_seconds() <= max_backoff:
|
||||
return
|
||||
summary = _summarize_status_error(e)
|
||||
raise ProviderRateLimitResetError(
|
||||
retry_at=retry_at,
|
||||
message=(
|
||||
f"Provider quota exhausted ({provider}/{model}, scope={scope}); retry at {retry_at.isoformat()}: {summary}"
|
||||
),
|
||||
) from e
|
||||
|
||||
|
||||
class OpenAICompatibleLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider for OpenAI-compatible APIs.
|
||||
@@ -269,7 +446,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
base_url: Base URL for the API (uses defaults for groq/ollama/lmstudio if empty).
|
||||
model: Model name.
|
||||
reasoning_effort: Reasoning effort level for supported models ("low", "medium", "high").
|
||||
timeout: Request timeout in seconds (uses env var or 300s default).
|
||||
timeout: Request timeout in seconds (uses env var or 120s default).
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto").
|
||||
extra_body: Extra body params merged into every API call.
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
@@ -288,8 +465,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"deepseek",
|
||||
"volcano",
|
||||
"openrouter",
|
||||
"requesty",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
"fireworks",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
@@ -311,10 +490,14 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
self.base_url = "https://api.deepseek.com"
|
||||
elif self.provider == "openrouter":
|
||||
self.base_url = "https://openrouter.ai/api/v1"
|
||||
elif self.provider == "requesty":
|
||||
self.base_url = "https://router.requesty.ai/v1"
|
||||
elif self.provider == "zai":
|
||||
self.base_url = "https://api.z.ai/api/coding/paas/v4"
|
||||
elif self.provider == "opencode-go":
|
||||
self.base_url = "https://opencode.ai/zen/go/v1"
|
||||
elif self.provider == "atlas":
|
||||
self.base_url = "https://api.atlascloud.ai/v1"
|
||||
elif self.provider == "fireworks":
|
||||
# OpenAI-compatible inference host (online path). The batch API
|
||||
# lives on a separate control-plane host — see FireworksLLM.
|
||||
@@ -333,8 +516,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"minimax",
|
||||
"deepseek",
|
||||
"openrouter",
|
||||
"requesty",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
"ollama-cloud",
|
||||
)
|
||||
and not self.api_key
|
||||
@@ -609,6 +794,9 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
try:
|
||||
if response_format is not None:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
# Stash usage before parse/validate, which may raise locally
|
||||
# even though the provider charged for these tokens (#2387).
|
||||
stash_response_usage(_usage_from_openai_response(response))
|
||||
|
||||
content, first_choice = _content_or_error(
|
||||
response,
|
||||
@@ -617,15 +805,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
scope=scope,
|
||||
)
|
||||
|
||||
# Strip reasoning model thinking tags
|
||||
# Strip reasoning model thinking tags (closed and unclosed).
|
||||
# Supports: <think>, <thinking>, <thought>, <reasoning>, |startthink|/|endthink|
|
||||
original_len = len(content)
|
||||
content = re.sub(r"<think>.*?</think>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"<thinking>.*?</thinking>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"<thought>.*?</thought>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"<reasoning>.*?</reasoning>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"\|startthink\|.*?\|endthink\|", "", content, flags=re.DOTALL)
|
||||
content = content.strip()
|
||||
content = _strip_reasoning_tags(content)
|
||||
if len(content) < original_len:
|
||||
logger.debug(f"Stripped {original_len - len(content)} chars of reasoning tokens")
|
||||
|
||||
@@ -667,6 +850,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
result = response_format.model_validate(json_data)
|
||||
else:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
stash_response_usage(_usage_from_openai_response(response))
|
||||
result, first_choice = _content_or_error(
|
||||
response,
|
||||
provider=self.provider,
|
||||
@@ -674,15 +858,33 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
scope=scope,
|
||||
)
|
||||
|
||||
# Free-form (non-structured) output also leaks reasoning tags:
|
||||
# reasoning models like MiniMax-M3 wrap their chain-of-thought
|
||||
# in <think>...</think> in the response body. Without this strip
|
||||
# a mental-model markdown blob is stored verbatim with the raw
|
||||
# thinking tags. Mirrors the structured-output path above.
|
||||
result = _strip_reasoning_tags(result)
|
||||
|
||||
# Record token usage metrics
|
||||
duration = time.time() - start_time
|
||||
usage = response.usage
|
||||
input_tokens = usage.prompt_tokens or 0 if usage else 0
|
||||
output_tokens = usage.completion_tokens or 0 if usage else 0
|
||||
response_usage = _usage_from_openai_response(response)
|
||||
input_tokens = response_usage.input_tokens
|
||||
output_tokens = response_usage.output_tokens
|
||||
total_tokens = usage.total_tokens or 0 if usage else 0
|
||||
cached_tokens = 0
|
||||
if usage and getattr(usage, "prompt_tokens_details", None):
|
||||
cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
|
||||
cached_tokens = response_usage.cached_tokens
|
||||
thoughts_tokens = 0
|
||||
if usage and getattr(usage, "completion_tokens_details", None):
|
||||
thoughts_tokens = getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0
|
||||
# OpenAI-compatible providers fold reasoning tokens into
|
||||
# ``completion_tokens`` (and thus ``total_tokens``), but the
|
||||
# TokenUsage contract — and the Gemini provider — treat
|
||||
# ``output_tokens``/``total_tokens`` as visible-only, surfacing
|
||||
# reasoning separately in ``thoughts_tokens``. Subtract so the
|
||||
# two fields don't double-count reasoning (cost over-attribution).
|
||||
if thoughts_tokens:
|
||||
output_tokens = max(0, output_tokens - thoughts_tokens)
|
||||
total_tokens = max(0, total_tokens - thoughts_tokens)
|
||||
|
||||
# Record LLM metrics
|
||||
metrics = get_metrics_collector()
|
||||
@@ -731,6 +933,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=total_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
@@ -761,6 +964,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
logger.error(f"Auth error (HTTP {e.status_code}), not retrying: {str(e)}")
|
||||
raise
|
||||
|
||||
_raise_provider_quota_defer(
|
||||
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
|
||||
)
|
||||
|
||||
# Handle tool_use_failed error - model outputted in tool call format
|
||||
if e.status_code == 400 and response_format is not None:
|
||||
try:
|
||||
@@ -814,7 +1021,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
f"scope={scope}): {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except ProviderResponseError as e:
|
||||
last_exception = e
|
||||
if e.retryable and attempt < max_retries:
|
||||
@@ -978,6 +1184,17 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
usage = response.usage
|
||||
input_tokens = usage.prompt_tokens or 0 if usage else 0
|
||||
output_tokens = usage.completion_tokens or 0 if usage else 0
|
||||
cached_tokens = 0
|
||||
if usage and getattr(usage, "prompt_tokens_details", None):
|
||||
cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
|
||||
thoughts_tokens = 0
|
||||
if usage and getattr(usage, "completion_tokens_details", None):
|
||||
thoughts_tokens = getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0
|
||||
# See ``call()``: OpenAI-compatible ``completion_tokens`` includes
|
||||
# reasoning, so make ``output_tokens`` visible-only to avoid
|
||||
# double-counting it against ``thoughts_tokens``.
|
||||
if thoughts_tokens:
|
||||
output_tokens = max(0, output_tokens - thoughts_tokens)
|
||||
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
@@ -1020,6 +1237,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
finish_reason=finish_reason,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
|
||||
except APIConnectionError as e:
|
||||
@@ -1047,6 +1266,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
f"not retrying: {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
_raise_provider_quota_defer(
|
||||
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
|
||||
)
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
@@ -1060,7 +1283,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
f"({self.provider}/{self.model}, scope={scope}): {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@ import time
|
||||
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||
|
||||
from ...config import get_config
|
||||
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, TokenUsageSummary, ToolCall
|
||||
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, StructuredOutputResult, TokenUsageSummary, ToolCall
|
||||
from .prompts import (
|
||||
_extract_directive_rules,
|
||||
build_final_prompt,
|
||||
@@ -90,12 +90,87 @@ _LEAKED_JSON_SUFFIX = re.compile(
|
||||
r'\s*```(?:json)?\s*\{[^}]*(?:"(?:observation_ids|memory_ids|mental_model_ids)"|\})\s*```\s*$',
|
||||
re.DOTALL | re.IGNORECASE,
|
||||
)
|
||||
_LEAKED_JSON_OBJECT = re.compile(
|
||||
r'\s*\{[^{]*"(?:observation_ids|memory_ids|mental_model_ids|answer)"[^}]*\}\s*$', re.DOTALL
|
||||
)
|
||||
_TRAILING_IDS_PATTERN = re.compile(
|
||||
r"\s*(?:observation_ids|memory_ids|mental_model_ids)\s*[=:]\s*\[.*?\]\s*$", re.DOTALL | re.IGNORECASE
|
||||
)
|
||||
_JSON_CODE_FENCE_PATTERN = re.compile(r"^\s*```(?:json)?\s*(\{.*\})\s*```\s*$", re.DOTALL | re.IGNORECASE)
|
||||
|
||||
_DONE_ARGUMENT_KEYS = frozenset(
|
||||
{
|
||||
"answer",
|
||||
"directive_compliance",
|
||||
"memory_ids",
|
||||
"mental_model_ids",
|
||||
"observation_ids",
|
||||
"model_ids",
|
||||
}
|
||||
)
|
||||
_DONE_ARGUMENT_MARKER_KEYS = _DONE_ARGUMENT_KEYS - {"answer"}
|
||||
_LEAKED_JSON_ID_KEYS = frozenset({"memory_ids", "mental_model_ids", "observation_ids", "model_ids"})
|
||||
|
||||
|
||||
def _unwrap_leaked_done_arguments(text: str) -> str | None:
|
||||
"""Return the answer when a done tool call was rendered as JSON text.
|
||||
|
||||
Some providers leak the done tool's argument object instead of surfacing it
|
||||
as a native tool call, e.g. {"answer": "...", "memory_ids": [...]}. Only
|
||||
unwrap objects that match the done argument shape so normal JSON answers
|
||||
stay intact.
|
||||
"""
|
||||
candidate = text.strip()
|
||||
if not candidate:
|
||||
return None
|
||||
|
||||
fenced = _JSON_CODE_FENCE_PATTERN.match(candidate)
|
||||
if fenced:
|
||||
candidate = fenced.group(1).strip()
|
||||
|
||||
try:
|
||||
payload = json.loads(candidate)
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
answer = payload.get("answer")
|
||||
if not isinstance(answer, str) or not answer.strip():
|
||||
return None
|
||||
|
||||
keys = set(payload)
|
||||
if not keys.intersection(_DONE_ARGUMENT_MARKER_KEYS):
|
||||
return None
|
||||
if not keys.issubset(_DONE_ARGUMENT_KEYS):
|
||||
return None
|
||||
|
||||
for key in ("memory_ids", "mental_model_ids", "observation_ids", "model_ids"):
|
||||
value = payload.get(key)
|
||||
if value is not None and not isinstance(value, list):
|
||||
return None
|
||||
|
||||
return answer.strip()
|
||||
|
||||
|
||||
def _strip_trailing_id_json_object(text: str) -> str:
|
||||
stripped = text.rstrip()
|
||||
if not stripped.endswith("}"):
|
||||
return text.strip()
|
||||
|
||||
start = stripped.rfind("{")
|
||||
if start < 0:
|
||||
return text.strip()
|
||||
|
||||
try:
|
||||
payload = json.loads(stripped[start:])
|
||||
except json.JSONDecodeError:
|
||||
return text.strip()
|
||||
|
||||
if not isinstance(payload, dict) or not payload:
|
||||
return text.strip()
|
||||
keys = set(payload)
|
||||
if not keys.issubset(_LEAKED_JSON_ID_KEYS):
|
||||
return text.strip()
|
||||
|
||||
return stripped[:start].strip()
|
||||
|
||||
|
||||
def _clean_answer_text(text: str) -> str:
|
||||
@@ -104,6 +179,10 @@ def _clean_answer_text(text: str) -> str:
|
||||
Some LLMs output the done() call as text instead of a proper tool call.
|
||||
This strips out patterns like: done({"answer": "...", ...})
|
||||
"""
|
||||
unwrapped = _unwrap_leaked_done_arguments(text)
|
||||
if unwrapped is not None:
|
||||
return unwrapped
|
||||
|
||||
# Remove done() call pattern from the end of the text
|
||||
cleaned = _DONE_CALL_PATTERN.sub("", text).strip()
|
||||
return cleaned if cleaned else text
|
||||
@@ -122,13 +201,17 @@ def _clean_done_answer(text: str) -> str:
|
||||
if not text:
|
||||
return text
|
||||
|
||||
unwrapped = _unwrap_leaked_done_arguments(text)
|
||||
if unwrapped is not None:
|
||||
return unwrapped
|
||||
|
||||
cleaned = text
|
||||
|
||||
# Remove leaked JSON in code blocks at the end
|
||||
cleaned = _LEAKED_JSON_SUFFIX.sub("", cleaned).strip()
|
||||
|
||||
# Remove leaked raw JSON objects at the end
|
||||
cleaned = _LEAKED_JSON_OBJECT.sub("", cleaned).strip()
|
||||
cleaned = _strip_trailing_id_json_object(cleaned)
|
||||
|
||||
# Remove trailing ID patterns
|
||||
cleaned = _TRAILING_IDS_PATTERN.sub("", cleaned).strip()
|
||||
@@ -141,7 +224,7 @@ async def _generate_structured_output(
|
||||
response_schema: dict,
|
||||
llm_config: "LLMProvider",
|
||||
reflect_id: str,
|
||||
) -> tuple[dict[str, Any] | None, int, int]:
|
||||
) -> StructuredOutputResult:
|
||||
"""Generate structured output from an answer using the provided JSON schema.
|
||||
|
||||
Args:
|
||||
@@ -151,8 +234,8 @@ async def _generate_structured_output(
|
||||
reflect_id: Reflect ID for logging
|
||||
|
||||
Returns:
|
||||
Tuple of (structured_output, input_tokens, output_tokens).
|
||||
structured_output is None if generation fails.
|
||||
A StructuredOutputResult carrying the structured output (None if
|
||||
generation fails) and the call's token usage.
|
||||
"""
|
||||
try:
|
||||
from typing import Any as TypingAny
|
||||
@@ -186,7 +269,7 @@ async def _generate_structured_output(
|
||||
|
||||
if not fields:
|
||||
logger.warning(f"[REFLECT {reflect_id}] No fields found in response_schema, skipping structured output")
|
||||
return None, 0, 0
|
||||
return StructuredOutputResult()
|
||||
|
||||
DynamicModel = create_model("StructuredResponse", **fields)
|
||||
|
||||
@@ -239,6 +322,9 @@ OUTPUT:"""
|
||||
],
|
||||
response_format=DynamicModel,
|
||||
scope="reflect_structured",
|
||||
max_retries=1,
|
||||
initial_backoff=0.25,
|
||||
max_backoff=1.0,
|
||||
skip_validation=True, # We'll handle the dict ourselves
|
||||
return_usage=True,
|
||||
)
|
||||
@@ -259,11 +345,17 @@ OUTPUT:"""
|
||||
logger.warning(f"[REFLECT {reflect_id}] Required field '{field_name}' is empty in structured output")
|
||||
|
||||
logger.info(f"[REFLECT {reflect_id}] Generated structured output with {len(structured_output)} fields")
|
||||
return structured_output, usage.input_tokens, usage.output_tokens
|
||||
return StructuredOutputResult(
|
||||
structured_output=structured_output,
|
||||
input_tokens=usage.input_tokens,
|
||||
output_tokens=usage.output_tokens,
|
||||
cached_tokens=usage.cached_tokens,
|
||||
thoughts_tokens=usage.thoughts_tokens,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"[REFLECT {reflect_id}] Failed to generate structured output: {e}")
|
||||
return None, 0, 0
|
||||
return StructuredOutputResult()
|
||||
|
||||
|
||||
def _count_messages_tokens(messages: list[dict[str, Any]]) -> int:
|
||||
@@ -435,9 +527,14 @@ async def run_reflect_agent(
|
||||
llm_trace: list[dict[str, Any]] = []
|
||||
context_history: list[dict[str, Any]] = [] # For final prompt fallback
|
||||
|
||||
# Token usage tracking - accumulate across all LLM calls
|
||||
# Token usage tracking - accumulate across all LLM calls.
|
||||
# cached_tokens and thoughts_tokens are surfaced for cost attribution
|
||||
# and prompt-cache tuning. Both are subsets of (or parallel to) the
|
||||
# input/output counts and are NOT double-counted in total_tokens.
|
||||
total_input_tokens = 0
|
||||
total_output_tokens = 0
|
||||
total_cached_tokens = 0
|
||||
total_thoughts_tokens = 0
|
||||
|
||||
# Track available IDs for validation (prevents hallucinated citations)
|
||||
available_memory_ids: set[str] = set()
|
||||
@@ -460,6 +557,8 @@ async def run_reflect_agent(
|
||||
input_tokens=total_input_tokens,
|
||||
output_tokens=total_output_tokens,
|
||||
total_tokens=total_input_tokens + total_output_tokens,
|
||||
cached_tokens=total_cached_tokens,
|
||||
thoughts_tokens=total_thoughts_tokens,
|
||||
)
|
||||
|
||||
def _log_completion(answer: str, iterations: int, forced: bool = False):
|
||||
@@ -526,6 +625,8 @@ async def run_reflect_agent(
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
@@ -539,11 +640,12 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
@@ -588,6 +690,8 @@ async def run_reflect_agent(
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
@@ -600,11 +704,12 @@ async def run_reflect_agent(
|
||||
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
@@ -661,6 +766,8 @@ async def run_reflect_agent(
|
||||
consecutive_errors = 0
|
||||
total_input_tokens += result.input_tokens
|
||||
total_output_tokens += result.output_tokens
|
||||
total_cached_tokens += getattr(result, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(result, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": f"agent_{iteration + 1}",
|
||||
@@ -709,6 +816,8 @@ async def run_reflect_agent(
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
@@ -722,11 +831,12 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
@@ -783,6 +893,8 @@ async def run_reflect_agent(
|
||||
)
|
||||
total_input_tokens += rewrite_usage.input_tokens
|
||||
total_output_tokens += rewrite_usage.output_tokens
|
||||
total_cached_tokens += getattr(rewrite_usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(rewrite_usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final_rewrite",
|
||||
@@ -796,11 +908,12 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
|
||||
_log_completion(answer, iteration + 1)
|
||||
return ReflectAgentResult(
|
||||
@@ -835,6 +948,8 @@ async def run_reflect_agent(
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
@@ -848,11 +963,12 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
@@ -1147,14 +1263,15 @@ async def _process_done_tool(
|
||||
structured_output = None
|
||||
final_usage = usage
|
||||
if response_schema and llm_config and answer:
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
structured_output = struct.structured_output
|
||||
# Add structured output tokens to usage
|
||||
final_usage = TokenUsageSummary(
|
||||
input_tokens=usage.input_tokens + struct_in,
|
||||
output_tokens=usage.output_tokens + struct_out,
|
||||
total_tokens=usage.total_tokens + struct_in + struct_out,
|
||||
input_tokens=usage.input_tokens + struct.input_tokens,
|
||||
output_tokens=usage.output_tokens + struct.output_tokens,
|
||||
total_tokens=usage.total_tokens + struct.input_tokens + struct.output_tokens,
|
||||
cached_tokens=usage.cached_tokens + struct.cached_tokens,
|
||||
thoughts_tokens=usage.thoughts_tokens + struct.thoughts_tokens,
|
||||
)
|
||||
|
||||
log_completion(answer, iterations)
|
||||
|
||||
@@ -78,9 +78,32 @@ class DirectiveInfo(BaseModel):
|
||||
class TokenUsageSummary(BaseModel):
|
||||
"""Total token usage across all LLM calls."""
|
||||
|
||||
input_tokens: int = Field(default=0, description="Total input tokens used")
|
||||
output_tokens: int = Field(default=0, description="Total output tokens used")
|
||||
total_tokens: int = Field(default=0, description="Total tokens (input + output)")
|
||||
input_tokens: int = Field(default=0, description="Total input tokens used (includes any cached prefix tokens)")
|
||||
output_tokens: int = Field(default=0, description="Total visible output tokens used (excludes reasoning/thoughts)")
|
||||
total_tokens: int = Field(default=0, description="Total tokens (input + output, excludes thoughts)")
|
||||
cached_tokens: int = Field(
|
||||
default=0,
|
||||
description="Cached/cache-read prompt tokens summed across calls. Subset of input_tokens.",
|
||||
)
|
||||
thoughts_tokens: int = Field(
|
||||
default=0,
|
||||
description=(
|
||||
"Reasoning/thinking tokens summed across calls. Billed at the output rate by some providers "
|
||||
"but not part of visible output."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class StructuredOutputResult(BaseModel):
|
||||
"""Result of structured-output generation, including token usage for the call."""
|
||||
|
||||
structured_output: dict[str, Any] | None = Field(
|
||||
default=None, description="Generated structured output, or None if generation failed"
|
||||
)
|
||||
input_tokens: int = Field(default=0, description="Input tokens used")
|
||||
output_tokens: int = Field(default=0, description="Visible output tokens used")
|
||||
cached_tokens: int = Field(default=0, description="Cached prefix tokens. Subset of input_tokens.")
|
||||
thoughts_tokens: int = Field(default=0, description="Reasoning/thinking tokens, when reported by the provider")
|
||||
|
||||
|
||||
class ReflectAgentResult(BaseModel):
|
||||
|
||||
@@ -31,8 +31,20 @@ class LLMToolCallResult(BaseModel):
|
||||
content: str | None = Field(default=None, description="Text content if any")
|
||||
tool_calls: list[LLMToolCall] = Field(default_factory=list, description="Tool calls requested by the LLM")
|
||||
finish_reason: str | None = Field(default=None, description="Reason the LLM stopped: 'stop', 'tool_calls', etc.")
|
||||
input_tokens: int = Field(default=0, description="Input tokens used in this call")
|
||||
output_tokens: int = Field(default=0, description="Output tokens used in this call")
|
||||
input_tokens: int = Field(
|
||||
default=0,
|
||||
description="Input tokens used in this call (includes any cached prefix tokens reported by the provider)",
|
||||
)
|
||||
output_tokens: int = Field(
|
||||
default=0, description="Visible output tokens used in this call (excludes reasoning/thoughts)"
|
||||
)
|
||||
cached_tokens: int = Field(
|
||||
default=0, description="Cached prefix tokens, when reported by the provider. Subset of input_tokens."
|
||||
)
|
||||
thoughts_tokens: int = Field(
|
||||
default=0,
|
||||
description="Reasoning/thinking tokens. Billed at the output rate by some providers but not part of visible output.",
|
||||
)
|
||||
|
||||
|
||||
class ToolCallTrace(BaseModel):
|
||||
@@ -91,9 +103,18 @@ class TokenUsage(BaseModel):
|
||||
)
|
||||
|
||||
input_tokens: int = Field(default=0, description="Number of input/prompt tokens consumed")
|
||||
output_tokens: int = Field(default=0, description="Number of output/completion tokens generated")
|
||||
total_tokens: int = Field(default=0, description="Total tokens (input + output)")
|
||||
output_tokens: int = Field(
|
||||
default=0, description="Number of visible output/completion tokens generated (excludes reasoning/thoughts)"
|
||||
)
|
||||
total_tokens: int = Field(default=0, description="Total tokens (input + output, excludes thoughts)")
|
||||
cached_tokens: int = Field(default=0, description="Cached/cache-read prompt tokens, when reported by the provider")
|
||||
thoughts_tokens: int = Field(
|
||||
default=0,
|
||||
description=(
|
||||
"Reasoning/thinking tokens generated by the model. Billed at the output rate by some providers "
|
||||
"(e.g. Gemini 2.5+ family) but not surfaced in the visible response."
|
||||
),
|
||||
)
|
||||
|
||||
def __add__(self, other: "TokenUsage") -> "TokenUsage":
|
||||
"""Allow aggregating token usage from multiple calls."""
|
||||
@@ -102,6 +123,7 @@ class TokenUsage(BaseModel):
|
||||
output_tokens=self.output_tokens + other.output_tokens,
|
||||
total_tokens=self.total_tokens + other.total_tokens,
|
||||
cached_tokens=self.cached_tokens + other.cached_tokens,
|
||||
thoughts_tokens=self.thoughts_tokens + other.thoughts_tokens,
|
||||
)
|
||||
|
||||
|
||||
@@ -150,6 +172,47 @@ class DispositionTraits(BaseModel):
|
||||
model_config = ConfigDict(json_schema_extra={"example": {"skepticism": 3, "literalism": 3, "empathy": 3}})
|
||||
|
||||
|
||||
class RecallScores(BaseModel):
|
||||
"""Per-result recall scores from different stages of the pipeline.
|
||||
|
||||
``final`` is the value results are ranked by. The others are diagnostic and
|
||||
can be filtered on via the recall ``min_scores`` request parameter. ``semantic``
|
||||
and ``keyword`` are the raw per-strategy retrieval scores (``None`` when that
|
||||
strategy did not surface this result); ``reranker`` is the cross-encoder's
|
||||
normalized relevance.
|
||||
"""
|
||||
|
||||
final: float = Field(description="Final ranking score (combined reranker + recency/temporal/proof boosts)")
|
||||
reranker: float | None = Field(
|
||||
default=None,
|
||||
description="Cross-encoder relevance, normalized 0-1. None when the reranker is a passthrough (rrf/interleave modes).",
|
||||
)
|
||||
semantic: float | None = Field(
|
||||
default=None, description="Vector cosine similarity (0-1). None if this result was not surfaced semantically."
|
||||
)
|
||||
keyword: float | None = Field(
|
||||
default=None,
|
||||
description="Keyword/full-text (BM25) score (>= 0, unbounded). None if this result was not surfaced by keyword search.",
|
||||
)
|
||||
|
||||
|
||||
class MinScores(BaseModel):
|
||||
"""Optional per-stage score floors for recall (all inclusive, AND-ed).
|
||||
|
||||
``semantic`` and ``keyword`` are **retrieval-level** cutoffs pushed into the SQL
|
||||
arms (overriding the global ``semantic_min_similarity`` / ``bm25_min_score``
|
||||
config for this request), so they prune weak matches before fusion. ``reranker``
|
||||
and ``final`` are **post-query** filters applied to the scored results after
|
||||
reranking. Any field left None imposes no floor; all-None (the default) means
|
||||
no score filtering.
|
||||
"""
|
||||
|
||||
semantic: float | None = Field(default=None, description="Retrieval-level: minimum vector similarity (0-1).")
|
||||
keyword: float | None = Field(default=None, description="Retrieval-level: minimum keyword/full-text (BM25) score.")
|
||||
reranker: float | None = Field(default=None, description="Post-query: minimum normalized reranker score (0-1).")
|
||||
final: float | None = Field(default=None, description="Post-query: minimum final ranking score.")
|
||||
|
||||
|
||||
class MemoryFact(BaseModel):
|
||||
"""
|
||||
A single memory fact returned by search or think operations.
|
||||
@@ -180,7 +243,7 @@ class MemoryFact(BaseModel):
|
||||
|
||||
id: str = Field(description="Unique identifier for the memory fact")
|
||||
text: str = Field(description="The actual text content of the memory")
|
||||
fact_type: str = Field(description="Type of fact: 'world', 'experience', 'opinion', or 'observation'")
|
||||
fact_type: str = Field(description="Type of fact: 'world', 'experience', or 'observation'")
|
||||
entities: list[str] | None = Field(None, description="Entity names mentioned in this fact")
|
||||
context: str | None = Field(None, description="Additional context for the memory")
|
||||
occurred_start: str | None = Field(None, description="ISO format date when the event started occurring")
|
||||
@@ -209,6 +272,10 @@ class MemoryFact(BaseModel):
|
||||
None,
|
||||
description="IDs of source facts this observation was derived from (observation type only, when source_facts is enabled)",
|
||||
)
|
||||
scores: RecallScores | None = Field(
|
||||
None,
|
||||
description="Recall scores from each pipeline stage (final/reranker/semantic/keyword). Not returned for source facts.",
|
||||
)
|
||||
|
||||
|
||||
class ChunkInfo(BaseModel):
|
||||
@@ -307,7 +374,8 @@ class ReflectResult(BaseModel):
|
||||
],
|
||||
"experience": [],
|
||||
"opinion": [],
|
||||
"mental_models": [],
|
||||
"observation": [],
|
||||
"mental-models": [],
|
||||
"directives": [
|
||||
{
|
||||
"id": "directive-123",
|
||||
@@ -324,7 +392,7 @@ class ReflectResult(BaseModel):
|
||||
|
||||
text: str = Field(description="The formulated answer text")
|
||||
based_on: dict[str, Any] = Field(
|
||||
description="Facts used to formulate the answer, organized by type (world, experience, mental_models, directives)"
|
||||
description="Facts used to formulate the answer, organized by type (world, experience, observation, mental-models, directives)"
|
||||
)
|
||||
structured_output: dict[str, Any] | None = Field(
|
||||
default=None,
|
||||
|
||||
@@ -14,6 +14,7 @@ from typing import Any, Literal, cast
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, create_model, field_validator
|
||||
|
||||
from ..llm_interface import ProviderRateLimitResetError
|
||||
from ..llm_wrapper import LLMConfig, OutputTooLongError, sanitize_llm_output
|
||||
from ..operation_metadata import RetainExtractionErrors
|
||||
from ..response_models import TokenUsage
|
||||
@@ -192,7 +193,7 @@ class ExtractedFact(BaseModel):
|
||||
occurred_start: str | None = Field(default=None, description="ISO timestamp for events")
|
||||
occurred_end: str | None = Field(default=None, description="ISO timestamp for event end")
|
||||
fact_type: Literal["world", "assistant"] = Field(
|
||||
description="'world' = objective/external facts. 'assistant' = first-person actions, experiences, or observations by the speaker."
|
||||
description="'world' = objective/external facts, including user preferences, rules, corrections, and constraints even when stated during a conversation. 'assistant' = actions, experiences, or observations the assistant/agent actually performed."
|
||||
)
|
||||
entities: list[Entity] | None = Field(default=None, description="People, places, concepts")
|
||||
causal_relations: list[FactCausalRelation] | None = Field(
|
||||
@@ -295,7 +296,7 @@ class ExtractedFactVerbose(BaseModel):
|
||||
)
|
||||
|
||||
fact_type: Literal["world", "assistant"] = Field(
|
||||
description="'world' = objective/external facts about other people, events, general knowledge. 'assistant' = first-person actions, experiences, or observations by the speaker (e.g., 'I changed X', 'I discovered Y')."
|
||||
description="'world' = objective/external facts about the user, other people, events, general knowledge, preferences, rules, corrections, or constraints. 'assistant' = actions, experiences, or observations the assistant/agent actually performed (e.g., 'I changed X', 'I discovered Y')."
|
||||
)
|
||||
|
||||
entities: list[Entity] | None = Field(
|
||||
@@ -345,7 +346,7 @@ class ExtractedFactNoCausal(BaseModel):
|
||||
occurred_start: str | None = Field(default=None, description="WHEN the event happened (ISO timestamp).")
|
||||
occurred_end: str | None = Field(default=None, description="WHEN the event ended (ISO timestamp).")
|
||||
fact_type: Literal["world", "assistant"] = Field(
|
||||
description="'world' = about the user/others. 'assistant' = experience with assistant."
|
||||
description="'world' = about the user/others, including user preferences, rules, corrections, and constraints. 'assistant' = actions or experiences the assistant/agent actually performed."
|
||||
)
|
||||
entities: list[Entity] | None = Field(
|
||||
default=None,
|
||||
@@ -451,6 +452,11 @@ def chunk_text(text: str, max_chars: int, structured_chunk_size: int | None = No
|
||||
``structured_chunk_size``. When unset, that limit defaults to ``max_chars``.
|
||||
For plain text, uses sentence-aware splitting.
|
||||
|
||||
The result is idempotent: re-chunking any chunk this returns yields that chunk
|
||||
unchanged. The streaming retain pipeline pre-chunks each document once and then
|
||||
re-chunks every piece during extraction; if a piece re-split, its sub-chunks
|
||||
would inherit one chunk_index and collide on ``chunk_id`` (issue #2301).
|
||||
|
||||
Args:
|
||||
text: Input text to chunk (plain text, JSON conversation, or JSONL)
|
||||
max_chars: Target maximum characters per chunk
|
||||
@@ -469,11 +475,23 @@ def chunk_text(text: str, max_chars: int, structured_chunk_size: int | None = No
|
||||
# Try to parse as JSON conversation array
|
||||
try:
|
||||
parsed = json.loads(text)
|
||||
if isinstance(parsed, list) and all(isinstance(turn, dict) for turn in parsed):
|
||||
# This looks like a conversation - chunk at turn boundaries
|
||||
return _chunk_conversation(parsed, max_chars, structured_limit)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
parsed = None
|
||||
|
||||
if isinstance(parsed, list) and all(isinstance(turn, dict) for turn in parsed):
|
||||
# This looks like a conversation - chunk at turn boundaries
|
||||
return _chunk_conversation(parsed, max_chars, structured_limit)
|
||||
|
||||
if isinstance(parsed, dict):
|
||||
# A single JSON object — e.g. one JSONL line handed back to the extractor
|
||||
# after the producer already pre-chunked it. It is one structured unit:
|
||||
# keep it whole up to the structured limit, else split it as text within
|
||||
# the chunk budget. Without this, a lone object (one line, so _chunk_jsonl
|
||||
# declines) would fall through to plain-text splitting and re-split a chunk
|
||||
# the producer deliberately kept whole — breaking idempotency (issue #2301).
|
||||
if len(text) <= structured_limit:
|
||||
return [text]
|
||||
return _split_oversized_unit(text, max_chars)
|
||||
|
||||
# Try to parse as JSONL (newline-delimited JSON objects, e.g. session logs)
|
||||
jsonl_chunks = _chunk_jsonl(text, max_chars, structured_limit)
|
||||
@@ -515,10 +533,12 @@ def _chunk_conversation(turns: list[dict], max_chars: int, structured_limit: int
|
||||
turn_size = turn_unit_size + 1 # +1 for comma
|
||||
|
||||
# A turn too large to keep whole even alone: flush, then split it as
|
||||
# text so no chunk runs far over budget (the extractor won't re-chunk).
|
||||
# text. Fragment within min(structured_limit, max_chars) so no fragment
|
||||
# exceeds the chunk budget — otherwise a downstream re-chunk would split
|
||||
# it again and collide on chunk_id (issue #2301).
|
||||
if turn_unit_size > structured_limit:
|
||||
_flush()
|
||||
chunks.extend(_split_oversized_unit(turn_json, structured_limit))
|
||||
chunks.extend(_split_oversized_unit(turn_json, min(structured_limit, max_chars)))
|
||||
continue
|
||||
|
||||
# If adding this turn would exceed limit and we have turns, save current chunk
|
||||
@@ -581,10 +601,12 @@ def _chunk_jsonl(text: str, max_chars: int, structured_limit: int) -> list[str]
|
||||
line_size = len(line) + 1 # +1 for the joining newline
|
||||
|
||||
# A line too large to keep whole even alone: flush, then split it as
|
||||
# text so no chunk runs far over budget (the extractor won't re-chunk).
|
||||
# text. Fragment within min(structured_limit, max_chars) so no fragment
|
||||
# exceeds the chunk budget — otherwise a downstream re-chunk would split
|
||||
# it again and collide on chunk_id (issue #2301).
|
||||
if line_unit_size > structured_limit:
|
||||
_flush()
|
||||
chunks.extend(_split_oversized_unit(line, structured_limit))
|
||||
chunks.extend(_split_oversized_unit(line, min(structured_limit, max_chars)))
|
||||
continue
|
||||
|
||||
# If adding this line would exceed the limit and we have lines, flush.
|
||||
@@ -641,8 +663,8 @@ fact_kind:
|
||||
- "conversation": Ongoing state, preference, trait (no dates)
|
||||
|
||||
fact_type:
|
||||
- "world": About other people, external events, general knowledge, objective facts
|
||||
- "assistant": First-person actions, experiences, or observations by the speaker/author (e.g., "I changed X", "I discovered Y", "I debugged Z"). Also includes interactions with the user (requests, recommendations). If the narrator describes something they did, tried, learned, or decided — use "assistant".
|
||||
- "world": Objective/external facts, including the user's preferences, rules, corrections, constraints, plans, traits, or context. These stay "world" even when the user states them during an assistant interaction (e.g., "User prefers browser_navigate over web_search", "User corrected the project deadline").
|
||||
- "assistant": Actions, experiences, or observations the assistant/agent actually performed (e.g., "I changed X", "I discovered Y", "I debugged Z"). Use this for the assistant/agent doing, trying, learning, deciding, recommending, or responding — not merely for user facts mentioned in conversation.
|
||||
|
||||
══════════════════════════════════════════════════════════════════════════
|
||||
TEMPORAL HANDLING
|
||||
@@ -744,7 +766,7 @@ RULES:
|
||||
- Extract all entities (people, places, organizations, objects, concepts).
|
||||
- Extract temporal information (occurred_start, occurred_end, fact_kind, when).
|
||||
- Extract location (where) and people (who).
|
||||
- fact_type: use "world" unless the content is clearly an interaction with the assistant."""
|
||||
- fact_type: use "world" for user preferences, rules, corrections, constraints, traits, and other objective facts, even when stated during an assistant interaction. Use "assistant" only for actions or experiences the assistant/agent actually performed."""
|
||||
|
||||
VERBATIM_FACT_EXTRACTION_PROMPT = _BASE_FACT_EXTRACTION_PROMPT.format(
|
||||
retain_mission_section="{retain_mission_section}",
|
||||
@@ -845,8 +867,8 @@ For CONVERSATIONS (fact_kind="conversation"):
|
||||
FACT TYPE
|
||||
══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
- **world**: User's life, other people, events (would exist without this conversation)
|
||||
- **assistant**: Interactions with assistant (requests, recommendations, help)
|
||||
- **world**: User's life, preferences, rules, corrections, constraints, other people, and events (facts that would exist without this conversation)
|
||||
- **assistant**: Actions or experiences the assistant/agent actually performed while helping the user (requests, recommendations, help)
|
||||
⚠️ CRITICAL for assistant facts: ALWAYS capture the user's request/question in the fact!
|
||||
Include: what the user asked, what problem they wanted solved, what context they provided
|
||||
|
||||
@@ -1181,9 +1203,17 @@ def _build_request_body(llm_config, config, prompt: str, user_message: str, resp
|
||||
request_body = {
|
||||
"model": llm_config.model,
|
||||
"messages": [{"role": "system", "content": prompt}, {"role": "user", "content": user_message}],
|
||||
"temperature": 0.1,
|
||||
}
|
||||
|
||||
# Honour the configured retain temperature. ``None`` omits the parameter
|
||||
# entirely (for models like Azure GPT-5.5 that reject explicit temperatures),
|
||||
# mirroring LLMProvider.call, which drops temperature when it is None. The
|
||||
# batch path builds the request body directly instead of going through
|
||||
# LLMProvider.call (#2469 only de-hardcoded the streaming path), so it must
|
||||
# apply the same rule here.
|
||||
if config.llm_temperature_retain is not None:
|
||||
request_body["temperature"] = config.llm_temperature_retain
|
||||
|
||||
# Add max_completion_tokens if configured
|
||||
if config.retain_max_completion_tokens:
|
||||
request_body["max_completion_tokens"] = config.retain_max_completion_tokens
|
||||
@@ -1292,7 +1322,7 @@ async def _extract_facts_from_chunk(
|
||||
messages=[{"role": "system", "content": prompt}, {"role": "user", "content": user_message}],
|
||||
response_format=response_schema,
|
||||
scope="retain_extract_facts",
|
||||
temperature=0.1,
|
||||
temperature=config.llm_temperature_retain,
|
||||
max_completion_tokens=config.retain_max_completion_tokens,
|
||||
max_retries=llm_max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
@@ -1792,10 +1822,28 @@ async def extract_facts_from_text(
|
||||
total_usage = total_usage + chunk_usage
|
||||
|
||||
if failed_chunks:
|
||||
# Include the exception message — not just the type — so operators
|
||||
# can tell a structured-JSON parse failure apart from a rate limit
|
||||
# apart from a network 5xx, all of which can surface as the same
|
||||
# exception types. The error_message we propagate to the
|
||||
# async_operations row is the only inspection surface a worker-side
|
||||
# failure leaves behind, and a bare "chunk 0: RuntimeError" is not
|
||||
# actionable.
|
||||
failed_summary = ", ".join(f"chunk {idx}: {type(err).__name__}: {err}" for idx, err in failed_chunks[:5])
|
||||
quota_errors = [err for _, err in failed_chunks if isinstance(err, ProviderRateLimitResetError)]
|
||||
if quota_errors and len(quota_errors) == len(failed_chunks):
|
||||
retry_at = max(err.retry_at for err in quota_errors)
|
||||
raise ProviderRateLimitResetError(
|
||||
retry_at=retry_at,
|
||||
message=(
|
||||
f"Fact extraction deferred by provider quota: {len(failed_chunks)}/{len(chunks)} chunks failed. "
|
||||
f"First failures: {failed_summary}. Provider detail: {quota_errors[0]}"
|
||||
),
|
||||
) from quota_errors[0]
|
||||
|
||||
# Fail the entire retain — partial extraction is not acceptable.
|
||||
# All successfully extracted facts are discarded because the transaction
|
||||
# hasn't committed yet. The worker poller will retry the entire task.
|
||||
failed_summary = ", ".join(f"chunk {idx}: {type(err).__name__}" for idx, err in failed_chunks[:5])
|
||||
raise RuntimeError(
|
||||
f"Fact extraction failed: {len(failed_chunks)}/{len(chunks)} chunks failed. "
|
||||
f"First failures: {failed_summary}"
|
||||
|
||||
@@ -834,6 +834,28 @@ async def retain_batch(
|
||||
if first.get("tags"):
|
||||
existing_content["tags"] = first["tags"]
|
||||
contents_dicts = [existing_content, *contents_dicts]
|
||||
# Merge JSON arrays to keep original_text valid (#2409).
|
||||
# Without this, combined_content joins items with "\n", producing
|
||||
# "[...]\n[...]" which is not valid JSON. On the next append cycle
|
||||
# chunk_text() fails to parse it and falls through to sentence-
|
||||
# boundary text splitting, breaking speaker attribution.
|
||||
try:
|
||||
_merged = []
|
||||
for _item in contents_dicts:
|
||||
_parsed = json.loads(_item.get("content", ""))
|
||||
if isinstance(_parsed, list) and all(isinstance(_e, dict) for _e in _parsed):
|
||||
_merged.extend(_parsed)
|
||||
else:
|
||||
_merged = None
|
||||
break
|
||||
if _merged is not None:
|
||||
contents_dicts = [{"content": json.dumps(_merged, ensure_ascii=False)}]
|
||||
if first.get("context"):
|
||||
contents_dicts[0]["context"] = first["context"]
|
||||
if first.get("tags"):
|
||||
contents_dicts[0]["tags"] = first["tags"]
|
||||
except (json.JSONDecodeError, ValueError, TypeError):
|
||||
pass
|
||||
# Rebuild contents list to match
|
||||
contents = _build_contents(contents_dicts, document_tags)
|
||||
log_buffer.append(
|
||||
@@ -1615,8 +1637,19 @@ async def _streaming_retain_batch(
|
||||
# Check if facts are already committed (recovery from previous crash).
|
||||
# If so, skip extraction+writes and jump straight to final ANN pass.
|
||||
# ---------------------------------------------------------------------------
|
||||
# Only the call that starts a document at chunk 0 may take the whole-document
|
||||
# skip. When an oversized single item is split into several sequential
|
||||
# sub-batches that SHARE one document_id AND one operation_id (see
|
||||
# _split_contents_into_sub_batches), the first sub-batch commits its chunks
|
||||
# and stamps effective_doc_id into result_metadata.facts_committed_document_ids.
|
||||
# Without the offset gate, every later sub-batch (chunk_index_offset > 0) would
|
||||
# then see its own document already "committed" and skip extraction, dropping
|
||||
# all chunks past the first slice. A non-zero offset inherently means this call
|
||||
# continues a document another sub-batch already started, so it must always do
|
||||
# its work — crash-safety for those chunks still comes from the per-chunk hash
|
||||
# recovery (existing_chunk_hashes) below.
|
||||
facts_already_committed = False
|
||||
if operation_id:
|
||||
if operation_id and chunk_index_offset == 0:
|
||||
try:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
row = await conn.fetchrow(
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
Helper functions for hybrid search (semantic + BM25 + graph).
|
||||
"""
|
||||
|
||||
from .types import MergedCandidate, RetrievalResult
|
||||
from .types import ArmScores, MergedCandidate, RetrievalResult
|
||||
|
||||
|
||||
def cap_per_source(results: list[RetrievalResult], cap: int) -> list[RetrievalResult]:
|
||||
@@ -51,6 +51,7 @@ def reciprocal_rank_fusion(result_lists: list[list[RetrievalResult]], k: int = 6
|
||||
rrf_scores = {}
|
||||
source_ranks = {} # Track rank from each source for each doc_id
|
||||
all_retrievals = {} # Store the actual RetrievalResult (use first occurrence)
|
||||
arm_scores: dict[str, ArmScores] = {} # doc_id -> raw per-strategy scores across arms
|
||||
|
||||
source_names = ["semantic", "bm25", "graph", "temporal"]
|
||||
|
||||
@@ -79,17 +80,29 @@ def reciprocal_rank_fusion(result_lists: list[list[RetrievalResult]], k: int = 6
|
||||
if doc_id not in rrf_scores:
|
||||
rrf_scores[doc_id] = 0.0
|
||||
source_ranks[doc_id] = {}
|
||||
arm_scores[doc_id] = ArmScores()
|
||||
|
||||
rrf_scores[doc_id] += 1.0 / (k + rank)
|
||||
source_ranks[doc_id][f"{source_name}_rank"] = rank
|
||||
|
||||
# Capture this arm's raw score for the doc (the merged RetrievalResult
|
||||
# below keeps only the first arm's score, so record each arm here).
|
||||
if source_name == "semantic" and retrieval.similarity is not None:
|
||||
arm_scores[doc_id].semantic = retrieval.similarity
|
||||
elif source_name == "bm25" and retrieval.bm25_score is not None:
|
||||
arm_scores[doc_id].keyword = retrieval.bm25_score
|
||||
|
||||
# Combine into final results with metadata
|
||||
merged_results = []
|
||||
for rrf_rank, (doc_id, rrf_score) in enumerate(
|
||||
sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True), start=1
|
||||
):
|
||||
merged_candidate = MergedCandidate(
|
||||
retrieval=all_retrievals[doc_id], rrf_score=rrf_score, rrf_rank=rrf_rank, source_ranks=source_ranks[doc_id]
|
||||
retrieval=all_retrievals[doc_id],
|
||||
rrf_score=rrf_score,
|
||||
rrf_rank=rrf_rank,
|
||||
source_ranks=source_ranks[doc_id],
|
||||
arm_scores=arm_scores[doc_id],
|
||||
)
|
||||
merged_results.append(merged_candidate)
|
||||
|
||||
@@ -118,6 +131,7 @@ def interleave_fusion(result_lists: list[list[RetrievalResult]]) -> list[MergedC
|
||||
source_names = ["semantic", "bm25", "graph", "temporal"]
|
||||
source_ranks: dict[str, dict[str, int]] = {}
|
||||
all_retrievals: dict[str, RetrievalResult] = {}
|
||||
arm_scores: dict[str, ArmScores] = {}
|
||||
|
||||
for source_idx, results in enumerate(result_lists):
|
||||
source_name = source_names[source_idx] if source_idx < len(source_names) else f"source_{source_idx}"
|
||||
@@ -129,6 +143,11 @@ def interleave_fusion(result_lists: list[list[RetrievalResult]]) -> list[MergedC
|
||||
doc_id = retrieval.id
|
||||
all_retrievals.setdefault(doc_id, retrieval)
|
||||
source_ranks.setdefault(doc_id, {})[f"{source_name}_rank"] = rank
|
||||
arm = arm_scores.setdefault(doc_id, ArmScores())
|
||||
if source_name == "semantic" and retrieval.similarity is not None:
|
||||
arm.semantic = retrieval.similarity
|
||||
elif source_name == "bm25" and retrieval.bm25_score is not None:
|
||||
arm.keyword = retrieval.bm25_score
|
||||
|
||||
# Round-robin pick across arms in priority order: all #1s, then all #2s, ...
|
||||
ordered_ids: list[str] = []
|
||||
@@ -151,6 +170,7 @@ def interleave_fusion(result_lists: list[list[RetrievalResult]]) -> list[MergedC
|
||||
rrf_score=float(n - pos),
|
||||
rrf_rank=pos + 1,
|
||||
source_ranks=source_ranks[doc_id],
|
||||
arm_scores=arm_scores[doc_id],
|
||||
)
|
||||
for pos, doc_id in enumerate(ordered_ids)
|
||||
]
|
||||
|
||||
@@ -251,8 +251,10 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
result.activation = row["score"]
|
||||
results.append(result)
|
||||
|
||||
if tags:
|
||||
results = filter_results_by_tags(results, tags, match=tags_match)
|
||||
# filter_results_by_tags is a no-op when no filter applies (tags falsy and not
|
||||
# the exact-empty/global scope), so call it unconditionally — gating on `if tags:`
|
||||
# would skip the untagged-only filter for tags=[] + tags_match="exact".
|
||||
results = filter_results_by_tags(results, tags, match=tags_match)
|
||||
|
||||
if tag_groups:
|
||||
results = filter_results_by_tag_groups(results, tag_groups)
|
||||
|
||||
@@ -16,6 +16,44 @@ _RECENCY_ALPHA: float = 0.2
|
||||
_TEMPORAL_ALPHA: float = 0.2
|
||||
_PROOF_COUNT_ALPHA: float = 0.1 # Conservative: max ±5% for evidence strength
|
||||
|
||||
# Recency decay: maps a memory's age (days) onto a freshness signal in [0, 1]
|
||||
# where 0.5 is neutral (no boost). The signal is then folded into the
|
||||
# multiplicative recency_boost via `1 + recency_alpha * (recency - 0.5)`.
|
||||
#
|
||||
# "linear" — straight line from 1.0 (today) to a floor of 0.1, reaching
|
||||
# the floor at `linear_window_days`. The historical default.
|
||||
# "exponential" — 0.5 ** (days_ago / halflife_days). The half-life is the age
|
||||
# at which the signal is exactly neutral (0.5): younger
|
||||
# memories are boosted, older ones penalised, with a smooth
|
||||
# asymptote toward 0 (no hard cutoff).
|
||||
# "none" — always neutral (0.5), disabling the recency boost entirely.
|
||||
# The validated set of names lives in config.RECENCY_DECAY_FUNCTIONS.
|
||||
_RECENCY_DECAY_FUNCTION: str = "linear"
|
||||
_RECENCY_DECAY_LINEAR_WINDOW_DAYS: float = 365.0
|
||||
_RECENCY_DECAY_HALFLIFE_DAYS: float = 90.0
|
||||
|
||||
|
||||
def compute_recency_decay(
|
||||
days_ago: float,
|
||||
function: str = _RECENCY_DECAY_FUNCTION,
|
||||
linear_window_days: float = _RECENCY_DECAY_LINEAR_WINDOW_DAYS,
|
||||
halflife_days: float = _RECENCY_DECAY_HALFLIFE_DAYS,
|
||||
) -> float:
|
||||
"""Map a memory's age in days to a freshness signal in [0, 1] (neutral 0.5).
|
||||
|
||||
Future-dated memories (negative ``days_ago``) clamp to the maximum freshness
|
||||
so they are never penalised. See ``RECENCY_DECAY_FUNCTIONS`` for the shapes.
|
||||
"""
|
||||
if function == "none":
|
||||
return 0.5
|
||||
if function == "exponential":
|
||||
if halflife_days <= 0:
|
||||
return 0.5
|
||||
return min(1.0, 0.5 ** (days_ago / halflife_days))
|
||||
# "linear" (default): straight decay to a 0.1 floor over the window.
|
||||
window = linear_window_days if linear_window_days > 0 else _RECENCY_DECAY_LINEAR_WINDOW_DAYS
|
||||
return max(0.1, min(1.0, 1.0 - (days_ago / window)))
|
||||
|
||||
|
||||
def apply_combined_scoring(
|
||||
scored_results: list[ScoredResult],
|
||||
@@ -24,6 +62,9 @@ def apply_combined_scoring(
|
||||
temporal_alpha: float = _TEMPORAL_ALPHA,
|
||||
proof_count_alpha: float = _PROOF_COUNT_ALPHA,
|
||||
is_passthrough_reranker: bool = False,
|
||||
recency_decay_function: str = _RECENCY_DECAY_FUNCTION,
|
||||
recency_decay_linear_window_days: float = _RECENCY_DECAY_LINEAR_WINDOW_DAYS,
|
||||
recency_decay_halflife_days: float = _RECENCY_DECAY_HALFLIFE_DAYS,
|
||||
) -> None:
|
||||
"""Apply combined scoring to a list of ScoredResults in-place.
|
||||
|
||||
@@ -57,6 +98,12 @@ def apply_combined_scoring(
|
||||
recency_alpha: Max relative recency adjustment (default 0.2 → ±10%).
|
||||
temporal_alpha: Max relative temporal adjustment (default 0.2 → ±10%).
|
||||
proof_count_alpha: Max relative proof count adjustment (default 0.1 → ±5%).
|
||||
recency_decay_function: Age→freshness curve — "linear" (default),
|
||||
"exponential", or "none". See compute_recency_decay.
|
||||
recency_decay_linear_window_days: Days over which the linear curve
|
||||
decays to its floor (default 365).
|
||||
recency_decay_halflife_days: For the exponential curve, the age at which
|
||||
the recency signal is neutral (0.5) (default 90).
|
||||
"""
|
||||
if now.tzinfo is None:
|
||||
now = now.replace(tzinfo=UTC)
|
||||
@@ -98,7 +145,8 @@ def apply_combined_scoring(
|
||||
sr.cross_encoder_score_normalized = 1.0 - (0.9 * new_rank / denom)
|
||||
|
||||
for sr in scored_results:
|
||||
# Recency: linear decay over 365 days → [0.1, 1.0]; neutral 0.5 if no date.
|
||||
# Recency: configurable decay (linear default; see compute_recency_decay)
|
||||
# → [0.0, 1.0]; neutral 0.5 if no date.
|
||||
# Use the unit's effective time (occurred_start, then mentioned_at, then
|
||||
# occurred_end) — the same COALESCE order as retrieval._coalesce_date — so a
|
||||
# memory that carries only a mentioned_at / occurred_end (e.g. conversation
|
||||
@@ -111,7 +159,12 @@ def apply_combined_scoring(
|
||||
if occurred.tzinfo is None:
|
||||
occurred = occurred.replace(tzinfo=UTC)
|
||||
days_ago = (now - occurred).total_seconds() / 86400
|
||||
sr.recency = max(0.1, min(1.0, 1.0 - (days_ago / 365)))
|
||||
sr.recency = compute_recency_decay(
|
||||
days_ago,
|
||||
recency_decay_function,
|
||||
recency_decay_linear_window_days,
|
||||
recency_decay_halflife_days,
|
||||
)
|
||||
|
||||
# Temporal proximity: meaningful only for temporal queries; neutral otherwise.
|
||||
sr.temporal = sr.retrieval.temporal_proximity if sr.retrieval.temporal_proximity is not None else 0.5
|
||||
@@ -124,6 +177,9 @@ def apply_combined_scoring(
|
||||
else:
|
||||
# Neutral baseline is precisely 0.5, ensuring neutral multiplier (1.0)
|
||||
proof_norm = 0.5
|
||||
# Surface the proof signal so the trace can show the proof_count_boost
|
||||
# factor (otherwise the reranked breakdown can't reconcile CE × boosts).
|
||||
sr.proof_norm = proof_norm
|
||||
|
||||
# RRF: kept at 0.0 for trace continuity but excluded from scoring.
|
||||
# RRF is batch-relative (min-max normalised) and redundant after reranking.
|
||||
|
||||
@@ -104,6 +104,8 @@ async def retrieve_semantic_bm25_combined(
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
min_semantic: float | None = None,
|
||||
min_keyword: float | None = None,
|
||||
) -> dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]]:
|
||||
"""
|
||||
Combined semantic + BM25 retrieval for multiple fact types in a single query.
|
||||
@@ -143,6 +145,12 @@ async def retrieve_semantic_bm25_combined(
|
||||
config = get_config()
|
||||
tokens = tokenize_query(query_text)
|
||||
|
||||
# Per-request retrieval-level score floors (recall min_scores.semantic / .keyword)
|
||||
# override the global config defaults for this query, pruning weak matches in
|
||||
# the SQL arms before fusion.
|
||||
sem_min = min_semantic if min_semantic is not None else config.semantic_min_similarity
|
||||
bm25_min = min_keyword if min_keyword is not None else config.bm25_min_score
|
||||
|
||||
# Over-fetch for HNSW approximation; semantic results trimmed to limit in Python.
|
||||
hnsw_fetch = max(limit * 5, 100)
|
||||
|
||||
@@ -203,7 +211,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
embedding_param="$1",
|
||||
bank_id_param="$2",
|
||||
fetch_limit=hnsw_fetch,
|
||||
min_similarity=config.semantic_min_similarity,
|
||||
min_similarity=sem_min,
|
||||
tags_clause=tags_clause,
|
||||
groups_clause=groups_clause,
|
||||
extra_where=created_range_clause,
|
||||
@@ -229,7 +237,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
arm_index=i,
|
||||
text_search_extension=text_ext,
|
||||
bm25_language=config.text_search_extension_native_language,
|
||||
bm25_min_score=config.bm25_min_score,
|
||||
bm25_min_score=bm25_min,
|
||||
extra_where=created_range_clause,
|
||||
)
|
||||
)
|
||||
@@ -277,7 +285,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
embedding_param="$1",
|
||||
bank_id_param="$2",
|
||||
fetch_limit=hnsw_fetch,
|
||||
min_similarity=config.semantic_min_similarity,
|
||||
min_similarity=sem_min,
|
||||
tags_clause=fb_tags_clause,
|
||||
groups_clause=fb_groups_clause,
|
||||
extra_where=fb_created_clause,
|
||||
@@ -706,6 +714,8 @@ async def retrieve_all_fact_types_parallel(
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
min_semantic: float | None = None,
|
||||
min_keyword: float | None = None,
|
||||
) -> MultiFactTypeRetrievalResult:
|
||||
"""
|
||||
Optimized retrieval for multiple fact types using batched queries.
|
||||
@@ -766,6 +776,8 @@ async def retrieve_all_fact_types_parallel(
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
min_semantic=min_semantic,
|
||||
min_keyword=min_keyword,
|
||||
)
|
||||
semantic_bm25_time = time.time() - semantic_bm25_start
|
||||
|
||||
@@ -781,7 +793,7 @@ async def retrieve_all_fact_types_parallel(
|
||||
tc_start,
|
||||
tc_end,
|
||||
budget=thinking_budget,
|
||||
semantic_threshold=0.1,
|
||||
semantic_threshold=min_semantic if min_semantic is not None else 0.1,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
|
||||
@@ -14,6 +14,12 @@ AND matching (all/all_strict): Memory matches if ALL request tags are present in
|
||||
EXACT matching: Memory matches only if its tag set EQUALS the request tag set (order-
|
||||
independent). Used for observation "scope" filtering, where each observation lives
|
||||
under exactly one scope (its full tag set) and "scope [a]" must not match "[a, b]".
|
||||
An EMPTY request scope (no tags — ``[]`` or ``None``) is the global/untagged scope and
|
||||
matches only untagged memories — the scope that ``observation_scopes="shared"``
|
||||
consolidation writes to. This is the one mode where absent tags filter rather than
|
||||
meaning "no filter"; all other modes treat empty/absent tags as "no filtering". This
|
||||
mirrors the ``GET .../graph`` endpoint, where ``tags_match="exact"`` with no tags also
|
||||
selects the global scope.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -82,11 +88,16 @@ def build_tags_where_clause(
|
||||
>>> clause, params, next_offset = build_tags_where_clause(['user_a'], 3, 'mu.', 'any_strict')
|
||||
>>> print(clause) # "AND mu.tags IS NOT NULL AND mu.tags != '{}' AND mu.tags && $3"
|
||||
"""
|
||||
column = f"{table_alias}tags" if table_alias else "tags"
|
||||
|
||||
if match == "exact" and not tags:
|
||||
# Empty/absent scope = global/untagged: match only untagged rows. No bind param
|
||||
# needed (callers gate the param on truthy `tags`, so none is appended).
|
||||
return f"AND ({column} IS NULL OR {column} = '{{}}')", [], param_offset
|
||||
|
||||
if not tags:
|
||||
return "", [], param_offset
|
||||
|
||||
column = f"{table_alias}tags" if table_alias else "tags"
|
||||
|
||||
if match == "exact":
|
||||
# Set equality (order-independent): superset AND subset. Untagged rows
|
||||
# (empty array) never satisfy `@>` of a non-empty scope, so they're excluded.
|
||||
@@ -126,11 +137,16 @@ def build_tags_where_clause_simple(
|
||||
Returns:
|
||||
SQL clause string or empty string.
|
||||
"""
|
||||
column = f"{table_alias}tags" if table_alias else "tags"
|
||||
|
||||
if match == "exact" and not tags:
|
||||
# Empty/absent scope = global/untagged: match only untagged rows. No bind param
|
||||
# needed (callers gate the param on truthy `tags`, so none is appended).
|
||||
return f"AND ({column} IS NULL OR {column} = '{{}}')"
|
||||
|
||||
if not tags:
|
||||
return ""
|
||||
|
||||
column = f"{table_alias}tags" if table_alias else "tags"
|
||||
|
||||
if match == "exact":
|
||||
# Set equality (order-independent): superset AND subset. Untagged rows
|
||||
# (empty array) never satisfy `@>` of a non-empty scope, so they're excluded.
|
||||
@@ -164,6 +180,10 @@ def filter_results_by_tags(
|
||||
Returns:
|
||||
Filtered list of results.
|
||||
"""
|
||||
if match == "exact" and not tags:
|
||||
# Empty/absent scope = global/untagged: keep only untagged results.
|
||||
return [r for r in results if not getattr(r, "tags", None)]
|
||||
|
||||
if not tags:
|
||||
return results
|
||||
|
||||
@@ -267,6 +287,9 @@ def _build_group_clause(
|
||||
if isinstance(group, TagGroupLeaf):
|
||||
column = f"{table_alias}tags" if table_alias else "tags"
|
||||
if group.match == "exact":
|
||||
if len(group.tags) == 0:
|
||||
# Empty scope = global/untagged: match only untagged rows (no bind param).
|
||||
return f"({column} IS NULL OR {column} = '{{}}')", [], param_offset
|
||||
clause = f"({column} @> ${param_offset} AND {column} <@ ${param_offset})"
|
||||
return clause, [group.tags], param_offset + 1
|
||||
operator, include_untagged = _parse_tags_match(group.match)
|
||||
@@ -369,6 +392,9 @@ def _match_group(result: object, group: TagGroup) -> bool:
|
||||
if isinstance(group, TagGroupLeaf):
|
||||
result_tags = getattr(result, "tags", None)
|
||||
is_untagged = result_tags is None or len(result_tags) == 0
|
||||
if group.match == "exact" and len(group.tags) == 0:
|
||||
# Empty scope = global/untagged: match only untagged results.
|
||||
return is_untagged
|
||||
_, include_untagged = _parse_tags_match(group.match)
|
||||
is_any_match = group.match in ("any", "any_strict")
|
||||
tags_set = set(group.tags)
|
||||
|
||||
@@ -5,6 +5,7 @@ Think operation utilities for formulating answers based on agent and world facts
|
||||
import logging
|
||||
from datetime import datetime
|
||||
|
||||
from ...config import get_config
|
||||
from ..response_models import DispositionTraits, MemoryFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -251,7 +252,7 @@ async def reflect(
|
||||
answer_text = await llm_config.call(
|
||||
messages=[{"role": "system", "content": system_message}, {"role": "user", "content": prompt}],
|
||||
scope="memory_think",
|
||||
temperature=0.9,
|
||||
temperature=get_config().llm_temperature_reflect,
|
||||
max_completion_tokens=1000,
|
||||
)
|
||||
|
||||
|
||||
@@ -392,7 +392,7 @@ class SearchTracer:
|
||||
|
||||
# Extract score components (only include non-None values)
|
||||
# Keys from ScoredResult.to_dict(): cross_encoder_score, cross_encoder_score_normalized,
|
||||
# rrf_normalized, temporal, recency, combined_score, weight
|
||||
# rrf_normalized, temporal, recency, proof_norm, combined_score, weight
|
||||
score_components = {}
|
||||
for key in [
|
||||
"cross_encoder_score",
|
||||
@@ -401,6 +401,7 @@ class SearchTracer:
|
||||
"rrf_normalized",
|
||||
"temporal",
|
||||
"recency",
|
||||
"proof_norm",
|
||||
"combined_score",
|
||||
]:
|
||||
if key in result and result[key] is not None:
|
||||
|
||||
@@ -82,6 +82,20 @@ class RetrievalResult:
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ArmScores:
|
||||
"""Raw per-strategy retrieval scores for a single doc, aggregated across arms.
|
||||
|
||||
Fusion keeps only the first-seen RetrievalResult per doc, so its per-arm score
|
||||
fields reflect just one arm. This captures each arm's raw score for the same doc
|
||||
so the recall response can report them (and ``min_scores`` can filter on them).
|
||||
``None`` means the doc was not surfaced by that arm.
|
||||
"""
|
||||
|
||||
semantic: float | None = None # cosine similarity from the semantic arm
|
||||
keyword: float | None = None # BM25 / full-text score from the keyword arm
|
||||
|
||||
|
||||
@dataclass
|
||||
class MergedCandidate:
|
||||
"""
|
||||
@@ -97,6 +111,7 @@ class MergedCandidate:
|
||||
rrf_score: float
|
||||
rrf_rank: int = 0
|
||||
source_ranks: dict[str, int] = field(default_factory=dict) # method_name -> rank
|
||||
arm_scores: "ArmScores" = field(default_factory=lambda: ArmScores()) # raw per-strategy scores
|
||||
|
||||
@property
|
||||
def id(self) -> str:
|
||||
@@ -123,6 +138,7 @@ class ScoredResult:
|
||||
rrf_normalized: float = 0.0
|
||||
recency: float = 0.5
|
||||
temporal: float = 0.5
|
||||
proof_norm: float = 0.5 # log-normalized proof count (neutral 0.5); drives proof_count_boost
|
||||
|
||||
# Final combined score
|
||||
combined_score: float = 0.0
|
||||
@@ -179,6 +195,7 @@ class ScoredResult:
|
||||
result["rrf_normalized"] = self.rrf_normalized
|
||||
result["temporal"] = self.temporal
|
||||
result["recency"] = self.recency
|
||||
result["proof_norm"] = self.proof_norm
|
||||
result["combined_score"] = self.combined_score
|
||||
result["weight"] = self.weight
|
||||
result["activation"] = self.weight # Legacy field
|
||||
|
||||
@@ -39,7 +39,9 @@ from hindsight_api.extensions.operation_validator import (
|
||||
BankListContext,
|
||||
BankListResult,
|
||||
BankReadContext,
|
||||
BankReadOperation,
|
||||
BankWriteContext,
|
||||
BankWriteOperation,
|
||||
# Consolidation operation
|
||||
ConsolidateContext,
|
||||
ConsolidateResult,
|
||||
@@ -54,6 +56,7 @@ from hindsight_api.extensions.operation_validator import (
|
||||
OperationValidationError,
|
||||
OperationValidatorExtension,
|
||||
PrecheckContext,
|
||||
PrecheckOperation,
|
||||
RecallContext,
|
||||
RecallResult,
|
||||
ReflectContext,
|
||||
@@ -87,6 +90,7 @@ __all__ = [
|
||||
"OperationValidationError",
|
||||
"OperationValidatorExtension",
|
||||
"PrecheckContext",
|
||||
"PrecheckOperation",
|
||||
"RecallContext",
|
||||
"RecallResult",
|
||||
"ReflectContext",
|
||||
@@ -98,7 +102,9 @@ __all__ = [
|
||||
"BankListContext",
|
||||
"BankListResult",
|
||||
"BankReadContext",
|
||||
"BankReadOperation",
|
||||
"BankWriteContext",
|
||||
"BankWriteOperation",
|
||||
# Operation Validator - Consolidation
|
||||
"ConsolidateContext",
|
||||
"ConsolidateResult",
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from enum import StrEnum
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
@@ -82,6 +83,18 @@ class ValidationResult:
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class PrecheckOperation(StrEnum):
|
||||
"""Route operation names passed to the pre-body-parse precheck hook."""
|
||||
|
||||
DRY_RUN_EXTRACT = "dry_run_extract"
|
||||
FILES_RETAIN = "files_retain"
|
||||
MENTAL_MODEL_CREATE = "mental_model_create"
|
||||
MENTAL_MODEL_REFRESH = "mental_model_refresh"
|
||||
RECALL = "recall"
|
||||
REFLECT = "reflect"
|
||||
RETAIN = "retain"
|
||||
|
||||
|
||||
@dataclass
|
||||
class PrecheckContext:
|
||||
"""Context for a pre-body-parse precheck on an operation.
|
||||
@@ -91,12 +104,14 @@ class PrecheckContext:
|
||||
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"``.
|
||||
- ``operation``: a short string-compatible enum identifying the route.
|
||||
- ``bank_id``: parsed from the URL path.
|
||||
- ``request_context``: the authenticated :class:`RequestContext` (tenant
|
||||
already resolved by the tenant extension).
|
||||
- ``content_length``: value of the ``Content-Length`` request header as an
|
||||
int, or ``None`` when the header is absent or unparseable (e.g. chunked
|
||||
transfer encoding). Lets a precheck make size-aware decisions — such as
|
||||
an upper-bound cost estimate — without reading or deserialising the body.
|
||||
|
||||
Implementations should keep precheck cheap and side-effect-free. The
|
||||
full per-request validators (``validate_retain`` / ``validate_recall``
|
||||
@@ -104,9 +119,10 @@ class PrecheckContext:
|
||||
the source of truth for the precise per-call cost / quota arithmetic.
|
||||
"""
|
||||
|
||||
operation: str
|
||||
operation: PrecheckOperation
|
||||
bank_id: str
|
||||
request_context: "RequestContext"
|
||||
content_length: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -203,6 +219,16 @@ class RetainResult:
|
||||
llm_input_tokens: int | None = None
|
||||
llm_output_tokens: int | None = None
|
||||
llm_total_tokens: int | None = None
|
||||
# Diagnostic token splits surfaced for cost attribution and prompt-cache
|
||||
# tuning. ``llm_cached_input_tokens`` is the subset of llm_input_tokens
|
||||
# served from the provider's prompt cache (e.g. Gemini's
|
||||
# cached_content_token_count). ``llm_thoughts_tokens`` is reasoning tokens
|
||||
# that are billed at the output rate by some providers (Gemini 2.5+) but
|
||||
# are not part of the visible response. Both default to None when the
|
||||
# engine/provider didn't report them; downstream metering extensions
|
||||
# should treat None as 0.
|
||||
llm_cached_input_tokens: int | None = None
|
||||
llm_thoughts_tokens: int | None = None
|
||||
# Content tokens the retain pipeline actually processed, after
|
||||
# chunk-level content-hash deduplication. Semantics:
|
||||
# None — no dedup signal available (e.g. a first-time retain or a
|
||||
@@ -288,12 +314,77 @@ class ConsolidateResult:
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class BankReadOperation(StrEnum):
|
||||
"""Bank-scoped read operation names passed to validate_bank_read."""
|
||||
|
||||
GET_BANK_CONFIG = "get_bank_config"
|
||||
GET_BANK_PROFILE = "get_bank_profile"
|
||||
GET_BANK_STATS = "get_bank_stats"
|
||||
GET_CHUNK = "get_chunk"
|
||||
GET_DIRECTIVE = "get_directive"
|
||||
GET_DOCUMENT = "get_document"
|
||||
GET_ENTITY = "get_entity"
|
||||
GET_ENTITY_GRAPH = "get_entity_graph"
|
||||
GET_ENTITY_STATE = "get_entity_state"
|
||||
GET_GRAPH_DATA = "get_graph_data"
|
||||
GET_MEMORIES_TIMESERIES = "get_memories_timeseries"
|
||||
GET_MEMORY_UNIT = "get_memory_unit"
|
||||
GET_OBSERVATION_HISTORY = "get_observation_history"
|
||||
GET_OPERATION_STATUS = "get_operation_status"
|
||||
LIST_DIRECTIVES = "list_directives"
|
||||
LIST_DOCUMENT_CHUNKS = "list_document_chunks"
|
||||
LIST_DOCUMENTS = "list_documents"
|
||||
LIST_ENTITIES = "list_entities"
|
||||
LIST_MEMORY_UNITS = "list_memory_units"
|
||||
LIST_MENTAL_MODEL_TAGS = "list_mental_model_tags"
|
||||
LIST_MENTAL_MODELS = "list_mental_models"
|
||||
LIST_OBSERVATION_SCOPES = "list_observation_scopes"
|
||||
LIST_OPERATIONS = "list_operations"
|
||||
LIST_TAGS = "list_tags"
|
||||
LIST_WEBHOOK_DELIVERIES = "list_webhook_deliveries"
|
||||
LIST_WEBHOOKS = "list_webhooks"
|
||||
|
||||
|
||||
class BankWriteOperation(StrEnum):
|
||||
"""Bank-scoped write operation names passed to validate_bank_write."""
|
||||
|
||||
CANCEL_OPERATION = "cancel_operation"
|
||||
CLEAR_MENTAL_MODEL = "clear_mental_model"
|
||||
CLEAR_OBSERVATIONS = "clear_observations"
|
||||
CLEAR_OBSERVATIONS_FOR_MEMORY = "clear_observations_for_memory"
|
||||
CREATE_DIRECTIVE = "create_directive"
|
||||
CREATE_MENTAL_MODEL = "create_mental_model"
|
||||
CREATE_WEBHOOK = "create_webhook"
|
||||
DELETE_BANK = "delete_bank"
|
||||
DELETE_DIRECTIVE = "delete_directive"
|
||||
DELETE_DOCUMENT = "delete_document"
|
||||
DELETE_MENTAL_MODEL = "delete_mental_model"
|
||||
DELETE_WEBHOOK = "delete_webhook"
|
||||
MERGE_BANK_MISSION = "merge_bank_mission"
|
||||
REPROCESS_DOCUMENT = "reprocess_document"
|
||||
RESET_BANK_CONFIG = "reset_bank_config"
|
||||
RETRY_FAILED_CONSOLIDATION = "retry_failed_consolidation"
|
||||
RETRY_OPERATION = "retry_operation"
|
||||
RUN_CONSOLIDATION = "run_consolidation"
|
||||
SET_BANK_MISSION = "set_bank_mission"
|
||||
SUBMIT_ASYNC_CONSOLIDATION = "submit_async_consolidation"
|
||||
SUBMIT_ASYNC_GRAPH_MAINTENANCE = "submit_async_graph_maintenance"
|
||||
UPDATE_BANK = "update_bank"
|
||||
UPDATE_BANK_CONFIG = "update_bank_config"
|
||||
UPDATE_BANK_DISPOSITION = "update_bank_disposition"
|
||||
UPDATE_DIRECTIVE = "update_directive"
|
||||
UPDATE_DOCUMENT = "update_document"
|
||||
UPDATE_MEMORY_UNIT = "update_memory_unit"
|
||||
UPDATE_MENTAL_MODEL = "update_mental_model"
|
||||
UPDATE_WEBHOOK = "update_webhook"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BankReadContext:
|
||||
"""Context for a bank read operation validation (pre-operation)."""
|
||||
|
||||
bank_id: str
|
||||
operation: str # "get_bank_profile", "get_bank_stats"
|
||||
operation: BankReadOperation
|
||||
request_context: "RequestContext"
|
||||
|
||||
|
||||
@@ -302,7 +393,7 @@ class BankWriteContext:
|
||||
"""Context for a bank write operation validation (pre-operation)."""
|
||||
|
||||
bank_id: str
|
||||
operation: str # "delete_bank", "update_bank", "update_bank_disposition", "set_bank_mission", "merge_bank_mission", "clear_observations", "clear_observations_for_memory"
|
||||
operation: BankWriteOperation
|
||||
request_context: "RequestContext"
|
||||
|
||||
|
||||
|
||||
@@ -12,6 +12,7 @@ from datetime import datetime, timezone
|
||||
from typing import Any, Callable
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from mcp.types import ToolAnnotations
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
@@ -21,7 +22,7 @@ from hindsight_api.config import (
|
||||
)
|
||||
from hindsight_api.engine.audit import AuditEntry, AuditLogger
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES, MinScores
|
||||
from hindsight_api.engine.search.tags import TagGroup
|
||||
from hindsight_api.extensions import OperationValidationError
|
||||
from hindsight_api.models import RequestContext
|
||||
@@ -199,6 +200,47 @@ def build_content_dict(
|
||||
return content_dict, None
|
||||
|
||||
|
||||
# MCP tool annotations. Hindsight is a closed memory store (no open-world / internet
|
||||
# access), so openWorldHint=False throughout. readOnlyHint lets clients group and
|
||||
# auto-approve safe reads; destructiveHint flags tools that delete or clear memory.
|
||||
_READ_ONLY_TOOLS = {
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_banks",
|
||||
"get_bank",
|
||||
"get_bank_stats",
|
||||
"list_mental_models",
|
||||
"get_mental_model",
|
||||
"list_directives",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"list_operations",
|
||||
"get_operation",
|
||||
"list_tags",
|
||||
}
|
||||
_DESTRUCTIVE_TOOLS = {
|
||||
"delete_bank",
|
||||
"clear_memories",
|
||||
"clear_mental_model",
|
||||
"delete_mental_model",
|
||||
"delete_directive",
|
||||
"delete_document",
|
||||
"invalidate_memory",
|
||||
}
|
||||
|
||||
|
||||
def _tool_annotations(name: str) -> ToolAnnotations:
|
||||
if name in _READ_ONLY_TOOLS:
|
||||
return ToolAnnotations(readOnlyHint=True, openWorldHint=False)
|
||||
if name in _DESTRUCTIVE_TOOLS:
|
||||
return ToolAnnotations(readOnlyHint=False, destructiveHint=True, openWorldHint=False)
|
||||
# Everything else writes but does not destructively delete/clear memory
|
||||
# (retain, create_*, update_*, refresh_mental_model, cancel_operation).
|
||||
return ToolAnnotations(readOnlyHint=False, destructiveHint=False, openWorldHint=False)
|
||||
|
||||
|
||||
def register_mcp_tools(
|
||||
mcp: FastMCP,
|
||||
memory: MemoryEngine,
|
||||
@@ -552,7 +594,7 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
@mcp.tool(description=description, annotations=_tool_annotations("retain"))
|
||||
async def retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
@@ -608,7 +650,7 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
@mcp.tool(description=description, annotations=_tool_annotations("retain"))
|
||||
async def retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
@@ -666,7 +708,7 @@ def _register_sync_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("sync_retain"))
|
||||
async def sync_retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
@@ -724,7 +766,7 @@ def _register_sync_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("sync_retain"))
|
||||
async def sync_retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
@@ -785,16 +827,18 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
@mcp.tool(description=description, annotations=_tool_annotations("recall"))
|
||||
async def recall(
|
||||
query: str,
|
||||
max_tokens: int = 4096,
|
||||
budget: str = "high",
|
||||
types: list[str] | None = None,
|
||||
prefer_observations: bool = False,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list[dict] | None = None,
|
||||
query_timestamp: str | None = None,
|
||||
min_scores: dict | None = None,
|
||||
bank_id: str | None = None,
|
||||
) -> str | dict:
|
||||
"""
|
||||
@@ -803,6 +847,10 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
max_tokens: Maximum tokens to return in results (default: 4096)
|
||||
budget: Search budget - 'low', 'mid', or 'high' (default: 'high'). Higher budgets search more thoroughly.
|
||||
types: Fact types to include (e.g., ['world', 'experience']). Default: all types.
|
||||
prefer_observations: When recalling raw facts together with 'observation', drop any raw fact
|
||||
that a returned observation was consolidated from, so the observation supersedes it (no
|
||||
duplicate content). Disabled by default; set true to enable. No effect unless
|
||||
'observation' and a raw type are both in types. Default: False.
|
||||
tags: Optional tags to filter results by (e.g., ['project:alpha']). Mutually exclusive with tag_groups.
|
||||
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
|
||||
tag_groups: Compound tag filter using boolean groups (AND-ed together). Each group is a leaf
|
||||
@@ -811,6 +859,11 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
Mutually exclusive with tags.
|
||||
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z').
|
||||
Anchors relative temporal expressions and recency scoring.
|
||||
min_scores: Optional per-stage score floors as an object with any of: "semantic", "keyword"
|
||||
(retrieval-level cutoffs), "reranker", "final" (post-ranking). E.g. {"reranker": 0.5}.
|
||||
All inclusive and AND-ed; omit for no score filtering. The reranker's absolute scores are
|
||||
not calibrated across queries, so only threshold against scores you've calibrated for your
|
||||
own data.
|
||||
bank_id: Optional bank to search in (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
@@ -831,6 +884,7 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
"bank_id": target_bank,
|
||||
"query": query,
|
||||
"fact_type": fact_types,
|
||||
"prefer_observations": prefer_observations,
|
||||
"budget": budget_enum,
|
||||
"max_tokens": max_tokens,
|
||||
"request_context": _get_request_context(config),
|
||||
@@ -842,6 +896,8 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
recall_kwargs["tag_groups"] = _TAG_GROUP_LIST_ADAPTER.validate_python(tag_groups)
|
||||
if query_timestamp is not None:
|
||||
recall_kwargs["question_date"] = parse_timestamp(query_timestamp)
|
||||
if min_scores is not None:
|
||||
recall_kwargs["min_scores"] = MinScores.model_validate(min_scores)
|
||||
|
||||
recall_result = await memory.recall_async(**recall_kwargs)
|
||||
|
||||
@@ -857,16 +913,18 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
@mcp.tool(description=description, annotations=_tool_annotations("recall"))
|
||||
async def recall(
|
||||
query: str,
|
||||
max_tokens: int = 4096,
|
||||
budget: str = "high",
|
||||
types: list[str] | None = None,
|
||||
prefer_observations: bool = False,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list[dict] | None = None,
|
||||
query_timestamp: str | None = None,
|
||||
min_scores: dict | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Args:
|
||||
@@ -874,6 +932,10 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
max_tokens: Maximum tokens to return in results (default: 4096)
|
||||
budget: Search budget - 'low', 'mid', or 'high' (default: 'high'). Higher budgets search more thoroughly.
|
||||
types: Fact types to include (e.g., ['world', 'experience']). Default: all types.
|
||||
prefer_observations: When recalling raw facts together with 'observation', drop any raw fact
|
||||
that a returned observation was consolidated from, so the observation supersedes it (no
|
||||
duplicate content). Disabled by default; set true to enable. No effect unless
|
||||
'observation' and a raw type are both in types. Default: False.
|
||||
tags: Optional tags to filter results by (e.g., ['project:alpha']). Mutually exclusive with tag_groups.
|
||||
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
|
||||
tag_groups: Compound tag filter using boolean groups (AND-ed together). Each group is a leaf
|
||||
@@ -882,6 +944,11 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
Mutually exclusive with tags.
|
||||
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z').
|
||||
Anchors relative temporal expressions and recency scoring.
|
||||
min_scores: Optional per-stage score floors as an object with any of: "semantic", "keyword"
|
||||
(retrieval-level cutoffs), "reranker", "final" (post-ranking). E.g. {"reranker": 0.5}.
|
||||
All inclusive and AND-ed; omit for no score filtering. The reranker's absolute scores are
|
||||
not calibrated across queries, so only threshold against scores you've calibrated for your
|
||||
own data.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
@@ -901,6 +968,7 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
"bank_id": target_bank,
|
||||
"query": query,
|
||||
"fact_type": fact_types,
|
||||
"prefer_observations": prefer_observations,
|
||||
"budget": budget_enum,
|
||||
"max_tokens": max_tokens,
|
||||
"request_context": _get_request_context(config),
|
||||
@@ -912,6 +980,8 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
recall_kwargs["tag_groups"] = _TAG_GROUP_LIST_ADAPTER.validate_python(tag_groups)
|
||||
if query_timestamp is not None:
|
||||
recall_kwargs["question_date"] = parse_timestamp(query_timestamp)
|
||||
if min_scores is not None:
|
||||
recall_kwargs["min_scores"] = MinScores.model_validate(min_scores)
|
||||
|
||||
recall_result = await memory.recall_async(**recall_kwargs)
|
||||
|
||||
@@ -931,7 +1001,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("reflect"))
|
||||
async def reflect(
|
||||
query: str,
|
||||
context: str | None = None,
|
||||
@@ -941,6 +1011,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
include_based_on: bool = False,
|
||||
include_trace: bool = False,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -971,6 +1042,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
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.
|
||||
include_trace: Include the reflection's internal trace fields (tool_trace/llm_trace and directives_applied). Defaults to false because the trace can be tens of KB and overflow MCP client context; enable only for debugging.
|
||||
bank_id: Optional bank to reflect in (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
@@ -1000,6 +1072,15 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
result_data = json.loads(reflect_result.model_dump_json(indent=2))
|
||||
if not include_based_on:
|
||||
result_data.pop("based_on", None)
|
||||
if not include_trace:
|
||||
# The agentic reflect loop's trace fields can be tens of KB (full
|
||||
# mental-model text) and silently overflow MCP client context; the
|
||||
# REST API omits them by default too. directives_applied is built by
|
||||
# the engine "for the trace" and carries full directive content, so it
|
||||
# belongs with tool_trace/llm_trace here. Opt in via include_trace.
|
||||
result_data.pop("tool_trace", None)
|
||||
result_data.pop("llm_trace", None)
|
||||
result_data.pop("directives_applied", 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)
|
||||
@@ -1012,7 +1093,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("reflect"))
|
||||
async def reflect(
|
||||
query: str,
|
||||
context: str | None = None,
|
||||
@@ -1022,6 +1103,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
include_based_on: bool = False,
|
||||
include_trace: bool = False,
|
||||
) -> dict:
|
||||
"""
|
||||
Generate thoughtful analysis by synthesizing stored memories with the bank's personality.
|
||||
@@ -1051,6 +1133,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
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.
|
||||
include_trace: Include the reflection's internal trace fields (tool_trace/llm_trace and directives_applied). Defaults to false because the trace can be tens of KB and overflow MCP client context; enable only for debugging.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
@@ -1079,6 +1162,15 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
result_data = reflect_result.model_dump()
|
||||
if not include_based_on:
|
||||
result_data.pop("based_on", None)
|
||||
if not include_trace:
|
||||
# The agentic reflect loop's trace fields can be tens of KB (full
|
||||
# mental-model text) and silently overflow MCP client context; the
|
||||
# REST API omits them by default too. directives_applied is built by
|
||||
# the engine "for the trace" and carries full directive content, so it
|
||||
# belongs with tool_trace/llm_trace here. Opt in via include_trace.
|
||||
result_data.pop("tool_trace", None)
|
||||
result_data.pop("llm_trace", None)
|
||||
result_data.pop("directives_applied", None)
|
||||
if response_schema is not None and hasattr(reflect_result, "structured_output"):
|
||||
result_data["structured_output"] = reflect_result.structured_output
|
||||
return result_data
|
||||
@@ -1093,7 +1185,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
def _register_list_banks(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the list_banks tool."""
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_banks"))
|
||||
async def list_banks() -> str:
|
||||
"""
|
||||
List all available memory banks.
|
||||
@@ -1118,7 +1210,7 @@ def _register_list_banks(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
def _register_create_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the create_bank tool."""
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_bank"))
|
||||
async def create_bank(bank_id: str, name: str | None = None, mission: str | None = None) -> str:
|
||||
"""
|
||||
Create a new memory bank or get an existing one.
|
||||
@@ -1182,7 +1274,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_mental_models"))
|
||||
async def list_mental_models(
|
||||
tags: list[str] | None = None,
|
||||
detail: str = "full",
|
||||
@@ -1221,7 +1313,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_mental_models"))
|
||||
async def list_mental_models(
|
||||
tags: list[str] | None = None,
|
||||
detail: str = "full",
|
||||
@@ -1262,7 +1354,7 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_mental_model"))
|
||||
async def get_mental_model(
|
||||
mental_model_id: str,
|
||||
detail: str = "full",
|
||||
@@ -1302,7 +1394,7 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_mental_model"))
|
||||
async def get_mental_model(
|
||||
mental_model_id: str,
|
||||
detail: str = "full",
|
||||
@@ -1344,7 +1436,7 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_mental_model"))
|
||||
async def create_mental_model(
|
||||
name: str,
|
||||
source_query: str,
|
||||
@@ -1428,7 +1520,7 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_mental_model"))
|
||||
async def create_mental_model(
|
||||
name: str,
|
||||
source_query: str,
|
||||
@@ -1510,7 +1602,7 @@ def _register_update_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("update_mental_model"))
|
||||
async def update_mental_model(
|
||||
mental_model_id: str,
|
||||
name: str | None = None,
|
||||
@@ -1571,7 +1663,7 @@ def _register_update_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("update_mental_model"))
|
||||
async def update_mental_model(
|
||||
mental_model_id: str,
|
||||
name: str | None = None,
|
||||
@@ -1634,7 +1726,7 @@ def _register_delete_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_mental_model"))
|
||||
async def delete_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -1670,7 +1762,7 @@ def _register_delete_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_mental_model"))
|
||||
async def delete_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
@@ -1708,7 +1800,7 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("refresh_mental_model"))
|
||||
async def refresh_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -1752,7 +1844,7 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("refresh_mental_model"))
|
||||
async def refresh_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
@@ -1796,7 +1888,7 @@ def _register_clear_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("clear_mental_model"))
|
||||
async def clear_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -1842,7 +1934,7 @@ def _register_clear_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("clear_mental_model"))
|
||||
async def clear_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
@@ -1893,7 +1985,7 @@ def _register_list_directives(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_directives"))
|
||||
async def list_directives(
|
||||
tags: list[str] | None = None,
|
||||
active_only: bool = True,
|
||||
@@ -1931,7 +2023,7 @@ def _register_list_directives(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_directives"))
|
||||
async def list_directives(
|
||||
tags: list[str] | None = None,
|
||||
active_only: bool = True,
|
||||
@@ -1971,7 +2063,7 @@ def _register_create_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_directive"))
|
||||
async def create_directive(
|
||||
name: str,
|
||||
content: str,
|
||||
@@ -2017,7 +2109,7 @@ def _register_create_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_directive"))
|
||||
async def create_directive(
|
||||
name: str,
|
||||
content: str,
|
||||
@@ -2065,7 +2157,7 @@ def _register_delete_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_directive"))
|
||||
async def delete_directive(
|
||||
directive_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2101,7 +2193,7 @@ def _register_delete_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_directive"))
|
||||
async def delete_directive(
|
||||
directive_id: str,
|
||||
) -> dict:
|
||||
@@ -2144,7 +2236,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_memories"))
|
||||
async def list_memories(
|
||||
type: str | None = None,
|
||||
q: str | None = None,
|
||||
@@ -2159,7 +2251,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
browse/search without relevance ranking.
|
||||
|
||||
Args:
|
||||
type: Filter by fact type: 'world', 'experience', or 'opinion'
|
||||
type: Filter by fact type: 'world', 'experience', or 'observation'
|
||||
q: Optional text search query to filter memories
|
||||
limit: Maximum number of results (default: 100)
|
||||
offset: Pagination offset (default: 0)
|
||||
@@ -2188,7 +2280,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_memories"))
|
||||
async def list_memories(
|
||||
type: str | None = None,
|
||||
q: str | None = None,
|
||||
@@ -2202,7 +2294,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
browse/search without relevance ranking.
|
||||
|
||||
Args:
|
||||
type: Filter by fact type: 'world', 'experience', or 'opinion'
|
||||
type: Filter by fact type: 'world', 'experience', or 'observation'
|
||||
q: Optional text search query to filter memories
|
||||
limit: Maximum number of results (default: 100)
|
||||
offset: Pagination offset (default: 0)
|
||||
@@ -2234,7 +2326,7 @@ def _register_get_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_memory"))
|
||||
async def get_memory(
|
||||
memory_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2270,7 +2362,7 @@ def _register_get_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_memory"))
|
||||
async def get_memory(
|
||||
memory_id: str,
|
||||
) -> dict:
|
||||
@@ -2321,7 +2413,7 @@ def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(description=_EDIT_DOC)
|
||||
@mcp.tool(description=_EDIT_DOC, annotations=_tool_annotations("update_memory"))
|
||||
async def update_memory(
|
||||
memory_id: str,
|
||||
text: str | None = None,
|
||||
@@ -2367,7 +2459,7 @@ def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(description=_EDIT_DOC)
|
||||
@mcp.tool(description=_EDIT_DOC, annotations=_tool_annotations("update_memory"))
|
||||
async def update_memory(
|
||||
memory_id: str,
|
||||
text: str | None = None,
|
||||
@@ -2426,7 +2518,7 @@ def _register_invalidate_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPT
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(description=_INVALIDATE_DOC)
|
||||
@mcp.tool(description=_INVALIDATE_DOC, annotations=_tool_annotations("invalidate_memory"))
|
||||
async def invalidate_memory(
|
||||
memory_id: str,
|
||||
reason: str | None = None,
|
||||
@@ -2466,7 +2558,7 @@ def _register_invalidate_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPT
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(description=_INVALIDATE_DOC)
|
||||
@mcp.tool(description=_INVALIDATE_DOC, annotations=_tool_annotations("invalidate_memory"))
|
||||
async def invalidate_memory(
|
||||
memory_id: str,
|
||||
reason: str | None = None,
|
||||
@@ -2513,7 +2605,7 @@ def _register_list_documents(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_documents"))
|
||||
async def list_documents(
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
@@ -2551,7 +2643,7 @@ def _register_list_documents(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_documents"))
|
||||
async def list_documents(
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
@@ -2591,7 +2683,7 @@ def _register_get_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsC
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_document"))
|
||||
async def get_document(
|
||||
document_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2627,7 +2719,7 @@ def _register_get_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsC
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_document"))
|
||||
async def get_document(
|
||||
document_id: str,
|
||||
) -> dict:
|
||||
@@ -2665,7 +2757,7 @@ def _register_delete_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_document"))
|
||||
async def delete_document(
|
||||
document_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2699,7 +2791,7 @@ def _register_delete_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_document"))
|
||||
async def delete_document(
|
||||
document_id: str,
|
||||
) -> dict:
|
||||
@@ -2740,7 +2832,7 @@ def _register_list_operations(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_operations"))
|
||||
async def list_operations(
|
||||
status: str | None = None,
|
||||
limit: int = 20,
|
||||
@@ -2777,7 +2869,7 @@ def _register_list_operations(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_operations"))
|
||||
async def list_operations(
|
||||
status: str | None = None,
|
||||
limit: int = 20,
|
||||
@@ -2816,7 +2908,7 @@ def _register_get_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_operation"))
|
||||
async def get_operation(
|
||||
operation_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2850,7 +2942,7 @@ def _register_get_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_operation"))
|
||||
async def get_operation(
|
||||
operation_id: str,
|
||||
) -> dict:
|
||||
@@ -2886,7 +2978,7 @@ def _register_cancel_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("cancel_operation"))
|
||||
async def cancel_operation(
|
||||
operation_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2918,7 +3010,7 @@ def _register_cancel_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("cancel_operation"))
|
||||
async def cancel_operation(
|
||||
operation_id: str,
|
||||
) -> dict:
|
||||
@@ -2957,7 +3049,7 @@ def _register_list_tags(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConf
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_tags"))
|
||||
async def list_tags(
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
@@ -2994,7 +3086,7 @@ def _register_list_tags(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConf
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_tags"))
|
||||
async def list_tags(
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
@@ -3033,7 +3125,7 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_bank"))
|
||||
async def get_bank(
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
@@ -3066,7 +3158,7 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_bank"))
|
||||
async def get_bank() -> dict:
|
||||
"""
|
||||
Get the profile of this memory bank.
|
||||
@@ -3096,7 +3188,7 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
|
||||
def _register_get_bank_stats(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the get_bank_stats tool (multi-bank only)."""
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_bank_stats"))
|
||||
async def get_bank_stats(
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
@@ -3169,7 +3261,7 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("update_bank"))
|
||||
async def update_bank(
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
@@ -3230,7 +3322,7 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("update_bank"))
|
||||
async def update_bank(
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
@@ -3293,7 +3385,7 @@ def _register_delete_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_bank"))
|
||||
async def delete_bank(
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
@@ -3325,7 +3417,7 @@ def _register_delete_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_bank"))
|
||||
async def delete_bank() -> dict:
|
||||
"""
|
||||
Delete this memory bank and all its data.
|
||||
@@ -3356,7 +3448,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("clear_memories"))
|
||||
async def clear_memories(
|
||||
type: str | None = None,
|
||||
bank_id: str | None = None,
|
||||
@@ -3367,7 +3459,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
Optionally filter by fact type to only clear specific kinds of memories.
|
||||
|
||||
Args:
|
||||
type: Optional fact type filter: 'world', 'experience', or 'opinion'. If not specified, clears all.
|
||||
type: Optional fact type filter: 'world', 'experience', or 'observation'. If not specified, clears all.
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
@@ -3391,7 +3483,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("clear_memories"))
|
||||
async def clear_memories(
|
||||
type: str | None = None,
|
||||
) -> dict:
|
||||
@@ -3401,7 +3493,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
Optionally filter by fact type to only clear specific kinds of memories.
|
||||
|
||||
Args:
|
||||
type: Optional fact type filter: 'world', 'experience', or 'opinion'. If not specified, clears all.
|
||||
type: Optional fact type filter: 'world', 'experience', or 'observation'. If not specified, clears all.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
|
||||
@@ -11,15 +11,17 @@ This module provides metrics for:
|
||||
- Database connection pool metrics
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import importlib
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
|
||||
_resource_mod = importlib.import_module("resource") if importlib.util.find_spec("resource") else None
|
||||
import threading
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from typing import TYPE_CHECKING, Callable
|
||||
from typing import TYPE_CHECKING, Callable, NamedTuple
|
||||
|
||||
from opentelemetry import metrics
|
||||
from opentelemetry.exporter.prometheus import PrometheusMetricReader
|
||||
@@ -75,6 +77,28 @@ LLM_DURATION_BUCKETS = (0.1, 0.25, 0.5, 1.0, 2.0, 3.0, 5.0, 10.0, 15.0, 30.0, 60
|
||||
# HTTP request duration buckets (millisecond-level for fast endpoints)
|
||||
HTTP_DURATION_BUCKETS = (0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0)
|
||||
|
||||
# How often the backlog / queue-depth gauge caches are refreshed (seconds).
|
||||
# The counts are aggregate COUNT queries, so a background task refreshes a
|
||||
# cache and the observable gauges read from it — keeping the /metrics scrape
|
||||
# path synchronous (the same reason the db-pool gauges read cached state).
|
||||
BACKLOG_METRICS_REFRESH_SECONDS = 30
|
||||
|
||||
|
||||
class _AsyncOpKey(NamedTuple):
|
||||
"""Cache / label key for the async-operation queue gauge."""
|
||||
|
||||
tenant: str
|
||||
operation_type: str
|
||||
status: str
|
||||
bank_id: str | None
|
||||
|
||||
|
||||
class _BacklogKey(NamedTuple):
|
||||
"""Cache / label key for the consolidation backlog and failed gauges."""
|
||||
|
||||
tenant: str
|
||||
bank_id: str | None
|
||||
|
||||
|
||||
def get_token_bucket(token_count: int) -> str:
|
||||
"""
|
||||
@@ -113,6 +137,27 @@ def get_token_bucket(token_count: int) -> str:
|
||||
return "50k+"
|
||||
|
||||
|
||||
# Template unbounded id segments before a path is used as the low-cardinality
|
||||
# "endpoint" metric label. A raw per-bank path segment (e.g. user-123) would
|
||||
# otherwise create one never-evicted OTel series per bank.
|
||||
_METRIC_BANK_SEGMENT_RE = re.compile(r"(/banks/)[^/]+")
|
||||
_METRIC_UUID_RE = re.compile(r"/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}")
|
||||
_METRIC_NUMERIC_ID_RE = re.compile(r"/\d+(?=/|$)")
|
||||
|
||||
|
||||
def normalize_http_endpoint(path: str) -> str:
|
||||
"""Template high-cardinality id segments in an HTTP path for safe metric labeling.
|
||||
|
||||
Collapses the "/banks/<id>" segment (any bank id, including non-numeric ones like
|
||||
"user-123"), UUIDs, and numeric ids to placeholders so the "endpoint" metric label
|
||||
has bounded cardinality. Analogous to get_token_bucket for token counts.
|
||||
"""
|
||||
path = _METRIC_BANK_SEGMENT_RE.sub(r"\g<1>{bank_id}", path)
|
||||
path = _METRIC_UUID_RE.sub("/{id}", path)
|
||||
path = _METRIC_NUMERIC_ID_RE.sub("/{id}", path)
|
||||
return path
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Global meter instance
|
||||
@@ -201,6 +246,19 @@ class MetricsCollectorBase:
|
||||
"""Context manager to record operation duration and status."""
|
||||
raise NotImplementedError
|
||||
|
||||
def record_operation_result(
|
||||
self,
|
||||
operation: str,
|
||||
bank_id: str,
|
||||
success: bool,
|
||||
duration: float,
|
||||
source: str = "api",
|
||||
budget: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
):
|
||||
"""Record a single completed operation with an explicit success label."""
|
||||
raise NotImplementedError
|
||||
|
||||
def record_llm_call(
|
||||
self,
|
||||
provider: str,
|
||||
@@ -254,6 +312,19 @@ class NoOpMetricsCollector(MetricsCollectorBase):
|
||||
"""No-op context manager."""
|
||||
yield
|
||||
|
||||
def record_operation_result(
|
||||
self,
|
||||
operation: str,
|
||||
bank_id: str,
|
||||
success: bool,
|
||||
duration: float,
|
||||
source: str = "api",
|
||||
budget: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
):
|
||||
"""No-op operation result recording."""
|
||||
pass
|
||||
|
||||
def record_llm_call(
|
||||
self,
|
||||
provider: str,
|
||||
@@ -361,6 +432,13 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
# DB pool metrics holder (set via set_db_pool)
|
||||
self._db_pool: "asyncpg.Pool | None" = None
|
||||
|
||||
# Backlog / queue-depth gauge caches, refreshed by a background task
|
||||
# (see _setup_backlog_metrics) so the scrape path stays synchronous.
|
||||
self._async_ops_counts: dict[_AsyncOpKey, int] = {}
|
||||
self._consolidation_backlog: dict[_BacklogKey, int] = {}
|
||||
self._consolidation_failed: dict[_BacklogKey, int] = {}
|
||||
self._backlog_task: "asyncio.Task | None" = None
|
||||
|
||||
@contextmanager
|
||||
def record_operation(
|
||||
self,
|
||||
@@ -386,18 +464,6 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
max_tokens: Optional max tokens for the operation
|
||||
"""
|
||||
start_time = time.time()
|
||||
attributes = {
|
||||
"operation": operation,
|
||||
"source": source,
|
||||
"tenant": _get_tenant(),
|
||||
}
|
||||
if self._include_bank_id:
|
||||
attributes["bank_id"] = bank_id
|
||||
if budget:
|
||||
attributes["budget"] = budget
|
||||
if max_tokens:
|
||||
attributes["max_tokens"] = str(max_tokens)
|
||||
|
||||
success = True
|
||||
cancelled = False
|
||||
try:
|
||||
@@ -416,14 +482,51 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
raise
|
||||
finally:
|
||||
if not cancelled:
|
||||
duration = time.time() - start_time
|
||||
attributes["success"] = str(success).lower()
|
||||
self.record_operation_result(
|
||||
operation,
|
||||
bank_id,
|
||||
success=success,
|
||||
duration=time.time() - start_time,
|
||||
source=source,
|
||||
budget=budget,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
|
||||
# Record duration
|
||||
self.operation_duration.record(duration, attributes)
|
||||
def record_operation_result(
|
||||
self,
|
||||
operation: str,
|
||||
bank_id: str,
|
||||
success: bool,
|
||||
duration: float,
|
||||
source: str = "api",
|
||||
budget: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
):
|
||||
"""Record a single completed operation (duration + count) with a success label.
|
||||
|
||||
# Record operation count
|
||||
self.operation_total.add(1, attributes)
|
||||
Direct (non-context-manager) recording for code paths that need explicit
|
||||
success control rather than the exception-based ``record_operation`` — e.g.
|
||||
the async worker, where deferrals/retries are not terminal outcomes and must
|
||||
not be counted as completions.
|
||||
"""
|
||||
attributes = {
|
||||
"operation": operation,
|
||||
"source": source,
|
||||
"tenant": _get_tenant(),
|
||||
}
|
||||
if self._include_bank_id:
|
||||
attributes["bank_id"] = bank_id
|
||||
if budget:
|
||||
attributes["budget"] = budget
|
||||
if max_tokens:
|
||||
attributes["max_tokens"] = str(max_tokens)
|
||||
attributes["success"] = str(success).lower()
|
||||
|
||||
# Record duration
|
||||
self.operation_duration.record(duration, attributes)
|
||||
|
||||
# Record operation count
|
||||
self.operation_total.add(1, attributes)
|
||||
|
||||
def record_llm_call(
|
||||
self,
|
||||
@@ -628,6 +731,10 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
"""
|
||||
self._db_pool = pool
|
||||
self._setup_db_pool_metrics()
|
||||
from .config import get_config
|
||||
|
||||
if get_config().metrics_backlog_enabled:
|
||||
self._setup_backlog_metrics()
|
||||
|
||||
def _setup_db_pool_metrics(self):
|
||||
"""Set up observable gauges for database pool metrics."""
|
||||
@@ -693,6 +800,192 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
unit="{connections}",
|
||||
)
|
||||
|
||||
def _setup_backlog_metrics(self):
|
||||
"""Observable gauges for the async-operation queue and the
|
||||
consolidation backlog.
|
||||
|
||||
These mirror fields the bank-stats endpoint already computes
|
||||
(``operations_by_status``, ``pending_consolidation``,
|
||||
``failed_consolidation``) but expose them as scrapable gauges, so
|
||||
queue depth and backlog can be trended and alerted on instead of only
|
||||
polled per-bank over HTTP. The two motivating questions both come for
|
||||
free here: "is the worker keeping up?" (async-op queue) and "is the
|
||||
knowledge base caught up?" (consolidation backlog) — including the
|
||||
``processing`` state, which is the only signal that surfaces a hung
|
||||
operation stuck holding a worker slot.
|
||||
|
||||
Counts are aggregate ``COUNT`` queries, so a background task refreshes
|
||||
a cache every ``BACKLOG_METRICS_REFRESH_SECONDS`` and these callbacks
|
||||
read it — keeping the scrape path synchronous, the same approach as
|
||||
the db-pool gauges above.
|
||||
"""
|
||||
if self._backlog_task is not None:
|
||||
return # already started for this collector
|
||||
|
||||
def get_async_operations(_options):
|
||||
for key, value in list(self._async_ops_counts.items()):
|
||||
attrs = {"tenant": key.tenant, "operation_type": key.operation_type, "status": key.status}
|
||||
if key.bank_id is not None:
|
||||
attrs["bank_id"] = key.bank_id
|
||||
yield metrics.Observation(value, attrs)
|
||||
|
||||
def get_consolidation_backlog(_options):
|
||||
for key, value in list(self._consolidation_backlog.items()):
|
||||
attrs = {"tenant": key.tenant}
|
||||
if key.bank_id is not None:
|
||||
attrs["bank_id"] = key.bank_id
|
||||
yield metrics.Observation(value, attrs)
|
||||
|
||||
def get_consolidation_failed(_options):
|
||||
for key, value in list(self._consolidation_failed.items()):
|
||||
attrs = {"tenant": key.tenant}
|
||||
if key.bank_id is not None:
|
||||
attrs["bank_id"] = key.bank_id
|
||||
yield metrics.Observation(value, attrs)
|
||||
|
||||
self.meter.create_observable_gauge(
|
||||
name="hindsight.async_operations",
|
||||
callbacks=[get_async_operations],
|
||||
description="Async operations in a non-terminal state, by operation_type and status "
|
||||
"(pending=queued backlog, processing=in-flight, failed=stranded)",
|
||||
unit="{operations}",
|
||||
)
|
||||
self.meter.create_observable_gauge(
|
||||
name="hindsight.consolidation.backlog",
|
||||
callbacks=[get_consolidation_backlog],
|
||||
description="Source memories (experience/world) not yet consolidated into observations",
|
||||
unit="{memories}",
|
||||
)
|
||||
self.meter.create_observable_gauge(
|
||||
name="hindsight.consolidation.failed",
|
||||
callbacks=[get_consolidation_failed],
|
||||
description="Source memories whose consolidation permanently failed "
|
||||
"(recoverable via the consolidation recovery endpoint)",
|
||||
unit="{memories}",
|
||||
)
|
||||
|
||||
# Drive the caches from a background task on the running loop.
|
||||
# set_db_pool runs during async startup, so a loop is normally present;
|
||||
# if not, the gauges simply stay empty rather than crashing collection.
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
logger.warning("No running event loop; backlog metrics disabled")
|
||||
return
|
||||
# Process-lifetime task: there is no collector teardown hook to cancel it
|
||||
# on, so it's torn down with the event loop at process shutdown. If a
|
||||
# shutdown path is ever added, cancel self._backlog_task there.
|
||||
self._backlog_task = loop.create_task(self._backlog_refresh_loop())
|
||||
|
||||
async def _backlog_refresh_loop(self):
|
||||
"""Periodically refresh the backlog / queue-depth caches."""
|
||||
while True:
|
||||
try:
|
||||
await self._refresh_backlog()
|
||||
except Exception:
|
||||
logger.debug("Backlog metrics refresh failed", exc_info=True)
|
||||
await asyncio.sleep(BACKLOG_METRICS_REFRESH_SECONDS)
|
||||
|
||||
async def _refresh_backlog(self):
|
||||
"""Recount the async-operation queue and consolidation backlog across
|
||||
every provisioned Hindsight schema.
|
||||
|
||||
Per-bank labels are gated behind ``metrics_include_bank_id`` (off by
|
||||
default) to keep cardinality bounded; when off, counts are aggregated
|
||||
per tenant/schema. All SQL here is PostgreSQL-specific (``FILTER``,
|
||||
``information_schema``), which is consistent with this collector
|
||||
already being bound to an asyncpg pool.
|
||||
"""
|
||||
if self._db_pool is None:
|
||||
return
|
||||
|
||||
async_ops: dict[_AsyncOpKey, int] = {}
|
||||
backlog: dict[_BacklogKey, int] = {}
|
||||
failed: dict[_BacklogKey, int] = {}
|
||||
per_bank = self._include_bank_id
|
||||
bank_sel = "bank_id, " if per_bank else ""
|
||||
bank_grp = " GROUP BY bank_id" if per_bank else ""
|
||||
|
||||
async with self._db_pool.acquire() as conn:
|
||||
# memory_units is the central per-tenant table; its presence marks a
|
||||
# provisioned Hindsight schema.
|
||||
schema_rows = await conn.fetch(
|
||||
"SELECT table_schema FROM information_schema.tables WHERE table_name = 'memory_units'"
|
||||
)
|
||||
for schema_row in schema_rows:
|
||||
schema = schema_row["table_schema"]
|
||||
|
||||
# Worker queue depth — mirrors operations_by_status, split by
|
||||
# operation_type. Terminal states (completed/cancelled) are
|
||||
# excluded on purpose: a gauge of finished work grows without
|
||||
# bound and says nothing about current load.
|
||||
# Index: idx_async_operations_status.
|
||||
ops_grp = "operation_type, status" + (", bank_id" if per_bank else "")
|
||||
try:
|
||||
rows = await conn.fetch(
|
||||
f"SELECT operation_type, status, {bank_sel}COUNT(*) AS count "
|
||||
f'FROM "{schema}".async_operations '
|
||||
"WHERE status IN ('pending', 'processing', 'failed') "
|
||||
f"GROUP BY {ops_grp}"
|
||||
)
|
||||
for row in rows:
|
||||
bank = row["bank_id"] if per_bank else None
|
||||
key = _AsyncOpKey(schema, row["operation_type"] or "unknown", row["status"], bank)
|
||||
async_ops[key] = async_ops.get(key, 0) + int(row["count"])
|
||||
except Exception:
|
||||
logger.debug("Async-ops queue query failed for schema %s", schema, exc_info=True)
|
||||
|
||||
# Consolidation backlog + stranded counts. Two separate COUNT(*)
|
||||
# queries rather than one with two FILTERs — each WHERE matches a
|
||||
# partial-index predicate exactly:
|
||||
# idx_memory_units_unconsolidated WHERE consolidated_at IS NULL ...
|
||||
# idx_memory_units_consolidation_failed WHERE consolidation_failed_at IS NOT NULL ...
|
||||
# GROUP BY bank_id still composes — bank_id is each index's lead column.
|
||||
#
|
||||
# The backlog count runs with seqscan disabled in a scoped
|
||||
# transaction. The partial index matches its predicate, but
|
||||
# `consolidated_at IS NULL` is true for a large fraction of the
|
||||
# table (every observation has a null consolidated_at), so the
|
||||
# planner misjudges selectivity and otherwise seq-scans the whole
|
||||
# (largest) table on every refresh — verified on a 114k-row table
|
||||
# via EXPLAIN: seq scan ~92 ms vs index scan ~0.1 ms. SET LOCAL
|
||||
# forces the index path and resets at transaction end. The failed
|
||||
# count below needs no such nudge: `consolidation_failed_at IS NOT
|
||||
# NULL` is rare, so its index is chosen on cost.
|
||||
try:
|
||||
async with conn.transaction():
|
||||
await conn.execute("SET LOCAL enable_seqscan = off")
|
||||
rows = await conn.fetch(
|
||||
f"SELECT {bank_sel}COUNT(*) AS count "
|
||||
f'FROM "{schema}".memory_units '
|
||||
"WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')"
|
||||
f"{bank_grp}"
|
||||
)
|
||||
for row in rows:
|
||||
bank = row["bank_id"] if per_bank else None
|
||||
key = _BacklogKey(schema, bank)
|
||||
backlog[key] = backlog.get(key, 0) + int(row["count"])
|
||||
except Exception:
|
||||
logger.debug("Consolidation backlog query failed for schema %s", schema, exc_info=True)
|
||||
|
||||
try:
|
||||
rows = await conn.fetch(
|
||||
f"SELECT {bank_sel}COUNT(*) AS count "
|
||||
f'FROM "{schema}".memory_units '
|
||||
"WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')"
|
||||
f"{bank_grp}"
|
||||
)
|
||||
for row in rows:
|
||||
bank = row["bank_id"] if per_bank else None
|
||||
key = _BacklogKey(schema, bank)
|
||||
failed[key] = failed.get(key, 0) + int(row["count"])
|
||||
except Exception:
|
||||
logger.debug("Consolidation failed query failed for schema %s", schema, exc_info=True)
|
||||
|
||||
self._async_ops_counts = async_ops
|
||||
self._consolidation_backlog = backlog
|
||||
self._consolidation_failed = failed
|
||||
|
||||
|
||||
# Global metrics collector instance (defaults to no-op)
|
||||
_metrics_collector: MetricsCollectorBase = NoOpMetricsCollector()
|
||||
|
||||
@@ -32,6 +32,7 @@ from sqlalchemy.pool import NullPool
|
||||
from ._pg_search import normalize_pg_search_tokenizer, pg_search_bm25_columns
|
||||
from ._vector_index import (
|
||||
bootstrap_extension,
|
||||
configured_vector_extension,
|
||||
detect_vector_extension,
|
||||
index_type_keyword,
|
||||
index_using_clause,
|
||||
@@ -60,6 +61,86 @@ def _detect_vector_extension(conn, vector_extension: str = "pgvector") -> str:
|
||||
return detect_vector_extension(conn, vector_extension)
|
||||
|
||||
|
||||
def _ensure_pgvector_extension_in_public(conn: Connection) -> None:
|
||||
"""Ensure pgvector is installed before pgvector-backed migrations run."""
|
||||
logger.debug("Checking pgvector extension availability...")
|
||||
|
||||
# First, check if extension already exists
|
||||
ext_check = conn.execute(
|
||||
text(
|
||||
"SELECT extname, nspname FROM pg_extension e "
|
||||
"JOIN pg_namespace n ON e.extnamespace = n.oid "
|
||||
"WHERE extname = 'vector'"
|
||||
)
|
||||
).fetchone()
|
||||
|
||||
if ext_check:
|
||||
# Extension exists - check if in correct schema
|
||||
ext_schema = ext_check[1]
|
||||
if ext_schema == "public":
|
||||
logger.info("pgvector extension found in public schema - ready to use")
|
||||
else:
|
||||
# Extension in wrong schema - try to fix if we have permissions
|
||||
logger.warning(
|
||||
f"pgvector extension found in schema '{ext_schema}' instead of 'public'. Attempting to relocate..."
|
||||
)
|
||||
try:
|
||||
conn.execute(text("DROP EXTENSION vector CASCADE"))
|
||||
conn.execute(text("SET search_path TO public"))
|
||||
conn.execute(text("CREATE EXTENSION vector"))
|
||||
conn.commit()
|
||||
logger.info("pgvector extension relocated to public schema")
|
||||
except Exception as e:
|
||||
# Failed to relocate - log but don't fail if extension exists somewhere
|
||||
logger.warning(
|
||||
f"Could not relocate pgvector extension to public schema: {e}. "
|
||||
f"Continuing with extension in '{ext_schema}' schema."
|
||||
)
|
||||
conn.rollback()
|
||||
else:
|
||||
# Extension doesn't exist - try to install
|
||||
logger.info("pgvector extension not found, attempting to install...")
|
||||
try:
|
||||
conn.execute(text("SET search_path TO public"))
|
||||
conn.execute(text("CREATE EXTENSION vector"))
|
||||
conn.commit()
|
||||
logger.info("pgvector extension installed in public schema")
|
||||
except Exception as e:
|
||||
# Installation failed - this is only fatal if extension truly doesn't exist
|
||||
# Check one more time in case another process installed it
|
||||
conn.rollback()
|
||||
ext_recheck = conn.execute(
|
||||
text(
|
||||
"SELECT nspname FROM pg_extension e "
|
||||
"JOIN pg_namespace n ON e.extnamespace = n.oid "
|
||||
"WHERE extname = 'vector'"
|
||||
)
|
||||
).fetchone()
|
||||
|
||||
if ext_recheck:
|
||||
logger.warning(
|
||||
f"Could not install pgvector extension (permission denied?), "
|
||||
f"but extension exists in '{ext_recheck[0]}' schema. Continuing..."
|
||||
)
|
||||
else:
|
||||
# Extension truly doesn't exist and we can't install it
|
||||
logger.error(
|
||||
f"pgvector extension is not installed and cannot be installed: {e}. "
|
||||
f"Please ensure pgvector is installed by a database administrator. "
|
||||
f"See: https://github.com/pgvector/pgvector#installation"
|
||||
)
|
||||
raise RuntimeError(
|
||||
"pgvector extension is required but not installed. Please install it with: CREATE EXTENSION vector;"
|
||||
) from e
|
||||
|
||||
|
||||
def _bootstrap_vector_extension_for_migrations(conn: Connection, vector_extension: str) -> None:
|
||||
"""Bootstrap the configured vector backend before schema migrations run."""
|
||||
if vector_extension == "pgvector":
|
||||
_ensure_pgvector_extension_in_public(conn)
|
||||
bootstrap_extension(conn, vector_extension)
|
||||
|
||||
|
||||
def _drop_per_bank_vector_indexes(conn: Connection, schema_name: str) -> None:
|
||||
"""Drop per-bank partial memory_units vector indexes after global ScaNN is ready."""
|
||||
rows = conn.execute(
|
||||
@@ -275,83 +356,8 @@ def run_migrations(
|
||||
logger.debug("Migration advisory lock acquired")
|
||||
|
||||
try:
|
||||
# Ensure pgvector extension is installed globally BEFORE schema migrations
|
||||
# This is critical: the extension must exist database-wide before any schema
|
||||
# migrations run, otherwise custom schemas won't have access to vector types
|
||||
logger.debug("Checking pgvector extension availability...")
|
||||
|
||||
# First, check if extension already exists
|
||||
ext_check = conn.execute(
|
||||
text(
|
||||
"SELECT extname, nspname FROM pg_extension e "
|
||||
"JOIN pg_namespace n ON e.extnamespace = n.oid "
|
||||
"WHERE extname = 'vector'"
|
||||
)
|
||||
).fetchone()
|
||||
|
||||
if ext_check:
|
||||
# Extension exists - check if in correct schema
|
||||
ext_schema = ext_check[1]
|
||||
if ext_schema == "public":
|
||||
logger.info("pgvector extension found in public schema - ready to use")
|
||||
else:
|
||||
# Extension in wrong schema - try to fix if we have permissions
|
||||
logger.warning(
|
||||
f"pgvector extension found in schema '{ext_schema}' instead of 'public'. "
|
||||
f"Attempting to relocate..."
|
||||
)
|
||||
try:
|
||||
conn.execute(text("DROP EXTENSION vector CASCADE"))
|
||||
conn.execute(text("SET search_path TO public"))
|
||||
conn.execute(text("CREATE EXTENSION vector"))
|
||||
conn.commit()
|
||||
logger.info("pgvector extension relocated to public schema")
|
||||
except Exception as e:
|
||||
# Failed to relocate - log but don't fail if extension exists somewhere
|
||||
logger.warning(
|
||||
f"Could not relocate pgvector extension to public schema: {e}. "
|
||||
f"Continuing with extension in '{ext_schema}' schema."
|
||||
)
|
||||
conn.rollback()
|
||||
else:
|
||||
# Extension doesn't exist - try to install
|
||||
logger.info("pgvector extension not found, attempting to install...")
|
||||
try:
|
||||
conn.execute(text("SET search_path TO public"))
|
||||
conn.execute(text("CREATE EXTENSION vector"))
|
||||
conn.commit()
|
||||
logger.info("pgvector extension installed in public schema")
|
||||
except Exception as e:
|
||||
# Installation failed - this is only fatal if extension truly doesn't exist
|
||||
# Check one more time in case another process installed it
|
||||
conn.rollback()
|
||||
ext_recheck = conn.execute(
|
||||
text(
|
||||
"SELECT nspname FROM pg_extension e "
|
||||
"JOIN pg_namespace n ON e.extnamespace = n.oid "
|
||||
"WHERE extname = 'vector'"
|
||||
)
|
||||
).fetchone()
|
||||
|
||||
if ext_recheck:
|
||||
logger.warning(
|
||||
f"Could not install pgvector extension (permission denied?), "
|
||||
f"but extension exists in '{ext_recheck[0]}' schema. Continuing..."
|
||||
)
|
||||
else:
|
||||
# Extension truly doesn't exist and we can't install it
|
||||
logger.error(
|
||||
f"pgvector extension is not installed and cannot be installed: {e}. "
|
||||
f"Please ensure pgvector is installed by a database administrator. "
|
||||
f"See: https://github.com/pgvector/pgvector#installation"
|
||||
)
|
||||
raise RuntimeError(
|
||||
"pgvector extension is required but not installed. "
|
||||
"Please install it with: CREATE EXTENSION vector;"
|
||||
) from e
|
||||
|
||||
vector_extension = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
|
||||
bootstrap_extension(conn, vector_extension)
|
||||
vector_extension = configured_vector_extension()
|
||||
_bootstrap_vector_extension_for_migrations(conn, vector_extension)
|
||||
|
||||
# Commit any pending transaction on the advisory-lock connection
|
||||
# before running migrations. Some code paths above (e.g., the
|
||||
@@ -686,24 +692,20 @@ def ensure_vector_extension(
|
||||
|
||||
if not current_index_info:
|
||||
if table_name == "memory_units" and uses_per_bank_vector_indexes(target_ext):
|
||||
# Check whether per-bank partial vector indexes already cover this table
|
||||
# (created by the bank_utils lifecycle — no global index needed in that case)
|
||||
per_bank_index_count = conn.execute(
|
||||
text("""
|
||||
SELECT COUNT(*)
|
||||
FROM pg_indexes
|
||||
WHERE schemaname = :schema
|
||||
AND tablename = :table_name
|
||||
AND indexname LIKE 'idx_mu_emb_%'
|
||||
"""),
|
||||
{"schema": schema_name, "table_name": table_name},
|
||||
).scalar()
|
||||
if per_bank_index_count and per_bank_index_count > 0:
|
||||
logger.debug(
|
||||
f"No global embedding index on {table_name}, but {per_bank_index_count} "
|
||||
f"per-bank partial vector indexes exist — skipping global index creation"
|
||||
)
|
||||
continue
|
||||
# Per-bank backends never use a GLOBAL memory_units vector index.
|
||||
# Every vector search is bank + fact_type scoped and served by the
|
||||
# per-(bank, fact_type) partial indexes created at bank-creation time
|
||||
# (bank_utils.create_bank_vector_indexes); the planner never picks a
|
||||
# global index when bank_id is in the WHERE clause, which is exactly
|
||||
# why migration d5e6f7a8b9c0 drops it for these backends. So don't
|
||||
# create one here either — not even on an empty schema with no per-bank
|
||||
# indexes yet (those are built when the first bank is created). Verified
|
||||
# via EXPLAIN: the query uses idx_mu_emb_* whether or not the global
|
||||
# index exists, so creating it is dead weight.
|
||||
logger.debug(
|
||||
f"Per-bank vector backend ({target_ext}); skipping global {index_name} creation on {table_name}"
|
||||
)
|
||||
continue
|
||||
logger.warning(f"No embedding index found for {table_name}, will create it if safe")
|
||||
mismatched_tables.append((table_name, index_name, None, row_count))
|
||||
continue
|
||||
|
||||
@@ -1,6 +1,58 @@
|
||||
import logging
|
||||
import os
|
||||
from urllib.parse import urlparse, urlunparse
|
||||
|
||||
|
||||
def detect_container_runtime() -> str | None:
|
||||
"""Detect whether the process is running inside a container.
|
||||
|
||||
Returns "kubernetes", "docker", or None. Used to warn operators that the
|
||||
default ``socket.gethostname()`` worker id is unstable across container
|
||||
recreation (the random container id changes on restart, so tasks stuck in
|
||||
'processing' under the old id are never recovered).
|
||||
"""
|
||||
if os.getenv("KUBERNETES_SERVICE_HOST"):
|
||||
return "kubernetes"
|
||||
# Docker (and most OCI runtimes) create this marker file in every container.
|
||||
if os.path.exists("/.dockerenv"):
|
||||
return "docker"
|
||||
# cgroup v1 fallback for runtimes that don't write /.dockerenv.
|
||||
try:
|
||||
with open("/proc/1/cgroup", encoding="utf-8") as f:
|
||||
if any(token in f.read() for token in ("docker", "containerd", "kubepods")):
|
||||
return "docker"
|
||||
except OSError:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def warn_if_container_default_worker_id(worker_id: str | None) -> None:
|
||||
"""Warn when worker id will fall back to an unstable container hostname."""
|
||||
if worker_id:
|
||||
return
|
||||
|
||||
runtime = detect_container_runtime()
|
||||
if not runtime:
|
||||
return
|
||||
|
||||
logging.warning(
|
||||
"\n"
|
||||
"============================================================\n"
|
||||
" WARNING: HINDSIGHT_API_WORKER_ID is not set and Hindsight\n"
|
||||
f" appears to be running inside {runtime}.\n"
|
||||
"\n"
|
||||
" The worker id is defaulting to the container hostname,\n"
|
||||
" which CHANGES every time the container is recreated.\n"
|
||||
" When that happens, tasks left in 'processing' under the\n"
|
||||
" old hostname are never recovered — consolidation and other\n"
|
||||
" async operations can get stuck indefinitely.\n"
|
||||
"\n"
|
||||
" Set HINDSIGHT_API_WORKER_ID to a STABLE value (e.g. the\n"
|
||||
" compose service name or StatefulSet pod name) to avoid this.\n"
|
||||
"============================================================"
|
||||
)
|
||||
|
||||
|
||||
def mask_network_location(url):
|
||||
if not url:
|
||||
return url
|
||||
|
||||
@@ -136,7 +136,7 @@ def main():
|
||||
# Worker options
|
||||
parser.add_argument(
|
||||
"--worker-id",
|
||||
default=config.worker_id or socket.gethostname(),
|
||||
default=config.worker_id,
|
||||
help="Worker identifier (default: hostname, env: HINDSIGHT_API_WORKER_ID)",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -178,10 +178,17 @@ def main():
|
||||
# Configure logging
|
||||
config.configure_logging()
|
||||
|
||||
from ..utils import warn_if_container_default_worker_id
|
||||
|
||||
warn_if_container_default_worker_id(args.worker_id)
|
||||
worker_id = args.worker_id or socket.gethostname()
|
||||
worker_id_source = "HINDSIGHT_API_WORKER_ID/--worker-id" if args.worker_id else "hostname (default)"
|
||||
logger.info(f"Worker id: {worker_id} (source: {worker_id_source})")
|
||||
|
||||
# Import MemoryEngine here to avoid circular imports
|
||||
from .. import MemoryEngine
|
||||
|
||||
print(f"Starting Hindsight Worker: {args.worker_id}")
|
||||
print(f"Starting Hindsight Worker: {worker_id}")
|
||||
print(f" Poll interval: {args.poll_interval}ms")
|
||||
print(f" Max retries: {args.max_retries}")
|
||||
print(f" Max slots: {config.worker_max_slots}")
|
||||
@@ -249,7 +256,7 @@ def main():
|
||||
schema = None if config.database_schema == DEFAULT_DATABASE_SCHEMA else config.database_schema
|
||||
poller = WorkerPoller(
|
||||
backend=memory._backend,
|
||||
worker_id=args.worker_id,
|
||||
worker_id=worker_id,
|
||||
executor=memory.execute_task,
|
||||
poll_interval_ms=args.poll_interval,
|
||||
schema=schema,
|
||||
|
||||
@@ -20,9 +20,23 @@ from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..engine.schema import fq_table_explicit as fq_table
|
||||
from ..metrics import get_metrics_collector
|
||||
from .exceptions import DeferOperation, RetryTaskAt
|
||||
from .stage import StageHolder, bind_holder
|
||||
|
||||
# Map DB operation_type -> metric `operation` label, collapsing the retain
|
||||
# variants onto "retain" so async worker completions land on the same
|
||||
# operation="retain" series the synchronous API path emits. Unknown types
|
||||
# pass through unchanged.
|
||||
_RETAIN_OP_TYPES = {"retain", "batch_retain", "file_convert_retain"}
|
||||
|
||||
|
||||
def _metric_operation_label(operation_type: str | None) -> str:
|
||||
if operation_type in _RETAIN_OP_TYPES:
|
||||
return "retain"
|
||||
return operation_type or "unknown"
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.db.base import DatabaseBackend, DatabaseConnection
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
@@ -701,6 +715,24 @@ class WorkerPoller:
|
||||
"""
|
||||
task_type = task.task_dict.get("type", "unknown")
|
||||
bank_id = task.task_dict.get("bank_id", "unknown")
|
||||
# Operation metric (source="worker"): record on terminal outcomes only, so
|
||||
# async worker throughput and latency (retain, consolidation and the other
|
||||
# worker task types) are visible in Prometheus. Prefer the DB-authoritative
|
||||
# operation_type.
|
||||
#
|
||||
# success semantics are deliberately narrow: success=false means the task
|
||||
# raised out to the poller (an unexpected error, or retry-exhausted). It does
|
||||
# NOT capture deterministic failures that the executor handles itself and
|
||||
# returns from normally (file_convert_retain, non-retryable errors via
|
||||
# memory_engine.execute_task) — those record success=true here. Treat this as
|
||||
# a completion-throughput signal, not a failure-rate one: for authoritative
|
||||
# failure visibility use the hindsight_async_operations{status="failed"} gauge,
|
||||
# which reads each operation's final DB status.
|
||||
op_label = _metric_operation_label(task.task_dict.get("operation_type") or task_type)
|
||||
op_start = time.time()
|
||||
metrics = get_metrics_collector()
|
||||
# None = not a terminal outcome (deferred/retried) → no metric.
|
||||
terminal_success: bool | None = None
|
||||
|
||||
# Bind the stage holder in this task's own contextvar scope so engine
|
||||
# code running under us can update it via stage.set_stage(). If holder
|
||||
@@ -717,14 +749,28 @@ class WorkerPoller:
|
||||
task.task_dict["_schema"] = task.schema
|
||||
await self._executor(task.task_dict)
|
||||
logger.debug(f"Task {task.operation_id} execution finished")
|
||||
terminal_success = True
|
||||
except DeferOperation as e:
|
||||
# Deferral is not a terminal outcome — do not record a completion.
|
||||
await self._defer_operation(task.operation_id, e.exec_date, e.reason, task.schema)
|
||||
except RetryTaskAt as e:
|
||||
# Retry is not a terminal outcome — do not record a completion.
|
||||
await self._schedule_retry(task.operation_id, e.retry_at, str(e), task.schema)
|
||||
except Exception as e:
|
||||
logger.error(f"Task {task.operation_id} failed: {e}")
|
||||
traceback.print_exc()
|
||||
await self._mark_failed(task.operation_id, str(e), task.schema)
|
||||
terminal_success = False
|
||||
|
||||
# Record the metric outside the executor's exception scope so a metrics
|
||||
# reporting failure can never be mistaken for a task failure and flip terminal state.
|
||||
if terminal_success is not None:
|
||||
try:
|
||||
metrics.record_operation_result(
|
||||
op_label, bank_id, success=terminal_success, duration=time.time() - op_start, source="worker"
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(f"Failed to record worker operation metric for {task.operation_id}", exc_info=True)
|
||||
|
||||
async def recover_own_tasks(self) -> int:
|
||||
"""
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api-slim"
|
||||
version = "0.8.2"
|
||||
version = "0.8.4"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -60,10 +60,10 @@ dependencies = [
|
||||
"pyasn1>=0.6.3", # DoS vulnerability fix
|
||||
"urllib3>=2.7.0", # Decompression-bomb safeguards bypass + sensitive header forwarding fixes
|
||||
"langchain-core>=1.2.22", # Path traversal in legacy load_prompt functions fix
|
||||
"langsmith>=0.6.3", # SSRF via tracing header injection fix
|
||||
"langsmith>=0.8.18", # GHSA-f4xh-w4cj-qxq8: arbitrary server-side file read in TracingMiddleware fix (supersedes >=0.6.3 SSRF tracing-header-injection floor)
|
||||
"protobuf>=6.33.5", # JSON recursion depth bypass fix
|
||||
"pillow>=12.1.1", # Out-of-bounds write in PSD image loading fix
|
||||
"cryptography>=46.0.6,<47", # Incomplete DNS name constraint enforcement fix; cap <47 — 47.0.0 SIGILLs on some ARM64 Linux VMs (Docker/Podman on Apple Silicon), pyca/cryptography#14733
|
||||
"cryptography>=48.0.1", # GHSA-537c-gmf6-5ccf: bundled-OpenSSL OOB read fix needs >=48.0.1. Prior <47 cap (47.0.0 SIGILL on ARM64 Docker/Podman, pyca/cryptography#14733) lifted — 47/48/49 verified importing + RSA sign/verify cleanly on linux/arm64 (Docker on Apple Silicon) and native arm64 macOS; upstream issue closed unconfirmed.
|
||||
"filelock>=3.20.1", # TOCTOU race condition fix
|
||||
"authlib>=1.6.9", # Account takeover/JWS header injection vulnerability fix
|
||||
"pyjwt>=2.12.0", # Accepts unknown crit header extensions fix
|
||||
@@ -74,6 +74,7 @@ dependencies = [
|
||||
"pygments>=2.20.0", # ReDoS via inefficient GUID regex fix
|
||||
"claude-agent-sdk>=0.2.82",
|
||||
"boto3>=1.42.74",
|
||||
"croniter>=2.0.0", # Cron parsing for scheduled mental model refresh
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
|
||||
@@ -11,6 +11,22 @@ import pytest
|
||||
import pytest_asyncio
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Force torch to initialize exactly once, in the main thread, at conftest import
|
||||
# time — before any fixture spins up an event loop or sentence-transformers'
|
||||
# thread pools. torch's C-level `_add_docstr(_has_torch_function, ...)` in
|
||||
# torch/overrides.py is not re-entrancy-safe: when the first `import torch`
|
||||
# happens lazily from inside concurrent/async code (e.g.
|
||||
# embeddings.initialize() -> sentence_transformers -> transformers -> torch, or
|
||||
# cross_encoder's ThreadPoolExecutor), torch/overrides.py can execute twice and
|
||||
# raise "RuntimeError: function '_has_torch_function' already has a docstring",
|
||||
# failing collection of every test on the pytest-xdist shard. Importing it here
|
||||
# (single-threaded, before any concurrency) makes that registration happen once
|
||||
# per worker process. Guarded so slim/no-torch environments still collect.
|
||||
try:
|
||||
import torch # noqa: F401 # eager one-time init; see comment above
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
from hindsight_api import LLMConfig, LocalSTEmbeddings, MemoryEngine, RequestContext
|
||||
from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder
|
||||
from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer
|
||||
@@ -38,6 +54,29 @@ async def _teardown_memory_engine(mem: MemoryEngine) -> None:
|
||||
unregister_span_recorder(mem._llm_recorder)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _cleanup_leaked_span_recorders():
|
||||
"""Fail-safe for the process-global LLM-trace recorder registry (#2229).
|
||||
|
||||
``MemoryEngine.__init__`` registers its recorder in the shared registry, and
|
||||
only ``close()`` removes it. Tests that construct an engine directly (without
|
||||
``_teardown_memory_engine``/``close()``) leak an *enabled* recorder; a later
|
||||
test's LLM calls then get recorded into the shared DB, flaking
|
||||
``test_llm_trace::test_disabled_writes_no_rows`` (it observes rows for its
|
||||
bank even though its own recorder is disabled). ``_teardown_memory_engine``
|
||||
guards the fixtures; this guards everything else by dropping any recorder a
|
||||
test added to the registry.
|
||||
"""
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
|
||||
recorders = get_span_recorder()._recorders
|
||||
before = {id(r) for r in recorders}
|
||||
yield
|
||||
for recorder in list(recorders):
|
||||
if id(recorder) not in before:
|
||||
recorders.remove(recorder)
|
||||
|
||||
|
||||
# Default pg0 instance configuration for tests
|
||||
DEFAULT_PG0_INSTANCE_NAME = "hindsight-test"
|
||||
DEFAULT_PG0_PORT = int(os.environ.get("HINDSIGHT_TEST_PG_PORT", "5556"))
|
||||
@@ -45,11 +84,13 @@ DEFAULT_PG0_PORT = int(os.environ.get("HINDSIGHT_TEST_PG_PORT", "5556"))
|
||||
# Keep the background MaintenanceLoop from auto-starting during tests. In
|
||||
# production it sweeps retention and re-schedules consolidation, but its timers
|
||||
# would race shared-pg0 test data (e.g. delete llm_requests/audit_log rows a test
|
||||
# just inserted). Disabling the reconcile interval and llm-trace retention — with
|
||||
# audit retention already off by default — leaves no job enabled, so the loop
|
||||
# never starts. Tests that exercise it call MaintenanceLoop methods
|
||||
# (_run_reconcile / _purge_expired) directly.
|
||||
# just inserted). Disabling the reconcile interval, the mental-model refresh tick
|
||||
# and llm-trace retention — with audit retention already off by default — leaves
|
||||
# no job enabled, so the loop never starts. Tests that exercise it call
|
||||
# MaintenanceLoop methods (_run_reconcile / _run_scheduled_mm_refresh /
|
||||
# _purge_expired) directly.
|
||||
os.environ.setdefault("HINDSIGHT_API_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS", "0")
|
||||
os.environ.setdefault("HINDSIGHT_API_MENTAL_MODEL_REFRESH_TICK_SECONDS", "0")
|
||||
os.environ.setdefault("HINDSIGHT_API_LLM_TRACE_RETENTION_DAYS", "-1")
|
||||
|
||||
|
||||
|
||||
@@ -88,7 +88,10 @@ async def test_backup_tables_covers_entire_schema(backup_test_schema):
|
||||
await conn.close()
|
||||
|
||||
# alembic_version is migration bookkeeping, not data — never backed up.
|
||||
schema_tables = {r["table_name"] for r in rows} - {"alembic_version"}
|
||||
# bank_stats_cache is a derived TTL cache of get_bank_stats results: it has no
|
||||
# FK to banks (so the restore cascade never touches it) and repopulates itself
|
||||
# on demand, so it is deliberately not backed up — a restore starts it cold.
|
||||
schema_tables = {r["table_name"] for r in rows} - {"alembic_version", "bank_stats_cache"}
|
||||
backup_tables = set(BACKUP_TABLES)
|
||||
|
||||
missing = schema_tables - backup_tables
|
||||
@@ -547,3 +550,36 @@ async def test_run_migration_with_schema_only_runs_requested_schema(monkeypatch)
|
||||
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")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("ensure_extensions", "expected"),
|
||||
[(True, True), (False, False)],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_migration_threads_ensure_extensions_flag(monkeypatch, ensure_extensions, expected):
|
||||
"""The --skip-extension-reconcile flag (ensure_extensions=False) must reach run_migrations_for_schemas.
|
||||
|
||||
The post-migration vector/text-search reconcile only does work on a backend change, so operators
|
||||
can skip it on a no-change re-migration over many tenant schemas. Verify the flag is threaded through
|
||||
rather than silently dropped.
|
||||
"""
|
||||
monkeypatch.setenv("HINDSIGHT_API_DATABASE_URL", "postgresql://test")
|
||||
captured: dict = {}
|
||||
|
||||
async def fake_resolve_database_url(db_url: str) -> str:
|
||||
return f"resolved::{db_url}"
|
||||
|
||||
def fake_run_migrations_for_schemas(database_url, schemas, **kwargs):
|
||||
captured["ensure_extensions"] = kwargs.get("ensure_extensions")
|
||||
|
||||
monkeypatch.setattr(admin_cli, "load_extension", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(admin_cli, "resolve_database_url", fake_resolve_database_url)
|
||||
|
||||
from hindsight_api import migrations as migrations_module
|
||||
|
||||
monkeypatch.setattr(migrations_module, "run_migrations_for_schemas", fake_run_migrations_for_schemas)
|
||||
|
||||
await admin_cli._run_migration("postgresql://test", schema="tenant_demo", ensure_extensions=ensure_extensions)
|
||||
|
||||
assert captured["ensure_extensions"] is expected
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
"""Regression tests for issue #1002 — Anthropic structured output via forced tool_use.
|
||||
|
||||
When strict_schema=True, AnthropicLLM.call() must request the schema through a single
|
||||
forced tool_use tool (tool_choice={"type":"tool",...}) and read the validated args from
|
||||
the tool_use block, NOT inject the schema as text and json.loads() the reply (which caused
|
||||
a ~1:1 invalid-JSON retry storm / OOM in production).
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class _Decision(BaseModel):
|
||||
action: str
|
||||
reason: str
|
||||
|
||||
|
||||
def _make_anthropic_provider():
|
||||
with patch("anthropic.AsyncAnthropic") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
from hindsight_api.engine.providers.anthropic_llm import AnthropicLLM
|
||||
|
||||
provider = AnthropicLLM(
|
||||
provider="anthropic",
|
||||
api_key="fake-key",
|
||||
base_url="",
|
||||
model="claude-sonnet-4-20250514",
|
||||
)
|
||||
provider._client = MagicMock()
|
||||
return provider
|
||||
|
||||
|
||||
def _tool_use_response(args: dict):
|
||||
block = MagicMock()
|
||||
block.type = "tool_use"
|
||||
block.name = "structured_response"
|
||||
block.input = args
|
||||
resp = MagicMock()
|
||||
resp.content = [block]
|
||||
resp.usage = MagicMock(input_tokens=5, output_tokens=2, cache_read_input_tokens=0)
|
||||
resp.stop_reason = "tool_use"
|
||||
return resp
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strict_schema_uses_forced_tool_choice():
|
||||
"""strict_schema=True ⇒ a single tool is defined and tool_choice forces it (no schema text-injection)."""
|
||||
provider = _make_anthropic_provider()
|
||||
provider._client.messages.create = AsyncMock(return_value=_tool_use_response({"action": "skip", "reason": "dup"}))
|
||||
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
|
||||
result = await provider.call(
|
||||
messages=[{"role": "user", "content": "decide"}],
|
||||
response_format=_Decision,
|
||||
strict_schema=True,
|
||||
scope="test",
|
||||
max_retries=0,
|
||||
)
|
||||
kwargs = provider._client.messages.create.call_args.kwargs
|
||||
# forced tool_use requested
|
||||
assert "tools" in kwargs and len(kwargs["tools"]) == 1
|
||||
assert kwargs["tool_choice"] == {"type": "tool", "name": "structured_response"}
|
||||
# schema NOT injected as text into the system prompt
|
||||
assert "valid JSON matching this schema" not in (kwargs.get("system") or "")
|
||||
# validated model returned straight from tool_use.input
|
||||
assert isinstance(result, _Decision)
|
||||
assert result.action == "skip"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strict_schema_tool_use_never_hits_json_retry_loop():
|
||||
"""A tool_use response is structurally valid → no second messages.create call (no retry storm)."""
|
||||
provider = _make_anthropic_provider()
|
||||
create = AsyncMock(return_value=_tool_use_response({"action": "keep", "reason": "novel"}))
|
||||
provider._client.messages.create = create
|
||||
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
|
||||
await provider.call(
|
||||
messages=[{"role": "user", "content": "x"}],
|
||||
response_format=_Decision,
|
||||
strict_schema=True,
|
||||
scope="test",
|
||||
max_retries=10, # would allow 11 attempts on the old text-parse path
|
||||
)
|
||||
assert create.await_count == 1 # exactly one call — the bug was N retries on malformed text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_strict_keeps_text_injection_fallback():
|
||||
"""strict_schema=False (default) preserves the legacy schema-in-prompt behavior."""
|
||||
provider = _make_anthropic_provider()
|
||||
block = MagicMock()
|
||||
block.type = "text"
|
||||
block.text = '{"action":"skip","reason":"d"}'
|
||||
resp = MagicMock()
|
||||
resp.content = [block]
|
||||
resp.usage = MagicMock(input_tokens=5, output_tokens=2, cache_read_input_tokens=0)
|
||||
resp.stop_reason = "end_turn"
|
||||
provider._client.messages.create = AsyncMock(return_value=resp)
|
||||
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
|
||||
result = await provider.call(
|
||||
messages=[{"role": "user", "content": "decide"}],
|
||||
response_format=_Decision,
|
||||
strict_schema=False,
|
||||
scope="test",
|
||||
max_retries=0,
|
||||
)
|
||||
kwargs = provider._client.messages.create.call_args.kwargs
|
||||
assert "tools" not in kwargs # no forced tool when not strict
|
||||
assert "valid JSON matching this schema" in (kwargs.get("system") or "")
|
||||
assert isinstance(result, _Decision)
|
||||
@@ -0,0 +1,90 @@
|
||||
"""Regression test: submitting an async op for a bank that doesn't exist must
|
||||
raise a clean validation error, not a raw asyncpg `ForeignKeyViolationError`.
|
||||
|
||||
`_submit_async_operation` inserts into `async_operations`, which has an FK to
|
||||
`banks.bank_id`. If a caller submits for a missing bank (typo, race against a
|
||||
deletion, integration that derives bank IDs before the bank is created), the
|
||||
INSERT raises `asyncpg.exceptions.ForeignKeyViolationError`. The FastAPI
|
||||
endpoint's broad `except Exception` then surfaces it as a 500 — but this is
|
||||
a client error, not a server error, and should be a 404.
|
||||
|
||||
This test exercises the call directly via `MemoryEngine.submit_async_*` so
|
||||
the failure mode is observable without spinning up the HTTP layer.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.extensions.operation_validator import OperationValidationError
|
||||
|
||||
pytestmark = pytest.mark.xdist_group("async_submit_bank_not_found_tests")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_inline_execution(memory):
|
||||
"""Prevent SyncTaskBackend from running the submitted op inline so we
|
||||
only test the submit-path failure, not downstream execution."""
|
||||
|
||||
async def _noop(_payload):
|
||||
return None
|
||||
|
||||
original = memory._task_backend.submit_task
|
||||
memory._task_backend.submit_task = _noop
|
||||
yield
|
||||
memory._task_backend.submit_task = original
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consolidation_submit_on_missing_bank_raises_validation_error(
|
||||
memory, request_context, no_inline_execution
|
||||
):
|
||||
"""A `/consolidate` submit against a bank that doesn't exist must raise
|
||||
OperationValidationError(404), not a raw asyncpg FK violation that bubbles
|
||||
out as a 500 from the API."""
|
||||
missing_bank = f"does-not-exist-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
with pytest.raises(OperationValidationError) as exc_info:
|
||||
await memory.submit_async_consolidation(
|
||||
bank_id=missing_bank,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert missing_bank in exc_info.value.reason
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_consolidation_submit_on_missing_bank_raises_validation_error(
|
||||
memory, request_context, no_inline_execution
|
||||
):
|
||||
"""Scoped consolidates (with `observation_scopes`) take the
|
||||
`dedupe_by_bank=False` branch, which historically skipped the bank lock
|
||||
entirely and went straight to the FK-violating INSERT. Same 404 contract."""
|
||||
missing_bank = f"does-not-exist-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
with pytest.raises(OperationValidationError) as exc_info:
|
||||
await memory.submit_async_consolidation(
|
||||
bank_id=missing_bank,
|
||||
request_context=request_context,
|
||||
observation_scopes=[{"tag": "anything"}],
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert missing_bank in exc_info.value.reason
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_maintenance_on_missing_bank_short_circuits(memory, request_context, no_inline_execution):
|
||||
"""`submit_async_graph_maintenance` has its own short-circuit that checks
|
||||
the per-bank queue before calling `_submit_async_operation`. A missing
|
||||
bank means an empty queue, so it returns `no_work=True` without reaching
|
||||
the FK-violating INSERT. This test pins that behaviour."""
|
||||
missing_bank = f"does-not-exist-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
result = await memory.submit_async_graph_maintenance(
|
||||
bank_id=missing_bank,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result == {"operation_id": None, "no_work": True}
|
||||
@@ -0,0 +1,244 @@
|
||||
"""
|
||||
Tests for the async-operation queue and consolidation backlog gauges
|
||||
(``_setup_backlog_metrics`` / ``_refresh_backlog`` in metrics.py).
|
||||
|
||||
These gauges expose, as scrapable time-series, the same counts the bank-stats
|
||||
endpoint already returns per bank (``operations_by_status``,
|
||||
``pending_consolidation``, ``failed_consolidation``):
|
||||
|
||||
- ``hindsight_async_operations{operation_type,status}`` — worker queue depth
|
||||
(pending=backlog, processing=in-flight, failed=stranded)
|
||||
- ``hindsight_consolidation_backlog`` — source memories not yet consolidated
|
||||
- ``hindsight_consolidation_failed`` — source memories permanently failed
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.metrics import MetricsCollector, _AsyncOpKey, _BacklogKey
|
||||
|
||||
|
||||
class _FakeTxn:
|
||||
async def __aenter__(self):
|
||||
return None
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
return False
|
||||
|
||||
|
||||
class _FakeConn:
|
||||
"""asyncpg-like connection whose fetch() is dispatched by SQL substring."""
|
||||
|
||||
def __init__(self, fetch_fn):
|
||||
self._fetch_fn = fetch_fn
|
||||
self.executed = []
|
||||
|
||||
async def fetch(self, sql, *args):
|
||||
return self._fetch_fn(sql, *args)
|
||||
|
||||
async def execute(self, sql, *args):
|
||||
self.executed.append(sql)
|
||||
|
||||
def transaction(self):
|
||||
return _FakeTxn()
|
||||
|
||||
|
||||
class _FakeAcquire:
|
||||
def __init__(self, conn):
|
||||
self._conn = conn
|
||||
|
||||
async def __aenter__(self):
|
||||
return self._conn
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
return False
|
||||
|
||||
|
||||
class _FakePool:
|
||||
def __init__(self, fetch_fn):
|
||||
self._conn = _FakeConn(fetch_fn)
|
||||
|
||||
def acquire(self):
|
||||
return _FakeAcquire(self._conn)
|
||||
|
||||
|
||||
def _collector(include_bank_id=False):
|
||||
mock_config = MagicMock()
|
||||
mock_config.metrics_include_bank_id = include_bank_id
|
||||
with (
|
||||
patch("hindsight_api.metrics.get_meter", return_value=MagicMock()),
|
||||
patch("hindsight_api.config.get_config", return_value=mock_config),
|
||||
):
|
||||
return MetricsCollector()
|
||||
|
||||
|
||||
def _set_db_pool_with_backlog_enabled(collector, pool):
|
||||
"""Call set_db_pool with the backlog flag forced on (it's off by default)."""
|
||||
mock_config = MagicMock()
|
||||
mock_config.metrics_backlog_enabled = True
|
||||
with patch("hindsight_api.config.get_config", return_value=mock_config):
|
||||
collector.set_db_pool(pool)
|
||||
|
||||
|
||||
def _rows_for(sql):
|
||||
"""Canned results, keyed off distinctive substrings of each query."""
|
||||
if "information_schema.tables" in sql:
|
||||
return [{"table_schema": "public"}]
|
||||
if "async_operations" in sql:
|
||||
return [
|
||||
{"operation_type": "retain", "status": "pending", "count": 5},
|
||||
{"operation_type": "consolidation", "status": "pending", "count": 12},
|
||||
{"operation_type": "consolidation", "status": "processing", "count": 1},
|
||||
{"operation_type": "consolidation", "status": "failed", "count": 2},
|
||||
]
|
||||
if "memory_units" in sql and "consolidated_at IS NULL" in sql:
|
||||
return [{"count": 42}]
|
||||
if "memory_units" in sql and "consolidation_failed_at IS NOT NULL" in sql:
|
||||
return [{"count": 3}]
|
||||
return []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_backlog_aggregates_queue_and_consolidation():
|
||||
collector = _collector(include_bank_id=False)
|
||||
collector._db_pool = _FakePool(lambda sql, *a: _rows_for(sql))
|
||||
|
||||
await collector._refresh_backlog()
|
||||
|
||||
# Worker queue depth keyed by (schema, operation_type, status, bank=None)
|
||||
assert collector._async_ops_counts[("public", "retain", "pending", None)] == 5
|
||||
assert collector._async_ops_counts[("public", "consolidation", "pending", None)] == 12
|
||||
assert collector._async_ops_counts[("public", "consolidation", "processing", None)] == 1
|
||||
assert collector._async_ops_counts[("public", "consolidation", "failed", None)] == 2
|
||||
# Consolidation backlog (source memories), keyed by (schema, bank=None)
|
||||
assert collector._consolidation_backlog[("public", None)] == 42
|
||||
assert collector._consolidation_failed[("public", None)] == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_backlog_uses_index_matched_predicates_not_filter_scan():
|
||||
"""Backlog/failed must be two separate COUNT(*) queries whose WHERE matches
|
||||
a partial-index predicate exactly (no FILTER over a full-table scan), and
|
||||
the queue query must exclude terminal statuses."""
|
||||
captured = []
|
||||
collector = _collector()
|
||||
collector._db_pool = _FakePool(lambda sql, *a: (captured.append(sql), _rows_for(sql))[1])
|
||||
await collector._refresh_backlog()
|
||||
|
||||
mem_queries = [s for s in captured if "memory_units" in s and "COUNT(*)" in s]
|
||||
assert len(mem_queries) == 2 # split, not a single two-FILTER aggregate
|
||||
assert all("FILTER" not in s for s in mem_queries)
|
||||
assert any("consolidated_at IS NULL AND fact_type IN ('experience', 'world')" in s for s in mem_queries)
|
||||
assert any("consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')" in s for s in mem_queries)
|
||||
|
||||
ops_sql = next(s for s in captured if "async_operations" in s and "GROUP BY" in s)
|
||||
assert "status IN ('pending', 'processing', 'failed')" in ops_sql
|
||||
assert "completed" not in ops_sql and "cancelled" not in ops_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_backlog_count_runs_with_seqscan_disabled():
|
||||
"""`consolidated_at IS NULL` is true for a large fraction of the table, so
|
||||
the planner misjudges selectivity and won't use the partial index without a
|
||||
nudge — the backlog count must issue SET LOCAL enable_seqscan=off."""
|
||||
collector = _collector()
|
||||
pool = _FakePool(lambda sql, *a: _rows_for(sql))
|
||||
collector._db_pool = pool
|
||||
await collector._refresh_backlog()
|
||||
|
||||
assert any("enable_seqscan" in s.lower() and "off" in s.lower() for s in pool._conn.executed)
|
||||
# the result is still correct under the nudge
|
||||
assert collector._consolidation_backlog[("public", None)] == 42
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_backlog_per_bank_labels_and_group_by_when_enabled():
|
||||
"""With metrics_include_bank_id on, bank_id enters the cache key and the
|
||||
SQL switches to GROUP BY bank_id."""
|
||||
captured = []
|
||||
|
||||
def fetch(sql, *a):
|
||||
captured.append(sql)
|
||||
if "information_schema.tables" in sql:
|
||||
return [{"table_schema": "public"}]
|
||||
if "async_operations" in sql:
|
||||
return [{"operation_type": "retain", "status": "pending", "bank_id": "bankA", "count": 4}]
|
||||
if "memory_units" in sql and "consolidated_at IS NULL" in sql:
|
||||
return [{"bank_id": "bankA", "count": 11}]
|
||||
if "memory_units" in sql and "consolidation_failed_at IS NOT NULL" in sql:
|
||||
return [{"bank_id": "bankA", "count": 2}]
|
||||
return []
|
||||
|
||||
collector = _collector(include_bank_id=True)
|
||||
collector._db_pool = _FakePool(fetch)
|
||||
await collector._refresh_backlog()
|
||||
|
||||
assert collector._async_ops_counts[("public", "retain", "pending", "bankA")] == 4
|
||||
assert collector._consolidation_backlog[("public", "bankA")] == 11
|
||||
assert collector._consolidation_failed[("public", "bankA")] == 2
|
||||
# bank_id must be grouped in every per-bank count query
|
||||
assert all("GROUP BY bank_id" in s for s in captured if "memory_units" in s and "COUNT(*)" in s)
|
||||
|
||||
|
||||
def test_gauges_register_and_emit_cached_values_without_bank_id():
|
||||
collector = _collector(include_bank_id=False)
|
||||
# Sync call: no running loop, so gauges register but no background task spawns.
|
||||
_set_db_pool_with_backlog_enabled(collector, MagicMock())
|
||||
|
||||
gauges = {
|
||||
c.kwargs["name"]: c.kwargs["callbacks"][0]
|
||||
for c in collector.meter.create_observable_gauge.call_args_list
|
||||
if "callbacks" in c.kwargs
|
||||
}
|
||||
assert "hindsight.async_operations" in gauges
|
||||
assert "hindsight.consolidation.backlog" in gauges
|
||||
assert "hindsight.consolidation.failed" in gauges
|
||||
|
||||
collector._async_ops_counts = {
|
||||
_AsyncOpKey("public", "retain", "pending", None): 7,
|
||||
_AsyncOpKey("public", "consolidation", "processing", None): 1,
|
||||
}
|
||||
collector._consolidation_backlog = {_BacklogKey("public", None): 9}
|
||||
|
||||
obs = list(gauges["hindsight.async_operations"](None))
|
||||
by_label = {(o.attributes["operation_type"], o.attributes["status"]): o.value for o in obs}
|
||||
assert by_label[("retain", "pending")] == 7
|
||||
assert by_label[("consolidation", "processing")] == 1
|
||||
assert all("bank_id" not in o.attributes for o in obs) # cardinality guard
|
||||
|
||||
backlog_obs = list(gauges["hindsight.consolidation.backlog"](None))
|
||||
assert backlog_obs[0].value == 9
|
||||
assert backlog_obs[0].attributes["tenant"] == "public"
|
||||
|
||||
|
||||
def test_gauge_emits_bank_id_attribute_when_present():
|
||||
collector = _collector(include_bank_id=True)
|
||||
_set_db_pool_with_backlog_enabled(collector, MagicMock())
|
||||
gauges = {
|
||||
c.kwargs["name"]: c.kwargs["callbacks"][0]
|
||||
for c in collector.meter.create_observable_gauge.call_args_list
|
||||
if "callbacks" in c.kwargs
|
||||
}
|
||||
collector._consolidation_backlog = {_BacklogKey("public", "bankA"): 4}
|
||||
obs = list(gauges["hindsight.consolidation.backlog"](None))
|
||||
assert obs[0].value == 4
|
||||
assert obs[0].attributes["bank_id"] == "bankA"
|
||||
|
||||
|
||||
def test_backlog_gauges_not_registered_when_flag_disabled():
|
||||
"""Backlog metrics are off by default: set_db_pool must not register the
|
||||
gauges unless metrics_backlog_enabled is set."""
|
||||
collector = _collector()
|
||||
mock_config = MagicMock()
|
||||
mock_config.metrics_backlog_enabled = False
|
||||
with patch("hindsight_api.config.get_config", return_value=mock_config):
|
||||
collector.set_db_pool(MagicMock())
|
||||
|
||||
names = [
|
||||
c.kwargs.get("name") for c in collector.meter.create_observable_gauge.call_args_list if "callbacks" in c.kwargs
|
||||
]
|
||||
assert "hindsight.async_operations" not in names
|
||||
assert "hindsight.consolidation.backlog" not in names
|
||||
assert "hindsight.consolidation.failed" not in names
|
||||
assert collector._backlog_task is None
|
||||
@@ -186,6 +186,43 @@ async def test_invalidate_drops_entry() -> None:
|
||||
assert calls[0] == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalidate_detaches_in_flight_loader() -> None:
|
||||
cache = BankStatsCache(ttl_seconds=60, max_entries=100)
|
||||
stale_started = asyncio.Event()
|
||||
release_stale = asyncio.Event()
|
||||
fresh_started = asyncio.Event()
|
||||
|
||||
async def stale_loader() -> dict[str, Any]:
|
||||
stale_started.set()
|
||||
await release_stale.wait()
|
||||
return {"v": "stale"}
|
||||
|
||||
async def fresh_loader() -> dict[str, Any]:
|
||||
fresh_started.set()
|
||||
return {"v": "fresh"}
|
||||
|
||||
stale_task = asyncio.create_task(cache.get_or_load("schema", "bank", stale_loader))
|
||||
await stale_started.wait()
|
||||
await cache.invalidate("schema", "bank")
|
||||
|
||||
# A request after invalidation must start a new load instead of joining the
|
||||
# pre-invalidation query, which may contain data from before a bank write.
|
||||
fresh_result = await asyncio.wait_for(cache.get_or_load("schema", "bank", fresh_loader), timeout=1)
|
||||
assert fresh_started.is_set()
|
||||
assert fresh_result == {"v": "fresh"}
|
||||
|
||||
release_stale.set()
|
||||
assert await stale_task == {"v": "stale"}
|
||||
|
||||
# The stale loader completed last, but must not overwrite the fresh value.
|
||||
async def should_not_run() -> dict[str, Any]:
|
||||
raise AssertionError("fresh value was not cached")
|
||||
|
||||
cached = await cache.get_or_load("schema", "bank", should_not_run)
|
||||
assert cached == {"v": "fresh"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_drops_all_entries() -> None:
|
||||
cache = BankStatsCache(ttl_seconds=60, max_entries=100)
|
||||
@@ -199,3 +236,27 @@ async def test_clear_drops_all_entries() -> None:
|
||||
await cache.get_or_load("s", "a", loader)
|
||||
await cache.get_or_load("s", "b", loader)
|
||||
assert calls[0] == 4
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_detaches_in_flight_loaders() -> None:
|
||||
cache = BankStatsCache(ttl_seconds=60, max_entries=100)
|
||||
stale_started = asyncio.Event()
|
||||
release_stale = asyncio.Event()
|
||||
|
||||
async def stale_loader() -> dict[str, Any]:
|
||||
stale_started.set()
|
||||
await release_stale.wait()
|
||||
return {"v": "stale"}
|
||||
|
||||
async def fresh_loader() -> dict[str, Any]:
|
||||
return {"v": "fresh"}
|
||||
|
||||
stale_task = asyncio.create_task(cache.get_or_load("schema", "bank", stale_loader))
|
||||
await stale_started.wait()
|
||||
await cache.clear()
|
||||
assert await cache.get_or_load("schema", "bank", fresh_loader) == {"v": "fresh"}
|
||||
|
||||
release_stale.set()
|
||||
assert await stale_task == {"v": "stale"}
|
||||
assert await cache.get_or_load("schema", "bank", fresh_loader) == {"v": "fresh"}
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
"""Tests for the table-backed (cross-process) get_bank_stats cache.
|
||||
|
||||
On PostgreSQL the engine backs `get_bank_stats` with the `bank_stats_cache`
|
||||
table (`DistributedBankStatsCache`) instead of a per-process dict, so one
|
||||
worker's computation is shared with every other worker. These tests verify:
|
||||
|
||||
* the PG engine actually selects the distributed cache,
|
||||
* a computed result is written to the table and served from it on the next call,
|
||||
* invalidation deletes the row so the next call recomputes, and
|
||||
* an unreachable cache table degrades to computing without caching rather than
|
||||
failing the endpoint.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.bank_stats_cache import DistributedBankStatsCache
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine, get_current_schema
|
||||
|
||||
_PINNED_TTL_SECONDS = 300.0
|
||||
|
||||
|
||||
async def _insert_memory(conn, bank_id: str, text: str, fact_type: str = "experience") -> uuid.UUID:
|
||||
mem_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, created_at, updated_at, consolidated_at)
|
||||
VALUES ($1, $2, $3, $4, NOW(), NOW(), NOW(), NOW())
|
||||
""",
|
||||
mem_id,
|
||||
bank_id,
|
||||
text,
|
||||
fact_type,
|
||||
)
|
||||
return mem_id
|
||||
|
||||
|
||||
async def _ensure_bank(memory: MemoryEngine, bank_id: str, request_context: RequestContext) -> None:
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
|
||||
def _pin_distributed_cache(memory: MemoryEngine) -> DistributedBankStatsCache:
|
||||
cache = DistributedBankStatsCache(backend=memory._backend, ttl_seconds=_PINNED_TTL_SECONDS)
|
||||
memory._bank_stats_cache = cache
|
||||
return cache
|
||||
|
||||
|
||||
class TestDistributedBankStatsCache:
|
||||
@pytest.mark.asyncio
|
||||
async def test_pg_engine_selects_distributed_cache(self, memory: MemoryEngine):
|
||||
if memory._database_backend_type != "postgresql":
|
||||
pytest.skip("distributed cache is PostgreSQL-only")
|
||||
assert isinstance(memory._bank_stats_cache, DistributedBankStatsCache)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_result_is_written_and_served_from_table(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
if memory._database_backend_type != "postgresql":
|
||||
pytest.skip("distributed cache is PostgreSQL-only")
|
||||
|
||||
bank_id = f"test-dist-stats-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_memory(conn, bank_id, "Alice loves hiking.")
|
||||
|
||||
_pin_distributed_cache(memory)
|
||||
try:
|
||||
first = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert first["node_counts"].get("experience") == 1
|
||||
|
||||
# The computed result was persisted to the shared table.
|
||||
async with pool.acquire() as conn:
|
||||
rows = await conn.fetchval("SELECT count(*) FROM bank_stats_cache WHERE bank_id = $1", bank_id)
|
||||
assert rows == 1
|
||||
|
||||
# Mutate the underlying data WITHOUT going through an invalidating
|
||||
# engine method — the long-TTL cache must serve the stale row.
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_memory(conn, bank_id, "Bob enjoys cycling.")
|
||||
served = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert served["node_counts"].get("experience") == 1 # still cached
|
||||
|
||||
# Invalidating drops the row → next call recomputes the true count.
|
||||
await memory._bank_stats_cache.invalidate(get_current_schema(), bank_id)
|
||||
async with pool.acquire() as conn:
|
||||
rows = await conn.fetchval("SELECT count(*) FROM bank_stats_cache WHERE bank_id = $1", bank_id)
|
||||
assert rows == 0
|
||||
fresh = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert fresh["node_counts"].get("experience") == 2
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_force_refresh_bypasses_and_updates_cache(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
if memory._database_backend_type != "postgresql":
|
||||
pytest.skip("distributed cache is PostgreSQL-only")
|
||||
|
||||
bank_id = f"test-dist-stats-fresh-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_memory(conn, bank_id, "Alice loves hiking.")
|
||||
|
||||
_pin_distributed_cache(memory)
|
||||
try:
|
||||
# Warm the cache, then mutate the data without invalidation.
|
||||
assert (await memory.get_bank_stats(bank_id, request_context=request_context))["node_counts"][
|
||||
"experience"
|
||||
] == 1
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_memory(conn, bank_id, "Bob enjoys cycling.")
|
||||
|
||||
# A normal read is served the stale cached count...
|
||||
stale = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert stale["node_counts"]["experience"] == 1
|
||||
# ...but force_refresh recomputes the true count.
|
||||
fresh = await memory.get_bank_stats(bank_id, request_context=request_context, force_refresh=True)
|
||||
assert fresh["node_counts"]["experience"] == 2
|
||||
# The forced result also refreshed the cache for the next caller.
|
||||
served = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert served["node_counts"]["experience"] == 2
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_degrades_when_cache_table_unreachable(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
if memory._database_backend_type != "postgresql":
|
||||
pytest.skip("distributed cache is PostgreSQL-only")
|
||||
|
||||
bank_id = f"test-dist-stats-degrade-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_memory(conn, bank_id, "Alice loves hiking.")
|
||||
|
||||
# Point the cache at a table that does not exist: reads and writes fail,
|
||||
# so it must fall back to computing the real result (no real table touched).
|
||||
cache = _pin_distributed_cache(memory)
|
||||
cache._qualified = lambda schema: '"public".bank_stats_cache_does_not_exist' # type: ignore[method-assign]
|
||||
try:
|
||||
stats = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert stats["node_counts"].get("experience") == 1
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -0,0 +1,164 @@
|
||||
"""Regression tests: get_bank_stats cache must be invalidated by mutations.
|
||||
|
||||
`get_bank_stats` is served from a short-TTL per-process cache (`BankStatsCache`).
|
||||
`delete_bank` already invalidates that cache after it mutates counts, but the
|
||||
other operations that change the same counts — `delete_memory_unit`,
|
||||
`delete_document`, `clear_observations`, and `update_document` (when a tag
|
||||
change deletes observations) — did not, so a client polling stats right after a
|
||||
deletion would see pre-mutation counts until the TTL expired (up to a minute).
|
||||
|
||||
Each test pins a long TTL on the engine's stats cache so that, *without* the
|
||||
invalidation fix, the second `get_bank_stats` call would be served the stale
|
||||
cached value and the assertion would fail.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.bank_stats_cache import BankStatsCache
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
# A TTL long enough that, absent invalidation, the warmed cache would still be
|
||||
# served on the post-mutation read within the same test.
|
||||
_PINNED_TTL_SECONDS = 300.0
|
||||
|
||||
|
||||
async def _insert_memory(conn, bank_id: str, text: str, fact_type: str = "experience") -> uuid.UUID:
|
||||
"""Insert a memory unit directly, bypassing the LLM retain pipeline."""
|
||||
mem_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, created_at, updated_at, consolidated_at)
|
||||
VALUES ($1, $2, $3, $4, NOW(), NOW(), NOW(), NOW())
|
||||
""",
|
||||
mem_id,
|
||||
bank_id,
|
||||
text,
|
||||
fact_type,
|
||||
)
|
||||
return mem_id
|
||||
|
||||
|
||||
async def _insert_observation(conn, bank_id: str, text: str, source_memory_ids: list[uuid.UUID]) -> uuid.UUID:
|
||||
"""Insert an observation unit directly."""
|
||||
obs_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_units (
|
||||
id, bank_id, text, fact_type, event_date, source_memory_ids, proof_count, created_at, updated_at
|
||||
) VALUES ($1, $2, $3, 'observation', NOW(), $4, $5, NOW(), NOW())
|
||||
""",
|
||||
obs_id,
|
||||
bank_id,
|
||||
text,
|
||||
source_memory_ids,
|
||||
len(source_memory_ids),
|
||||
)
|
||||
return obs_id
|
||||
|
||||
|
||||
async def _insert_document(conn, bank_id: str, doc_id: str) -> None:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO documents (id, bank_id, original_text, content_hash)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
""",
|
||||
doc_id,
|
||||
bank_id,
|
||||
f"text-for-{doc_id}",
|
||||
doc_id,
|
||||
)
|
||||
|
||||
|
||||
async def _attach_unit_to_doc(conn, unit_id: uuid.UUID, doc_id: str) -> None:
|
||||
await conn.execute("UPDATE memory_units SET document_id = $1 WHERE id = $2", doc_id, unit_id)
|
||||
|
||||
|
||||
async def _ensure_bank(memory: MemoryEngine, bank_id: str, request_context: RequestContext) -> None:
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
|
||||
def _pin_cache(memory: MemoryEngine) -> None:
|
||||
"""Replace the stats cache with one that has a deterministic long TTL."""
|
||||
memory._bank_stats_cache = BankStatsCache(ttl_seconds=_PINNED_TTL_SECONDS, max_entries=128)
|
||||
|
||||
|
||||
class TestBankStatsCacheInvalidation:
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_memory_unit_invalidates_stats_cache(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
bank_id = f"test-stats-cache-delunit-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
|
||||
await _insert_memory(conn, bank_id, "Bob enjoys cycling.")
|
||||
|
||||
_pin_cache(memory)
|
||||
try:
|
||||
before = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert before["node_counts"].get("experience") == 2
|
||||
|
||||
await memory.delete_memory_unit(str(m1), request_context=request_context)
|
||||
|
||||
after = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
# Without invalidation the long-TTL cache would still report 2.
|
||||
assert after["node_counts"].get("experience") == 1
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_document_invalidates_stats_cache(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
bank_id = f"test-stats-cache-deldoc-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
document_id = f"doc-{uuid.uuid4().hex[:8]}"
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_document(conn, bank_id, document_id)
|
||||
unit_id = await _insert_memory(conn, bank_id, "Alice works at Acme.")
|
||||
await _attach_unit_to_doc(conn, unit_id, document_id)
|
||||
|
||||
_pin_cache(memory)
|
||||
try:
|
||||
before = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert before["total_documents"] == 1
|
||||
assert before["node_counts"].get("experience") == 1
|
||||
|
||||
await memory.delete_document(document_id, bank_id, request_context=request_context)
|
||||
|
||||
after = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
# Without invalidation the long-TTL cache would still report 1 document.
|
||||
assert after["total_documents"] == 0
|
||||
assert after["node_counts"].get("experience", 0) == 0
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_observations_invalidates_stats_cache(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
bank_id = f"test-stats-cache-clearobs-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
|
||||
await _insert_observation(conn, bank_id, "Alice enjoys hiking regularly.", [m1])
|
||||
|
||||
_pin_cache(memory)
|
||||
try:
|
||||
before = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert before["total_observations"] == 1
|
||||
|
||||
await memory.clear_observations(bank_id, request_context=request_context)
|
||||
|
||||
after = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
# Without invalidation the long-TTL cache would still report 1 observation.
|
||||
assert after["total_observations"] == 0
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -322,3 +322,136 @@ def test_plain_text_lines_not_treated_as_jsonl():
|
||||
# Sanity: these are not JSON objects (so the JSONL path correctly declined).
|
||||
with pytest.raises(json.JSONDecodeError):
|
||||
json.loads(chunks[0])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Idempotency — re-chunking a produced chunk must be a no-op (issue #2301)
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# The streaming retain pipeline pre-chunks each document once (producer) and then
|
||||
# re-chunks every piece during extraction (consumer), stamping all sub-chunks of
|
||||
# one piece with that piece's single chunk_index. If a piece re-split, its
|
||||
# sub-chunks would derive the same chunk_id = {bank}_{doc}_{index} and the
|
||||
# ON CONFLICT upsert would fail with CardinalityViolationError. This can only
|
||||
# happen when structured_chunk_size > max_chars (a chunk legitimately exceeds the
|
||||
# re-chunk budget); the defaults (structured == max_chars) never trip it.
|
||||
|
||||
|
||||
def _assert_idempotent(text: str, *, max_chars: int, structured_chunk_size: int) -> list[str]:
|
||||
chunks = chunk_text(text, max_chars=max_chars, structured_chunk_size=structured_chunk_size)
|
||||
for chunk in chunks:
|
||||
rechunked = chunk_text(chunk, max_chars=max_chars, structured_chunk_size=structured_chunk_size)
|
||||
assert rechunked == [chunk], (
|
||||
f"re-chunking a produced chunk split it again ({len(chunk)} chars -> "
|
||||
f"{len(rechunked)} pieces) — not idempotent (issue #2301)"
|
||||
)
|
||||
return chunks
|
||||
|
||||
|
||||
def test_conversation_turn_over_chunk_size_is_rechunk_stable():
|
||||
"""A conversation turn larger than max_chars but kept whole by the larger
|
||||
structured cap must survive a re-chunk unchanged (issue #2301)."""
|
||||
content = json.dumps([{"role": "assistant", "content": "x" * 6000}])
|
||||
|
||||
_assert_idempotent(content, max_chars=3000, structured_chunk_size=5000)
|
||||
|
||||
|
||||
def test_jsonl_line_over_chunk_size_is_rechunk_stable():
|
||||
"""A single oversized JSONL line, kept whole within the structured cap, must
|
||||
not be re-split when handed back through chunk_text (issue #2301)."""
|
||||
text = "\n".join([json.dumps({"event": "x" * 3800}), json.dumps({"event": "small"})])
|
||||
|
||||
_assert_idempotent(text, max_chars=3000, structured_chunk_size=4500)
|
||||
|
||||
|
||||
def test_oversized_unit_fragments_stay_within_chunk_budget():
|
||||
"""A unit past even the structured cap is fragmented as text; no fragment may
|
||||
exceed max_chars, so a re-chunk leaves the fragments intact (issue #2301)."""
|
||||
text = "\n".join([json.dumps({"event": "z" * 9000}), json.dumps({"e": "s"})])
|
||||
|
||||
chunks = _assert_idempotent(text, max_chars=3000, structured_chunk_size=4500)
|
||||
|
||||
assert all(len(c) <= 3000 for c in chunks)
|
||||
|
||||
|
||||
def test_single_json_object_kept_whole_within_structured_cap():
|
||||
"""A lone JSON object over max_chars but within the structured cap is returned
|
||||
whole rather than plain-text-split (the basis of re-chunk stability)."""
|
||||
obj = json.dumps({"role": "assistant", "content": "x" * 4000})
|
||||
|
||||
assert chunk_text(obj, max_chars=3000, structured_chunk_size=5000) == [obj]
|
||||
|
||||
|
||||
def test_rechunk_preserves_one_chunk_id_per_pre_chunk():
|
||||
"""End-to-end of the producer/consumer chunk_id derivation: each pre-chunk
|
||||
(one global index) must re-chunk to exactly one piece, so the derived
|
||||
chunk_ids stay unique within an upsert batch (issue #2301)."""
|
||||
content = json.dumps([{"role": "assistant", "content": "x" * 6000}])
|
||||
|
||||
pre_chunks = chunk_text(content, max_chars=3000, structured_chunk_size=5000)
|
||||
chunk_ids = []
|
||||
for global_idx, pre in enumerate(pre_chunks):
|
||||
for _ in chunk_text(pre, max_chars=3000, structured_chunk_size=5000):
|
||||
chunk_ids.append(f"bank_doc_{global_idx}")
|
||||
|
||||
assert len(chunk_ids) == len(set(chunk_ids)), f"duplicate chunk_ids in one batch: {chunk_ids}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Append-mode JSON array merge simulation (issue #2409)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_newline_joined_json_arrays_bypass_conversation_chunking():
|
||||
"""Newline-joined JSON arrays (the pre-fix append-mode storage format)
|
||||
fail both the conversation and JSONL detection paths and fall through
|
||||
to sentence-boundary text splitting.
|
||||
|
||||
This test documents the broken state that issue #2409 fixes at the
|
||||
orchestrator level. chunk_text() itself is not changed; the fix
|
||||
merges the arrays before they reach chunk_text().
|
||||
"""
|
||||
turn1 = json.dumps([{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there"}])
|
||||
turn2 = json.dumps([{"role": "user", "content": "How are you"}, {"role": "assistant", "content": "Fine"}])
|
||||
corrupted = turn1 + "\n" + turn2
|
||||
|
||||
chunks = chunk_text(corrupted, max_chars=80)
|
||||
|
||||
# The corrupted format does NOT route through _chunk_conversation.
|
||||
# At least one chunk will not be a valid JSON array of dicts.
|
||||
has_non_json_chunk = False
|
||||
for chunk in chunks:
|
||||
try:
|
||||
parsed = json.loads(chunk)
|
||||
if not (isinstance(parsed, list) and all(isinstance(e, dict) for e in parsed)):
|
||||
has_non_json_chunk = True
|
||||
except json.JSONDecodeError:
|
||||
has_non_json_chunk = True
|
||||
assert has_non_json_chunk, (
|
||||
"Newline-joined JSON arrays should NOT produce valid conversation chunks. "
|
||||
"If this fails, chunk_text() learned to handle the format and the "
|
||||
"orchestrator-level merge in #2409 may be redundant."
|
||||
)
|
||||
|
||||
|
||||
def test_merged_json_array_routes_to_conversation_chunking():
|
||||
"""A properly merged flat JSON array (the post-fix format) routes
|
||||
through _chunk_conversation and produces chunks that are each valid
|
||||
JSON arrays of complete message dicts.
|
||||
"""
|
||||
messages = [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hi there"},
|
||||
{"role": "user", "content": "How are you"},
|
||||
{"role": "assistant", "content": "Fine, thanks for asking"},
|
||||
]
|
||||
text = json.dumps(messages)
|
||||
|
||||
chunks = chunk_text(text, max_chars=120)
|
||||
|
||||
assert len(chunks) > 1, "Should produce multiple chunks at this budget"
|
||||
for chunk in chunks:
|
||||
parsed = json.loads(chunk)
|
||||
assert isinstance(parsed, list), f"Chunk must be a JSON array: {chunk[:60]}"
|
||||
assert all(isinstance(e, dict) for e in parsed), f"Every element must be a dict: {chunk[:60]}"
|
||||
assert all("role" in e for e in parsed), f"Every element must have a role key: {chunk[:60]}"
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Tests for ``CODEX_HOME`` resolution of the Codex ``auth.json`` location.
|
||||
|
||||
Codex stores its OAuth credentials under a configurable home directory. The
|
||||
canonical ``@openai/codex`` CLI honors the ``CODEX_HOME`` environment variable
|
||||
and falls back to ``~/.codex``. Hindsight's Codex auth/LLM/embeddings paths
|
||||
must resolve the same way so that a user who relocates ``CODEX_HOME`` is still
|
||||
authenticated.
|
||||
"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from hindsight_api.engine.providers.codex_auth import (
|
||||
CodexAuthManager,
|
||||
default_codex_auth_file,
|
||||
)
|
||||
from hindsight_api.engine.providers.codex_llm import CodexLLM
|
||||
|
||||
|
||||
def _write_auth(auth_dir: Path, access_token: str = "at-test") -> Path:
|
||||
"""Write a minimal chatgpt-mode auth.json under ``auth_dir``."""
|
||||
auth_dir.mkdir(parents=True, exist_ok=True)
|
||||
auth_file = auth_dir / "auth.json"
|
||||
auth_file.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"auth_mode": "chatgpt",
|
||||
"tokens": {
|
||||
"access_token": access_token,
|
||||
"refresh_token": "rt-test",
|
||||
"account_id": "acct-test",
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
return auth_file
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# default_codex_auth_file()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_default_auth_file_falls_back_to_home_codex_when_unset(tmp_path, monkeypatch):
|
||||
monkeypatch.delenv("CODEX_HOME", raising=False)
|
||||
monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path))
|
||||
|
||||
assert default_codex_auth_file() == tmp_path / ".codex" / "auth.json"
|
||||
|
||||
|
||||
def test_default_auth_file_honors_codex_home_when_set(tmp_path, monkeypatch):
|
||||
codex_home = tmp_path / "custom-codex"
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
|
||||
assert default_codex_auth_file() == codex_home / "auth.json"
|
||||
|
||||
|
||||
def test_default_auth_file_empty_codex_home_falls_back(tmp_path, monkeypatch):
|
||||
"""An empty ``CODEX_HOME`` is treated as unset (matches shell semantics)."""
|
||||
monkeypatch.setenv("CODEX_HOME", "")
|
||||
monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path))
|
||||
|
||||
assert default_codex_auth_file() == tmp_path / ".codex" / "auth.json"
|
||||
|
||||
|
||||
def test_default_auth_file_resolved_lazily(tmp_path, monkeypatch):
|
||||
"""The env var is read on each call, not cached at import time."""
|
||||
monkeypatch.setenv("CODEX_HOME", str(tmp_path / "a"))
|
||||
assert default_codex_auth_file() == tmp_path / "a" / "auth.json"
|
||||
|
||||
monkeypatch.setenv("CODEX_HOME", str(tmp_path / "b"))
|
||||
assert default_codex_auth_file() == tmp_path / "b" / "auth.json"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CodexAuthManager.from_file() — honors CODEX_HOME by default
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_auth_manager_from_file_uses_codex_home(tmp_path, monkeypatch):
|
||||
codex_home = tmp_path / "custom-codex"
|
||||
_write_auth(codex_home, access_token="at-from-codex-home")
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
|
||||
mgr = CodexAuthManager.from_file()
|
||||
|
||||
assert mgr.access_token == "at-from-codex-home"
|
||||
assert mgr._auth_file == codex_home / "auth.json"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CodexLLM — loads credentials from CODEX_HOME
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_codex_llm_loads_from_codex_home(tmp_path, monkeypatch):
|
||||
codex_home = tmp_path / "custom-codex"
|
||||
_write_auth(codex_home, access_token="at-llm")
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
|
||||
llm = CodexLLM(
|
||||
provider="codex",
|
||||
api_key="ignored",
|
||||
base_url="",
|
||||
model="gpt-5-codex",
|
||||
)
|
||||
|
||||
assert llm.access_token == "at-llm"
|
||||
assert llm._auth_file == codex_home / "auth.json"
|
||||
@@ -8,9 +8,12 @@ relevance score, independent of the cross-encoder model's score calibration.
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.search.reranking import apply_combined_scoring, _RECENCY_ALPHA, _TEMPORAL_ALPHA
|
||||
from hindsight_api.engine.search.reranking import (
|
||||
_RECENCY_ALPHA,
|
||||
_TEMPORAL_ALPHA,
|
||||
apply_combined_scoring,
|
||||
compute_recency_decay,
|
||||
)
|
||||
from hindsight_api.engine.search.types import MergedCandidate, RetrievalResult, ScoredResult
|
||||
|
||||
UTC = timezone.utc
|
||||
@@ -200,3 +203,52 @@ class TestBoostFormula:
|
||||
|
||||
def test_empty_list_is_noop(self):
|
||||
apply_combined_scoring([], now=NOW) # must not raise
|
||||
|
||||
|
||||
class TestRecencyDecayFunction:
|
||||
"""The configurable age→freshness curve (compute_recency_decay)."""
|
||||
|
||||
def test_linear_is_default_and_unchanged(self):
|
||||
"""Default function reproduces the historical linear decay over 365 days."""
|
||||
assert compute_recency_decay(0) == 1.0
|
||||
assert abs(compute_recency_decay(182.5) - 0.5) < 1e-6 # neutral at half the window
|
||||
assert compute_recency_decay(400) == 0.1 # floored past the window
|
||||
|
||||
def test_linear_window_is_configurable(self):
|
||||
"""A custom window moves the neutral crossing; 730d window → neutral at 365d."""
|
||||
assert abs(compute_recency_decay(365, "linear", linear_window_days=730) - 0.5) < 1e-6
|
||||
|
||||
def test_exponential_neutral_at_halflife(self):
|
||||
"""Exponential decay is exactly neutral (0.5) at the configured half-life."""
|
||||
assert compute_recency_decay(0, "exponential", halflife_days=90) == 1.0
|
||||
assert abs(compute_recency_decay(90, "exponential", halflife_days=90) - 0.5) < 1e-9
|
||||
assert abs(compute_recency_decay(180, "exponential", halflife_days=90) - 0.25) < 1e-9
|
||||
|
||||
def test_exponential_penalises_old_less_harshly_than_linear(self):
|
||||
"""A 1-year-old memory keeps more freshness under a 90d-halflife exponential
|
||||
than under the linear floor — the curve never hard-cuts to 0.1."""
|
||||
lin = compute_recency_decay(365, "linear")
|
||||
exp = compute_recency_decay(365, "exponential", halflife_days=180)
|
||||
assert exp > lin
|
||||
|
||||
def test_none_is_always_neutral(self):
|
||||
"""'none' disables the recency signal — always neutral, no boost."""
|
||||
assert compute_recency_decay(0, "none") == 0.5
|
||||
assert compute_recency_decay(10_000, "none") == 0.5
|
||||
|
||||
def test_future_dates_clamp_to_max(self):
|
||||
"""Negative ages (future-dated memories) never exceed full freshness."""
|
||||
assert compute_recency_decay(-100, "linear") == 1.0
|
||||
assert compute_recency_decay(-100, "exponential", halflife_days=90) == 1.0
|
||||
|
||||
def test_nonpositive_halflife_falls_back_to_neutral(self):
|
||||
"""A misconfigured (<=0) half-life degrades to neutral rather than dividing by zero."""
|
||||
assert compute_recency_decay(30, "exponential", halflife_days=0) == 0.5
|
||||
|
||||
def test_function_threads_through_apply_combined_scoring(self):
|
||||
"""The decay function chosen at the call site is what scores sr.recency."""
|
||||
old = NOW - timedelta(days=180)
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=old)
|
||||
apply_combined_scoring([sr], now=NOW, recency_decay_function="none")
|
||||
assert sr.recency == 0.5
|
||||
assert abs(sr.weight - 0.5) < 1e-9 # neutral → no recency boost
|
||||
|
||||
@@ -25,6 +25,7 @@ def setup_test_env():
|
||||
"HINDSIGHT_API_LLM_MODEL",
|
||||
"HINDSIGHT_API_LLM_REASONING_EFFORT",
|
||||
"HINDSIGHT_API_LLM_BEDROCK_SERVICE_TIER",
|
||||
"HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER",
|
||||
"HINDSIGHT_API_SEMANTIC_MIN_SIMILARITY",
|
||||
"HINDSIGHT_API_DATABASE_URL",
|
||||
"HINDSIGHT_API_MIGRATION_DATABASE_URL",
|
||||
@@ -452,6 +453,53 @@ def test_llm_output_language_empty_string_is_unset(monkeypatch):
|
||||
assert config.llm_output_language is None
|
||||
|
||||
|
||||
def test_markitdown_ocr_defaults_disabled(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.file_parser_markitdown_ocr_enabled is False
|
||||
|
||||
|
||||
def test_markitdown_ocr_does_not_fall_back_to_main_llm_config(monkeypatch):
|
||||
from hindsight_api.config import DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT, HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED", "true")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "anthropic")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "main-key")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_BASE_URL", "https://main.example/v1")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_MODEL", "main-vision-model")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.file_parser_markitdown_ocr_enabled is True
|
||||
assert config.file_parser_markitdown_ocr_api_key is None
|
||||
assert config.file_parser_markitdown_ocr_base_url is None
|
||||
assert config.file_parser_markitdown_ocr_model is None
|
||||
assert config.file_parser_markitdown_ocr_prompt == DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
|
||||
|
||||
|
||||
def test_markitdown_ocr_uses_explicit_config(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED", "true")
|
||||
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY", "parser-key")
|
||||
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL", "https://parser.example/v1")
|
||||
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL", "parser-vision-model")
|
||||
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT", "Extract this document exactly.")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "main-key")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_BASE_URL", "https://main.example/v1")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_MODEL", "main-vision-model")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.file_parser_markitdown_ocr_enabled is True
|
||||
assert config.file_parser_markitdown_ocr_api_key == "parser-key"
|
||||
assert config.file_parser_markitdown_ocr_base_url == "https://parser.example/v1"
|
||||
assert config.file_parser_markitdown_ocr_model == "parser-vision-model"
|
||||
assert config.file_parser_markitdown_ocr_prompt == "Extract this document exactly."
|
||||
|
||||
|
||||
def test_llm_reasoning_effort_defaults_to_low(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
@@ -580,3 +628,81 @@ def test_bedrock_service_tier_rejects_invalid_value(monkeypatch):
|
||||
assert "HINDSIGHT_API_LLM_BEDROCK_SERVICE_TIER" in error_message
|
||||
assert "standard" in error_message
|
||||
assert "'standard' is not a valid Bedrock service tier" in error_message
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Gemini service tier (HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_gemini_service_tier_defaults_to_none(monkeypatch):
|
||||
"""Gemini service tier defaults to None (standard tier) when unset."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.delenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", raising=False)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_service_tier is None
|
||||
|
||||
|
||||
def test_gemini_service_tier_flex(monkeypatch):
|
||||
"""Flex tier is accepted for Gemini."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "flex")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "gemini")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "fake-key")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_service_tier == "flex"
|
||||
|
||||
|
||||
def test_gemini_service_tier_accepts_mixed_case_provider(monkeypatch):
|
||||
"""Gemini tier parsing follows provider's case-insensitive handling."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "flex")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "Gemini")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "fake-key")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_service_tier == "flex"
|
||||
|
||||
|
||||
def test_gemini_service_tier_rejects_invalid_value(monkeypatch):
|
||||
"""Unknown Gemini service tiers are rejected early."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "standard")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "gemini")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "fake-key")
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
HindsightConfig.from_env()
|
||||
|
||||
error_message = str(exc_info.value)
|
||||
assert "HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER" in error_message
|
||||
assert "standard" in error_message
|
||||
|
||||
|
||||
def test_gemini_service_tier_ignored_for_non_gemini_provider(monkeypatch):
|
||||
"""Invalid Gemini-only tiers do not break unrelated providers."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "standard")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_service_tier is None
|
||||
|
||||
|
||||
def test_gemini_service_tier_empty_env_is_unset(monkeypatch):
|
||||
"""Empty env values are treated as unset for templated deployments."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_service_tier is None
|
||||
|
||||
@@ -3619,3 +3619,39 @@ def test_consolidation_prompt_split_is_cacheable_and_complete():
|
||||
)
|
||||
assert "OBSERVATION LIMIT REACHED" in capped
|
||||
assert "OBSERVATION LIMIT REACHED" not in sys_prompt
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_observation_populates_search_vector_native(memory, request_context):
|
||||
"""Observations created via consolidation must have search_vector populated
|
||||
when text_search_extension == 'native', so BM25 retrieval finds them."""
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
config = get_config()
|
||||
if config.text_search_extension != "native":
|
||||
pytest.skip("Only applies to native text search backend")
|
||||
|
||||
bank_id = f"test-search-vector-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Django uses middleware for request processing.",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT search_vector
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
LIMIT 1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
assert row is not None, "Consolidation should have created an observation"
|
||||
assert row["search_vector"] is not None, "search_vector must be populated for BM25 retrieval under native backend"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -84,7 +84,13 @@ def _ctx(threshold: float = 0.97):
|
||||
conn=conn,
|
||||
memory_engine=types.SimpleNamespace(embeddings=object()),
|
||||
bank_id="bank1",
|
||||
config=types.SimpleNamespace(consolidation_dedup_threshold=threshold),
|
||||
# The merge path builds a search_vector UPDATE clause from the text-search
|
||||
# config, so these must be present (production defaults: native/english).
|
||||
config=types.SimpleNamespace(
|
||||
consolidation_dedup_threshold=threshold,
|
||||
text_search_extension="native",
|
||||
text_search_extension_native_language="english",
|
||||
),
|
||||
dedup_llm_config=llm,
|
||||
create_text="YouTube content in Uzbek is very rich.",
|
||||
create_source_ids=[uuid.uuid4()],
|
||||
@@ -126,6 +132,16 @@ async def test_dedup_llm_keep_does_not_merge() -> None:
|
||||
conn.execute.assert_not_called() # kept distinct → no merge
|
||||
|
||||
|
||||
async def test_dedup_llm_missing_action_defaults_to_keep() -> None:
|
||||
kwargs, conn, llm = _ctx()
|
||||
llm.call.return_value = _DedupDecision(reason="underfilled structured response")
|
||||
with _patch_embed(), _patch_probe([_obs("Uzbek content on YouTube is described as very rich.", 0.98)]):
|
||||
result = await _dedup_reconcile_create(**kwargs)
|
||||
assert result is None
|
||||
llm.call.assert_awaited_once()
|
||||
conn.execute.assert_not_called() # missing action is a conservative no-merge
|
||||
|
||||
|
||||
async def test_dedup_llm_merge_folds_into_twin() -> None:
|
||||
kwargs, conn, llm = _ctx()
|
||||
kwargs["create_source_ids"] = [uuid.uuid4(), uuid.uuid4()]
|
||||
@@ -169,7 +185,13 @@ def _update_ctx(threshold: float = 0.97):
|
||||
conn=conn,
|
||||
memory_engine=types.SimpleNamespace(embeddings=object()),
|
||||
bank_id="bank1",
|
||||
config=types.SimpleNamespace(consolidation_dedup_threshold=threshold),
|
||||
# The merge path builds a search_vector UPDATE clause from the text-search
|
||||
# config, so these must be present (production defaults: native/english).
|
||||
config=types.SimpleNamespace(
|
||||
consolidation_dedup_threshold=threshold,
|
||||
text_search_extension="native",
|
||||
text_search_extension_native_language="english",
|
||||
),
|
||||
dedup_llm_config=llm,
|
||||
updated_id=_UPDATED_ID,
|
||||
updated_text="Uzbek content on YouTube is very rich and growing.",
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
"""Tests for container-runtime detection used to warn about unstable worker ids."""
|
||||
|
||||
import builtins
|
||||
|
||||
from hindsight_api.utils import detect_container_runtime, warn_if_container_default_worker_id
|
||||
|
||||
|
||||
def test_detects_kubernetes_via_env(monkeypatch):
|
||||
monkeypatch.setenv("KUBERNETES_SERVICE_HOST", "10.0.0.1")
|
||||
assert detect_container_runtime() == "kubernetes"
|
||||
|
||||
|
||||
def test_detects_docker_via_dockerenv(monkeypatch):
|
||||
monkeypatch.delenv("KUBERNETES_SERVICE_HOST", raising=False)
|
||||
monkeypatch.setattr("os.path.exists", lambda p: p == "/.dockerenv")
|
||||
assert detect_container_runtime() == "docker"
|
||||
|
||||
|
||||
def test_detects_docker_via_cgroup(monkeypatch):
|
||||
monkeypatch.delenv("KUBERNETES_SERVICE_HOST", raising=False)
|
||||
monkeypatch.setattr("os.path.exists", lambda p: False)
|
||||
|
||||
real_open = builtins.open
|
||||
|
||||
def fake_open(path, *args, **kwargs):
|
||||
if path == "/proc/1/cgroup":
|
||||
import io
|
||||
|
||||
return io.StringIO("12:devices:/docker/abcdef123456\n")
|
||||
return real_open(path, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr("builtins.open", fake_open)
|
||||
assert detect_container_runtime() == "docker"
|
||||
|
||||
|
||||
def test_returns_none_when_not_containerized(monkeypatch):
|
||||
monkeypatch.delenv("KUBERNETES_SERVICE_HOST", raising=False)
|
||||
monkeypatch.setattr("os.path.exists", lambda p: False)
|
||||
|
||||
def fake_open(path, *args, **kwargs):
|
||||
raise OSError("no such file")
|
||||
|
||||
monkeypatch.setattr("builtins.open", fake_open)
|
||||
assert detect_container_runtime() is None
|
||||
|
||||
|
||||
def test_warns_when_default_worker_id_is_used_in_container(monkeypatch, caplog):
|
||||
monkeypatch.setattr("hindsight_api.utils.detect_container_runtime", lambda: "docker")
|
||||
|
||||
warn_if_container_default_worker_id(None)
|
||||
|
||||
assert "HINDSIGHT_API_WORKER_ID is not set" in caplog.text
|
||||
assert "appears to be running inside docker" in caplog.text
|
||||
|
||||
|
||||
def test_skips_warning_when_worker_id_is_explicit(monkeypatch, caplog):
|
||||
monkeypatch.setattr("hindsight_api.utils.detect_container_runtime", lambda: "docker")
|
||||
|
||||
warn_if_container_default_worker_id("worker-1")
|
||||
|
||||
assert caplog.text == ""
|
||||
|
||||
|
||||
def test_skips_warning_outside_containers(monkeypatch, caplog):
|
||||
monkeypatch.setattr("hindsight_api.utils.detect_container_runtime", lambda: None)
|
||||
|
||||
warn_if_container_default_worker_id(None)
|
||||
|
||||
assert caplog.text == ""
|
||||
@@ -77,6 +77,9 @@ def _make_tool_call_response(tool_name: str = "search_observations") -> MagicMoc
|
||||
mock_response.usage.prompt_tokens = 100
|
||||
mock_response.usage.completion_tokens = 20
|
||||
mock_response.usage.total_tokens = 120
|
||||
# Explicit None: an auto-MagicMock here is truthy, so the reasoning-token
|
||||
# accounting (#2378) would do arithmetic on a MagicMock and crash.
|
||||
mock_response.usage.completion_tokens_details = None
|
||||
mock_response.choices[0].finish_reason = "tool_calls"
|
||||
mock_response.choices[0].message.content = None
|
||||
mock_response.choices[0].message.tool_calls = [mock_tc]
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
"""Tests for write-side validation of bank disposition config overrides.
|
||||
|
||||
Disposition traits (skepticism / literalism / empathy) are integers on a 1-5
|
||||
scale. The ``PATCH /v1/{tenant}/banks/{id}/config`` write path must reject
|
||||
out-of-contract values (floats, 0-1 scales, ints outside 1-5) at write time;
|
||||
otherwise a single malformed bank 500s the entire bank list because the read
|
||||
overlay injects the stored value verbatim into a strict
|
||||
``DispositionTraits(int, ge=1, le=5)``. See issue #2348.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config_resolver import _validate_disposition_updates
|
||||
|
||||
|
||||
_DISPOSITION_FIELD_NAMES = (
|
||||
"disposition_skepticism",
|
||||
"disposition_literalism",
|
||||
"disposition_empathy",
|
||||
)
|
||||
|
||||
|
||||
class TestValidateDispositionUpdates:
|
||||
def test_no_op_passes(self):
|
||||
_validate_disposition_updates({})
|
||||
_validate_disposition_updates({"unrelated_field": 123})
|
||||
|
||||
def test_valid_in_range_integers_pass(self):
|
||||
for key in _DISPOSITION_FIELD_NAMES:
|
||||
for value in (1, 2, 3, 4, 5):
|
||||
_validate_disposition_updates({key: value})
|
||||
|
||||
def test_none_clears_override(self):
|
||||
# None is the "unset this per-bank override" sentinel (field is int | None).
|
||||
for key in _DISPOSITION_FIELD_NAMES:
|
||||
_validate_disposition_updates({key: None})
|
||||
|
||||
def test_out_of_range_integer_raises(self):
|
||||
for key in _DISPOSITION_FIELD_NAMES:
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: 0})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: 6})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: -1})
|
||||
|
||||
def test_float_raises(self):
|
||||
# The reported v0.8.3 case: a 0-1 scale used by mistake.
|
||||
for key in _DISPOSITION_FIELD_NAMES:
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: 0.7})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: 3.0}) # float, even if in 1-5 range
|
||||
|
||||
def test_bool_raises(self):
|
||||
# bool is an int subclass and would sneak past a naive isinstance(int) check.
|
||||
for key in _DISPOSITION_FIELD_NAMES:
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: True})
|
||||
|
||||
def test_string_raises(self):
|
||||
for key in _DISPOSITION_FIELD_NAMES:
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: "3"})
|
||||
@@ -124,6 +124,10 @@ def test_openai_codex_provider_uses_codex_oauth_token_and_configured_batch_size(
|
||||
)
|
||||
|
||||
monkeypatch.setenv("HOME", str(tmp_path))
|
||||
# Codex auth resolves via CODEX_HOME first (falling back to ~/.codex), so a
|
||||
# CODEX_HOME leaking in from the runner's environment would point auth.json
|
||||
# away from the tmp_path fixture. Pin resolution to the patched HOME.
|
||||
monkeypatch.delenv("CODEX_HOME", raising=False)
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_PROVIDER"] = "openai-codex"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL"] = "text-embedding-3-small"
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
"""Regression test: `enqueue_graph_maintenance` must insert unit_ids in a
|
||||
deterministic sorted order so concurrent transactions can't deadlock on the
|
||||
graph_maintenance_queue unique-key check.
|
||||
|
||||
Symptom (production): under load, concurrent `PATCH /memories/{id}` requests
|
||||
on the same bank generate overlapping `victim_ids` sets (the surviving units
|
||||
whose outgoing links pointed at the updated unit). Each transaction inserts
|
||||
those victims into `graph_maintenance_queue` with
|
||||
`ON CONFLICT (bank_id, unit_id) DO NOTHING`. The conflict check takes a
|
||||
short-lived row-level lock per (bank_id, unit_id) being inserted, and when
|
||||
two transactions insert overlapping sets in different orders Postgres
|
||||
detects a deadlock and aborts one of them — surfacing as
|
||||
`asyncpg.exceptions.DeadlockDetectedError` from the API, which becomes a 500.
|
||||
|
||||
The fix sorts the input list inside both `ops_postgresql` and `ops_oracle`
|
||||
before passing it to the INSERT, so every transaction acquires the per-row
|
||||
locks in the same global (sorted-UUID) order. With a total order over the
|
||||
lock set, deadlock is mathematically impossible — Postgres still serializes
|
||||
the conflicting inserts but they queue cleanly instead of cycling.
|
||||
|
||||
This test pins that post-condition by capturing the array passed to the
|
||||
underlying `conn.execute` (PG path) / `conn.executemany` (Oracle path) and
|
||||
asserting it's sorted.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.db.ops_oracle import OracleOps
|
||||
from hindsight_api.engine.db.ops_postgresql import PostgreSQLOps
|
||||
|
||||
|
||||
def _shuffled_uuids(n: int) -> list[uuid.UUID]:
|
||||
"""Generate n UUIDs in a deliberately non-monotonic order. Hex literals
|
||||
avoid `uuid.uuid4()` because uuid4 is random and we want determinism."""
|
||||
raw = [
|
||||
"ffffffff-ffff-4fff-8fff-ffffffffffff",
|
||||
"00000000-0000-4000-8000-000000000001",
|
||||
"88888888-8888-4888-8888-888888888888",
|
||||
"11111111-1111-4111-8111-111111111111",
|
||||
"ccccccccc-cccc-4ccc-8ccc-cccccccccccc"[:36],
|
||||
"44444444-4444-4444-8444-444444444444",
|
||||
]
|
||||
return [uuid.UUID(s) for s in raw[:n]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pg_enqueue_graph_maintenance_inserts_in_sorted_order():
|
||||
"""The PostgreSQL ops impl must pass the unit_ids to the INSERT in
|
||||
sorted order, regardless of how the caller ordered them."""
|
||||
ops = PostgreSQLOps()
|
||||
conn = AsyncMock()
|
||||
|
||||
unit_ids = _shuffled_uuids(6)
|
||||
assert unit_ids != sorted(unit_ids), "test inputs must be unsorted"
|
||||
|
||||
await ops.enqueue_graph_maintenance(
|
||||
conn=conn,
|
||||
table="graph_maintenance_queue",
|
||||
bank_id="test-bank",
|
||||
unit_ids=unit_ids,
|
||||
)
|
||||
|
||||
assert conn.execute.await_count == 1
|
||||
_sql, bank_id_arg, ids_arg = conn.execute.await_args.args
|
||||
assert bank_id_arg == "test-bank"
|
||||
assert ids_arg == sorted(unit_ids), f"expected sorted unit_ids for deadlock-free concurrent inserts, got {ids_arg}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oracle_enqueue_graph_maintenance_inserts_in_sorted_order():
|
||||
"""The Oracle ops impl applies the same sort. `executemany` receives a
|
||||
list of (bank_id, unit_id) tuples; the unit_id projection must be
|
||||
sorted."""
|
||||
ops = OracleOps()
|
||||
conn = AsyncMock()
|
||||
|
||||
unit_ids = _shuffled_uuids(6)
|
||||
assert unit_ids != sorted(unit_ids), "test inputs must be unsorted"
|
||||
|
||||
await ops.enqueue_graph_maintenance(
|
||||
conn=conn,
|
||||
table="graph_maintenance_queue",
|
||||
bank_id="test-bank",
|
||||
unit_ids=unit_ids,
|
||||
)
|
||||
|
||||
assert conn.executemany.await_count == 1
|
||||
_sql, rows = conn.executemany.await_args.args
|
||||
assert [r[0] for r in rows] == ["test-bank"] * len(unit_ids)
|
||||
assert [r[1] for r in rows] == sorted(unit_ids), (
|
||||
f"expected sorted unit_ids for deadlock-free concurrent inserts, got {[r[1] for r in rows]}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pg_empty_unit_ids_short_circuits():
|
||||
"""Empty input must remain a no-op — the early return predates this fix
|
||||
and must continue to skip the INSERT entirely."""
|
||||
ops = PostgreSQLOps()
|
||||
conn = AsyncMock()
|
||||
await ops.enqueue_graph_maintenance(conn, "graph_maintenance_queue", "b", [])
|
||||
conn.execute.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oracle_empty_unit_ids_short_circuits():
|
||||
ops = OracleOps()
|
||||
conn = AsyncMock()
|
||||
await ops.enqueue_graph_maintenance(conn, "graph_maintenance_queue", "b", [])
|
||||
conn.executemany.assert_not_awaited()
|
||||
@@ -0,0 +1,70 @@
|
||||
"""ensure_vector_extension must not create the (unused) global memory_units index.
|
||||
|
||||
For per-bank backends (pgvector / pgvectorscale / vchord) every vector search is
|
||||
bank + fact_type scoped and served by the per-(bank, fact_type) partial indexes
|
||||
created at bank-creation time. The global `idx_memory_units_embedding` is never
|
||||
chosen by the planner (migration d5e6f7a8b9c0 drops it for exactly this reason),
|
||||
so the post-migration reconcile must not recreate it on a fresh schema.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import create_engine, text
|
||||
|
||||
from hindsight_api._vector_index import uses_per_bank_vector_indexes
|
||||
from hindsight_api.config import HindsightConfig
|
||||
from hindsight_api.migrations import ensure_vector_extension, run_migrations
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def vec_db_url():
|
||||
"""A dedicated pg0 instance so the test owns its schema/index state."""
|
||||
from hindsight_api.pg0 import EmbeddedPostgres
|
||||
|
||||
pg0 = EmbeddedPostgres(name="hindsight-vecidx-test", port=5570)
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
return loop.run_until_complete(pg0.ensure_running())
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
|
||||
def test_per_bank_backend_does_not_create_global_memory_units_index(vec_db_url):
|
||||
config = HindsightConfig.from_env()
|
||||
vec = config.vector_extension
|
||||
if not uses_per_bank_vector_indexes(vec):
|
||||
pytest.skip(f"backend {vec!r} uses a global vector index by design (no per-bank indexes)")
|
||||
|
||||
schema = "vecidx_fresh"
|
||||
|
||||
engine = create_engine(vec_db_url)
|
||||
try:
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
conn.commit()
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
run_migrations(vec_db_url, schema=schema)
|
||||
# Fresh, empty schema (no banks yet) → the reconcile must be a no-op for the
|
||||
# global index, not recreate it.
|
||||
ensure_vector_extension(vec_db_url, vector_extension=vec, schema=schema)
|
||||
|
||||
engine = create_engine(vec_db_url)
|
||||
try:
|
||||
with engine.connect() as conn:
|
||||
global_index_count = conn.execute(
|
||||
text(
|
||||
"SELECT COUNT(*) FROM pg_indexes "
|
||||
"WHERE schemaname = :schema AND tablename = 'memory_units' "
|
||||
"AND indexname = 'idx_memory_units_embedding'"
|
||||
),
|
||||
{"schema": schema},
|
||||
).scalar()
|
||||
conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
conn.commit()
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
assert global_index_count == 0
|
||||
@@ -9,11 +9,19 @@ from fastapi.testclient import TestClient
|
||||
from hindsight_api.extensions import (
|
||||
ApiKeyTenantExtension,
|
||||
AuthenticationError,
|
||||
BankReadContext,
|
||||
BankReadOperation,
|
||||
BankWriteContext,
|
||||
BankWriteOperation,
|
||||
# Consolidation operation
|
||||
ConsolidateContext,
|
||||
ConsolidateResult,
|
||||
Extension,
|
||||
HttpExtension,
|
||||
OperationValidationError,
|
||||
OperationValidatorExtension,
|
||||
PrecheckContext,
|
||||
PrecheckOperation,
|
||||
RecallContext,
|
||||
RecallResult,
|
||||
ReflectContext,
|
||||
@@ -21,13 +29,8 @@ from hindsight_api.extensions import (
|
||||
RequestContext,
|
||||
RetainContext,
|
||||
RetainResult,
|
||||
TenantContext,
|
||||
TenantExtension,
|
||||
ValidationResult,
|
||||
load_extension,
|
||||
# Consolidation operation
|
||||
ConsolidateContext,
|
||||
ConsolidateResult,
|
||||
)
|
||||
|
||||
|
||||
@@ -68,6 +71,36 @@ class TestExtensionLoader:
|
||||
await ext.on_shutdown()
|
||||
assert ext.stopped
|
||||
|
||||
def test_operation_enums_remain_string_compatible(self):
|
||||
"""Operation enums centralize names without breaking string comparisons."""
|
||||
request_context = RequestContext(tenant_id="tenant-1")
|
||||
|
||||
precheck_ctx = PrecheckContext(
|
||||
bank_id="bank-1",
|
||||
operation=PrecheckOperation.RETAIN,
|
||||
request_context=request_context,
|
||||
)
|
||||
read_ctx = BankReadContext(
|
||||
bank_id="bank-1",
|
||||
operation=BankReadOperation.GET_BANK_STATS,
|
||||
request_context=request_context,
|
||||
)
|
||||
write_ctx = BankWriteContext(
|
||||
bank_id="bank-1",
|
||||
operation=BankWriteOperation.UPDATE_BANK_CONFIG,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert precheck_ctx.operation is PrecheckOperation.RETAIN
|
||||
assert read_ctx.operation is BankReadOperation.GET_BANK_STATS
|
||||
assert write_ctx.operation is BankWriteOperation.UPDATE_BANK_CONFIG
|
||||
assert precheck_ctx.operation == "retain"
|
||||
assert read_ctx.operation == "get_bank_stats"
|
||||
assert write_ctx.operation == "update_bank_config"
|
||||
assert isinstance(precheck_ctx.operation, str)
|
||||
assert isinstance(read_ctx.operation, str)
|
||||
assert isinstance(write_ctx.operation, str)
|
||||
|
||||
|
||||
class LifecycleTestExtension(Extension):
|
||||
"""Test extension for config and lifecycle tests."""
|
||||
@@ -356,6 +389,7 @@ class TestOperationHooksParameters:
|
||||
async def test_recall_pre_hook_receives_all_parameters(self, memory_with_tracking_validator):
|
||||
"""Pre-recall hook receives all user-provided parameters."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
memory, validator = memory_with_tracking_validator
|
||||
@@ -882,7 +916,7 @@ class TestPrecheckDefault:
|
||||
validator = RecordingPrecheckValidator(reject=False)
|
||||
# Bypass our override by calling the base implementation directly.
|
||||
ctx = PrecheckContext(
|
||||
operation="retain",
|
||||
operation=PrecheckOperation.RETAIN,
|
||||
bank_id="bank-x",
|
||||
request_context=RequestContext(),
|
||||
)
|
||||
@@ -911,10 +945,10 @@ class TestPrecheckHttpWiring:
|
||||
def _build_app(validator):
|
||||
"""Mirror the precheck wiring from ``hindsight_api.api.http`` in a
|
||||
standalone FastAPI app."""
|
||||
from fastapi import Depends, FastAPI, HTTPException
|
||||
from fastapi import Depends, FastAPI, HTTPException, Request
|
||||
from pydantic import BaseModel, model_validator
|
||||
|
||||
from hindsight_api.extensions import PrecheckContext
|
||||
from hindsight_api.extensions import PrecheckContext, PrecheckOperation
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
body_parses: list[str] = []
|
||||
@@ -949,15 +983,26 @@ class TestPrecheckHttpWiring:
|
||||
async def _request_context() -> RequestContext:
|
||||
return RequestContext()
|
||||
|
||||
def _precheck_for(operation: str):
|
||||
def _precheck_for(operation: PrecheckOperation):
|
||||
async def _dep(
|
||||
bank_id: str,
|
||||
request: Request,
|
||||
request_context: RequestContext = Depends(_request_context),
|
||||
) -> None:
|
||||
cl_header = request.headers.get("content-length")
|
||||
content_length: int | None = None
|
||||
if cl_header is not None:
|
||||
try:
|
||||
parsed = int(cl_header)
|
||||
except ValueError:
|
||||
parsed = -1
|
||||
if parsed >= 0:
|
||||
content_length = parsed
|
||||
ctx = PrecheckContext(
|
||||
operation=operation,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
content_length=content_length,
|
||||
)
|
||||
result = await validator.precheck(ctx)
|
||||
if not result.allowed:
|
||||
@@ -974,7 +1019,7 @@ class TestPrecheckHttpWiring:
|
||||
async def retain(
|
||||
bank_id: str,
|
||||
body: _RetainBody,
|
||||
_: None = Depends(_precheck_for("retain")),
|
||||
_: None = Depends(_precheck_for(PrecheckOperation.RETAIN)),
|
||||
):
|
||||
return {"ok": True, "bank_id": bank_id, "n": len(body.items)}
|
||||
|
||||
@@ -982,7 +1027,7 @@ class TestPrecheckHttpWiring:
|
||||
async def recall(
|
||||
bank_id: str,
|
||||
body: _RecallBody,
|
||||
_: None = Depends(_precheck_for("recall")),
|
||||
_: None = Depends(_precheck_for(PrecheckOperation.RECALL)),
|
||||
):
|
||||
return {"ok": True}
|
||||
|
||||
@@ -990,7 +1035,7 @@ class TestPrecheckHttpWiring:
|
||||
async def reflect(
|
||||
bank_id: str,
|
||||
body: _ReflectBody,
|
||||
_: None = Depends(_precheck_for("reflect")),
|
||||
_: None = Depends(_precheck_for(PrecheckOperation.REFLECT)),
|
||||
):
|
||||
return {"ok": True}
|
||||
|
||||
@@ -1081,3 +1126,86 @@ class TestPrecheckHttpWiring:
|
||||
resp = client.get("/v1/default/banks/precheck-bank/memories/list")
|
||||
assert resp.status_code == 200
|
||||
assert len(validator.precheck_calls) == 0
|
||||
|
||||
def test_precheck_context_carries_content_length(self):
|
||||
"""Content-Length header is exposed to the precheck so a validator
|
||||
can make size-aware decisions (e.g. upper-bound cost estimate)
|
||||
before the body is deserialised."""
|
||||
validator = RecordingPrecheckValidator(reject=False)
|
||||
app, _ = self._build_app(validator)
|
||||
client = TestClient(app)
|
||||
|
||||
# Body must contain at least 500 'x' bytes; check the surfaced
|
||||
# Content-Length is within a tight band around that floor (allows
|
||||
# for JSON envelope + httpx's serialisation choices without
|
||||
# depending on exact byte counts).
|
||||
payload = {"items": [{"content": "x" * 500}]}
|
||||
resp = client.post(
|
||||
"/v1/default/banks/precheck-bank/memories",
|
||||
json=payload,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert len(validator.precheck_calls) == 1
|
||||
ctx = validator.precheck_calls[0]
|
||||
assert ctx.content_length is not None
|
||||
assert 500 <= ctx.content_length <= 600
|
||||
|
||||
def test_precheck_context_content_length_zero_is_not_none(self):
|
||||
"""An empty POST body has Content-Length: 0. That should surface
|
||||
as the int 0, not None — None means 'unknown', 0 means 'known to
|
||||
be empty'."""
|
||||
validator = RecordingPrecheckValidator(reject=False)
|
||||
app, _ = self._build_app(validator)
|
||||
client = TestClient(app)
|
||||
|
||||
# Empty body fails Pydantic parse (422), but precheck runs first
|
||||
# and records the Content-Length.
|
||||
client.post(
|
||||
"/v1/default/banks/precheck-bank/memories",
|
||||
content=b"",
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
assert len(validator.precheck_calls) >= 1
|
||||
ctx = validator.precheck_calls[-1]
|
||||
assert ctx.content_length == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_precheck_context_content_length_none_when_header_missing(self):
|
||||
"""When the Content-Length header isn't set (e.g. chunked transfer
|
||||
encoding) the validator sees None, not a crash and not a default 0."""
|
||||
from starlette.requests import Request as _StarletteRequest
|
||||
|
||||
from hindsight_api.extensions import PrecheckContext
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
validator = RecordingPrecheckValidator(reject=False)
|
||||
|
||||
# Replicate the wiring's parse step inline so the test exercises
|
||||
# the same code-path semantics introduced in
|
||||
# ``hindsight_api.api.http._precheck_dep``.
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/v1/default/banks/bank-x/memories",
|
||||
"headers": [], # no content-length
|
||||
"query_string": b"",
|
||||
}
|
||||
req = _StarletteRequest(scope)
|
||||
cl_header = req.headers.get("content-length")
|
||||
content_length: int | None = None
|
||||
if cl_header is not None:
|
||||
try:
|
||||
parsed = int(cl_header)
|
||||
except ValueError:
|
||||
parsed = -1
|
||||
if parsed >= 0:
|
||||
content_length = parsed
|
||||
|
||||
ctx = PrecheckContext(
|
||||
operation=PrecheckOperation.RETAIN,
|
||||
bank_id="bank-x",
|
||||
request_context=RequestContext(),
|
||||
content_length=content_length,
|
||||
)
|
||||
await validator.precheck(ctx)
|
||||
assert validator.precheck_calls[-1].content_length is None
|
||||
|
||||
@@ -77,6 +77,58 @@ async def test_dry_run_extracts_without_persisting(api_client, memory):
|
||||
assert after["total"] == before["total"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dry_run_does_not_create_missing_bank(api_client, memory):
|
||||
bank_id = f"dryrun-missing-{uuid.uuid4().hex[:8]}"
|
||||
request_context = RequestContext()
|
||||
|
||||
assert (
|
||||
await memory.get_bank_profile(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
create_if_missing=False,
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
resp = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/memories/dry-run-extract",
|
||||
json={"content": "Alice moved to Berlin in 2021."},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert resp.json()["facts"]
|
||||
|
||||
assert (
|
||||
await memory.get_bank_profile(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
create_if_missing=False,
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dry_run_rejects_empty_content(api_client, memory):
|
||||
"""Empty/whitespace-only content is rejected by request validation (422) before the
|
||||
billable LLM extraction call runs — matching retain (RetainItem.content) and recall
|
||||
(RecallRequest.query), which already reject empty input."""
|
||||
bank_id = f"dryrun-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=RequestContext())
|
||||
|
||||
before = await memory.list_memory_units(bank_id=bank_id, request_context=RequestContext())
|
||||
for content in ("", " ", "\n\t "):
|
||||
resp = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/memories/dry-run-extract",
|
||||
json={"content": content},
|
||||
)
|
||||
assert resp.status_code == 422, resp.text
|
||||
|
||||
# Rejected before extraction: nothing was persisted.
|
||||
after = await memory.list_memory_units(bank_id=bank_id, request_context=RequestContext())
|
||||
assert after["total"] == before["total"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dry_run_disabled_returns_404(api_client, memory):
|
||||
"""With HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT=false the endpoint is removed (returns 404)."""
|
||||
|
||||
@@ -27,6 +27,13 @@ def llm_config():
|
||||
api_key=config.retain_llm_api_key or config.llm_api_key,
|
||||
model=config.retain_llm_model or config.llm_model,
|
||||
base_url=config.retain_llm_base_url or config.llm_base_url,
|
||||
# LLMConfig uses these as-passed and no longer reads them from global config,
|
||||
# so the caller must forward the Vertex AI settings (mirrors MemoryEngine's
|
||||
# own LLMConfig construction). Without this, provider=vertexai raises
|
||||
# "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID is required" even when it is set.
|
||||
vertexai_project_id=config.llm_vertexai_project_id,
|
||||
vertexai_region=config.llm_vertexai_region,
|
||||
vertexai_service_account_key=config.llm_vertexai_service_account_key,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from hindsight_api.engine.retain.fact_extraction import (
|
||||
ExtractedFact,
|
||||
ExtractedFactNoCausal,
|
||||
ExtractedFactVerbose,
|
||||
_build_extraction_prompt_and_schema,
|
||||
)
|
||||
|
||||
|
||||
def _baseline_config() -> MagicMock:
|
||||
config = MagicMock()
|
||||
config.entity_labels = None
|
||||
config.entities_allow_free_form = True
|
||||
config.retain_extraction_mode = "concise"
|
||||
config.retain_extract_causal_links = False
|
||||
config.retain_mission = None
|
||||
config.retain_custom_instructions = None
|
||||
config.llm_output_language = None
|
||||
return config
|
||||
|
||||
|
||||
def test_concise_prompt_keeps_user_preferences_rules_and_corrections_world():
|
||||
prompt, _ = _build_extraction_prompt_and_schema(_baseline_config())
|
||||
|
||||
assert '"world": Objective/external facts' in prompt
|
||||
assert "user's preferences, rules, corrections, constraints" in prompt
|
||||
assert 'These stay "world" even when the user states them during an assistant interaction' in prompt
|
||||
assert "Use this for the assistant/agent doing" in prompt
|
||||
assert "not merely for user facts mentioned in conversation" in prompt
|
||||
|
||||
|
||||
def test_fact_type_schema_descriptions_distinguish_user_facts_from_agent_actions():
|
||||
for model in (ExtractedFact, ExtractedFactVerbose, ExtractedFactNoCausal):
|
||||
description = model.model_fields["fact_type"].description
|
||||
|
||||
assert description is not None
|
||||
assert "preferences" in description
|
||||
assert "rules" in description
|
||||
assert "corrections" in description
|
||||
assert "assistant/agent actually performed" in description
|
||||
@@ -30,6 +30,7 @@ def _make_config(llm_max_retries: int = 3, retain_llm_max_retries: int | None =
|
||||
cfg.retain_extraction_mode = "concise"
|
||||
cfg.retain_extract_causal_links = False
|
||||
cfg.retain_mission = None
|
||||
cfg.llm_temperature_retain = 0.1
|
||||
return cfg
|
||||
|
||||
|
||||
@@ -216,3 +217,43 @@ async def test_none_event_date_with_valid_facts_no_crash():
|
||||
|
||||
assert len(facts) == 1
|
||||
assert "Alice visited Paris" in facts[0].fact
|
||||
|
||||
|
||||
def _make_batch_temp_config(temperature):
|
||||
"""Minimal config for _build_request_body temperature tests."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
cfg = MagicMock(spec=HindsightConfig)
|
||||
cfg.llm_temperature_retain = temperature
|
||||
cfg.retain_max_completion_tokens = None
|
||||
cfg.llm_strict_schema = False
|
||||
return cfg
|
||||
|
||||
|
||||
def _make_batch_llm_config():
|
||||
"""Minimal LLMProvider mock for _build_request_body (non-openai skips service_tier)."""
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
llm = MagicMock(spec=LLMProvider)
|
||||
llm.model = "gpt-test"
|
||||
llm.provider = "mock"
|
||||
return llm
|
||||
|
||||
|
||||
def test_build_request_body_forwards_configured_temperature():
|
||||
"""Batch retain path must send the configured retain temperature."""
|
||||
from hindsight_api.engine.retain.fact_extraction import _build_request_body
|
||||
|
||||
body = _build_request_body(_make_batch_llm_config(), _make_batch_temp_config(0.7), "sys", "user", dict)
|
||||
assert body["temperature"] == 0.7
|
||||
|
||||
|
||||
def test_build_request_body_omits_temperature_when_none():
|
||||
"""HINDSIGHT_API_LLM_TEMPERATURE=none must drop temperature from the batch
|
||||
request body too (Azure GPT-5.5 rejects explicit temperatures). Follow-up to
|
||||
#2469, which only de-hardcoded the streaming path and left the batch
|
||||
_build_request_body hardcoding temperature=0.1."""
|
||||
from hindsight_api.engine.retain.fact_extraction import _build_request_body
|
||||
|
||||
body = _build_request_body(_make_batch_llm_config(), _make_batch_temp_config(None), "sys", "user", dict)
|
||||
assert "temperature" not in body
|
||||
|
||||
@@ -346,6 +346,149 @@ async def test_markitdown_converter():
|
||||
assert "test document" in result.lower() or "multiple lines" in result.lower()
|
||||
|
||||
|
||||
def test_markitdown_converter_does_not_enable_ocr_by_default(monkeypatch):
|
||||
"""Markitdown should keep its local/default behavior unless OCR is explicitly enabled."""
|
||||
import markitdown
|
||||
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
calls = []
|
||||
|
||||
class FakeMarkItDown:
|
||||
def __init__(self, **kwargs):
|
||||
calls.append(kwargs)
|
||||
|
||||
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
|
||||
|
||||
MarkitdownParser()
|
||||
|
||||
assert calls == [{}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_markitdown_image_without_ocr_has_actionable_error(monkeypatch):
|
||||
"""Image uploads should explain that MarkItDown OCR is disabled instead of surfacing a low-level error."""
|
||||
import markitdown
|
||||
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
class FakeMarkItDown:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def convert(self, path):
|
||||
raise AssertionError("MarkItDown should not be called when image OCR is disabled")
|
||||
|
||||
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
|
||||
|
||||
parser = MarkitdownParser()
|
||||
with pytest.raises(RuntimeError, match="Image OCR is not enabled for the markitdown parser"):
|
||||
await parser.convert(b"\x89PNG\r\n\x1a\n", "screenshot.png")
|
||||
|
||||
|
||||
def test_markitdown_converter_can_enable_ocr(monkeypatch):
|
||||
"""When enabled, Markitdown receives an OpenAI-compatible client, model, and OCR prompt."""
|
||||
import markitdown
|
||||
import openai
|
||||
|
||||
from hindsight_api.config import DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
markitdown_calls = []
|
||||
openai_calls = []
|
||||
|
||||
class FakeMarkItDown:
|
||||
def __init__(self, **kwargs):
|
||||
markitdown_calls.append(kwargs)
|
||||
|
||||
class FakeOpenAI:
|
||||
def __init__(self, **kwargs):
|
||||
openai_calls.append(kwargs)
|
||||
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
|
||||
monkeypatch.setattr(openai, "OpenAI", FakeOpenAI)
|
||||
|
||||
MarkitdownParser(
|
||||
ocr_enabled=True,
|
||||
ocr_api_key="parser-key",
|
||||
ocr_base_url="https://vision.example/v1",
|
||||
ocr_model="vision-model",
|
||||
)
|
||||
|
||||
assert openai_calls == [
|
||||
{
|
||||
"api_key": "parser-key",
|
||||
"base_url": "https://vision.example/v1",
|
||||
}
|
||||
]
|
||||
assert markitdown_calls[0]["llm_client"].__class__ is FakeOpenAI
|
||||
assert markitdown_calls[0]["llm_model"] == "vision-model"
|
||||
assert markitdown_calls[0]["llm_prompt"] == DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
|
||||
|
||||
|
||||
def test_markitdown_converter_requires_model_when_ocr_enabled(monkeypatch):
|
||||
"""OCR should fail fast when enabled without a model."""
|
||||
import markitdown
|
||||
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
class FakeMarkItDown:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
|
||||
|
||||
with pytest.raises(ValueError, match="no model"):
|
||||
MarkitdownParser(ocr_enabled=True, ocr_api_key="parser-key")
|
||||
|
||||
|
||||
def test_markitdown_converter_requires_base_url_when_ocr_enabled(monkeypatch):
|
||||
"""OCR should fail fast when enabled without a dedicated OpenAI-compatible endpoint."""
|
||||
import markitdown
|
||||
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
class FakeMarkItDown:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
|
||||
|
||||
with pytest.raises(ValueError, match="no base URL"):
|
||||
MarkitdownParser(ocr_enabled=True, ocr_api_key="parser-key", ocr_model="vision-model")
|
||||
|
||||
|
||||
def test_markitdown_converter_reports_missing_openai_when_ocr_enabled(monkeypatch):
|
||||
"""Missing OpenAI SDK should not be reported as missing MarkItDown."""
|
||||
import builtins
|
||||
import markitdown
|
||||
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
real_import = builtins.__import__
|
||||
|
||||
class FakeMarkItDown:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def fake_import(name, globals=None, locals=None, fromlist=(), level=0):
|
||||
if name == "openai":
|
||||
raise ImportError("no openai")
|
||||
return real_import(name, globals, locals, fromlist, level)
|
||||
|
||||
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
|
||||
monkeypatch.setattr(builtins, "__import__", fake_import)
|
||||
|
||||
with pytest.raises(RuntimeError, match="openai package is required"):
|
||||
MarkitdownParser(
|
||||
ocr_enabled=True,
|
||||
ocr_api_key="parser-key",
|
||||
ocr_base_url="https://vision.example/v1",
|
||||
ocr_model="vision-model",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_converter_registry():
|
||||
"""Test file parser registry."""
|
||||
@@ -474,6 +617,59 @@ async def test_file_conversion_creates_separate_retain_operation(memory_no_llm_v
|
||||
assert len(doc["original_text"]) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_operations_surfaces_file_document_id_and_filename(memory_no_llm_verify, sample_txt_content):
|
||||
"""list_operations must expose document_id + filename for file_convert_retain ops.
|
||||
|
||||
The control plane derives its pending-upload rows from these fields (it
|
||||
matches an in-flight operation to the real document via document_id and
|
||||
labels the row with the original filename), so both must round-trip from
|
||||
the operation's result_metadata into the list response.
|
||||
"""
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
bank_id = "test_file_op_fields_bank"
|
||||
context = RequestContext(internal=True)
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id, request_context=context)
|
||||
|
||||
class MockFile:
|
||||
def __init__(self, content, filename, content_type):
|
||||
self.content = content
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
|
||||
async def read(self):
|
||||
return self.content
|
||||
|
||||
file_items = [
|
||||
{
|
||||
"file": MockFile(sample_txt_content, "report.txt", "text/plain"),
|
||||
"document_id": "doc_op_fields",
|
||||
"context": None,
|
||||
"metadata": {},
|
||||
"tags": [],
|
||||
"timestamp": None,
|
||||
"parser": ["markitdown"],
|
||||
}
|
||||
]
|
||||
|
||||
await memory_no_llm_verify.submit_async_file_retain(
|
||||
bank_id=bank_id,
|
||||
file_items=file_items,
|
||||
document_tags=None,
|
||||
request_context=context,
|
||||
)
|
||||
|
||||
result = await memory_no_llm_verify.list_operations(
|
||||
bank_id, task_type="file_convert_retain", request_context=context
|
||||
)
|
||||
|
||||
file_ops = [op for op in result["operations"] if op["task_type"] == "file_convert_retain"]
|
||||
assert len(file_ops) == 1
|
||||
assert file_ops[0]["document_id"] == "doc_op_fields"
|
||||
assert file_ops[0]["filename"] == "report.txt"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_file_retain_serializes_datetime_timestamp(memory_no_llm_verify, sample_txt_content):
|
||||
"""Async file retain should accept Python datetimes in task payloads."""
|
||||
|
||||
@@ -12,6 +12,8 @@ import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
@@ -21,6 +23,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
from testcontainers.core.container import DockerContainer
|
||||
from testcontainers.core.docker_client import DockerClient as _DockerClient
|
||||
|
||||
_has_testcontainers = True
|
||||
except ImportError:
|
||||
@@ -38,6 +41,8 @@ SEAWEEDFS_S3_PORT = 8333
|
||||
TEST_BUCKET = "hindsight-test"
|
||||
ACCESS_KEY = "test_access_key"
|
||||
SECRET_KEY = "test_secret_key"
|
||||
_PORT_MAPPING_RETRY_TIMEOUT_SECONDS = 10.0
|
||||
_PORT_MAPPING_RETRY_INTERVAL_SECONDS = 0.1
|
||||
|
||||
# SeaweedFS S3 IAM config granting full access to our test credentials
|
||||
_S3_CONFIG = {
|
||||
@@ -64,6 +69,33 @@ def _docker_available() -> bool:
|
||||
return False
|
||||
|
||||
|
||||
if _has_testcontainers:
|
||||
|
||||
@contextmanager
|
||||
def _retry_testcontainers_port_mapping() -> Iterator[None]:
|
||||
original_port = _DockerClient.port
|
||||
|
||||
def port_with_retry(self: _DockerClient, container_id: str, port: int) -> str:
|
||||
deadline = time.monotonic() + _PORT_MAPPING_RETRY_TIMEOUT_SECONDS
|
||||
while True:
|
||||
try:
|
||||
return original_port(self, container_id, port)
|
||||
except ConnectionError:
|
||||
# Docker Desktop can report a container as running before its
|
||||
# published port appears in NetworkSettings.Ports. This affects
|
||||
# both Ryuk's 8080 lookup inside testcontainers and the
|
||||
# SeaweedFS S3 port lookup below.
|
||||
if time.monotonic() >= deadline:
|
||||
raise
|
||||
time.sleep(_PORT_MAPPING_RETRY_INTERVAL_SECONDS)
|
||||
|
||||
_DockerClient.port = port_with_retry
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_DockerClient.port = original_port
|
||||
|
||||
|
||||
def _wait_for_seaweedfs(endpoint: str, timeout: int = 30) -> None:
|
||||
"""Poll SeaweedFS S3 endpoint until ready."""
|
||||
deadline = time.time() + timeout
|
||||
@@ -101,11 +133,11 @@ def seaweedfs_container():
|
||||
.with_command(f"server -s3 -s3.port={SEAWEEDFS_S3_PORT} -s3.config=/etc/seaweedfs/s3.json -ip.bind=0.0.0.0")
|
||||
)
|
||||
|
||||
container.start()
|
||||
|
||||
try:
|
||||
host = container.get_container_host_ip()
|
||||
port = container.get_exposed_port(SEAWEEDFS_S3_PORT)
|
||||
with _retry_testcontainers_port_mapping():
|
||||
container.start()
|
||||
host = container.get_container_host_ip()
|
||||
port = container.get_exposed_port(SEAWEEDFS_S3_PORT)
|
||||
endpoint = f"http://{host}:{port}"
|
||||
|
||||
_wait_for_seaweedfs(endpoint, timeout=240)
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user