Compare commits
128
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 |
@@ -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"
|
||||
|
||||
+30
-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
|
||||
|
||||
@@ -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 }}
|
||||
@@ -50,6 +51,7 @@ jobs:
|
||||
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 }}
|
||||
@@ -59,6 +61,7 @@ jobs:
|
||||
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 }}
|
||||
@@ -143,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:
|
||||
@@ -169,6 +174,8 @@ jobs:
|
||||
- 'hindsight-integrations/paperclip/**'
|
||||
integrations-opencode:
|
||||
- 'hindsight-integrations/opencode/**'
|
||||
integrations-eve:
|
||||
- 'hindsight-integrations/eve/**'
|
||||
integrations-cursor:
|
||||
- 'hindsight-integrations/cursor/**'
|
||||
integrations-zed:
|
||||
@@ -189,6 +196,8 @@ jobs:
|
||||
- 'hindsight-integrations/openai-agents/**'
|
||||
integrations-openhands:
|
||||
- 'hindsight-integrations/openhands/**'
|
||||
integrations-devin-desktop:
|
||||
- 'hindsight-integrations/devin-desktop/**'
|
||||
integrations-pipecat:
|
||||
- 'hindsight-integrations/pipecat/**'
|
||||
integrations-agentcore:
|
||||
@@ -591,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: >-
|
||||
@@ -748,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: >-
|
||||
@@ -3787,6 +3866,45 @@ jobs:
|
||||
# (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: >-
|
||||
@@ -4788,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
|
||||
@@ -4837,6 +4957,7 @@ jobs:
|
||||
- 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.3
|
||||
appVersion: "0.8.3"
|
||||
version: 0.8.4
|
||||
appVersion: "0.8.4"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.8.3",
|
||||
"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.3"
|
||||
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.3",
|
||||
"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.3"
|
||||
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.3",
|
||||
"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.3",
|
||||
"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.3"
|
||||
__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,6 +155,8 @@ 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
|
||||
@@ -270,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
|
||||
@@ -286,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
|
||||
@@ -353,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):
|
||||
@@ -1953,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).",
|
||||
@@ -2011,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)."""
|
||||
@@ -2048,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."""
|
||||
|
||||
@@ -2546,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,
|
||||
@@ -3055,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(
|
||||
@@ -3310,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``.
|
||||
|
||||
@@ -3465,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",
|
||||
@@ -3513,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),
|
||||
):
|
||||
"""
|
||||
@@ -3561,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.
|
||||
@@ -3591,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 = (
|
||||
@@ -3759,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
|
||||
@@ -3826,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,
|
||||
@@ -3838,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,
|
||||
@@ -3859,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]
|
||||
@@ -3960,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()
|
||||
|
||||
@@ -4118,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"]
|
||||
@@ -4248,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."""
|
||||
@@ -4284,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),
|
||||
):
|
||||
@@ -4507,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:
|
||||
@@ -4556,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:
|
||||
@@ -4702,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
|
||||
# =========================================================================
|
||||
@@ -4911,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),
|
||||
):
|
||||
"""
|
||||
@@ -5096,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),
|
||||
):
|
||||
"""
|
||||
@@ -5625,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:
|
||||
@@ -5642,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")
|
||||
@@ -6133,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)
|
||||
)
|
||||
@@ -6181,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)
|
||||
)
|
||||
@@ -6240,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)
|
||||
)
|
||||
@@ -6613,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()
|
||||
@@ -6796,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
|
||||
|
||||
@@ -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)
|
||||
@@ -148,6 +148,29 @@ 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).
|
||||
@@ -156,6 +179,12 @@ 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)
|
||||
@@ -179,6 +208,46 @@ def parse_gemini_service_tier(value: str | None) -> str | 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"
|
||||
@@ -272,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"
|
||||
@@ -499,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"
|
||||
@@ -575,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"
|
||||
@@ -588,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"
|
||||
@@ -610,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",
|
||||
@@ -623,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",
|
||||
}
|
||||
@@ -713,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
|
||||
@@ -771,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.
|
||||
@@ -1011,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.
|
||||
|
||||
@@ -1208,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}``.
|
||||
|
||||
@@ -1263,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.
|
||||
@@ -1340,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}.
|
||||
@@ -1429,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
|
||||
@@ -1464,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
|
||||
@@ -1634,7 +1885,6 @@ class HindsightConfig:
|
||||
|
||||
# Optimization flags
|
||||
skip_llm_verification: bool
|
||||
lazy_reranker: bool
|
||||
|
||||
# Database migrations
|
||||
run_migrations_on_startup: bool
|
||||
@@ -1689,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)
|
||||
@@ -1713,6 +1966,20 @@ class HindsightConfig:
|
||||
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
|
||||
|
||||
# CREDENTIAL_FIELDS: Never exposed via API, never configurable per-tenant/bank
|
||||
@@ -1727,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",
|
||||
@@ -2047,6 +2319,18 @@ class HindsightConfig:
|
||||
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,
|
||||
@@ -2146,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),
|
||||
@@ -2210,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"),
|
||||
@@ -2316,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),
|
||||
@@ -2421,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))
|
||||
@@ -2699,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,
|
||||
|
||||
@@ -128,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,
|
||||
@@ -402,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
|
||||
@@ -516,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
|
||||
""",
|
||||
@@ -1845,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"""
|
||||
@@ -1857,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,
|
||||
@@ -2333,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(
|
||||
|
||||
@@ -1638,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:
|
||||
@@ -1701,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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -253,6 +253,7 @@ def create_llm_provider(
|
||||
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.
|
||||
@@ -272,12 +273,20 @@ def create_llm_provider(
|
||||
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.
|
||||
@@ -375,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":
|
||||
@@ -393,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":
|
||||
@@ -405,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":
|
||||
@@ -452,6 +467,7 @@ def create_llm_provider(
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
elif provider_lower in (
|
||||
@@ -464,8 +480,10 @@ def create_llm_provider(
|
||||
"deepseek",
|
||||
"volcano",
|
||||
"openrouter",
|
||||
"requesty",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
):
|
||||
return OpenAICompatibleLLM(
|
||||
provider=provider,
|
||||
@@ -476,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:
|
||||
@@ -505,6 +524,13 @@ class LLMProvider:
|
||||
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.
|
||||
@@ -523,20 +549,49 @@ class LLMProvider:
|
||||
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
|
||||
@@ -553,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 = [
|
||||
@@ -586,8 +634,10 @@ class LLMProvider:
|
||||
"bedrock",
|
||||
"volcano",
|
||||
"openrouter",
|
||||
"requesty",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
"fireworks",
|
||||
"nous",
|
||||
]
|
||||
@@ -610,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:
|
||||
@@ -659,61 +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
|
||||
|
||||
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
|
||||
|
||||
# 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
|
||||
|
||||
self.gemini_service_tier = parse_gemini_service_tier(self.gemini_service_tier)
|
||||
|
||||
if self.provider == "gemini" and self.gemini_service_tier is None:
|
||||
from ..config import _get_raw_config
|
||||
|
||||
try:
|
||||
raw_config = _get_raw_config()
|
||||
self.gemini_service_tier = raw_config.llm_gemini_service_tier
|
||||
except Exception:
|
||||
pass # Config may not be initialized in test environments
|
||||
elif self.provider != "gemini":
|
||||
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(
|
||||
@@ -734,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
|
||||
@@ -790,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,
|
||||
@@ -807,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
|
||||
@@ -834,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,
|
||||
@@ -850,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(
|
||||
@@ -861,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):
|
||||
@@ -888,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,
|
||||
)
|
||||
@@ -911,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
|
||||
|
||||
@@ -921,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":
|
||||
@@ -936,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:
|
||||
@@ -948,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(
|
||||
@@ -961,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):
|
||||
@@ -985,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,
|
||||
)
|
||||
@@ -1008,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
|
||||
|
||||
@@ -1166,19 +1239,37 @@ 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,
|
||||
)
|
||||
|
||||
@@ -1196,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,
|
||||
@@ -1205,12 +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)
|
||||
@@ -5,13 +5,35 @@ 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:
|
||||
@@ -134,8 +156,9 @@ class MarkitdownParser(FileParser):
|
||||
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}'")
|
||||
@@ -153,6 +176,23 @@ 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."""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -26,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
|
||||
|
||||
@@ -414,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.
|
||||
@@ -258,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()
|
||||
@@ -287,9 +297,15 @@ class GeminiLLM(LLMInterface):
|
||||
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:
|
||||
@@ -307,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
|
||||
@@ -323,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
|
||||
|
||||
@@ -424,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)
|
||||
@@ -675,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
|
||||
@@ -762,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:
|
||||
|
||||
@@ -23,6 +23,7 @@ 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
|
||||
@@ -30,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.
|
||||
@@ -54,6 +71,7 @@ class LiteLLMLLM(LLMInterface):
|
||||
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)
|
||||
@@ -67,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:
|
||||
@@ -83,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,
|
||||
)
|
||||
@@ -127,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
|
||||
@@ -219,6 +253,10 @@ class LiteLLMLLM(LLMInterface):
|
||||
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
|
||||
@@ -249,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
|
||||
@@ -386,6 +425,14 @@ class LiteLLMLLM(LLMInterface):
|
||||
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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -37,6 +37,7 @@ 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, 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
|
||||
@@ -232,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."""
|
||||
|
||||
@@ -449,8 +465,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"deepseek",
|
||||
"volcano",
|
||||
"openrouter",
|
||||
"requesty",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
"fireworks",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
@@ -472,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.
|
||||
@@ -494,8 +516,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"minimax",
|
||||
"deepseek",
|
||||
"openrouter",
|
||||
"requesty",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
"ollama-cloud",
|
||||
)
|
||||
and not self.api_key
|
||||
@@ -770,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,
|
||||
@@ -823,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,
|
||||
@@ -840,12 +868,23 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
# 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()
|
||||
@@ -894,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
|
||||
@@ -1144,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(
|
||||
@@ -1186,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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -193,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(
|
||||
@@ -296,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(
|
||||
@@ -346,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,
|
||||
@@ -663,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
|
||||
@@ -766,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}",
|
||||
@@ -867,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
|
||||
|
||||
@@ -1203,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
|
||||
@@ -1314,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,
|
||||
@@ -1814,7 +1822,14 @@ async def extract_facts_from_text(
|
||||
total_usage = total_usage + chunk_usage
|
||||
|
||||
if failed_chunks:
|
||||
failed_summary = ", ".join(f"chunk {idx}: {type(err).__name__}" for idx, err in failed_chunks[:5])
|
||||
# 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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,9 +104,7 @@ 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).
|
||||
@@ -108,7 +119,7 @@ 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
|
||||
@@ -208,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
|
||||
@@ -293,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"
|
||||
|
||||
|
||||
@@ -307,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"
|
||||
|
||||
|
||||
|
||||
@@ -22,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
|
||||
@@ -833,10 +833,12 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
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:
|
||||
"""
|
||||
@@ -845,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
|
||||
@@ -853,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:
|
||||
@@ -873,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),
|
||||
@@ -884,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)
|
||||
|
||||
@@ -905,10 +919,12 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
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:
|
||||
@@ -916,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
|
||||
@@ -924,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()
|
||||
@@ -943,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),
|
||||
@@ -954,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)
|
||||
|
||||
@@ -1014,7 +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 tool_trace/llm_trace. Defaults to false because the trace can be tens of KB and overflow MCP client context; enable only for debugging.
|
||||
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:
|
||||
@@ -1045,11 +1073,14 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
if not include_based_on:
|
||||
result_data.pop("based_on", None)
|
||||
if not include_trace:
|
||||
# The agentic reflect loop's tool_trace/llm_trace can be tens of KB
|
||||
# (full mental-model text) and silently overflow MCP client context;
|
||||
# the REST API omits it by default too. Opt in via 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)
|
||||
@@ -1102,7 +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 tool_trace/llm_trace. Defaults to false because the trace can be tens of KB and overflow MCP client context; enable only for debugging.
|
||||
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()
|
||||
@@ -1132,11 +1163,14 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
if not include_based_on:
|
||||
result_data.pop("based_on", None)
|
||||
if not include_trace:
|
||||
# The agentic reflect loop's tool_trace/llm_trace can be tens of KB
|
||||
# (full mental-model text) and silently overflow MCP client context;
|
||||
# the REST API omits it by default too. Opt in via 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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api-slim"
|
||||
version = "0.8.3"
|
||||
version = "0.8.4"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -60,7 +60,7 @@ 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>=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.
|
||||
@@ -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}
|
||||
@@ -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)
|
||||
@@ -395,3 +395,63 @@ def test_rechunk_preserves_one_chunk_id_per_pre_chunk():
|
||||
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]}"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"})
|
||||
@@ -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(),
|
||||
)
|
||||
@@ -914,7 +948,7 @@ class TestPrecheckHttpWiring:
|
||||
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,7 +983,7 @@ 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,
|
||||
@@ -985,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)}
|
||||
|
||||
@@ -993,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}
|
||||
|
||||
@@ -1001,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}
|
||||
|
||||
@@ -1168,7 +1202,7 @@ class TestPrecheckHttpWiring:
|
||||
content_length = parsed
|
||||
|
||||
ctx = PrecheckContext(
|
||||
operation="retain",
|
||||
operation=PrecheckOperation.RETAIN,
|
||||
bank_id="bank-x",
|
||||
request_context=RequestContext(),
|
||||
content_length=content_length,
|
||||
|
||||
@@ -77,6 +77,37 @@ 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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -617,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."""
|
||||
|
||||
@@ -106,8 +106,6 @@ def test_gemini_llm_no_safety_settings_is_none():
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_applies_safety_settings():
|
||||
"""call() includes safety_settings in GenerateContentConfig when configured."""
|
||||
from google.genai import types as genai_types
|
||||
|
||||
provider = _make_gemini_provider(safety_settings=SAMPLE_SAFETY_SETTINGS)
|
||||
|
||||
# Build a fake successful response
|
||||
@@ -319,8 +317,12 @@ async def test_with_config_resets_after_call():
|
||||
# ─── LLMProvider reads safety settings from config ────────────────────────────
|
||||
|
||||
|
||||
def test_llm_provider_reads_safety_settings_from_config():
|
||||
"""LLMProvider reads llm_gemini_safety_settings from global config for Gemini provider."""
|
||||
def test_llm_provider_from_env_reads_safety_settings():
|
||||
"""from_env() resolves llm_gemini_safety_settings from the environment.
|
||||
|
||||
The constructor itself is config-free; the env-reading factory supplies the
|
||||
server default (the engine builds do the same from config).
|
||||
"""
|
||||
import json
|
||||
|
||||
from hindsight_api.config import ENV_LLM_GEMINI_SAFETY_SETTINGS, clear_config_cache
|
||||
@@ -338,12 +340,7 @@ def test_llm_provider_reads_safety_settings_from_config():
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
provider = LLMProvider(
|
||||
provider="gemini",
|
||||
api_key="fake-key",
|
||||
base_url="",
|
||||
model="gemini-2.5-flash",
|
||||
)
|
||||
provider = LLMProvider.from_env()
|
||||
|
||||
assert provider.gemini_safety_settings == SAMPLE_SAFETY_SETTINGS
|
||||
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
"""Reproduces the concurrent-insert deadlock on ``graph_maintenance_queue``
|
||||
that PR #2353 targets, and demonstrates that a shared insertion order cures it.
|
||||
|
||||
A deadlock is a *database*-level phenomenon, so unlike the PR's own tests (which
|
||||
only assert that the Python list handed to ``conn.execute`` is sorted) these run
|
||||
against the real Postgres test DB and drive two genuinely-concurrent
|
||||
transactions, forcing the exact interleaving that produces a lock cycle.
|
||||
|
||||
Modelling note
|
||||
--------------
|
||||
Production enqueues a whole victim set in ONE statement::
|
||||
|
||||
INSERT INTO graph_maintenance_queue (bank_id, unit_id)
|
||||
SELECT $1, v FROM unnest($2::uuid[]) ON CONFLICT (bank_id, unit_id) DO NOTHING
|
||||
|
||||
That single statement still takes the per-row unique-key locks one row at a time,
|
||||
in the order ``unnest`` yields — we just can't pause *inside* a single statement.
|
||||
So each worker here issues the rows one at a time with a barrier between them.
|
||||
That makes the otherwise-racy interleaving deterministic while exercising the
|
||||
identical lock: ``ON CONFLICT`` on the ``(bank_id, unit_id)`` primary key.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
from asyncpg.exceptions import DeadlockDetectedError
|
||||
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
# Two keys with an unambiguous sort order (low < high as UUIDs / as text).
|
||||
K_LOW = uuid.UUID("00000000-0000-4000-8000-000000000001")
|
||||
K_HIGH = uuid.UUID("ffffffff-ffff-4fff-8fff-ffffffffffff")
|
||||
|
||||
|
||||
async def _insert_one(conn, bank_id: str, unit_id: uuid.UUID) -> None:
|
||||
"""One row of the production INSERT ... ON CONFLICT DO NOTHING."""
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO graph_maintenance_queue (bank_id, unit_id)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT (bank_id, unit_id) DO NOTHING
|
||||
""",
|
||||
bank_id,
|
||||
unit_id,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unordered_concurrent_enqueue_deadlocks(memory: MemoryEngine):
|
||||
"""Two transactions inserting the same two keys in OPPOSITE orders deadlock.
|
||||
|
||||
This is the pre-fix reality: ``enqueue_relink_victims`` feeds whatever order
|
||||
``SELECT DISTINCT`` returns, so two overlapping victim sets can acquire the
|
||||
unique-key locks in opposite orders and cycle. Postgres aborts one with
|
||||
``DeadlockDetectedError``, which the API surfaces as a 500.
|
||||
"""
|
||||
pool = await memory._get_pool()
|
||||
bank_id = f"dl-bug-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Both transactions hold their first lock before either takes its second,
|
||||
# so the cross-wait (and thus the cycle) is guaranteed rather than racy.
|
||||
barrier = asyncio.Barrier(2)
|
||||
|
||||
async def worker(order: list[uuid.UUID]) -> None:
|
||||
async with pool.acquire() as conn:
|
||||
async with conn.transaction():
|
||||
await _insert_one(conn, bank_id, order[0])
|
||||
await barrier.wait()
|
||||
await _insert_one(conn, bank_id, order[1])
|
||||
|
||||
results = await asyncio.wait_for(
|
||||
asyncio.gather(
|
||||
worker([K_LOW, K_HIGH]),
|
||||
worker([K_HIGH, K_LOW]),
|
||||
return_exceptions=True,
|
||||
),
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
deadlocks = [r for r in results if isinstance(r, DeadlockDetectedError)]
|
||||
assert deadlocks, f"expected one transaction aborted with DeadlockDetectedError, got {results!r}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ordered_concurrent_enqueue_does_not_deadlock(memory: MemoryEngine):
|
||||
"""With both transactions inserting in the SAME (sorted) order — exactly what
|
||||
PR #2353's ``sorted(unit_ids)`` guarantees per call — there is no cycle. The
|
||||
second transaction simply waits on the first shared key and proceeds once the
|
||||
first commits; both victim sets land in the queue.
|
||||
"""
|
||||
pool = await memory._get_pool()
|
||||
bank_id = f"dl-fix-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
order = sorted([K_LOW, K_HIGH]) # identical order for both workers
|
||||
|
||||
async def worker() -> None:
|
||||
async with pool.acquire() as conn:
|
||||
async with conn.transaction():
|
||||
for uid in order:
|
||||
await _insert_one(conn, bank_id, uid)
|
||||
|
||||
# Sorted order cannot cycle; the timeout only guards against an unexpected hang.
|
||||
results = await asyncio.wait_for(
|
||||
asyncio.gather(worker(), worker(), return_exceptions=True),
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
errors = [r for r in results if isinstance(r, BaseException)]
|
||||
assert not errors, f"sorted concurrent inserts must not deadlock, got {results!r}"
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT unit_id FROM graph_maintenance_queue WHERE bank_id = $1 ORDER BY unit_id",
|
||||
bank_id,
|
||||
)
|
||||
assert [r["unit_id"] for r in rows] == order
|
||||
@@ -165,6 +165,12 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
assert "total_nodes" in stats
|
||||
assert stats["total_nodes"] > 0
|
||||
|
||||
# ?refresh=true forces a fresh recompute, bypassing the cache; same shape.
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/stats?refresh=true")
|
||||
assert response.status_code == 200
|
||||
fresh_stats = response.json()
|
||||
assert fresh_stats["total_nodes"] == stats["total_nodes"]
|
||||
|
||||
# Verify bank list returns stats (fact_count, last_document_at)
|
||||
response = await api_client.get("/v1/default/banks")
|
||||
assert response.status_code == 200
|
||||
@@ -1354,6 +1360,26 @@ async def test_patch_config_persists_override_for_uncreated_bank(api_client, fie
|
||||
assert response.json()["name"] == test_bank_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_patch_bank_does_not_create_missing_bank(api_client):
|
||||
"""PATCH /banks/{bank_id} updates existing banks only."""
|
||||
test_bank_id = f"patch_missing_bank_{datetime.now().timestamp()}"
|
||||
|
||||
response = await api_client.patch(
|
||||
f"/v1/default/banks/{test_bank_id}",
|
||||
json={"name": "Should Not Exist"},
|
||||
)
|
||||
assert response.status_code == 404, response.text
|
||||
assert response.json()["detail"] == f"Bank '{test_bank_id}' not found"
|
||||
|
||||
profile = await api_client.get(f"/v1/default/banks/{test_bank_id}/profile")
|
||||
assert profile.status_code == 404, profile.text
|
||||
|
||||
banks = await api_client.get("/v1/default/banks")
|
||||
assert banks.status_code == 200, banks.text
|
||||
assert test_bank_id not in {bank["bank_id"] for bank in banks.json()["banks"]}
|
||||
|
||||
|
||||
@pytest.mark.hs_llm_core
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_api_workflow_llm_quality(api_client_real_llm):
|
||||
|
||||
@@ -0,0 +1,241 @@
|
||||
"""HTTP + engine integration tests for the knowledge base (folders + pages).
|
||||
|
||||
Pages are seeded directly via the engine (deterministic content, no LLM) so the
|
||||
tree, OKF projection, move/rename, and cascade-delete behaviour can be asserted
|
||||
without consolidation.
|
||||
"""
|
||||
|
||||
import urllib.parse
|
||||
import uuid
|
||||
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
|
||||
def _enc(bank_id: str) -> str:
|
||||
return urllib.parse.quote(bank_id, safe="")
|
||||
|
||||
|
||||
class _Seed:
|
||||
"""Holds the ids created by the seed fixture for assertions."""
|
||||
|
||||
def __init__(self, **ids):
|
||||
self.__dict__.update(ids)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def kb_bank(memory: MemoryEngine, request_context):
|
||||
"""A bank with folders, nested folders, and pages."""
|
||||
bank_id = f"test-kb-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
runbooks = await memory.create_knowledge_folder(bank_id, "Runbooks", request_context=request_context)
|
||||
policies = await memory.create_knowledge_folder(bank_id, "Policies", request_context=request_context)
|
||||
sub = await memory.create_knowledge_folder(
|
||||
bank_id, "Sub", parent_id=runbooks["id"], request_context=request_context
|
||||
)
|
||||
orders = await memory.create_knowledge_page(
|
||||
bank_id,
|
||||
"Orders",
|
||||
"What are the order facts?",
|
||||
"# Orders\n\nOne row per order.",
|
||||
parent_id=runbooks["id"],
|
||||
tags=["type:runbook", "sales", "revenue"],
|
||||
request_context=request_context,
|
||||
)
|
||||
billing = await memory.create_knowledge_page(
|
||||
bank_id,
|
||||
"Billing",
|
||||
"What is the billing policy?",
|
||||
"# Billing\n\nNet-30.",
|
||||
parent_id=policies["id"],
|
||||
tags=["type:policy", "revenue"],
|
||||
request_context=request_context,
|
||||
)
|
||||
loose = await memory.create_knowledge_page(
|
||||
bank_id,
|
||||
"Loose",
|
||||
"A root page.",
|
||||
"# Loose\n\nNo folder, no tags.",
|
||||
tags=[],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
yield (
|
||||
bank_id,
|
||||
_Seed(
|
||||
runbooks=runbooks["id"],
|
||||
policies=policies["id"],
|
||||
sub=sub["id"],
|
||||
orders=orders["id"],
|
||||
billing=billing["id"],
|
||||
loose=loose["id"],
|
||||
orders_mm=orders["mental_model_id"],
|
||||
),
|
||||
)
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
class TestTree:
|
||||
async def test_nested_tree(self, api_client, kb_bank):
|
||||
bank_id, ids = kb_bank
|
||||
resp = await api_client.get(f"/v1/default/banks/{_enc(bank_id)}/knowledge-base/tree")
|
||||
assert resp.status_code == 200, resp.text
|
||||
roots = {r["name"]: r for r in resp.json()["roots"]}
|
||||
assert set(roots) == {"Runbooks", "Policies", "Loose"}
|
||||
|
||||
runbooks = roots["Runbooks"]
|
||||
assert runbooks["kind"] == "folder"
|
||||
child_names = {c["name"] for c in runbooks["children"]}
|
||||
assert child_names == {"Sub", "Orders"}
|
||||
|
||||
orders = next(c for c in runbooks["children"] if c["name"] == "Orders")
|
||||
assert orders["kind"] == "page"
|
||||
# Human-created pages are pinned (not curator-managed).
|
||||
assert orders["managed"] is False
|
||||
assert "sales" in orders["tags"]
|
||||
assert roots["Loose"]["kind"] == "page"
|
||||
|
||||
|
||||
class TestPageDefaults:
|
||||
"""A knowledge page is a living document by default: observation-only, delta,
|
||||
auto-refreshing, with a larger token budget than a plain mental model."""
|
||||
|
||||
async def test_default_trigger_and_max_tokens(self, memory: MemoryEngine, request_context):
|
||||
bank_id = f"test-kb-def-{uuid.uuid4().hex[:8]}"
|
||||
page = await memory.create_knowledge_page(
|
||||
bank_id, "P", "What is P?", "seed", request_context=request_context
|
||||
)
|
||||
mm = await memory.get_mental_model(bank_id, page["mental_model_id"], request_context=request_context)
|
||||
assert mm["trigger"] == {
|
||||
"mode": "delta",
|
||||
"fact_types": ["observation"],
|
||||
"exclude_mental_models": True,
|
||||
"refresh_after_consolidation": True,
|
||||
}
|
||||
assert mm["max_tokens"] == 4096
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_client_trigger_and_max_tokens_override_defaults(self, memory: MemoryEngine, request_context):
|
||||
bank_id = f"test-kb-ovr-{uuid.uuid4().hex[:8]}"
|
||||
page = await memory.create_knowledge_page(
|
||||
bank_id,
|
||||
"P",
|
||||
"What is P?",
|
||||
"seed",
|
||||
trigger={"mode": "full", "refresh_after_consolidation": False},
|
||||
max_tokens=1024,
|
||||
request_context=request_context,
|
||||
)
|
||||
mm = await memory.get_mental_model(bank_id, page["mental_model_id"], request_context=request_context)
|
||||
assert mm["trigger"]["mode"] == "full"
|
||||
assert mm["trigger"].get("refresh_after_consolidation") is False
|
||||
assert mm["max_tokens"] == 1024
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
class TestGetPage:
|
||||
async def test_okf_document(self, api_client, kb_bank):
|
||||
bank_id, ids = kb_bank
|
||||
resp = await api_client.get(f"/v1/default/banks/{_enc(bank_id)}/knowledge-base/pages/{ids.orders}")
|
||||
assert resp.status_code == 200, resp.text
|
||||
page = resp.json()
|
||||
assert page["type"] == "runbook"
|
||||
assert page["body"].startswith("# Orders")
|
||||
assert page["markdown"].startswith("---\n")
|
||||
assert 'type: "runbook"' in page["markdown"]
|
||||
|
||||
async def test_missing_page_404(self, api_client, kb_bank):
|
||||
bank_id, ids = kb_bank
|
||||
resp = await api_client.get(f"/v1/default/banks/{_enc(bank_id)}/knowledge-base/pages/nope")
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
class TestCreate:
|
||||
async def test_create_folder(self, api_client, kb_bank):
|
||||
bank_id, ids = kb_bank
|
||||
resp = await api_client.post(
|
||||
f"/v1/default/banks/{_enc(bank_id)}/knowledge-base/folders",
|
||||
json={"name": "Guides", "parent_id": None},
|
||||
)
|
||||
assert resp.status_code == 201, resp.text
|
||||
assert resp.json()["kind"] == "folder"
|
||||
assert resp.json()["name"] == "Guides"
|
||||
|
||||
async def test_create_folder_bad_parent(self, api_client, kb_bank):
|
||||
bank_id, ids = kb_bank
|
||||
# parent that is a page, not a folder → 400
|
||||
resp = await api_client.post(
|
||||
f"/v1/default/banks/{_enc(bank_id)}/knowledge-base/folders",
|
||||
json={"name": "Nope", "parent_id": ids.orders},
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
class TestGraphAndExport:
|
||||
async def test_graph_shared_tag_edge(self, api_client, kb_bank):
|
||||
bank_id, ids = kb_bank
|
||||
resp = await api_client.get(f"/v1/default/banks/{_enc(bank_id)}/knowledge-base/graph")
|
||||
assert resp.status_code == 200, resp.text
|
||||
data = resp.json()
|
||||
assert data["total_pages"] == 3
|
||||
# orders & billing share "revenue"; loose has no tags
|
||||
assert data["total_edges"] == 1
|
||||
edge = data["edges"][0]["data"]
|
||||
assert {edge["source"], edge["target"]} == {ids.orders, ids.billing}
|
||||
assert edge["sharedTags"] == ["revenue"]
|
||||
|
||||
async def test_export_bundle_nested_index(self, api_client, kb_bank):
|
||||
bank_id, ids = kb_bank
|
||||
resp = await api_client.get(f"/v1/default/banks/{_enc(bank_id)}/knowledge-base/export")
|
||||
assert resp.status_code == 200, resp.text
|
||||
files = {f["path"]: f["content"] for f in resp.json()["files"]}
|
||||
assert "index.md" in files
|
||||
assert f"{ids.orders}.md" in files
|
||||
# index reflects the folder hierarchy
|
||||
assert "**Runbooks/**" in files["index.md"]
|
||||
assert "One row per order." in files[f"{ids.orders}.md"]
|
||||
|
||||
|
||||
class TestMoveRenameDelete:
|
||||
async def test_rename(self, api_client, kb_bank):
|
||||
bank_id, ids = kb_bank
|
||||
resp = await api_client.patch(
|
||||
f"/v1/default/banks/{_enc(bank_id)}/knowledge-base/nodes/{ids.policies}",
|
||||
json={"name": "Compliance"},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert resp.json()["name"] == "Compliance"
|
||||
|
||||
async def test_move_into_folder(self, api_client, kb_bank):
|
||||
bank_id, ids = kb_bank
|
||||
# move the Loose root page under Policies
|
||||
resp = await api_client.patch(
|
||||
f"/v1/default/banks/{_enc(bank_id)}/knowledge-base/nodes/{ids.loose}",
|
||||
json={"parent_id": ids.policies},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert resp.json()["parent_id"] == ids.policies
|
||||
|
||||
async def test_move_cycle_rejected(self, api_client, kb_bank):
|
||||
bank_id, ids = kb_bank
|
||||
# moving Runbooks under its own descendant Sub must fail
|
||||
resp = await api_client.patch(
|
||||
f"/v1/default/banks/{_enc(bank_id)}/knowledge-base/nodes/{ids.runbooks}",
|
||||
json={"parent_id": ids.sub},
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
async def test_delete_folder_cascades(self, api_client, kb_bank, memory, request_context):
|
||||
bank_id, ids = kb_bank
|
||||
# deleting Runbooks removes Sub + Orders (and Orders' mental model)
|
||||
resp = await api_client.delete(f"/v1/default/banks/{_enc(bank_id)}/knowledge-base/nodes/{ids.runbooks}")
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
||||
tree = (await api_client.get(f"/v1/default/banks/{_enc(bank_id)}/knowledge-base/tree")).json()
|
||||
root_names = {r["name"] for r in tree["roots"]}
|
||||
assert "Runbooks" not in root_names
|
||||
# the backing mental model is gone too
|
||||
mm = await memory.get_mental_model(bank_id, ids.orders_mm, request_context=request_context)
|
||||
assert mm is None
|
||||
@@ -5,6 +5,7 @@ import pytest
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from hindsight_api.config import clear_config_cache
|
||||
from hindsight_api.engine.retain.link_utils import (
|
||||
_normalize_datetime,
|
||||
_cap_links_per_unit,
|
||||
@@ -408,6 +409,19 @@ class TestComputeSemanticLinksAnnPgBouncerSafety:
|
||||
following the CREATE TEMP TABLE.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_config_cache(self):
|
||||
# Tests below monkeypatch HINDSIGHT_API_VECTOR_EXTENSION. The ANN code
|
||||
# path reads it through the process-global config cache, and monkeypatch
|
||||
# reverts only the env var — not the cache. Left uncleared, a leaked
|
||||
# "vchord" makes every later bank-creating test on the same xdist worker
|
||||
# emit `USING vchordrq` against the pgvector-only test DB and fail with
|
||||
# `access method "vchordrq" does not exist`. Clear before and after so
|
||||
# the cache is rebuilt from the current env for each test.
|
||||
clear_config_cache()
|
||||
yield
|
||||
clear_config_cache()
|
||||
|
||||
@pytest.fixture
|
||||
def mock_conn(self):
|
||||
"""An asyncpg-like connection mock with an async `transaction()`
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
"""Regression: user-facing GET list endpoints must reject negative limit/offset
|
||||
with a clean 422 at the FastAPI boundary instead of letting the value reach
|
||||
Postgres (``LIMIT/OFFSET must not be negative``) and surfacing as an opaque 500
|
||||
that also leaks the raw Postgres error string.
|
||||
|
||||
This makes pagination validation consistent with the sibling list endpoints in
|
||||
the same router (document-chunks / directives / async-ops / audit) that already
|
||||
declare ``Query(..., ge=...)``. The engine emits ``LIMIT $n OFFSET $n`` with no
|
||||
``max(0, ...)`` clamp, so the guard has to live at the request boundary.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client(memory):
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
def _url(bank_id: str, suffix: str) -> str:
|
||||
return f"/v1/default/banks/{bank_id}/{suffix}"
|
||||
|
||||
|
||||
# Endpoints that accept a ``limit`` query param.
|
||||
LIMIT_ENDPOINTS = [
|
||||
"graph",
|
||||
"memories/list",
|
||||
"entities",
|
||||
"entities/graph",
|
||||
"documents",
|
||||
"tags",
|
||||
]
|
||||
|
||||
# Subset that also accept an ``offset`` query param.
|
||||
OFFSET_ENDPOINTS = [
|
||||
"memories/list",
|
||||
"entities",
|
||||
"documents",
|
||||
"tags",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("suffix", LIMIT_ENDPOINTS)
|
||||
async def test_negative_limit_returns_422_not_500(api_client, suffix):
|
||||
bank_id = f"pag-{uuid.uuid4().hex[:8]}"
|
||||
resp = await api_client.get(_url(bank_id, suffix), params={"limit": -1})
|
||||
# FastAPI validation runs before the handler / DB, so a bad pagination input
|
||||
# is a clean 422 — never a 500 leaking the raw Postgres error.
|
||||
assert resp.status_code == 422, resp.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("suffix", OFFSET_ENDPOINTS)
|
||||
async def test_negative_offset_returns_422_not_500(api_client, suffix):
|
||||
bank_id = f"pag-{uuid.uuid4().hex[:8]}"
|
||||
resp = await api_client.get(_url(bank_id, suffix), params={"offset": -1})
|
||||
assert resp.status_code == 422, resp.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("limit", [0, 1, 100])
|
||||
async def test_valid_limit_is_accepted_including_zero(api_client, memory, limit):
|
||||
# Positive control: ge=0 rejects only NEGATIVE limits. A non-negative limit
|
||||
# — including limit=0 (a valid empty page, LIMIT 0) — must still be accepted,
|
||||
# so this fix does not change behavior for any previously-valid input.
|
||||
bank_id = f"pag-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=RequestContext())
|
||||
resp = await api_client.get(_url(bank_id, "memories/list"), params={"limit": limit, "offset": 0})
|
||||
assert resp.status_code == 200, resp.text
|
||||
@@ -22,6 +22,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
||||
import pytest
|
||||
|
||||
EXTRA_BODY = {"temperature": 0.2, "top_p": 0.9}
|
||||
DEFAULT_HEADERS = {"X-Component-Id": "hindsight", "X-Trace": "abc"}
|
||||
|
||||
|
||||
# ─── config / env parsing ─────────────────────────────────────────────────────
|
||||
@@ -205,10 +206,150 @@ async def test_gemini_extra_body_service_tier_takes_precedence():
|
||||
assert provider._extra_body["http_options"]["extra_body"]["service_tier"] == "standard"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_structured_call_uses_native_schema_without_prompt_duplicate():
|
||||
"""Structured Gemini calls send schema through response_schema only."""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class StructuredAnswer(BaseModel):
|
||||
answer: str
|
||||
|
||||
provider = _make_gemini_provider()
|
||||
response = _fake_gemini_response()
|
||||
response.text = '{"answer": "ok"}'
|
||||
provider._client.aio.models.generate_content = AsyncMock(return_value=response)
|
||||
|
||||
result = await provider.call(
|
||||
messages=[
|
||||
{"role": "system", "content": "Return concise JSON."},
|
||||
{"role": "user", "content": "hello"},
|
||||
],
|
||||
response_format=StructuredAnswer,
|
||||
scope="test",
|
||||
)
|
||||
|
||||
config_arg = provider._client.aio.models.generate_content.call_args.kwargs.get("config")
|
||||
assert result.answer == "ok"
|
||||
assert config_arg.response_mime_type == "application/json"
|
||||
assert config_arg.response_schema is StructuredAnswer
|
||||
assert config_arg.system_instruction == "Return concise JSON."
|
||||
assert "valid JSON matching this schema" not in config_arg.system_instruction
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_cached_structured_call_keeps_native_schema():
|
||||
"""Cached Gemini calls still send response_schema per request."""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class StructuredAnswer(BaseModel):
|
||||
answer: str
|
||||
|
||||
provider = _make_gemini_provider()
|
||||
response = _fake_gemini_response()
|
||||
response.text = '{"answer": "ok"}'
|
||||
provider._client.aio.models.generate_content = AsyncMock(return_value=response)
|
||||
|
||||
result = await provider.call(
|
||||
messages=[
|
||||
{"role": "system", "content": "Return concise JSON."},
|
||||
{"role": "user", "content": "hello"},
|
||||
],
|
||||
response_format=StructuredAnswer,
|
||||
cached_prefix="cachedContents/test",
|
||||
scope="test",
|
||||
)
|
||||
|
||||
config_arg = provider._client.aio.models.generate_content.call_args.kwargs.get("config")
|
||||
assert result.answer == "ok"
|
||||
assert config_arg.cached_content == "cachedContents/test"
|
||||
assert config_arg.system_instruction is None
|
||||
assert config_arg.response_mime_type == "application/json"
|
||||
assert config_arg.response_schema is StructuredAnswer
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_structured_parse_failure_falls_back_to_prompt_schema():
|
||||
"""Malformed native-schema output gets one prompt-schema compatibility retry."""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class StructuredAnswer(BaseModel):
|
||||
answer: str
|
||||
|
||||
provider = _make_gemini_provider()
|
||||
invalid = _fake_gemini_response()
|
||||
invalid.text = "not json"
|
||||
valid = _fake_gemini_response()
|
||||
valid.text = '{"answer": "ok"}'
|
||||
provider._client.aio.models.generate_content = AsyncMock(side_effect=[invalid, valid])
|
||||
|
||||
result = await provider.call(
|
||||
messages=[
|
||||
{"role": "system", "content": "Return concise JSON."},
|
||||
{"role": "user", "content": "hello"},
|
||||
],
|
||||
response_format=StructuredAnswer,
|
||||
scope="test",
|
||||
max_retries=1,
|
||||
initial_backoff=0,
|
||||
max_backoff=0,
|
||||
)
|
||||
|
||||
first_config = provider._client.aio.models.generate_content.call_args_list[0].kwargs["config"]
|
||||
fallback_config = provider._client.aio.models.generate_content.call_args_list[1].kwargs["config"]
|
||||
|
||||
assert result.answer == "ok"
|
||||
assert first_config.response_schema is StructuredAnswer
|
||||
assert first_config.system_instruction == "Return concise JSON."
|
||||
assert fallback_config.response_schema is None
|
||||
assert fallback_config.response_mime_type is None
|
||||
assert fallback_config.system_instruction.startswith("Return concise JSON.")
|
||||
assert "valid JSON matching this schema" in fallback_config.system_instruction
|
||||
assert '"answer"' in fallback_config.system_instruction
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_cached_parse_retry_keeps_cached_native_schema():
|
||||
"""Cached structured retries keep cache context instead of switching prompts."""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class StructuredAnswer(BaseModel):
|
||||
answer: str
|
||||
|
||||
provider = _make_gemini_provider()
|
||||
invalid = _fake_gemini_response()
|
||||
invalid.text = "not json"
|
||||
valid = _fake_gemini_response()
|
||||
valid.text = '{"answer": "ok"}'
|
||||
provider._client.aio.models.generate_content = AsyncMock(side_effect=[invalid, valid])
|
||||
|
||||
result = await provider.call(
|
||||
messages=[
|
||||
{"role": "system", "content": "Return concise JSON."},
|
||||
{"role": "user", "content": "hello"},
|
||||
],
|
||||
response_format=StructuredAnswer,
|
||||
cached_prefix="cachedContents/test",
|
||||
scope="test",
|
||||
max_retries=1,
|
||||
initial_backoff=0,
|
||||
max_backoff=0,
|
||||
)
|
||||
|
||||
first_config = provider._client.aio.models.generate_content.call_args_list[0].kwargs["config"]
|
||||
retry_config = provider._client.aio.models.generate_content.call_args_list[1].kwargs["config"]
|
||||
|
||||
assert result.answer == "ok"
|
||||
assert first_config.cached_content == "cachedContents/test"
|
||||
assert retry_config.cached_content == "cachedContents/test"
|
||||
assert retry_config.response_schema is StructuredAnswer
|
||||
assert retry_config.response_mime_type == "application/json"
|
||||
assert retry_config.system_instruction is None
|
||||
|
||||
|
||||
# ─── LiteLLM ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _make_litellm_provider(extra_body=None):
|
||||
def _make_litellm_provider(extra_body=None, default_headers=None):
|
||||
pytest.importorskip("litellm")
|
||||
from hindsight_api.engine.providers.litellm_llm import LiteLLMLLM
|
||||
|
||||
@@ -218,6 +359,7 @@ def _make_litellm_provider(extra_body=None):
|
||||
base_url="",
|
||||
model="gpt-4o",
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
|
||||
@@ -279,3 +421,107 @@ def test_litellm_router_forwards_extra_body():
|
||||
extra_body=EXTRA_BODY,
|
||||
)
|
||||
assert provider._extra_body == EXTRA_BODY
|
||||
|
||||
|
||||
def test_litellm_stores_default_headers():
|
||||
provider = _make_litellm_provider(default_headers=DEFAULT_HEADERS)
|
||||
assert provider._default_headers == DEFAULT_HEADERS
|
||||
|
||||
|
||||
def test_litellm_empty_default_headers_defaults_to_dict():
|
||||
provider = _make_litellm_provider(default_headers=None)
|
||||
assert provider._default_headers == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_call_passes_default_headers_as_extra_headers():
|
||||
"""``call()`` forwards default_headers to acompletion via ``extra_headers``."""
|
||||
provider = _make_litellm_provider(default_headers=DEFAULT_HEADERS)
|
||||
provider._acompletion = AsyncMock(return_value=_fake_litellm_response())
|
||||
|
||||
with patch("hindsight_api.engine.providers.litellm_llm.get_metrics_collector"):
|
||||
await provider.call(messages=[{"role": "user", "content": "hi"}], scope="test", max_retries=0)
|
||||
|
||||
assert provider._acompletion.call_args.kwargs.get("extra_headers") == DEFAULT_HEADERS
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_no_default_headers_omits_extra_headers():
|
||||
"""``call()`` does not pass ``extra_headers`` when none are configured."""
|
||||
provider = _make_litellm_provider(default_headers=None)
|
||||
provider._acompletion = AsyncMock(return_value=_fake_litellm_response())
|
||||
|
||||
with patch("hindsight_api.engine.providers.litellm_llm.get_metrics_collector"):
|
||||
await provider.call(messages=[{"role": "user", "content": "hi"}], scope="test", max_retries=0)
|
||||
|
||||
assert "extra_headers" not in provider._acompletion.call_args.kwargs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_default_headers_passed_as_fresh_copy():
|
||||
"""Each call gets its own ``extra_headers`` copy so downstream mutation can't
|
||||
contaminate the stored headers or other requests."""
|
||||
provider = _make_litellm_provider(default_headers=DEFAULT_HEADERS)
|
||||
provider._acompletion = AsyncMock(return_value=_fake_litellm_response())
|
||||
|
||||
with patch("hindsight_api.engine.providers.litellm_llm.get_metrics_collector"):
|
||||
await provider.call(messages=[{"role": "user", "content": "hi"}], scope="test", max_retries=0)
|
||||
|
||||
passed = provider._acompletion.call_args.kwargs["extra_headers"]
|
||||
assert passed == DEFAULT_HEADERS
|
||||
assert passed is not provider._default_headers
|
||||
passed["X-Injected"] = "1"
|
||||
assert "X-Injected" not in provider._default_headers
|
||||
|
||||
|
||||
def test_litellm_default_headers_copied_from_caller_dict():
|
||||
"""A caller-owned dict cannot be mutated through the provider."""
|
||||
caller_dict = {"X-Component-Id": "hindsight"}
|
||||
provider = _make_litellm_provider(default_headers=caller_dict)
|
||||
caller_dict["X-Mutated"] = "1"
|
||||
assert "X-Mutated" not in provider._default_headers
|
||||
|
||||
|
||||
def _make_litellm_router_provider(default_headers=None):
|
||||
pytest.importorskip("litellm")
|
||||
from hindsight_api.engine.providers.litellm_router_llm import LiteLLMRouterLLM
|
||||
|
||||
config = {"model_list": [{"model_name": "default", "litellm_params": {"model": "gpt-4o", "api_key": "x"}}]}
|
||||
return LiteLLMRouterLLM(
|
||||
provider="litellmrouter",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="default",
|
||||
config=config,
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
|
||||
def test_litellm_router_stores_default_headers():
|
||||
provider = _make_litellm_router_provider(default_headers=DEFAULT_HEADERS)
|
||||
assert provider._default_headers == DEFAULT_HEADERS
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_router_call_passes_default_headers_as_extra_headers():
|
||||
"""The Router's ``_build_common_kwargs`` override must also forward default_headers
|
||||
as ``extra_headers`` — storage alone doesn't reach the provider behind the Router."""
|
||||
provider = _make_litellm_router_provider(default_headers=DEFAULT_HEADERS)
|
||||
provider._acompletion = AsyncMock(return_value=_fake_litellm_response())
|
||||
|
||||
with patch("hindsight_api.engine.providers.litellm_llm.get_metrics_collector"):
|
||||
await provider.call(messages=[{"role": "user", "content": "hi"}], scope="test", max_retries=0)
|
||||
|
||||
assert provider._acompletion.call_args.kwargs.get("extra_headers") == DEFAULT_HEADERS
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_router_no_default_headers_omits_extra_headers():
|
||||
"""The Router omits ``extra_headers`` entirely when none are configured."""
|
||||
provider = _make_litellm_router_provider(default_headers=None)
|
||||
provider._acompletion = AsyncMock(return_value=_fake_litellm_response())
|
||||
|
||||
with patch("hindsight_api.engine.providers.litellm_llm.get_metrics_collector"):
|
||||
await provider.call(messages=[{"role": "user", "content": "hi"}], scope="test", max_retries=0)
|
||||
|
||||
assert "extra_headers" not in provider._acompletion.call_args.kwargs
|
||||
|
||||
@@ -45,11 +45,22 @@ def _get_api_key() -> str:
|
||||
|
||||
|
||||
def _make_llm() -> LLMProvider:
|
||||
# LLMProvider uses provider-specific settings as-passed (it does not resolve
|
||||
# them from global config), so forward the ones whose providers require them:
|
||||
# Vertex AI needs project/region, and litellmrouter needs its router config.
|
||||
# Without these, provider=vertexai/litellmrouter raise at construction.
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
config = get_config()
|
||||
return LLMProvider(
|
||||
provider=_PROVIDER,
|
||||
api_key=_get_api_key(),
|
||||
base_url=os.environ.get("HINDSIGHT_API_LLM_BASE_URL", ""),
|
||||
model=_MODEL,
|
||||
vertexai_project_id=config.llm_vertexai_project_id,
|
||||
vertexai_region=config.llm_vertexai_region,
|
||||
vertexai_service_account_key=config.llm_vertexai_service_account_key,
|
||||
litellmrouter_config=config.llm_litellmrouter_config,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -194,6 +194,7 @@ def _make_router_provider(config: dict[str, Any], mock_router: Any) -> LiteLLMRo
|
||||
provider.model = "unused"
|
||||
provider.reasoning_effort = "low"
|
||||
provider.timeout = 300.0
|
||||
provider._default_headers = {}
|
||||
provider.config = config
|
||||
provider._litellm = fake_litellm
|
||||
provider._router = mock_router
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Tests for per-operation LLM temperature configuration from environment variables.
|
||||
|
||||
Covers the resolution order (per-operation env -> global env -> built-in default)
|
||||
and the "omit" sentinels that drop the temperature parameter for models that reject
|
||||
explicit temperatures (e.g. Azure gpt-5.5 -- see issue #2459).
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config import HindsightConfig, _parse_temperature
|
||||
|
||||
_OP_FIELDS = {
|
||||
"HINDSIGHT_API_LLM_TEMPERATURE_VERIFICATION": ("llm_temperature_verification", 0.0),
|
||||
"HINDSIGHT_API_LLM_TEMPERATURE_RETAIN": ("llm_temperature_retain", 0.1),
|
||||
"HINDSIGHT_API_LLM_TEMPERATURE_REFLECT": ("llm_temperature_reflect", 0.9),
|
||||
"HINDSIGHT_API_LLM_TEMPERATURE_CONSOLIDATION": ("llm_temperature_consolidation", 0.0),
|
||||
}
|
||||
|
||||
|
||||
def _clear_temperature_env(monkeypatch) -> None:
|
||||
monkeypatch.delenv("HINDSIGHT_API_LLM_TEMPERATURE", raising=False)
|
||||
for env_name in _OP_FIELDS:
|
||||
monkeypatch.delenv(env_name, raising=False)
|
||||
|
||||
|
||||
def test_defaults_preserve_historical_values(monkeypatch):
|
||||
_clear_temperature_env(monkeypatch)
|
||||
config = HindsightConfig.from_env()
|
||||
for _, (field, default) in _OP_FIELDS.items():
|
||||
assert getattr(config, field) == default
|
||||
|
||||
|
||||
def test_global_override_applies_to_all_operations(monkeypatch):
|
||||
_clear_temperature_env(monkeypatch)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_TEMPERATURE", "0.2")
|
||||
config = HindsightConfig.from_env()
|
||||
for _, (field, _default) in _OP_FIELDS.items():
|
||||
assert getattr(config, field) == 0.2
|
||||
|
||||
|
||||
def test_global_none_omits_temperature_everywhere(monkeypatch):
|
||||
_clear_temperature_env(monkeypatch)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_TEMPERATURE", "none")
|
||||
config = HindsightConfig.from_env()
|
||||
for _, (field, _default) in _OP_FIELDS.items():
|
||||
assert getattr(config, field) is None
|
||||
|
||||
|
||||
def test_per_operation_override_beats_global(monkeypatch):
|
||||
_clear_temperature_env(monkeypatch)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_TEMPERATURE", "none")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_TEMPERATURE_RETAIN", "0.5")
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_temperature_retain == 0.5
|
||||
# Other operations still follow the global "none" (omit).
|
||||
assert config.llm_temperature_reflect is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sentinel", ["none", "NONE", "default", "off", "unset", "", " "])
|
||||
def test_omit_sentinels(sentinel):
|
||||
assert _parse_temperature(sentinel) is None
|
||||
|
||||
|
||||
def test_parse_temperature_rejects_out_of_range():
|
||||
with pytest.raises(ValueError):
|
||||
_parse_temperature("2.5")
|
||||
with pytest.raises(ValueError):
|
||||
_parse_temperature("-0.1")
|
||||
|
||||
|
||||
def test_parse_temperature_rejects_non_numeric():
|
||||
with pytest.raises(ValueError):
|
||||
_parse_temperature("warm")
|
||||
@@ -0,0 +1,99 @@
|
||||
"""End-to-end checks that per-operation temperature reaches the LLM call.
|
||||
|
||||
These drive the real pipeline with the mock LLM provider (which records the
|
||||
``temperature`` it receives) and assert that each operation forwards the
|
||||
configured value -- including ``None``, which omits the parameter for models
|
||||
that reject explicit temperatures (issue #2459).
|
||||
|
||||
Config resolution itself is unit-tested in ``test_llm_temperature_env.py``;
|
||||
here we verify the value is actually threaded through to ``provider.call()``.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config import clear_config_cache
|
||||
from hindsight_api.engine.search import think_utils
|
||||
|
||||
|
||||
def _calls_for_scope(memory, scope: str) -> list[dict]:
|
||||
"""Collect mock call records for a scope across the engine's LLM configs.
|
||||
|
||||
retain/reflect/consolidation each wrap a distinct provider instance, so a
|
||||
given scope only lands on one of them; gather from all and filter.
|
||||
"""
|
||||
seen_impls: dict[int, object] = {}
|
||||
for config in (
|
||||
memory._llm_config,
|
||||
memory._retain_llm_config,
|
||||
memory._reflect_llm_config,
|
||||
memory._consolidation_llm_config,
|
||||
):
|
||||
impl = config._provider_impl
|
||||
seen_impls[id(impl)] = impl
|
||||
|
||||
calls: list[dict] = []
|
||||
for impl in seen_impls.values():
|
||||
calls.extend(impl.get_mock_calls())
|
||||
return [c for c in calls if c.get("scope") == scope]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_forwards_configured_temperature(memory, request_context):
|
||||
"""Retain's fact extraction must call the LLM with the retain temperature (0.1 default)."""
|
||||
bank_id = f"test_temp_retain_{datetime.now(timezone.utc).timestamp()}"
|
||||
try:
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice is a senior engineer at TechCorp. She works on distributed systems.",
|
||||
context="team overview",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
extract_calls = _calls_for_scope(memory, "retain_extract_facts")
|
||||
assert extract_calls, "retain should have made a fact-extraction LLM call"
|
||||
assert all(c["temperature"] == 0.1 for c in extract_calls)
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_think_forwards_configured_temperature(memory):
|
||||
"""The reflect 'thinking' path must call the LLM with the reflect temperature (0.9 default)."""
|
||||
reflect_config = memory._reflect_llm_config
|
||||
reflect_config._provider_impl.clear_mock_calls()
|
||||
|
||||
await think_utils.reflect(
|
||||
llm_config=reflect_config,
|
||||
query="What does Alice work on?",
|
||||
world_facts=["Alice works on distributed systems."],
|
||||
)
|
||||
|
||||
think_calls = [c for c in reflect_config._provider_impl.get_mock_calls() if c["scope"] == "memory_think"]
|
||||
assert think_calls, "reflect should have made a memory_think LLM call"
|
||||
assert all(c["temperature"] == 0.9 for c in think_calls)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_none_omits_temperature_on_real_call(memory, monkeypatch):
|
||||
"""HINDSIGHT_API_LLM_TEMPERATURE=none must omit (None) the temperature on a live call."""
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_TEMPERATURE", "none")
|
||||
clear_config_cache()
|
||||
try:
|
||||
reflect_config = memory._reflect_llm_config
|
||||
reflect_config._provider_impl.clear_mock_calls()
|
||||
|
||||
await think_utils.reflect(
|
||||
llm_config=reflect_config,
|
||||
query="What does Alice work on?",
|
||||
world_facts=["Alice works on distributed systems."],
|
||||
)
|
||||
|
||||
think_calls = [c for c in reflect_config._provider_impl.get_mock_calls() if c["scope"] == "memory_think"]
|
||||
assert think_calls, "reflect should have made a memory_think LLM call"
|
||||
assert all(c["temperature"] is None for c in think_calls), "temperature should be omitted"
|
||||
finally:
|
||||
# Restore the cached config so later tests see default temperatures.
|
||||
clear_config_cache()
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user