Compare commits
92
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
24ee73a304 | ||
|
|
63a92bef5f | ||
|
|
1533c0915d | ||
|
|
8eb2937cdb | ||
|
|
d9a372a92e | ||
|
|
199ae146ab | ||
|
|
b4874672fa | ||
|
|
f8d277697d | ||
|
|
4db8a12362 | ||
|
|
cabcb3bb0b | ||
|
|
5f0b715517 | ||
|
|
0ba613c3ce | ||
|
|
04703d2153 | ||
|
|
20da6d7609 | ||
|
|
387c09e91e | ||
|
|
cbce937042 | ||
|
|
246803bcfe | ||
|
|
625c331e80 | ||
|
|
f21944d789 | ||
|
|
0672fba279 | ||
|
|
53a52afe8b | ||
|
|
a2166ee4ff | ||
|
|
26bfd2ece4 | ||
|
|
0d60f0c638 | ||
|
|
735172f806 | ||
|
|
2c2a20b290 | ||
|
|
1a09a9cccd | ||
|
|
f183b09b93 | ||
|
|
5543992d7d | ||
|
|
dcabd76911 | ||
|
|
1a51184a32 | ||
|
|
c7e5095a86 | ||
|
|
f187d32351 | ||
|
|
ee81c65e4b | ||
|
|
ccd3eb24c9 | ||
|
|
51cb32896f | ||
|
|
af42382983 | ||
|
|
955b0c523c | ||
|
|
80281e2543 | ||
|
|
bb4dd4f393 | ||
|
|
ba5ddd59af | ||
|
|
aab7032071 | ||
|
|
adb6dcd683 | ||
|
|
ae93c93182 | ||
|
|
1c0c53c062 | ||
|
|
731add1fdf | ||
|
|
a7f82453a3 | ||
|
|
2bd6be8d85 | ||
|
|
e1014cc790 | ||
|
|
da2125cf13 | ||
|
|
f4bac2d41d | ||
|
|
39abf0ad3f | ||
|
|
8426b0c359 | ||
|
|
f4a0a31f70 | ||
|
|
65862c4fef | ||
|
|
ef548833fd | ||
|
|
55f70e1d27 | ||
|
|
b8cfddd7b6 | ||
|
|
52cb9a2bae | ||
|
|
539101af38 | ||
|
|
efa37cb15f | ||
|
|
faaa97d4a0 | ||
|
|
70804fe2c6 | ||
|
|
ae2532b165 | ||
|
|
81865bf873 | ||
|
|
9e47759347 | ||
|
|
d8665d7ab0 | ||
|
|
9bde15331e | ||
|
|
2fb2de1aa8 | ||
|
|
d68bd07423 | ||
|
|
44972d3215 | ||
|
|
aa308ad201 | ||
|
|
cb73790c27 | ||
|
|
b1fe23fbe4 | ||
|
|
ce81217381 | ||
|
|
94619ce52b | ||
|
|
27aa6bbf46 | ||
|
|
acf4d5c860 | ||
|
|
551932991d | ||
|
|
5ee53c512f | ||
|
|
4efa204727 | ||
|
|
c3bb647640 | ||
|
|
12851bc7ee | ||
|
|
a5f4d30ea6 | ||
|
|
0c9bc765ce | ||
|
|
cd34efa596 | ||
|
|
2b521c3a09 | ||
|
|
9681d96195 | ||
|
|
a32ecfeb33 | ||
|
|
bf73a1dfbe | ||
|
|
ca2ce5c16d | ||
|
|
7b17da7a0c |
+22
-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
|
||||
@@ -37,6 +37,11 @@ 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
|
||||
@@ -87,6 +92,18 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# korean_lindera/lindera(korean), ngram(min,max), edge_ngram(min,max)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER=
|
||||
|
||||
# File Parser (Optional - uses markitdown by default)
|
||||
# HINDSIGHT_API_FILE_PARSER=markitdown
|
||||
# Enable image OCR for MarkItDown using an OpenAI-compatible OCR/vision endpoint.
|
||||
# These OCR settings are independent from HINDSIGHT_API_LLM_* because MarkItDown
|
||||
# uses the OpenAI SDK directly and requires Chat Completions image input support.
|
||||
# When OCR is enabled, API_KEY, BASE_URL, and MODEL are required.
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED=false
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT=
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default), "onnx", "tei", "openai", "cohere", "google", "openrouter", "zeroentropy", "litellm", or "litellm-sdk"
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
@@ -150,6 +167,10 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Custom service name and environment (optional, defaults: hindsight-api, development)
|
||||
# HINDSIGHT_API_OTEL_SERVICE_NAME=hindsight-production
|
||||
# HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT=production
|
||||
#
|
||||
# Expose async-operation queue + consolidation-backlog gauges on /metrics.
|
||||
# Runs periodic per-schema COUNT queries on a background task (disabled by default).
|
||||
# HINDSIGHT_API_METRICS_BACKLOG_ENABLED=true
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Control Plane (Optional)
|
||||
|
||||
@@ -1,6 +0,0 @@
|
||||
version: 2
|
||||
updates:
|
||||
- package-ecosystem: "github-actions"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
@@ -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
|
||||
|
||||
@@ -45,17 +45,20 @@ jobs:
|
||||
integrations-pydantic-ai: ${{ steps.filter.outputs.integrations-pydantic-ai }}
|
||||
integrations-ag2: ${{ steps.filter.outputs.integrations-ag2 }}
|
||||
integrations-autogen: ${{ steps.filter.outputs.integrations-autogen }}
|
||||
integrations-aider: ${{ steps.filter.outputs.integrations-aider }}
|
||||
integrations-langgraph: ${{ steps.filter.outputs.integrations-langgraph }}
|
||||
integrations-llamaindex: ${{ steps.filter.outputs.integrations-llamaindex }}
|
||||
integrations-paperclip: ${{ steps.filter.outputs.integrations-paperclip }}
|
||||
integrations-opencode: ${{ steps.filter.outputs.integrations-opencode }}
|
||||
integrations-cursor: ${{ steps.filter.outputs.integrations-cursor }}
|
||||
integrations-zed: ${{ steps.filter.outputs.integrations-zed }}
|
||||
integrations-n8n: ${{ steps.filter.outputs.integrations-n8n }}
|
||||
integrations-zapier: ${{ steps.filter.outputs.integrations-zapier }}
|
||||
integrations-cloudflare-oauth-proxy: ${{ steps.filter.outputs.integrations-cloudflare-oauth-proxy }}
|
||||
integrations-superagent: ${{ steps.filter.outputs.integrations-superagent }}
|
||||
integrations-lockfiles: ${{ steps.filter.outputs.integrations-lockfiles }}
|
||||
integrations-openai-agents: ${{ steps.filter.outputs.integrations-openai-agents }}
|
||||
integrations-openhands: ${{ steps.filter.outputs.integrations-openhands }}
|
||||
integrations-pipecat: ${{ steps.filter.outputs.integrations-pipecat }}
|
||||
integrations-agentcore: ${{ steps.filter.outputs.integrations-agentcore }}
|
||||
integrations-smolagents: ${{ steps.filter.outputs.integrations-smolagents }}
|
||||
@@ -154,6 +157,8 @@ jobs:
|
||||
- 'hindsight-integrations/ag2/**'
|
||||
integrations-autogen:
|
||||
- 'hindsight-integrations/autogen/**'
|
||||
integrations-aider:
|
||||
- 'hindsight-integrations/aider/**'
|
||||
integrations-langgraph:
|
||||
- 'hindsight-integrations/langgraph/**'
|
||||
integrations-llamaindex:
|
||||
@@ -166,6 +171,8 @@ jobs:
|
||||
- 'hindsight-integrations/opencode/**'
|
||||
integrations-cursor:
|
||||
- 'hindsight-integrations/cursor/**'
|
||||
integrations-zed:
|
||||
- 'hindsight-integrations/zed/**'
|
||||
integrations-n8n:
|
||||
- 'hindsight-integrations/n8n/**'
|
||||
integrations-zapier:
|
||||
@@ -180,6 +187,8 @@ jobs:
|
||||
- 'scripts/check-integration-lockfiles.sh'
|
||||
integrations-openai-agents:
|
||||
- 'hindsight-integrations/openai-agents/**'
|
||||
integrations-openhands:
|
||||
- 'hindsight-integrations/openhands/**'
|
||||
integrations-pipecat:
|
||||
- 'hindsight-integrations/pipecat/**'
|
||||
integrations-agentcore:
|
||||
@@ -488,6 +497,37 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/cursor
|
||||
run: python -m pytest tests/ -v
|
||||
|
||||
test-zed-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-zed == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.11'
|
||||
|
||||
- name: Install package and pytest
|
||||
working-directory: ./hindsight-integrations/zed
|
||||
# Installs the package (incl. the zstandard runtime dep) so the threads.db
|
||||
# reader tests can decompress Zed's zstd blobs.
|
||||
run: pip install -e . pytest
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/zed
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: python -m pytest tests/ -v -m "not requires_real_llm"
|
||||
|
||||
test-omo-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -3021,6 +3061,45 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/ag2
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-aider-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-aider == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build aider integration
|
||||
working-directory: ./hindsight-integrations/aider
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/aider
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/aider
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-autogen-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -3669,6 +3748,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-openhands-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-openhands == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build openhands integration
|
||||
working-directory: ./hindsight-integrations/openhands
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/openhands
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/openhands
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-claude-agent-sdk-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -4703,6 +4821,7 @@ jobs:
|
||||
- test-openclaw-integration
|
||||
- test-integration
|
||||
- test-ag2-integration
|
||||
- test-aider-integration
|
||||
- test-autogen-integration
|
||||
- test-continue-integration
|
||||
- test-smolagents-integration
|
||||
@@ -4717,6 +4836,7 @@ jobs:
|
||||
- test-pydantic-ai-integration
|
||||
- test-llamaindex-integration
|
||||
- test-openai-agents-integration
|
||||
- test-openhands-integration
|
||||
- test-agentcore-integration
|
||||
- test-haystack-integration
|
||||
- test-pip-slim
|
||||
|
||||
@@ -16,6 +16,42 @@
|
||||
|
||||
---
|
||||
|
||||
### Powered by Atlas Cloud (OpenAI-compatible)
|
||||
|
||||
<p align="center">
|
||||
<a href="https://www.atlascloud.ai/?utm_source=github&utm_medium=link&utm_campaign=hindsight">
|
||||
<img src="./hindsight-docs/static/img/atlas-cloud-logo.png" alt="Atlas Cloud" width="200">
|
||||
</a>
|
||||
</p>
|
||||
|
||||
> 🎁 **[Atlas Cloud](https://www.atlascloud.ai/?utm_source=github&utm_medium=link&utm_campaign=hindsight)** is a full-modal, OpenAI-compatible AI inference platform — plug it in as a drop-in LLM backend for Hindsight's fact extraction, reflection and consolidation, with one API for DeepSeek, Qwen, GLM, Kimi, MiniMax and more. No multi-vendor setup needed.
|
||||
> Budget-friendly: [coding plan](https://www.atlascloud.ai/console/coding-plan)
|
||||
|
||||
```bash
|
||||
export HINDSIGHT_API_LLM_PROVIDER=atlas
|
||||
export HINDSIGHT_API_LLM_API_KEY=your-atlascloud-api-key # base_url defaults to https://api.atlascloud.ai/v1
|
||||
export HINDSIGHT_API_LLM_MODEL=deepseek-ai/deepseek-v4-pro
|
||||
```
|
||||
|
||||
`deepseek-ai/deepseek-v4-pro` is a reasoning model — give it enough `max_tokens` (>= 512).
|
||||
|
||||
<details>
|
||||
<summary>All Atlas Cloud chat models (59)</summary>
|
||||
|
||||
- **Anthropic (Claude):** `anthropic/claude-haiku-4.5-20251001`, `anthropic/claude-opus-4.8`, `anthropic/claude-sonnet-4.6`
|
||||
- **OpenAI (GPT):** `openai/gpt-5.4`, `openai/gpt-5.5`
|
||||
- **Google (Gemini):** `google/gemini-3.1-flash-lite`, `google/gemini-3.1-pro-preview`, `google/gemini-3.5-flash`
|
||||
- **Alibaba Qwen:** `qwen/qwen2.5-7b-instruct`, `Qwen/Qwen3-235B-A22B-Instruct-2507`, `qwen/qwen3-235b-a22b-thinking-2507`, `qwen/qwen3-30b-a3b`, `Qwen/Qwen3-30B-A3B-Instruct-2507`, `qwen/qwen3-30b-a3b-thinking-2507`, `qwen/qwen3-32b`, `qwen/qwen3-8b`, `Qwen/Qwen3-Coder`, `qwen/qwen3-coder-next`, `qwen/qwen3-max-2026-01-23`, `Qwen/Qwen3-Next-80B-A3B-Instruct`, `Qwen/Qwen3-Next-80B-A3B-Thinking`, `Qwen/Qwen3-VL-235B-A22B-Instruct`, `qwen/qwen3-vl-235b-a22b-thinking`, `qwen/qwen3-vl-30b-a3b-instruct`, `qwen/qwen3-vl-30b-a3b-thinking`, `qwen/qwen3-vl-8b-instruct`, `qwen/qwen3.5-122b-a10b`, `qwen/qwen3.5-27b`, `qwen/qwen3.5-35b-a3b`, `qwen/qwen3.5-397b-a17b`, `qwen/qwen3.6-35b-a3b`, `qwen/qwen3.6-plus`
|
||||
- **DeepSeek:** `deepseek-ai/deepseek-ocr`, `deepseek-ai/deepseek-r1-0528`, `deepseek-ai/DeepSeek-V3-0324`, `deepseek-ai/DeepSeek-V3.1`, `deepseek-ai/DeepSeek-V3.1-Terminus`, `deepseek-ai/deepseek-v3.2`, `deepseek-ai/DeepSeek-V3.2-Exp`, `deepseek-ai/deepseek-v4-flash`, `deepseek-ai/deepseek-v4-pro`
|
||||
- **Moonshot (Kimi):** `moonshotai/Kimi-K2-Instruct`, `moonshotai/Kimi-K2-Instruct-0905`, `moonshotai/Kimi-K2-Thinking`, `moonshotai/kimi-k2.5`, `moonshotai/kimi-k2.6`
|
||||
- **Zhipu GLM:** `zai-org/GLM-4.6`, `zai-org/glm-4.7`, `zai-org/glm-5`, `zai-org/glm-5-turbo`, `zai-org/glm-5.1`, `zai-org/glm-5v-turbo`
|
||||
- **MiniMax:** `MiniMaxAI/MiniMax-M2`, `minimaxai/minimax-m2.1`, `minimaxai/minimax-m2.5`, `minimaxai/minimax-m2.7`
|
||||
- **xAI:** `xai/grok-4.3`
|
||||
- **Kuaishou KAT:** `kwaipilot/kat-coder-pro-v2`
|
||||
- **Other:** `owl`
|
||||
|
||||
</details>
|
||||
|
||||
## What is Hindsight?
|
||||
|
||||
Hindsight™ is an agent memory system built to create smarter agents that learn over time. Most agent memory systems focus on recalling conversation history. Hindsight is focused on making agents that learn, not just remember.
|
||||
@@ -70,7 +106,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).
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.8.2
|
||||
appVersion: "0.8.2"
|
||||
version: 0.8.3
|
||||
appVersion: "0.8.3"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.8.2",
|
||||
"version": "0.8.3",
|
||||
"description": "Node.js programmatic lifecycle manager for Hindsight — embeds a local hindsight daemon in a Node application. Pair with @vectorize-io/hindsight-client for memory operations.",
|
||||
"main": "dist/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.8.2"
|
||||
version = "0.8.3"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim==0.8.2",
|
||||
"hindsight-api-slim==0.8.3",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.8.2"
|
||||
version = "0.8.3"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim[all]==0.8.2",
|
||||
"hindsight-api-slim[all]==0.8.3",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
@@ -21,7 +21,7 @@ hindsight-embed = { workspace = true }
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-llm = [
|
||||
"hindsight-api-slim[local-llm]==0.8.2",
|
||||
"hindsight-api-slim[local-llm]==0.8.3",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -53,4 +53,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.8.2"
|
||||
__version__ = "0.8.3"
|
||||
|
||||
@@ -256,6 +256,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 +293,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 +312,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 +337,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 +346,7 @@ def run_db_migration(
|
||||
schema=schema,
|
||||
base_schema=config.database_schema,
|
||||
embedding_dimension=embedding_dimension,
|
||||
ensure_extensions=not skip_extension_reconcile,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
+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)
|
||||
@@ -158,7 +158,12 @@ from hindsight_api.engine.response_models import (
|
||||
)
|
||||
from hindsight_api.engine.search.tags import TagGroup, TagsMatch
|
||||
from hindsight_api.extensions import HttpExtension, OperationValidationError, load_extension
|
||||
from hindsight_api.metrics import create_metrics_collector, get_metrics_collector, initialize_metrics
|
||||
from hindsight_api.metrics import (
|
||||
create_metrics_collector,
|
||||
get_metrics_collector,
|
||||
initialize_metrics,
|
||||
normalize_http_endpoint,
|
||||
)
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -265,6 +270,16 @@ class RecallRequest(BaseModel):
|
||||
default=None,
|
||||
description="List of fact types to recall: 'world', 'experience', 'observation'. Defaults to world and experience if not specified.",
|
||||
)
|
||||
prefer_observations: bool = Field(
|
||||
default=False,
|
||||
description=(
|
||||
"When recalling raw facts ('world'/'experience') together with 'observation', drop any raw "
|
||||
"fact that an observation in the results was consolidated from, so the observation supersedes "
|
||||
"it and you don't get duplicate content. The freed slots are backfilled with the next results, "
|
||||
"keeping the result count at the requested budget. Disabled by default; set to true to enable. "
|
||||
"No effect unless 'observation' and at least one raw type are both requested."
|
||||
),
|
||||
)
|
||||
budget: Budget = Budget.MID
|
||||
max_tokens: int = 4096
|
||||
trace: bool = False
|
||||
@@ -281,12 +296,16 @@ 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,
|
||||
@@ -1442,6 +1461,13 @@ class DryRunExtractRequest(BaseModel):
|
||||
entities_allow_free_form: bool | None = None
|
||||
llm_output_language: str | None = None
|
||||
|
||||
@field_validator("content")
|
||||
@classmethod
|
||||
def validate_content(cls, v: str) -> str:
|
||||
if not v.strip():
|
||||
raise ValueError("content cannot be empty")
|
||||
return v
|
||||
|
||||
|
||||
class ListDocumentsResponse(BaseModel):
|
||||
"""Response model for list documents endpoint."""
|
||||
@@ -3237,15 +3263,9 @@ def create_app(
|
||||
@app.middleware("http")
|
||||
async def http_metrics_middleware(request, call_next):
|
||||
"""Record HTTP request metrics."""
|
||||
# Normalize endpoint path to reduce cardinality
|
||||
# Replace UUIDs and numeric IDs with placeholders
|
||||
import re
|
||||
|
||||
path = request.url.path
|
||||
# Replace UUIDs
|
||||
path = re.sub(r"/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", "/{id}", path)
|
||||
# Replace numeric IDs
|
||||
path = re.sub(r"/\d+(?=/|$)", "/{id}", path)
|
||||
# Template id segments (bank ids, UUIDs, numeric ids) so the endpoint
|
||||
# metric label stays bounded-cardinality.
|
||||
path = normalize_http_endpoint(request.url.path)
|
||||
|
||||
status_code = [500] # Default to 500, will be updated
|
||||
metrics_collector = get_metrics_collector()
|
||||
@@ -3333,6 +3353,7 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
async def _precheck_dep(
|
||||
bank_id: str,
|
||||
request: Request,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
) -> None:
|
||||
validator = getattr(app.state.memory, "_operation_validator", None)
|
||||
@@ -3341,10 +3362,20 @@ def _register_routes(app: FastAPI):
|
||||
from hindsight_api.extensions import PrecheckContext
|
||||
|
||||
await app.state.memory._authenticate_tenant(request_context)
|
||||
cl_header = request.headers.get("content-length")
|
||||
content_length: int | None = None
|
||||
if cl_header is not None:
|
||||
try:
|
||||
parsed = int(cl_header)
|
||||
except ValueError:
|
||||
parsed = -1
|
||||
if parsed >= 0:
|
||||
content_length = parsed
|
||||
ctx = PrecheckContext(
|
||||
operation=operation,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
content_length=content_length,
|
||||
)
|
||||
result = await validator.precheck(ctx)
|
||||
if not result.allowed:
|
||||
@@ -3448,7 +3479,7 @@ def _register_routes(app: FastAPI):
|
||||
async def api_graph(
|
||||
bank_id: str,
|
||||
type: str | None = None,
|
||||
limit: int = 1000,
|
||||
limit: int = Query(default=1000, ge=0),
|
||||
q: str | None = None,
|
||||
tags: list[str] | None = Query(None),
|
||||
tags_match: str = "all_strict",
|
||||
@@ -3496,8 +3527,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),
|
||||
):
|
||||
"""
|
||||
@@ -3809,6 +3840,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,
|
||||
@@ -4231,8 +4263,8 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
async def api_list_entities(
|
||||
bank_id: str,
|
||||
limit: int = Query(default=100, description="Maximum number of entities to return"),
|
||||
offset: int = Query(default=0, description="Offset for pagination"),
|
||||
limit: int = Query(default=100, ge=0, description="Maximum number of entities to return"),
|
||||
offset: int = Query(default=0, ge=0, description="Offset for pagination"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""List entities for a memory bank with pagination."""
|
||||
@@ -4267,7 +4299,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),
|
||||
):
|
||||
@@ -4894,8 +4926,8 @@ def _register_routes(app: FastAPI):
|
||||
tags_match: str = Query(
|
||||
"any_strict", description="How to match tags: 'any', 'all', 'any_strict', 'all_strict'"
|
||||
),
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
limit: int = Query(default=100, ge=0),
|
||||
offset: int = Query(default=0, ge=0),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""
|
||||
@@ -5079,8 +5111,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),
|
||||
):
|
||||
"""
|
||||
@@ -6749,7 +6781,7 @@ def _register_routes(app: FastAPI):
|
||||
description="Upload files (PDF, DOCX, etc.), convert them to markdown, and retain as memories.\n\n"
|
||||
"This endpoint handles file upload, conversion, and memory creation in a single operation.\n\n"
|
||||
"**Features:**\n"
|
||||
"- Supports PDF, DOCX, PPTX, XLSX, images (with OCR), audio (with transcription)\n"
|
||||
"- Supports PDF, DOCX, PPTX, XLSX, images (parser-dependent OCR), audio (with transcription)\n"
|
||||
"- Automatic file-to-markdown conversion using pluggable parsers\n"
|
||||
"- Files stored in object storage (PostgreSQL by default, S3 for production)\n"
|
||||
"- Each file becomes a separate document with optional metadata/tags\n"
|
||||
|
||||
@@ -9,7 +9,7 @@ from fastmcp import FastMCP
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api import __version__ as HINDSIGHT_VERSION
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.config import DEFAULT_MCP_RECALL_DESCRIPTION, DEFAULT_MCP_RETAIN_DESCRIPTION, _get_raw_config
|
||||
from hindsight_api.engine.memory_engine import _current_schema
|
||||
from hindsight_api.extensions import MCPExtension, load_extension
|
||||
from hindsight_api.extensions.tenant import AuthenticationError
|
||||
@@ -78,6 +78,19 @@ def get_current_mcp_authenticated() -> bool:
|
||||
return _current_mcp_authenticated.get()
|
||||
|
||||
|
||||
def _build_mcp_tool_descriptions(extra_instructions: str | None) -> tuple[str | None, str | None]:
|
||||
"""Return custom retain/recall descriptions when server-level MCP instructions are set."""
|
||||
if not isinstance(extra_instructions, str):
|
||||
return None, None
|
||||
|
||||
extra_instructions = extra_instructions.strip()
|
||||
if not extra_instructions:
|
||||
return None, None
|
||||
|
||||
suffix = f"\n\nAdditional instructions: {extra_instructions}"
|
||||
return DEFAULT_MCP_RETAIN_DESCRIPTION + suffix, DEFAULT_MCP_RECALL_DESCRIPTION + suffix
|
||||
|
||||
|
||||
def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"""
|
||||
Create and configure the Hindsight MCP server.
|
||||
@@ -135,6 +148,10 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
allowed = frozenset(global_config.mcp_enabled_tools)
|
||||
base_tools = (base_tools if base_tools is not None else _ALL_TOOLS) & allowed
|
||||
|
||||
retain_description, recall_description = _build_mcp_tool_descriptions(
|
||||
getattr(global_config, "mcp_instructions", None)
|
||||
)
|
||||
|
||||
# Configure and register tools using shared module
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=get_current_bank_id,
|
||||
@@ -144,6 +161,8 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
mcp_authenticated_resolver=get_current_mcp_authenticated, # Propagate MCP pre-auth flag
|
||||
include_bank_id_param=multi_bank,
|
||||
tools=base_tools,
|
||||
retain_description=retain_description,
|
||||
recall_description=recall_description,
|
||||
)
|
||||
|
||||
register_mcp_tools(mcp, memory, config)
|
||||
|
||||
@@ -142,6 +142,7 @@ ENV_LLM_REASONING_EFFORT = "HINDSIGHT_API_LLM_REASONING_EFFORT"
|
||||
ENV_LLM_GROQ_SERVICE_TIER = "HINDSIGHT_API_LLM_GROQ_SERVICE_TIER"
|
||||
ENV_LLM_OPENAI_SERVICE_TIER = "HINDSIGHT_API_LLM_OPENAI_SERVICE_TIER"
|
||||
ENV_LLM_BEDROCK_SERVICE_TIER = "HINDSIGHT_API_LLM_BEDROCK_SERVICE_TIER"
|
||||
ENV_LLM_GEMINI_SERVICE_TIER = "HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER"
|
||||
ENV_LLM_EXTRA_BODY = "HINDSIGHT_API_LLM_EXTRA_BODY"
|
||||
ENV_LLM_DEFAULT_HEADERS = "HINDSIGHT_API_LLM_DEFAULT_HEADERS"
|
||||
ENV_LLM_STRICT_SCHEMA = "HINDSIGHT_API_LLM_STRICT_SCHEMA"
|
||||
@@ -159,11 +160,25 @@ ENV_LLM_LITELLMROUTER_CONFIG = "HINDSIGHT_API_LLM_LITELLMROUTER_CONFIG"
|
||||
DEFAULT_LLM_GROQ_SERVICE_TIER = "auto" # "on_demand", "flex", or "auto"
|
||||
DEFAULT_LLM_OPENAI_SERVICE_TIER = None # None (default) or "flex" (50% cheaper)
|
||||
DEFAULT_LLM_BEDROCK_SERVICE_TIER = None # None (default), "flex", "priority", or "reserved"
|
||||
DEFAULT_LLM_GEMINI_SERVICE_TIER = None # None (default) or "flex" (50% cheaper best-effort tier)
|
||||
DEFAULT_LLM_EXTRA_BODY = None # None = no extra body params; JSON dict merged into OpenAI extra_body
|
||||
DEFAULT_LLM_DEFAULT_HEADERS = (
|
||||
None # None = no extra headers; JSON dict passed as default_headers to provider SDK clients
|
||||
)
|
||||
|
||||
|
||||
def parse_gemini_service_tier(value: str | None) -> str | None:
|
||||
"""Normalize and validate the Gemini service tier."""
|
||||
tier = value or None
|
||||
valid_tiers = (None, "flex")
|
||||
if tier not in valid_tiers:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER: "
|
||||
f"{tier!r}. Must be one of: {', '.join(t for t in valid_tiers if t is not None)}."
|
||||
)
|
||||
return tier
|
||||
|
||||
|
||||
# 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"
|
||||
@@ -354,6 +369,7 @@ ENV_ACCESS_LOG = "HINDSIGHT_API_ACCESS_LOG"
|
||||
ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
|
||||
ENV_MCP_ENABLED_TOOLS = "HINDSIGHT_API_MCP_ENABLED_TOOLS"
|
||||
ENV_MCP_STATELESS = "HINDSIGHT_API_MCP_STATELESS"
|
||||
ENV_MCP_INSTRUCTIONS = "HINDSIGHT_API_MCP_INSTRUCTIONS"
|
||||
ENV_ENABLE_BANK_CONFIG_API = "HINDSIGHT_API_ENABLE_BANK_CONFIG_API"
|
||||
ENV_ENABLE_BANK_LLM_HEALTH = "HINDSIGHT_API_ENABLE_BANK_LLM_HEALTH"
|
||||
ENV_ENABLE_DRY_RUN_EXTRACT = "HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT"
|
||||
@@ -375,6 +391,7 @@ ENV_OTEL_EXPORTER_OTLP_HEADERS = "HINDSIGHT_API_OTEL_EXPORTER_OTLP_HEADERS"
|
||||
ENV_OTEL_SERVICE_NAME = "HINDSIGHT_API_OTEL_SERVICE_NAME"
|
||||
ENV_OTEL_DEPLOYMENT_ENVIRONMENT = "HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT"
|
||||
ENV_METRICS_INCLUDE_BANK_ID = "HINDSIGHT_API_METRICS_INCLUDE_BANK_ID"
|
||||
ENV_METRICS_BACKLOG_ENABLED = "HINDSIGHT_API_METRICS_BACKLOG_ENABLED"
|
||||
|
||||
# Vertex AI configuration
|
||||
ENV_LLM_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID"
|
||||
@@ -424,6 +441,11 @@ ENV_FILE_STORAGE_AZURE_ACCOUNT_NAME = "HINDSIGHT_API_FILE_STORAGE_AZURE_ACCOUNT_
|
||||
ENV_FILE_STORAGE_AZURE_ACCOUNT_KEY = "HINDSIGHT_API_FILE_STORAGE_AZURE_ACCOUNT_KEY"
|
||||
ENV_FILE_PARSER = "HINDSIGHT_API_FILE_PARSER"
|
||||
ENV_FILE_PARSER_ALLOWLIST = "HINDSIGHT_API_FILE_PARSER_ALLOWLIST"
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_ENABLED = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED"
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_API_KEY = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY"
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_BASE_URL = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL"
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_MODEL = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL"
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_PROMPT = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT"
|
||||
ENV_FILE_PARSER_IRIS_TOKEN = "HINDSIGHT_API_FILE_PARSER_IRIS_TOKEN"
|
||||
ENV_FILE_PARSER_IRIS_ORG_ID = "HINDSIGHT_API_FILE_PARSER_IRIS_ORG_ID"
|
||||
ENV_FILE_PARSER_LLAMA_PARSE_API_KEY = "HINDSIGHT_API_FILE_PARSER_LLAMA_PARSE_API_KEY"
|
||||
@@ -553,6 +575,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 +618,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",
|
||||
@@ -691,6 +722,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
|
||||
@@ -804,6 +843,7 @@ DEFAULT_ACCESS_LOG = False
|
||||
DEFAULT_MCP_ENABLED = True
|
||||
DEFAULT_MCP_ENABLED_TOOLS: list[str] | None = None # None = all tools enabled
|
||||
DEFAULT_MCP_STATELESS = False # False = stateful (supports SSE/GET); True = stateless (POST-only)
|
||||
DEFAULT_MCP_INSTRUCTIONS = None
|
||||
DEFAULT_ENABLE_BANK_CONFIG_API = True
|
||||
# Dry-run extraction is a preview tool that makes a real LLM call but stores nothing. Enabled by
|
||||
# default; set HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT=false to remove the endpoint (e.g. to cap
|
||||
@@ -847,6 +887,10 @@ DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60 # Batch API polling interval in
|
||||
DEFAULT_FILE_STORAGE_TYPE = "native" # PostgreSQL BYTEA storage
|
||||
DEFAULT_FILE_PARSER = "markitdown" # Default parser fallback chain (comma-separated, e.g. "iris,markitdown")
|
||||
DEFAULT_FILE_PARSER_ALLOWLIST = None # Allowlist of parsers clients may request (None = all registered parsers)
|
||||
DEFAULT_FILE_PARSER_MARKITDOWN_OCR_ENABLED = False
|
||||
DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT = """You are a precise OCR transcription engine.
|
||||
|
||||
Transcribe only the visible text in the image. Do not describe the image, summarize it, translate it, infer missing content, or add commentary. Preserve the original language, wording, numbers, punctuation, capitalization, and reading order. Reconstruct headings, lists, key-value fields, stamps, and tables as clean Markdown when the layout is clear. If text is unreadable or uncertain, write [unclear] for that span. Return only the extracted Markdown."""
|
||||
DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE_MB = 100 # Max total batch size in MB (all files combined)
|
||||
DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE = 10 # Max files per batch upload
|
||||
DEFAULT_ENABLE_FILE_UPLOAD_API = True # Enable file upload endpoint
|
||||
@@ -965,6 +1009,7 @@ DEFAULT_OTEL_TRACES_ENABLED = False # Disabled by default for backward compatib
|
||||
DEFAULT_OTEL_SERVICE_NAME = "hindsight-api"
|
||||
DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT = "development"
|
||||
DEFAULT_METRICS_INCLUDE_BANK_ID = False # Disabled by default to avoid high-cardinality OTel metric growth
|
||||
DEFAULT_METRICS_BACKLOG_ENABLED = False # Disabled by default: runs periodic per-schema COUNT queries
|
||||
|
||||
# Audit log defaults
|
||||
DEFAULT_AUDIT_LOG_ENABLED = False # Disabled by default
|
||||
@@ -1180,6 +1225,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}``.
|
||||
|
||||
@@ -1298,6 +1355,7 @@ class HindsightConfig:
|
||||
llm_groq_service_tier: str # Groq: "on_demand", "flex", or "auto"
|
||||
llm_openai_service_tier: str | None # OpenAI: None (default) or "flex" (50% cheaper)
|
||||
llm_bedrock_service_tier: str | None # Bedrock: None (default), "flex", "priority", or "reserved"
|
||||
llm_gemini_service_tier: str | None # Gemini: None (default) or "flex" (50% cheaper)
|
||||
llm_extra_body: (
|
||||
dict | None
|
||||
) # Extra body params merged into OpenAI-compatible API calls (e.g. {"chat_template_kwargs": {"enable_thinking": true}})
|
||||
@@ -1435,6 +1493,9 @@ class HindsightConfig:
|
||||
bm25_min_score: float
|
||||
recall_max_candidates_per_source: int
|
||||
recall_strategy_boosts: dict[str, str]
|
||||
recency_decay_function: str
|
||||
recency_decay_linear_window_days: float
|
||||
recency_decay_halflife_days: float
|
||||
reranker_cohere_api_key: str | None
|
||||
reranker_cohere_model: str
|
||||
reranker_cohere_base_url: str | None
|
||||
@@ -1478,6 +1539,7 @@ class HindsightConfig:
|
||||
mcp_enabled: bool
|
||||
mcp_enabled_tools: list[str] | None # None = all tools; explicit list = allowlist
|
||||
mcp_stateless: bool # True = stateless HTTP (POST-only); False = stateful (supports GET/SSE)
|
||||
mcp_instructions: str | None # Additional instructions appended to retain/recall MCP tool descriptions
|
||||
enable_bank_config_api: bool
|
||||
enable_bank_llm_health: bool
|
||||
enable_dry_run_extract: bool
|
||||
@@ -1642,6 +1704,7 @@ class HindsightConfig:
|
||||
otel_service_name: str
|
||||
otel_deployment_environment: str
|
||||
metrics_include_bank_id: bool
|
||||
metrics_backlog_enabled: bool
|
||||
|
||||
# Audit log configuration (static - server-level only)
|
||||
audit_log_enabled: bool # Master switch for audit logging
|
||||
@@ -1676,6 +1739,11 @@ class HindsightConfig:
|
||||
embeddings_zeroentropy_encoding_format: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT
|
||||
embeddings_zeroentropy_batch_size: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE
|
||||
embeddings_zeroentropy_latency: str | None = DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY
|
||||
file_parser_markitdown_ocr_enabled: bool = DEFAULT_FILE_PARSER_MARKITDOWN_OCR_ENABLED
|
||||
file_parser_markitdown_ocr_api_key: str | None = None
|
||||
file_parser_markitdown_ocr_base_url: str | None = None
|
||||
file_parser_markitdown_ocr_model: str | None = None
|
||||
file_parser_markitdown_ocr_prompt: str = DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
|
||||
|
||||
# Class-level sets for configuration categorization
|
||||
|
||||
@@ -1716,6 +1784,8 @@ class HindsightConfig:
|
||||
"file_storage_gcs_service_account_key",
|
||||
"file_storage_azure_account_key",
|
||||
# File parser credentials
|
||||
"file_parser_markitdown_ocr_api_key",
|
||||
"file_parser_markitdown_ocr_base_url",
|
||||
"file_parser_iris_token",
|
||||
"file_parser_llama_parse_api_key",
|
||||
}
|
||||
@@ -1879,6 +1949,9 @@ class HindsightConfig:
|
||||
f"Note: 'standard' is not a valid Bedrock service tier -- use unset for default tier."
|
||||
)
|
||||
|
||||
# Validate gemini_service_tier
|
||||
self.llm_gemini_service_tier = parse_gemini_service_tier(self.llm_gemini_service_tier)
|
||||
|
||||
# When LLM provider is "none", force chunks-only mode and disable LLM-dependent features
|
||||
if self.llm_provider == "none":
|
||||
self.retain_extraction_mode = "chunks"
|
||||
@@ -1996,6 +2069,11 @@ class HindsightConfig:
|
||||
llm_groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
|
||||
llm_openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
|
||||
llm_bedrock_service_tier=os.getenv(ENV_LLM_BEDROCK_SERVICE_TIER) or None,
|
||||
llm_gemini_service_tier=(
|
||||
parse_gemini_service_tier(os.getenv(ENV_LLM_GEMINI_SERVICE_TIER) or DEFAULT_LLM_GEMINI_SERVICE_TIER)
|
||||
if llm_provider.lower() == "gemini"
|
||||
else None
|
||||
),
|
||||
llm_extra_body=json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null")),
|
||||
llm_default_headers=json.loads(os.getenv(ENV_LLM_DEFAULT_HEADERS, "null")),
|
||||
llm_strict_schema=os.getenv(ENV_LLM_STRICT_SCHEMA, str(DEFAULT_LLM_STRICT_SCHEMA)).lower() in ("true", "1"),
|
||||
@@ -2270,6 +2348,15 @@ class HindsightConfig:
|
||||
recall_strategy_boosts=_parse_strategy_boosts(
|
||||
os.getenv(ENV_RECALL_STRATEGY_BOOSTS, DEFAULT_RECALL_STRATEGY_BOOSTS)
|
||||
),
|
||||
recency_decay_function=_validate_recency_decay_function(
|
||||
os.getenv(ENV_RECENCY_DECAY_FUNCTION, DEFAULT_RECENCY_DECAY_FUNCTION)
|
||||
),
|
||||
recency_decay_linear_window_days=float(
|
||||
os.getenv(ENV_RECENCY_DECAY_LINEAR_WINDOW_DAYS, str(DEFAULT_RECENCY_DECAY_LINEAR_WINDOW_DAYS))
|
||||
),
|
||||
recency_decay_halflife_days=float(
|
||||
os.getenv(ENV_RECENCY_DECAY_HALFLIFE_DAYS, str(DEFAULT_RECENCY_DECAY_HALFLIFE_DAYS))
|
||||
),
|
||||
# Cohere reranker (with backward-compatible fallback to shared API key)
|
||||
reranker_cohere_api_key=os.getenv(ENV_RERANKER_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
|
||||
reranker_cohere_model=os.getenv(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL),
|
||||
@@ -2345,6 +2432,7 @@ class HindsightConfig:
|
||||
if os.getenv(ENV_MCP_ENABLED_TOOLS)
|
||||
else DEFAULT_MCP_ENABLED_TOOLS,
|
||||
mcp_stateless=os.getenv(ENV_MCP_STATELESS, str(DEFAULT_MCP_STATELESS)).lower() == "true",
|
||||
mcp_instructions=os.getenv(ENV_MCP_INSTRUCTIONS) or DEFAULT_MCP_INSTRUCTIONS,
|
||||
enable_bank_llm_health=os.getenv(ENV_ENABLE_BANK_LLM_HEALTH, str(DEFAULT_ENABLE_BANK_LLM_HEALTH)).lower()
|
||||
== "true",
|
||||
enable_bank_config_api=os.getenv(ENV_ENABLE_BANK_CONFIG_API, str(DEFAULT_ENABLE_BANK_CONFIG_API)).lower()
|
||||
@@ -2424,6 +2512,18 @@ class HindsightConfig:
|
||||
file_parser_allowlist=_parse_str_list(os.getenv(ENV_FILE_PARSER_ALLOWLIST))
|
||||
if os.getenv(ENV_FILE_PARSER_ALLOWLIST)
|
||||
else None,
|
||||
file_parser_markitdown_ocr_enabled=os.getenv(
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_ENABLED,
|
||||
str(DEFAULT_FILE_PARSER_MARKITDOWN_OCR_ENABLED),
|
||||
).lower()
|
||||
in ("1", "true", "yes", "on"),
|
||||
file_parser_markitdown_ocr_api_key=os.getenv(ENV_FILE_PARSER_MARKITDOWN_OCR_API_KEY) or None,
|
||||
file_parser_markitdown_ocr_base_url=os.getenv(ENV_FILE_PARSER_MARKITDOWN_OCR_BASE_URL) or None,
|
||||
file_parser_markitdown_ocr_model=os.getenv(ENV_FILE_PARSER_MARKITDOWN_OCR_MODEL) or None,
|
||||
file_parser_markitdown_ocr_prompt=os.getenv(
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
|
||||
DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
|
||||
),
|
||||
file_parser_iris_token=os.getenv(ENV_FILE_PARSER_IRIS_TOKEN) or None,
|
||||
file_parser_iris_org_id=os.getenv(ENV_FILE_PARSER_IRIS_ORG_ID) or None,
|
||||
file_parser_llama_parse_api_key=os.getenv(ENV_FILE_PARSER_LLAMA_PARSE_API_KEY) or None,
|
||||
@@ -2614,6 +2714,8 @@ class HindsightConfig:
|
||||
otel_deployment_environment=os.getenv(ENV_OTEL_DEPLOYMENT_ENVIRONMENT, DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT),
|
||||
metrics_include_bank_id=os.getenv(ENV_METRICS_INCLUDE_BANK_ID, str(DEFAULT_METRICS_INCLUDE_BANK_ID)).lower()
|
||||
in ("true", "1", "yes"),
|
||||
metrics_backlog_enabled=os.getenv(ENV_METRICS_BACKLOG_ENABLED, str(DEFAULT_METRICS_BACKLOG_ENABLED)).lower()
|
||||
in ("true", "1", "yes"),
|
||||
# Audit log configuration (static, server-level only)
|
||||
audit_log_enabled=os.getenv(ENV_AUDIT_LOG_ENABLED, str(DEFAULT_AUDIT_LOG_ENABLED)).lower() == "true",
|
||||
audit_log_actions=[
|
||||
|
||||
@@ -8,6 +8,7 @@ Config values are resolved on every request to ensure consistency across
|
||||
multiple API servers.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import asdict, replace
|
||||
@@ -161,26 +162,83 @@ class ConfigResolver:
|
||||
resolved_config = await self.resolve_full_config(bank_id, context)
|
||||
config_dict = asdict(resolved_config)
|
||||
|
||||
# SECURITY: Filter to only configurable fields (exclude static/infrastructure)
|
||||
filtered = {k: v for k, v in config_dict.items() if k in self._configurable_fields}
|
||||
# SECURITY: drop static/infrastructure + credential fields, then permission-filter.
|
||||
filtered = self._strip_static_and_credential_fields(config_dict)
|
||||
return await self._apply_permission_filter(filtered, bank_id, context)
|
||||
|
||||
# SECURITY: Remove ALL credential fields (API keys, base URLs, etc.)
|
||||
filtered = {k: v for k, v in filtered.items() if k not in self._credential_fields}
|
||||
def _strip_static_and_credential_fields(self, config_dict: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Keep only configurable, non-credential fields.
|
||||
|
||||
# PERMISSIONS: Further filter based on tenant/bank permissions
|
||||
SECURITY: excludes static/infrastructure fields and ALL credential fields
|
||||
(API keys, base URLs, etc.) so a resolved config is safe to return over the API.
|
||||
"""
|
||||
return {
|
||||
k: v for k, v in config_dict.items() if k in self._configurable_fields and k not in self._credential_fields
|
||||
}
|
||||
|
||||
async def _apply_permission_filter(
|
||||
self, filtered: dict[str, Any], bank_id: str, context: RequestContext | None
|
||||
) -> dict[str, Any]:
|
||||
"""Further restrict already-stripped config to the tenant/bank permission allow-list.
|
||||
|
||||
On extension error, leaves ``filtered`` unchanged (parity with the historical
|
||||
single-bank path: a permissions lookup failure must not leak or drop fields).
|
||||
"""
|
||||
if not (self.tenant_extension and context):
|
||||
return filtered
|
||||
try:
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
filtered = {k: v for k, v in filtered.items() if k in allowed_fields}
|
||||
logger.debug(
|
||||
f"Applied permission filter for bank {bank_id}: allowed={len(allowed_fields)} fields, "
|
||||
f"returned={len(filtered)} fields"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load permissions for bank {bank_id}: {e}")
|
||||
return filtered
|
||||
|
||||
async def get_bank_configs(
|
||||
self, bank_ids: list[str], context: RequestContext | None = None
|
||||
) -> dict[str, dict[str, Any]]:
|
||||
"""Batch variant of :meth:`get_bank_config` for many banks.
|
||||
|
||||
Equivalent to calling ``get_bank_config`` per bank, but resolves the
|
||||
global + tenant base once and loads every bank's ``banks.config`` JSONB
|
||||
in a single query, instead of one config round-trip per bank. Used by
|
||||
``list_banks`` to overlay disposition + mission without an N+1.
|
||||
|
||||
Returns a mapping of bank_id -> filtered configurable-field dict. A bank
|
||||
with no config row still appears, mapped to the global+tenant base.
|
||||
"""
|
||||
if not bank_ids:
|
||||
return {}
|
||||
|
||||
# Global + tenant base, resolved once (tenant override is per-request, not per-bank).
|
||||
base_dict = asdict(self._global_config)
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
filtered = {k: v for k, v in filtered.items() if k in allowed_fields}
|
||||
logger.debug(
|
||||
f"Applied permission filter for bank {bank_id}: allowed={len(allowed_fields)} fields, "
|
||||
f"returned={len(filtered)} fields"
|
||||
)
|
||||
tenant_overrides = await self.tenant_extension.get_tenant_config(context)
|
||||
if tenant_overrides:
|
||||
normalized_tenant = normalize_config_dict(tenant_overrides)
|
||||
base_dict.update({k: v for k, v in normalized_tenant.items() if k in self._configurable_fields})
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load permissions for bank {bank_id}: {e}")
|
||||
logger.warning(f"Failed to load tenant config for bulk resolve: {e}")
|
||||
|
||||
return filtered
|
||||
# All bank overrides in one query, then merge + strip per bank.
|
||||
bank_overrides = await self._load_bank_configs(bank_ids)
|
||||
stripped = {
|
||||
bank_id: self._strip_static_and_credential_fields({**base_dict, **bank_overrides.get(bank_id, {})})
|
||||
for bank_id in bank_ids
|
||||
}
|
||||
|
||||
# Permission filter is per-bank; resolve concurrently when an extension is present.
|
||||
if not (self.tenant_extension and context):
|
||||
return stripped
|
||||
permission_filtered = await asyncio.gather(
|
||||
*(self._apply_permission_filter(stripped[bank_id], bank_id, context) for bank_id in bank_ids)
|
||||
)
|
||||
return dict(zip(bank_ids, permission_filtered, strict=True))
|
||||
|
||||
async def _load_bank_config(self, bank_id: str) -> dict[str, Any]:
|
||||
"""
|
||||
@@ -219,6 +277,45 @@ class ConfigResolver:
|
||||
|
||||
return {}
|
||||
|
||||
async def _load_bank_configs(self, bank_ids: list[str]) -> dict[str, dict[str, Any]]:
|
||||
"""Bulk variant of :meth:`_load_bank_config`: load many banks' overrides in one query.
|
||||
|
||||
Returns a mapping of bank_id -> normalized active overrides. Banks with no row
|
||||
(or an empty/all-tombstone config) are simply absent from the mapping.
|
||||
"""
|
||||
result: dict[str, dict[str, Any]] = {}
|
||||
if not bank_ids:
|
||||
return result
|
||||
try:
|
||||
async with self._backend.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT bank_id, config FROM {fq_table("banks")} WHERE bank_id = ANY($1)
|
||||
""",
|
||||
bank_ids,
|
||||
)
|
||||
for row in rows:
|
||||
config_data = row["config"]
|
||||
if not config_data:
|
||||
continue
|
||||
# Handle case where JSONB is returned as JSON string
|
||||
if isinstance(config_data, str):
|
||||
config_data = json.loads(config_data)
|
||||
|
||||
# Normalize keys (handle both env var format and Python field format)
|
||||
normalized = normalize_config_dict(config_data)
|
||||
|
||||
# Only active overrides for configurable fields. JSON null is a tombstone
|
||||
# for "Server Default" in the bank-config UI and must not override defaults.
|
||||
overrides = {
|
||||
k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None
|
||||
}
|
||||
if overrides:
|
||||
result[row["bank_id"]] = overrides
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to bulk-load bank configs: {e}")
|
||||
return result
|
||||
|
||||
async def update_bank_config(
|
||||
self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None
|
||||
) -> None:
|
||||
@@ -305,6 +402,9 @@ class ConfigResolver:
|
||||
# Validate recall budget fields
|
||||
_validate_recall_budget_updates(normalized_updates)
|
||||
|
||||
# Validate disposition trait fields (1-5 integer scale)
|
||||
_validate_disposition_updates(normalized_updates)
|
||||
|
||||
chunking_fields_updated = (
|
||||
"retain_chunk_size" in normalized_updates
|
||||
or "retain_structured_chunk_size" in normalized_updates
|
||||
@@ -419,6 +519,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.
|
||||
|
||||
@@ -96,7 +96,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 +109,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 +122,13 @@ 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()
|
||||
|
||||
@@ -449,6 +449,13 @@ class _CreateAction(BaseModel):
|
||||
def sanitize_text(cls, v: str) -> str:
|
||||
return sanitize_llm_output(v) or ""
|
||||
|
||||
@field_validator("source_fact_ids", mode="before")
|
||||
@classmethod
|
||||
def ensure_list(cls, v: str | list[str]) -> list[str]:
|
||||
if isinstance(v, str):
|
||||
return [v]
|
||||
return v
|
||||
|
||||
|
||||
class _UpdateAction(BaseModel):
|
||||
text: str
|
||||
@@ -461,6 +468,13 @@ class _UpdateAction(BaseModel):
|
||||
def sanitize_text(cls, v: str) -> str:
|
||||
return sanitize_llm_output(v) or ""
|
||||
|
||||
@field_validator("source_fact_ids", mode="before")
|
||||
@classmethod
|
||||
def ensure_list(cls, v: str | list[str]) -> list[str]:
|
||||
if isinstance(v, str):
|
||||
return [v]
|
||||
return v
|
||||
|
||||
|
||||
class _DeleteAction(BaseModel):
|
||||
observation_id: str # UUID of the observation to remove
|
||||
@@ -640,6 +654,7 @@ class ConsolidationPerfLog:
|
||||
self.start_time = time.time()
|
||||
self.lines: list[str] = []
|
||||
self.timings: dict[str, float] = {}
|
||||
self.timing_counts: dict[str, int] = {}
|
||||
self.llm_calls: int = 0
|
||||
self.total_obs_in_context: int = 0
|
||||
self.total_prompt_chars: int = 0
|
||||
@@ -649,11 +664,13 @@ class ConsolidationPerfLog:
|
||||
self.lines.append(message)
|
||||
|
||||
def record_timing(self, key: str, duration: float) -> None:
|
||||
"""Record a timing measurement."""
|
||||
if key in self.timings:
|
||||
self.timings[key] += duration
|
||||
else:
|
||||
self.timings[key] = duration
|
||||
"""Record a timing measurement.
|
||||
|
||||
Tracks both total seconds and call count so the summary can
|
||||
distinguish one slow call from many fast calls in aggregate.
|
||||
"""
|
||||
self.timings[key] = self.timings.get(key, 0.0) + duration
|
||||
self.timing_counts[key] = self.timing_counts.get(key, 0) + 1
|
||||
|
||||
def record_llm_call(self, obs_count: int, prompt_chars: int) -> None:
|
||||
"""Record stats for a single LLM call."""
|
||||
@@ -676,6 +693,8 @@ class ConsolidationPerfLog:
|
||||
"""
|
||||
for key, value in other.timings.items():
|
||||
self.timings[key] = self.timings.get(key, 0.0) + value
|
||||
for key, count in other.timing_counts.items():
|
||||
self.timing_counts[key] = self.timing_counts.get(key, 0) + count
|
||||
self.llm_calls += other.llm_calls
|
||||
self.total_obs_in_context += other.total_obs_in_context
|
||||
self.total_prompt_chars += other.total_prompt_chars
|
||||
@@ -1276,16 +1295,22 @@ async def _run_consolidation_job(
|
||||
f"{stats['skipped']} skipped)"
|
||||
)
|
||||
|
||||
# Add timing breakdown
|
||||
# Add timing breakdown. Each phase is recorded once per call, so the count
|
||||
# disambiguates a single slow call from many fast calls — important for
|
||||
# operators triaging "the recall phase took 15s" log lines, where the
|
||||
# total is the sum of many serial sub-calls rather than one slow query.
|
||||
def _fmt(key: str) -> str:
|
||||
total = perf.timings[key]
|
||||
count = perf.timing_counts.get(key, 0)
|
||||
if count > 1:
|
||||
avg_ms = total * 1000.0 / count
|
||||
return f"{key}={total:.3f}s ({count} calls, avg={avg_ms:.0f}ms)"
|
||||
return f"{key}={total:.3f}s"
|
||||
|
||||
timing_parts = []
|
||||
if "recall" in perf.timings:
|
||||
timing_parts.append(f"recall={perf.timings['recall']:.3f}s")
|
||||
if "llm" in perf.timings:
|
||||
timing_parts.append(f"llm={perf.timings['llm']:.3f}s")
|
||||
if "embedding" in perf.timings:
|
||||
timing_parts.append(f"embedding={perf.timings['embedding']:.3f}s")
|
||||
if "db_write" in perf.timings:
|
||||
timing_parts.append(f"db_write={perf.timings['db_write']:.3f}s")
|
||||
for key in ("recall", "llm", "embedding", "db_write"):
|
||||
if key in perf.timings:
|
||||
timing_parts.append(_fmt(key))
|
||||
|
||||
if perf.llm_calls > 0:
|
||||
timing_parts.append(f"avg_obs={perf.total_obs_in_context / perf.llm_calls:.1f}")
|
||||
|
||||
@@ -212,7 +212,7 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
device = "cpu"
|
||||
logger.info("Reranker: forcing CPU mode (HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU=1)")
|
||||
else:
|
||||
# Check for GPU (CUDA) or Apple Silicon (MPS)
|
||||
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
|
||||
# Wrap in try-except to gracefully handle any device detection issues
|
||||
# (e.g., in CI environments or when PyTorch is built without GPU support)
|
||||
device = "cpu" # Default to CPU
|
||||
@@ -220,10 +220,13 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
has_gpu = torch.cuda.is_available() or (
|
||||
hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
)
|
||||
# Intel Arc XPU support — torch.xpu is available when the XPU build is loaded
|
||||
if not has_gpu and hasattr(torch, "xpu"):
|
||||
has_gpu = torch.xpu.is_available()
|
||||
if has_gpu:
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS/XPU
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}")
|
||||
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
|
||||
|
||||
# Patch transformers 5.x compatibility for models using XLM-RoBERTa
|
||||
# (e.g., jina-reranker-v2-base-multilingual). transformers 5.x removed
|
||||
|
||||
@@ -256,13 +256,21 @@ class OracleOps(DataAccessOps):
|
||||
# Oracle doesn't support ON CONFLICT; rely on the PK and the
|
||||
# IGNORE_ROW_ON_DUPKEY_INDEX hint to skip duplicates server-side.
|
||||
# The hint name must match the PK constraint exactly.
|
||||
#
|
||||
# Sort to enforce a global lock-acquisition order on the
|
||||
# (bank_id, unit_id) PK. Without this, two concurrent
|
||||
# transactions inserting overlapping unit_id sets in different
|
||||
# orders can deadlock on the unique-check row locks. Sorting
|
||||
# gives every concurrent caller the same lock order, so
|
||||
# conflicting inserts queue cleanly instead of cycling.
|
||||
sorted_unit_ids = sorted(unit_ids)
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT /*+ IGNORE_ROW_ON_DUPKEY_INDEX({table}, pk_graph_maintenance_queue) */
|
||||
INTO {table} (bank_id, unit_id)
|
||||
VALUES ($1, $2)
|
||||
""",
|
||||
[(bank_id, uid) for uid in unit_ids],
|
||||
[(bank_id, uid) for uid in sorted_unit_ids],
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
|
||||
@@ -348,6 +348,15 @@ class PostgreSQLOps(DataAccessOps):
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
# Sort to enforce a global lock-acquisition order on the
|
||||
# (bank_id, unit_id) unique-key. Without this, two concurrent
|
||||
# transactions inserting overlapping unit_id sets in different
|
||||
# orders can deadlock on the ON CONFLICT row locks — Postgres
|
||||
# acquires a short-lived lock per row being checked, and cycle
|
||||
# detection then aborts one transaction. Sorting gives every
|
||||
# concurrent caller the same lock order, so conflicting inserts
|
||||
# queue cleanly instead of cycling.
|
||||
sorted_unit_ids = sorted(unit_ids)
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (bank_id, unit_id)
|
||||
@@ -355,7 +364,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
ON CONFLICT (bank_id, unit_id) DO NOTHING
|
||||
""",
|
||||
bank_id,
|
||||
unit_ids,
|
||||
sorted_unit_ids,
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
|
||||
@@ -190,7 +190,7 @@ class LocalSTEmbeddings(Embeddings):
|
||||
device = "cpu"
|
||||
logger.info("Embeddings: forcing CPU mode")
|
||||
else:
|
||||
# Check for GPU (CUDA) or Apple Silicon (MPS)
|
||||
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
|
||||
# Wrap in try-except to gracefully handle any device detection issues
|
||||
# (e.g., in CI environments or when PyTorch is built without GPU support)
|
||||
device = "cpu" # Default to CPU
|
||||
@@ -198,10 +198,13 @@ class LocalSTEmbeddings(Embeddings):
|
||||
has_gpu = torch.cuda.is_available() or (
|
||||
hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
)
|
||||
# Intel Arc XPU support — torch.xpu is available when the XPU build is loaded
|
||||
if not has_gpu and hasattr(torch, "xpu"):
|
||||
has_gpu = torch.xpu.is_available()
|
||||
if has_gpu:
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS/XPU
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}")
|
||||
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
|
||||
|
||||
# Suppress verbose transformers warnings during model loading
|
||||
# This suppresses the "UNEXPECTED" warnings from BertModel which are harmless
|
||||
@@ -709,7 +712,8 @@ class OpenAIEmbeddings(Embeddings):
|
||||
|
||||
class CodexOAuthEmbeddings(OpenAIEmbeddings):
|
||||
"""
|
||||
OpenAI embeddings using the Codex/ChatGPT OAuth token from ``~/.codex/auth.json``.
|
||||
OpenAI embeddings using the Codex/ChatGPT OAuth token from the Codex
|
||||
``auth.json`` (``$CODEX_HOME/auth.json``, or ``~/.codex/auth.json`` when unset).
|
||||
|
||||
Codex OAuth is an LLM-provider auth path in Hindsight, but the same bearer token
|
||||
can also authenticate against the standard OpenAI embeddings endpoint. This keeps
|
||||
|
||||
@@ -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]],
|
||||
|
||||
@@ -6,6 +6,7 @@ enabling support for multiple LLM backends (OpenAI, Anthropic, Gemini, Codex, et
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from .response_models import LLMToolCallResult
|
||||
@@ -252,3 +253,11 @@ class OutputTooLongError(Exception):
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ProviderRateLimitResetError(Exception):
|
||||
"""Raised when an upstream provider says quota will reopen at a known time."""
|
||||
|
||||
def __init__(self, retry_at: datetime, message: str = "") -> None:
|
||||
self.retry_at = retry_at
|
||||
super().__init__(message)
|
||||
|
||||
@@ -10,7 +10,6 @@ import re
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import AsyncExitStack
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
# Vertex AI imports (conditional - for LLMProvider to pass credentials to GeminiLLM)
|
||||
@@ -253,6 +252,7 @@ def create_llm_provider(
|
||||
gemini_safety_settings: list | None = None,
|
||||
prompt_cache_enabled: bool = False,
|
||||
litellmrouter_config: dict[str, Any] | None = None,
|
||||
gemini_service_tier: str | None = None,
|
||||
) -> Any: # Returns LLMInterface
|
||||
"""
|
||||
Factory function to create the appropriate LLM provider implementation.
|
||||
@@ -266,6 +266,7 @@ def create_llm_provider(
|
||||
groq_service_tier: Groq service tier (for Groq provider) - "on_demand", "flex", or "auto".
|
||||
openai_service_tier: OpenAI service tier (for OpenAI provider) - None (default) or "flex" (50% cheaper).
|
||||
bedrock_service_tier: Bedrock service tier (for Bedrock provider) - None (default), "flex", "priority", or "reserved".
|
||||
gemini_service_tier: Gemini service tier (for Gemini provider) - None (default) or "flex" (50% cheaper).
|
||||
extra_body: Extra request-body params merged into the provider's native
|
||||
call. Threaded into OpenAI-compatible, Fireworks, Anthropic, Gemini/
|
||||
VertexAI and LiteLLM providers (each merges them in its own parameter
|
||||
@@ -296,6 +297,12 @@ def create_llm_provider(
|
||||
)
|
||||
|
||||
provider_lower = provider.lower()
|
||||
if provider_lower == "gemini":
|
||||
from ..config import parse_gemini_service_tier
|
||||
|
||||
gemini_service_tier = parse_gemini_service_tier(gemini_service_tier)
|
||||
else:
|
||||
gemini_service_tier = None
|
||||
|
||||
if provider_lower == "openai-codex":
|
||||
return CodexLLM(
|
||||
@@ -344,6 +351,7 @@ def create_llm_provider(
|
||||
vertexai_region=vertexai_region,
|
||||
vertexai_credentials=vertexai_credentials,
|
||||
gemini_safety_settings=gemini_safety_settings,
|
||||
gemini_service_tier=gemini_service_tier,
|
||||
prompt_cache_enabled=prompt_cache_enabled,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
@@ -458,6 +466,7 @@ def create_llm_provider(
|
||||
"openrouter",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
):
|
||||
return OpenAICompatibleLLM(
|
||||
provider=provider,
|
||||
@@ -496,6 +505,7 @@ class LLMProvider:
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
default_headers: dict[str, str] | None = None,
|
||||
litellmrouter_config: dict[str, Any] | None = None,
|
||||
gemini_service_tier: str | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize LLM provider.
|
||||
@@ -509,6 +519,7 @@ class LLMProvider:
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto") - from config.
|
||||
openai_service_tier: OpenAI service tier (None or "flex") - from config.
|
||||
bedrock_service_tier: Bedrock service tier (None, "flex", "priority", "reserved") - from config.
|
||||
gemini_service_tier: Gemini service tier (None or "flex") - from config.
|
||||
gemini_safety_settings: Safety settings for Gemini/VertexAI providers.
|
||||
extra_body: Extra request-body params merged into the provider's native call
|
||||
(OpenAI-compatible, Fireworks, Anthropic, Gemini/VertexAI, LiteLLM).
|
||||
@@ -532,6 +543,7 @@ class LLMProvider:
|
||||
self.groq_service_tier = groq_service_tier
|
||||
self.openai_service_tier = openai_service_tier
|
||||
self.bedrock_service_tier = bedrock_service_tier
|
||||
self.gemini_service_tier = gemini_service_tier
|
||||
# Gemini safety settings (instance default; can be overridden per-request via context var)
|
||||
self.gemini_safety_settings = gemini_safety_settings
|
||||
# Gemini prompt caching: when True, retain extraction (and any future
|
||||
@@ -577,6 +589,7 @@ class LLMProvider:
|
||||
"openrouter",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
"fireworks",
|
||||
"nous",
|
||||
]
|
||||
@@ -603,6 +616,8 @@ class LLMProvider:
|
||||
self.base_url = "https://api.z.ai/api/coding/paas/v4"
|
||||
elif self.provider == "opencode-go":
|
||||
self.base_url = "https://opencode.ai/zen/go/v1"
|
||||
elif self.provider == "atlas":
|
||||
self.base_url = "https://api.atlascloud.ai/v1"
|
||||
elif self.provider == "nous":
|
||||
self.base_url = "https://inference-api.nousresearch.com/v1"
|
||||
|
||||
@@ -660,6 +675,22 @@ class LLMProvider:
|
||||
except Exception:
|
||||
pass # Config may not be initialized in test environments
|
||||
|
||||
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":
|
||||
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
|
||||
@@ -698,6 +729,7 @@ class LLMProvider:
|
||||
groq_service_tier=self.groq_service_tier,
|
||||
openai_service_tier=self.openai_service_tier,
|
||||
bedrock_service_tier=self.bedrock_service_tier,
|
||||
gemini_service_tier=self.gemini_service_tier,
|
||||
extra_body=self.extra_body,
|
||||
default_headers=self.default_headers,
|
||||
vertexai_project_id=vertexai_project_id,
|
||||
@@ -1023,7 +1055,9 @@ class LLMProvider:
|
||||
|
||||
def _load_codex_auth(self) -> tuple[str, str]:
|
||||
"""
|
||||
Load OAuth credentials from ~/.codex/auth.json.
|
||||
Load OAuth credentials from the Codex ``auth.json``.
|
||||
|
||||
Honors ``CODEX_HOME`` (falling back to ``~/.codex``).
|
||||
|
||||
Returns:
|
||||
Tuple of (access_token, account_id).
|
||||
@@ -1032,7 +1066,9 @@ class LLMProvider:
|
||||
FileNotFoundError: If auth file doesn't exist.
|
||||
ValueError: If auth file is invalid.
|
||||
"""
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
from .providers.codex_auth import default_codex_auth_file
|
||||
|
||||
auth_file = default_codex_auth_file()
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(
|
||||
@@ -1142,10 +1178,12 @@ class LLMProvider:
|
||||
ENV_LLM_BEDROCK_SERVICE_TIER,
|
||||
ENV_LLM_DEFAULT_HEADERS,
|
||||
ENV_LLM_EXTRA_BODY,
|
||||
ENV_LLM_GEMINI_SERVICE_TIER,
|
||||
ENV_LLM_MODEL,
|
||||
ENV_LLM_PROVIDER,
|
||||
ENV_LLM_REASONING_EFFORT,
|
||||
_get_default_model_for_provider,
|
||||
parse_gemini_service_tier,
|
||||
)
|
||||
|
||||
provider = os.getenv(ENV_LLM_PROVIDER, DEFAULT_LLM_PROVIDER)
|
||||
@@ -1172,6 +1210,11 @@ class LLMProvider:
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
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
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -45,6 +45,7 @@ from .audit import AuditLogger, audit_context
|
||||
from .bank_stats_cache import BankStatsCache
|
||||
from .db import DatabaseBackend, create_database_backend
|
||||
from .db_budget import budgeted_operation
|
||||
from .llm_interface import ProviderRateLimitResetError
|
||||
from .llm_trace import (
|
||||
LLMRequestEntry,
|
||||
LLMRequestListResponse,
|
||||
@@ -744,6 +745,37 @@ def _resolve_refresh_tag_filtering(
|
||||
return RefreshTagFiltering(tags=model_tags, tags_match=tags_match, tag_groups=None)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResolvedDispositionMission:
|
||||
"""Disposition + mission after overlaying resolved bank config on the legacy columns."""
|
||||
|
||||
disposition: dict[str, int]
|
||||
mission: str
|
||||
|
||||
|
||||
def _overlay_bank_config_disposition_mission(
|
||||
disposition: dict[str, int], mission: str, config_dict: dict[str, Any]
|
||||
) -> ResolvedDispositionMission:
|
||||
"""Overlay resolved bank config on top of the legacy banks.disposition /
|
||||
banks.mission column values.
|
||||
|
||||
``reflect_mission`` and ``disposition_*`` in the resolved bank config take
|
||||
precedence over the legacy DB columns. Shared by ``get_bank_profile`` and
|
||||
``list_banks`` so the single-bank and list paths return identical
|
||||
disposition + mission for the same bank.
|
||||
"""
|
||||
resolved_mission = config_dict.get("reflect_mission") or mission
|
||||
cfg_skep = config_dict.get("disposition_skepticism")
|
||||
cfg_lit = config_dict.get("disposition_literalism")
|
||||
cfg_emp = config_dict.get("disposition_empathy")
|
||||
resolved_disposition = {
|
||||
"skepticism": cfg_skep if cfg_skep is not None else disposition["skepticism"],
|
||||
"literalism": cfg_lit if cfg_lit is not None else disposition["literalism"],
|
||||
"empathy": cfg_emp if cfg_emp is not None else disposition["empathy"],
|
||||
}
|
||||
return ResolvedDispositionMission(disposition=resolved_disposition, mission=resolved_mission)
|
||||
|
||||
|
||||
class MemoryEngine(MemoryEngineInterface):
|
||||
"""
|
||||
Advanced memory system using temporal and semantic linking with PostgreSQL.
|
||||
@@ -933,6 +965,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
default_headers=config.llm_default_headers,
|
||||
litellmrouter_config=config.llm_litellmrouter_config,
|
||||
bedrock_service_tier=config.llm_bedrock_service_tier,
|
||||
gemini_service_tier=config.llm_gemini_service_tier,
|
||||
)
|
||||
|
||||
# Store client and model for convenience (deprecated: use _llm_config.call() instead)
|
||||
@@ -966,6 +999,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
default_headers=config.llm_default_headers,
|
||||
litellmrouter_config=config.retain_llm_litellmrouter_config or config.llm_litellmrouter_config,
|
||||
bedrock_service_tier=config.llm_bedrock_service_tier,
|
||||
gemini_service_tier=config.llm_gemini_service_tier,
|
||||
)
|
||||
|
||||
# Reflect LLM config - for think/observe operations (can use lighter models)
|
||||
@@ -994,6 +1028,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
default_headers=config.llm_default_headers,
|
||||
litellmrouter_config=config.reflect_llm_litellmrouter_config or config.llm_litellmrouter_config,
|
||||
bedrock_service_tier=config.llm_bedrock_service_tier,
|
||||
gemini_service_tier=config.llm_gemini_service_tier,
|
||||
)
|
||||
|
||||
# Consolidation LLM config - for mental model consolidation (can use efficient models)
|
||||
@@ -1022,6 +1057,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
default_headers=config.llm_default_headers,
|
||||
litellmrouter_config=config.consolidation_llm_litellmrouter_config or config.llm_litellmrouter_config,
|
||||
bedrock_service_tier=config.llm_bedrock_service_tier,
|
||||
gemini_service_tier=config.llm_gemini_service_tier,
|
||||
)
|
||||
|
||||
# Initialize cross-encoder reranker (cached for performance)
|
||||
@@ -1752,6 +1788,9 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
audit_entry.response = {"status": "completed", "operation_id": operation_id}
|
||||
|
||||
except ProviderRateLimitResetError as e:
|
||||
logger.warning(f"Task deferred until provider quota resets at {e.retry_at}: {e}")
|
||||
raise DeferOperation(exec_date=e.retry_at, reason=str(e)) from e
|
||||
except RetryTaskAt:
|
||||
# Task-owned retry: let the poller handle scheduling
|
||||
raise
|
||||
@@ -2751,7 +2790,15 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
self._parser_registry = FileParserRegistry()
|
||||
try:
|
||||
self._parser_registry.register(MarkitdownParser())
|
||||
self._parser_registry.register(
|
||||
MarkitdownParser(
|
||||
ocr_enabled=config.file_parser_markitdown_ocr_enabled,
|
||||
ocr_api_key=config.file_parser_markitdown_ocr_api_key,
|
||||
ocr_base_url=config.file_parser_markitdown_ocr_base_url,
|
||||
ocr_model=config.file_parser_markitdown_ocr_model,
|
||||
ocr_prompt=config.file_parser_markitdown_ocr_prompt,
|
||||
)
|
||||
)
|
||||
logger.debug("Registered markitdown parser")
|
||||
except ImportError:
|
||||
logger.warning("markitdown not available - file parsing disabled")
|
||||
@@ -3287,6 +3334,28 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
sub_doc_id = document_id or (sub_batch[0].get("document_id") if len(sub_batch) == 1 else None)
|
||||
sub_offset = chunk_offsets.get(sub_doc_id, 0) if sub_doc_id else 0
|
||||
|
||||
# Count the chunks this sub-batch will produce BEFORE handing it
|
||||
# to the orchestrator. retain_batch consumes (pops) each item's
|
||||
# "content" while streaming, so reading it back after the call
|
||||
# yields "" — and chunk_text("") returns [""] (count 1),
|
||||
# advancing the per-document cursor by 1 regardless of the real
|
||||
# chunk count. For slices that each span several chunks the next
|
||||
# sub-batch then restarts ~1 slot in, colliding chunk_ids and
|
||||
# overwriting earlier chunks (only ~1 new chunk survives per
|
||||
# sub-batch). Capture it here while content is still present.
|
||||
sub_chunk_count = 0
|
||||
if sub_doc_id:
|
||||
sub_chunk_count = sum(
|
||||
len(
|
||||
fact_extraction.chunk_text(
|
||||
item.get("content", "") or "",
|
||||
chunking_config.chunk_size,
|
||||
structured_chunk_size=chunking_config.structured_chunk_size,
|
||||
)
|
||||
)
|
||||
for item in sub_batch
|
||||
)
|
||||
|
||||
sub_results, sub_usage, sub_processed = await self._retain_batch_async_internal(
|
||||
bank_id=bank_id,
|
||||
contents=sub_batch,
|
||||
@@ -3306,20 +3375,10 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
)
|
||||
|
||||
# Advance the document's chunk_index cursor by the number of
|
||||
# chunks this sub-batch produced (computed with the same chunk
|
||||
# size the orchestrator uses), so the next sub-batch sharing the
|
||||
# document continues the sequence.
|
||||
# chunks this sub-batch produced (counted above, before the
|
||||
# orchestrator consumed the content), so the next sub-batch
|
||||
# sharing the document continues the sequence.
|
||||
if sub_doc_id:
|
||||
sub_chunk_count = sum(
|
||||
len(
|
||||
fact_extraction.chunk_text(
|
||||
item.get("content", "") or "",
|
||||
chunking_config.chunk_size,
|
||||
structured_chunk_size=chunking_config.structured_chunk_size,
|
||||
)
|
||||
)
|
||||
for item in sub_batch
|
||||
)
|
||||
# retain_batch only prepends the existing body on the global
|
||||
# first sub-batch (is_first_batch == i == 1), so fold its chunk
|
||||
# count in only there.
|
||||
@@ -3383,6 +3442,8 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
llm_input_tokens=total_usage.input_tokens,
|
||||
llm_output_tokens=total_usage.output_tokens,
|
||||
llm_total_tokens=total_usage.total_tokens,
|
||||
llm_cached_input_tokens=getattr(total_usage, "cached_tokens", 0) or 0,
|
||||
llm_thoughts_tokens=getattr(total_usage, "thoughts_tokens", 0) or 0,
|
||||
processed_content_tokens=total_processed_content_tokens,
|
||||
)
|
||||
try:
|
||||
@@ -3784,6 +3845,10 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
max_tokens: int = 4096,
|
||||
enable_trace: bool = False,
|
||||
fact_type: list[str] | None = None,
|
||||
# Opt-in (default False). Internal callers that recall raw facts on purpose —
|
||||
# notably consolidation, which needs the raw facts it folds into observations —
|
||||
# must leave this off so they aren't silently deduped away.
|
||||
prefer_observations: bool = False,
|
||||
question_date: datetime | None = None,
|
||||
include_entities: bool = False,
|
||||
max_entity_tokens: int = 500,
|
||||
@@ -3816,6 +3881,10 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
bank_id: bank ID to recall for
|
||||
query: Recall query
|
||||
fact_type: List of fact types to recall (e.g., ['world', 'experience'])
|
||||
prefer_observations: When True and both 'observation' and a raw type ('world'/'experience')
|
||||
are requested, drop raw facts that a returned observation was consolidated from
|
||||
(deduplication by provenance). Freed slots backfill, keeping the result count at
|
||||
the budget. No-op unless both observation and raw types are requested.
|
||||
budget: Budget level for graph traversal (low=100, mid=300, high=600 units)
|
||||
max_tokens: Maximum tokens to return (counts only 'text' field, default 4096)
|
||||
Results are returned until token budget is reached, stopping before
|
||||
@@ -3953,6 +4022,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
max_chunk_tokens,
|
||||
request_context,
|
||||
semaphore_wait=semaphore_wait,
|
||||
prefer_observations=prefer_observations,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
@@ -4089,6 +4159,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
max_chunk_tokens: int = 8192,
|
||||
request_context: "RequestContext" = None,
|
||||
semaphore_wait: float = 0.0,
|
||||
prefer_observations: bool = False,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
@@ -4563,10 +4634,14 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
ce = reranker_instance.cross_encoder
|
||||
# "rrf" mode is passthrough by construction; so is a configured "rrf" CE.
|
||||
is_passthrough = (reranking == "rrf") or (ce is not None and ce.provider_name == "rrf")
|
||||
scoring_config = get_config()
|
||||
apply_combined_scoring(
|
||||
scored_results,
|
||||
now=_recall_scoring_now(question_date),
|
||||
is_passthrough_reranker=is_passthrough,
|
||||
recency_decay_function=scoring_config.recency_decay_function,
|
||||
recency_decay_linear_window_days=scoring_config.recency_decay_linear_window_days,
|
||||
recency_decay_halflife_days=scoring_config.recency_decay_halflife_days,
|
||||
)
|
||||
# Per-strategy additive boost: nudge candidates surfaced by a
|
||||
# prioritised retrieval arm up the final ordering.
|
||||
@@ -4601,6 +4676,48 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
if request_context is not None:
|
||||
request_context.raise_if_cancelled()
|
||||
|
||||
# Step 4.8: prefer-observations dedup. When the caller asked for observations
|
||||
# alongside raw facts, an observation supersedes the raw facts it was
|
||||
# consolidated from: drop those raw facts so the same content isn't returned
|
||||
# twice. Runs BEFORE the Step 5 truncation so the freed slots backfill with
|
||||
# the next-best results, keeping the result count at the budget. No-op unless
|
||||
# 'observation' and at least one raw type were both requested.
|
||||
raw_types_requested = {"world", "experience"} & set(fact_type)
|
||||
if prefer_observations and "observation" in fact_type and raw_types_requested:
|
||||
# "The observation list" = observations within the window we would return.
|
||||
# Only those can supersede a raw fact; a far-down observation should not
|
||||
# suppress a top raw fact it merely happens to reference.
|
||||
observation_ids = [
|
||||
uuid.UUID(sr.id)
|
||||
for sr in scored_results[: thinking_budget * 2]
|
||||
if sr.retrieval.fact_type == "observation"
|
||||
]
|
||||
if observation_ids:
|
||||
superseded_ids: set[str] = set()
|
||||
async with acquire_with_retry(backend) as dedup_conn:
|
||||
obs_rows = await dedup_conn.fetch(
|
||||
f"""
|
||||
SELECT source_memory_ids
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[]) AND fact_type = 'observation'
|
||||
""",
|
||||
observation_ids,
|
||||
)
|
||||
for obs_row in obs_rows:
|
||||
for sid in obs_row["source_memory_ids"] or []:
|
||||
superseded_ids.add(str(sid))
|
||||
if superseded_ids:
|
||||
before_count = len(scored_results)
|
||||
scored_results = [
|
||||
sr
|
||||
for sr in scored_results
|
||||
if not (sr.retrieval.fact_type in ("world", "experience") and sr.id in superseded_ids)
|
||||
]
|
||||
log_buffer.append(
|
||||
f" [4.8] prefer_observations: dropped {before_count - len(scored_results)} "
|
||||
f"raw fact(s) superseded by {len(observation_ids)} observation(s)"
|
||||
)
|
||||
|
||||
# Step 5: Truncate to thinking_budget * 2 for token filtering
|
||||
rerank_limit = thinking_budget * 2
|
||||
top_scored = scored_results[:rerank_limit]
|
||||
@@ -5275,6 +5392,17 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
"memory_units_deleted": units_count if deleted else 0,
|
||||
}
|
||||
|
||||
# Drop any cached stats for this bank — deleting the document changed
|
||||
# the document count and (via cascade) the memory-unit/link counts
|
||||
# get_bank_stats reports, which the TTL would otherwise serve at
|
||||
# pre-delete values for up to a minute (mirrors delete_bank). Best-effort:
|
||||
# a cache-eviction failure must not fail an already-committed delete.
|
||||
if deleted:
|
||||
try:
|
||||
await self._bank_stats_cache.invalidate(get_current_schema(), bank_id)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to invalidate bank stats cache after document deletion for bank {bank_id}: {e}")
|
||||
|
||||
if invalidated_obs > 0:
|
||||
config = await self._config_resolver.resolve_full_config(bank_id, request_context)
|
||||
if config.enable_auto_consolidation:
|
||||
@@ -5442,6 +5570,14 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
)
|
||||
|
||||
if invalidated_obs > 0:
|
||||
# Observation units were deleted, changing the counts get_bank_stats
|
||||
# reports — drop the cached stats so the TTL does not serve pre-update
|
||||
# values for up to a minute (mirrors delete_bank). Best-effort: a
|
||||
# cache-eviction failure must not fail an already-committed update.
|
||||
try:
|
||||
await self._bank_stats_cache.invalidate(get_current_schema(), bank_id)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to invalidate bank stats cache after document update for bank {bank_id}: {e}")
|
||||
config = await self._config_resolver.resolve_full_config(bank_id, request_context)
|
||||
if config.enable_auto_consolidation:
|
||||
try:
|
||||
@@ -5533,6 +5669,19 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
else "Memory unit not found",
|
||||
}
|
||||
|
||||
# Drop any cached stats for this bank — the deleted unit (and its
|
||||
# cascaded links/entities) changed the counts get_bank_stats reports,
|
||||
# which the TTL would otherwise serve at pre-delete values for up to a
|
||||
# minute (mirrors delete_bank). Best-effort: a cache-eviction failure
|
||||
# must not fail an already-committed delete.
|
||||
if deleted and bank_id:
|
||||
try:
|
||||
await self._bank_stats_cache.invalidate(get_current_schema(), bank_id)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to invalidate bank stats cache after memory unit deletion for bank {bank_id}: {e}"
|
||||
)
|
||||
|
||||
if bank_id_for_consolidation:
|
||||
config = await self._config_resolver.resolve_full_config(bank_id_for_consolidation, request_context)
|
||||
if config.enable_auto_consolidation:
|
||||
@@ -5755,7 +5904,17 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
bank_id,
|
||||
)
|
||||
|
||||
return {"deleted_count": count or 0}
|
||||
# Drop any cached stats for this bank — clearing observations changed
|
||||
# the memory-unit/observation counts and the consolidation timestamps
|
||||
# get_bank_stats reports, which the TTL would otherwise serve at stale
|
||||
# values for up to a minute (mirrors delete_bank). Best-effort: a
|
||||
# cache-eviction failure must not fail an already-committed clear.
|
||||
try:
|
||||
await self._bank_stats_cache.invalidate(get_current_schema(), bank_id)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to invalidate bank stats cache after clearing observations for bank {bank_id}: {e}")
|
||||
|
||||
return {"deleted_count": count or 0}
|
||||
|
||||
async def list_observation_scopes(
|
||||
self,
|
||||
@@ -8112,25 +8271,15 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
# reflect_mission and disposition in config take precedence over the legacy DB columns
|
||||
config_dict = await self._config_resolver.get_bank_config(bank_id, request_context)
|
||||
mission = config_dict.get("reflect_mission") or profile["mission"]
|
||||
|
||||
# Overlay disposition from config if explicitly set; fall back to DB values
|
||||
db_disp = profile["disposition"]
|
||||
db_disp_dict = db_disp.model_dump() if hasattr(db_disp, "model_dump") else dict(db_disp)
|
||||
cfg_skep = config_dict.get("disposition_skepticism")
|
||||
cfg_lit = config_dict.get("disposition_literalism")
|
||||
cfg_emp = config_dict.get("disposition_empathy")
|
||||
disposition = {
|
||||
"skepticism": cfg_skep if cfg_skep is not None else db_disp_dict["skepticism"],
|
||||
"literalism": cfg_lit if cfg_lit is not None else db_disp_dict["literalism"],
|
||||
"empathy": cfg_emp if cfg_emp is not None else db_disp_dict["empathy"],
|
||||
}
|
||||
resolved = _overlay_bank_config_disposition_mission(db_disp_dict, profile["mission"], config_dict)
|
||||
|
||||
return {
|
||||
"bank_id": bank_id,
|
||||
"name": profile["name"],
|
||||
"disposition": disposition,
|
||||
"mission": mission,
|
||||
"disposition": resolved.disposition,
|
||||
"mission": resolved.mission,
|
||||
}
|
||||
|
||||
async def _ensure_bank_exists(
|
||||
@@ -8345,6 +8494,17 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
BankListContext(banks=banks, request_context=request_context)
|
||||
)
|
||||
banks = result.banks
|
||||
# Overlay resolved bank config (reflect_mission + disposition_*) on top of the
|
||||
# legacy banks.disposition / banks.mission columns, mirroring get_bank_profile so
|
||||
# the list and get paths return identical disposition + mission for a bank.
|
||||
# Resolve every bank's config in one batch (single config-column query + a single
|
||||
# tenant-config resolve) rather than one round-trip per bank.
|
||||
configs = await self._config_resolver.get_bank_configs([bank["bank_id"] for bank in banks], request_context)
|
||||
for bank in banks:
|
||||
resolved = _overlay_bank_config_disposition_mission(
|
||||
bank["disposition"], bank["mission"], configs.get(bank["bank_id"], {})
|
||||
)
|
||||
bank["disposition"], bank["mission"] = resolved.disposition, resolved.mission
|
||||
return banks
|
||||
|
||||
# ==================== Reflect Methods ====================
|
||||
@@ -11695,6 +11855,8 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
**task_payload,
|
||||
}
|
||||
|
||||
from hindsight_api.extensions.operation_validator import OperationValidationError
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
if dedupe_by_bank:
|
||||
@@ -11716,10 +11878,20 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
# serialize) but not with FOR KEY SHARE (so those inserts proceed).
|
||||
# On Oracle this rewrites to FOR UPDATE, which there does not block
|
||||
# indexed-FK child inserts.
|
||||
await conn.execute(
|
||||
#
|
||||
# Use fetchval so we can also verify the bank actually exists.
|
||||
# Without this check, callers that race against bank deletion
|
||||
# or that derive bank IDs before creating the bank reach the
|
||||
# INSERT below and get an asyncpg.ForeignKeyViolationError, which
|
||||
# surfaces as an opaque 500 from the API. A clean
|
||||
# OperationValidationError(404) is the right shape — the FastAPI
|
||||
# handler already converts it via its existing except clause.
|
||||
bank_exists = await conn.fetchval(
|
||||
f"SELECT 1 FROM {fq_table('banks')} WHERE bank_id = $1 FOR NO KEY UPDATE",
|
||||
bank_id,
|
||||
)
|
||||
if bank_exists is None:
|
||||
raise OperationValidationError(f"Bank '{bank_id}' not found", status_code=404)
|
||||
# Only check 'pending', not 'processing': a processing task uses a
|
||||
# watermark from when it started, so memories added after that need
|
||||
# a fresh run regardless.
|
||||
@@ -11749,6 +11921,16 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
"operation_id": str(row["operation_id"]),
|
||||
"deduplicated": True,
|
||||
}
|
||||
else:
|
||||
# Scoped/non-dedupe submits skip the lock + dedup above.
|
||||
# Still verify the bank exists so an FK violation can't
|
||||
# escape as a 500.
|
||||
bank_exists = await conn.fetchval(
|
||||
f"SELECT 1 FROM {fq_table('banks')} WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
if bank_exists is None:
|
||||
raise OperationValidationError(f"Bank '{bank_id}' not found", status_code=404)
|
||||
|
||||
await conn.execute(
|
||||
f"""
|
||||
|
||||
@@ -3,43 +3,116 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from hindsight_api.config import DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
|
||||
|
||||
from .base import FileParser
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MarkitdownOcrOptions:
|
||||
"""OpenAI-compatible OCR options passed through to MarkItDown."""
|
||||
|
||||
# Keep this typed as object so the OpenAI SDK import stays lazy for non-OCR users.
|
||||
llm_client: object
|
||||
llm_model: str
|
||||
llm_prompt: str
|
||||
|
||||
|
||||
class MarkitdownParser(FileParser):
|
||||
"""
|
||||
Markitdown file parser.
|
||||
|
||||
Uses Microsoft's markitdown library to convert various file formats
|
||||
to markdown including PDF, Office docs, images (via OCR), audio, HTML.
|
||||
to markdown including PDF, Office docs, images with optional OCR,
|
||||
audio, HTML.
|
||||
|
||||
Supported formats:
|
||||
- PDF (.pdf)
|
||||
- Word (.docx, .doc)
|
||||
- PowerPoint (.pptx, .ppt)
|
||||
- Excel (.xlsx, .xls)
|
||||
- Images (.jpg, .jpeg, .png) - with OCR
|
||||
- Images (.jpg, .jpeg, .png) - optional OCR
|
||||
- HTML (.html, .htm)
|
||||
- Text (.txt, .md)
|
||||
- Audio (.mp3, .wav) - with transcription
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
ocr_enabled: bool = False,
|
||||
ocr_api_key: str | None = None,
|
||||
ocr_base_url: str | None = None,
|
||||
ocr_model: str | None = None,
|
||||
ocr_prompt: str | None = None,
|
||||
):
|
||||
"""Initialize markitdown parser."""
|
||||
# Lazy import to avoid requiring markitdown for all users
|
||||
try:
|
||||
from markitdown import MarkItDown
|
||||
|
||||
self._markitdown = MarkItDown()
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"markitdown package is required for file parsing. Install with: pip install markitdown"
|
||||
) from e
|
||||
|
||||
self._ocr_enabled = ocr_enabled
|
||||
if ocr_enabled:
|
||||
ocr_options = self._build_ocr_options(
|
||||
api_key=ocr_api_key,
|
||||
base_url=ocr_base_url,
|
||||
model=ocr_model,
|
||||
prompt=ocr_prompt,
|
||||
)
|
||||
self._markitdown = MarkItDown(
|
||||
llm_client=ocr_options.llm_client,
|
||||
llm_model=ocr_options.llm_model,
|
||||
llm_prompt=ocr_options.llm_prompt,
|
||||
)
|
||||
else:
|
||||
self._markitdown = MarkItDown()
|
||||
|
||||
def _build_ocr_options(
|
||||
self,
|
||||
*,
|
||||
api_key: str | None,
|
||||
base_url: str | None,
|
||||
model: str | None,
|
||||
prompt: str | None,
|
||||
) -> MarkitdownOcrOptions:
|
||||
"""Build MarkItDown options for OpenAI-compatible image OCR."""
|
||||
if not model or not model.strip():
|
||||
raise ValueError(
|
||||
"Markitdown OCR is enabled but no model is configured. "
|
||||
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL to an OpenAI-compatible OCR/vision model "
|
||||
"with image-input support."
|
||||
)
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"Markitdown OCR is enabled but no API key is configured. "
|
||||
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY."
|
||||
)
|
||||
if not base_url or not base_url.strip():
|
||||
raise ValueError(
|
||||
"Markitdown OCR is enabled but no base URL is configured. "
|
||||
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL to an OpenAI-compatible OCR/vision endpoint."
|
||||
)
|
||||
|
||||
try:
|
||||
from openai import OpenAI
|
||||
except ImportError as e:
|
||||
raise RuntimeError("openai package is required when Markitdown OCR is enabled.") from e
|
||||
|
||||
return MarkitdownOcrOptions(
|
||||
llm_client=OpenAI(api_key=api_key, base_url=base_url.strip()),
|
||||
llm_model=model.strip(),
|
||||
llm_prompt=prompt or DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
|
||||
)
|
||||
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
"""Parse file to markdown using markitdown."""
|
||||
# markitdown is synchronous, so we run it in executor to avoid blocking
|
||||
@@ -48,6 +121,13 @@ class MarkitdownParser(FileParser):
|
||||
|
||||
def _convert_sync(self, file_data: bytes, filename: str) -> str:
|
||||
"""Synchronous parsing (runs in thread pool)."""
|
||||
if self._is_image_file(filename) and not self._ocr_enabled:
|
||||
raise RuntimeError(
|
||||
"Image OCR is not enabled for the markitdown parser. "
|
||||
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED=true and configure an OpenAI-compatible "
|
||||
"OCR/vision endpoint with image-input support, or choose an OCR-capable parser."
|
||||
)
|
||||
|
||||
# Write to temp file (markitdown requires file path)
|
||||
with tempfile.NamedTemporaryFile(suffix=Path(filename).suffix, delete=False) as tmp:
|
||||
tmp.write(file_data)
|
||||
@@ -73,6 +153,11 @@ class MarkitdownParser(FileParser):
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _is_image_file(filename: str) -> bool:
|
||||
"""Return whether the file type needs OCR to extract useful text."""
|
||||
return Path(filename).suffix.lower() in {".jpg", ".jpeg", ".png"}
|
||||
|
||||
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||
"""Check if markitdown supports this file type."""
|
||||
# Supported extensions (from markitdown docs)
|
||||
@@ -85,7 +170,7 @@ class MarkitdownParser(FileParser):
|
||||
".ppt",
|
||||
".xlsx",
|
||||
".xls",
|
||||
# Images (with OCR)
|
||||
# Images (optional OCR)
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".png",
|
||||
|
||||
@@ -136,7 +136,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 +169,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 +195,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
|
||||
|
||||
@@ -195,32 +212,49 @@ class AnthropicLLM(LLMInterface):
|
||||
try:
|
||||
response = await self._client.messages.create(**call_params)
|
||||
|
||||
# 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
|
||||
|
||||
@@ -60,6 +60,22 @@ _CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
|
||||
)
|
||||
|
||||
|
||||
def default_codex_auth_file() -> Path:
|
||||
"""Return the path to Codex's ``auth.json``.
|
||||
|
||||
Honors the ``CODEX_HOME`` environment variable — the same variable the
|
||||
canonical ``@openai/codex`` CLI uses to relocate its config/credentials
|
||||
directory — and falls back to ``~/.codex`` when it is unset or empty.
|
||||
|
||||
Resolved lazily on each call (rather than cached at import time) so that
|
||||
the environment is read at the point of use.
|
||||
"""
|
||||
codex_home = os.environ.get("CODEX_HOME")
|
||||
if codex_home:
|
||||
return Path(codex_home) / "auth.json"
|
||||
return Path.home() / ".codex" / "auth.json"
|
||||
|
||||
|
||||
class CodexRefreshExpiredError(RuntimeError):
|
||||
"""Raised when the Codex refresh_token itself is no longer valid.
|
||||
|
||||
@@ -86,7 +102,7 @@ class CodexAuthManager:
|
||||
The OAuth refresh token. May be ``None`` when the auth file omits it;
|
||||
the provider still works as a one-shot loader in that case.
|
||||
auth_file:
|
||||
Path to ``~/.codex/auth.json``. Used for re-reading the refresh token
|
||||
Path to the Codex ``auth.json``. Used for re-reading the refresh token
|
||||
on demand and for atomic persistence of rotated credentials.
|
||||
"""
|
||||
|
||||
@@ -115,7 +131,8 @@ class CodexAuthManager:
|
||||
Parameters
|
||||
----------
|
||||
auth_file:
|
||||
Defaults to ``~/.codex/auth.json``.
|
||||
Defaults to ``$CODEX_HOME/auth.json`` (or ``~/.codex/auth.json``
|
||||
when ``CODEX_HOME`` is unset).
|
||||
|
||||
Raises
|
||||
------
|
||||
@@ -126,7 +143,7 @@ class CodexAuthManager:
|
||||
``auth_mode``.
|
||||
"""
|
||||
if auth_file is None:
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
auth_file = default_codex_auth_file()
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(f"Codex auth file not found: {auth_file}. Run 'codex auth login' to authenticate.")
|
||||
|
||||
@@ -2,8 +2,9 @@
|
||||
OpenAI Codex LLM provider using ChatGPT Plus/Pro OAuth authentication.
|
||||
|
||||
This provider enables using ChatGPT Plus/Pro subscriptions for API calls
|
||||
without separate OpenAI Platform API credits. It uses OAuth tokens from
|
||||
~/.codex/auth.json and communicates with the ChatGPT backend API.
|
||||
without separate OpenAI Platform API credits. It uses OAuth tokens from the
|
||||
Codex ``auth.json`` (``$CODEX_HOME/auth.json``, or ``~/.codex/auth.json`` when
|
||||
``CODEX_HOME`` is unset) and communicates with the ChatGPT backend API.
|
||||
|
||||
Tokens are refreshed automatically: the provider decodes the access_token
|
||||
JWT's ``exp`` claim and proactively refreshes via
|
||||
@@ -35,6 +36,7 @@ from .codex_auth import (
|
||||
_CODEX_TOKEN_REFRESH_SKEW_SECONDS,
|
||||
CodexAuthManager,
|
||||
CodexRefreshExpiredError,
|
||||
default_codex_auth_file,
|
||||
)
|
||||
|
||||
# Re-export for backward compatibility (tests import from this module).
|
||||
@@ -55,14 +57,15 @@ class CodexLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using OpenAI Codex OAuth authentication.
|
||||
|
||||
Authenticates using ChatGPT Plus/Pro credentials stored in ~/.codex/auth.json
|
||||
and makes API calls to chatgpt.com/backend-api/codex/responses.
|
||||
Authenticates using ChatGPT Plus/Pro credentials stored in the Codex
|
||||
``auth.json`` (honoring ``CODEX_HOME``, default ``~/.codex``) and makes API
|
||||
calls to chatgpt.com/backend-api/codex/responses.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str, # Will be ignored, reads from ~/.codex/auth.json
|
||||
api_key: str, # Will be ignored, reads from the Codex auth.json (CODEX_HOME or ~/.codex)
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
@@ -81,12 +84,14 @@ class CodexLLM(LLMInterface):
|
||||
refresh_token = self._load_codex_refresh_token()
|
||||
logger.info(f"Loaded Codex OAuth credentials for account: {account_id}")
|
||||
except Exception as e:
|
||||
auth_file = default_codex_auth_file()
|
||||
raise RuntimeError(
|
||||
f"Failed to load Codex OAuth credentials from ~/.codex/auth.json: {e}\n\n"
|
||||
f"Failed to load Codex OAuth credentials from {auth_file}: {e}\n\n"
|
||||
"To set up Codex authentication:\n"
|
||||
"1. Install Codex CLI: npm install -g @openai/codex\n"
|
||||
"2. Login: codex auth login\n"
|
||||
"3. Verify: ls ~/.codex/auth.json\n\n"
|
||||
f"3. Verify: ls {auth_file}\n\n"
|
||||
"(Set CODEX_HOME to use a credentials directory other than ~/.codex.)\n\n"
|
||||
"Or use a different provider (openai, anthropic, gemini) with API keys."
|
||||
) from e
|
||||
|
||||
@@ -94,7 +99,7 @@ class CodexLLM(LLMInterface):
|
||||
access_token=access_token,
|
||||
account_id=account_id,
|
||||
refresh_token=refresh_token,
|
||||
auth_file=Path.home() / ".codex" / "auth.json",
|
||||
auth_file=default_codex_auth_file(),
|
||||
)
|
||||
|
||||
# Use ChatGPT backend API endpoint. Codex auth is tied to
|
||||
@@ -156,7 +161,7 @@ class CodexLLM(LLMInterface):
|
||||
|
||||
def _load_codex_auth(self) -> tuple[str, str]:
|
||||
"""
|
||||
Load OAuth credentials from ~/.codex/auth.json.
|
||||
Load OAuth credentials from the Codex ``auth.json`` (CODEX_HOME or ~/.codex).
|
||||
|
||||
Returns:
|
||||
Tuple of (access_token, account_id).
|
||||
@@ -165,7 +170,7 @@ class CodexLLM(LLMInterface):
|
||||
FileNotFoundError: If auth file doesn't exist.
|
||||
ValueError: If auth file is invalid.
|
||||
"""
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
auth_file = default_codex_auth_file()
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(
|
||||
@@ -197,9 +202,7 @@ class CodexLLM(LLMInterface):
|
||||
pre- and post-``__init__`` because it does not depend on
|
||||
``_auth_manager`` being constructed yet.
|
||||
"""
|
||||
auth_file = (
|
||||
self._auth_manager._auth_file if hasattr(self, "_auth_manager") else Path.home() / ".codex" / "auth.json"
|
||||
)
|
||||
auth_file = self._auth_manager._auth_file if hasattr(self, "_auth_manager") else default_codex_auth_file()
|
||||
return CodexAuthManager.load_refresh_token_from_file(auth_file)
|
||||
|
||||
@staticmethod
|
||||
|
||||
@@ -76,6 +76,7 @@ class GeminiLLM(LLMInterface):
|
||||
|
||||
# Safety settings: None means use Gemini's defaults
|
||||
self._safety_settings: list | None = kwargs.get("gemini_safety_settings")
|
||||
self._service_tier: str | None = kwargs.get("gemini_service_tier")
|
||||
|
||||
# User-configured extra params merged into the GenerateContentConfig of
|
||||
# every call. Gemini's request body nests generation params, so we expose
|
||||
@@ -106,6 +107,16 @@ class GeminiLLM(LLMInterface):
|
||||
self._client = genai.Client(api_key=self.api_key)
|
||||
logger.info(f"Gemini API: model={self.model}")
|
||||
|
||||
def _apply_service_tier(self, config_kwargs: dict[str, Any]) -> None:
|
||||
if not self._service_tier:
|
||||
return
|
||||
|
||||
http_options = dict(config_kwargs.get("http_options") or {})
|
||||
extra_body = dict(http_options.get("extra_body") or {})
|
||||
extra_body.setdefault("service_tier", self._service_tier)
|
||||
http_options["extra_body"] = extra_body
|
||||
config_kwargs["http_options"] = http_options
|
||||
|
||||
def _init_vertexai(self, **kwargs: Any) -> None:
|
||||
"""Initialize Vertex AI client with project, region, and credentials."""
|
||||
# Extract Vertex AI config from kwargs
|
||||
@@ -247,16 +258,13 @@ class GeminiLLM(LLMInterface):
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
|
||||
|
||||
# Add the JSON schema as a textual hint in the system_instruction (matching
|
||||
# the normal uncached path). Structured output is still enforced via
|
||||
# response_schema regardless; this is just guidance text.
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
def _system_instruction_with_schema() -> str:
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
if system_instruction:
|
||||
system_instruction += schema_msg
|
||||
else:
|
||||
system_instruction = schema_msg
|
||||
schema_msg = (
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n"
|
||||
f"{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
)
|
||||
return (system_instruction + schema_msg) if system_instruction else schema_msg
|
||||
|
||||
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
|
||||
effective_safety_settings = _safety_settings_ctx.get()
|
||||
@@ -273,11 +281,18 @@ class GeminiLLM(LLMInterface):
|
||||
def _build_generation_config(use_cache: bool) -> "genai_types.GenerateContentConfig | None":
|
||||
# Seed with user-configured extra params; explicit settings below win.
|
||||
config_kwargs: dict[str, Any] = dict(self._extra_body)
|
||||
self._apply_service_tier(config_kwargs)
|
||||
if use_cache:
|
||||
config_kwargs["cached_content"] = cached_prefix
|
||||
elif (
|
||||
use_schema_prompt_fallback
|
||||
and response_format is not None
|
||||
and hasattr(response_format, "model_json_schema")
|
||||
):
|
||||
config_kwargs["system_instruction"] = _system_instruction_with_schema()
|
||||
elif system_instruction:
|
||||
config_kwargs["system_instruction"] = system_instruction
|
||||
if response_format is not None:
|
||||
if response_format is not None and not use_schema_prompt_fallback:
|
||||
config_kwargs["response_mime_type"] = "application/json"
|
||||
config_kwargs["response_schema"] = response_format
|
||||
if temperature is not None:
|
||||
@@ -295,6 +310,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
|
||||
@@ -412,12 +428,26 @@ class GeminiLLM(LLMInterface):
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=input_tokens + output_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
last_exception = e
|
||||
if (
|
||||
attempt < max_retries
|
||||
and response_format is not None
|
||||
and hasattr(response_format, "model_json_schema")
|
||||
and not cache_active
|
||||
and not use_schema_prompt_fallback
|
||||
):
|
||||
logger.warning("Gemini returned invalid JSON, retrying with prompt-side schema guidance...")
|
||||
cache_active = False
|
||||
use_schema_prompt_fallback = True
|
||||
generation_config = _build_generation_config(cache_active)
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
if attempt < max_retries:
|
||||
logger.warning("Gemini returned invalid JSON, retrying...")
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
@@ -604,6 +634,7 @@ class GeminiLLM(LLMInterface):
|
||||
def _build_tools_config(use_cache: bool) -> "genai_types.GenerateContentConfig":
|
||||
# Seed with user-configured extra params; explicit settings below win.
|
||||
config_kwargs: dict[str, Any] = dict(self._extra_body)
|
||||
self._apply_service_tier(config_kwargs)
|
||||
if use_cache:
|
||||
config_kwargs["cached_content"] = cached_prefix
|
||||
else:
|
||||
@@ -749,6 +780,8 @@ class GeminiLLM(LLMInterface):
|
||||
finish_reason=finish_reason,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cached_tokens=cached_input_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
|
||||
except genai_errors.APIError as e:
|
||||
|
||||
@@ -15,9 +15,13 @@ is handled automatically by LiteLLM.
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from litellm.exceptions import Timeout as LiteLLMTimeout
|
||||
|
||||
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
@@ -47,13 +51,15 @@ class LiteLLMLLM(LLMInterface):
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float = 300.0,
|
||||
timeout: float | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
bedrock_service_tier: str | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
self.timeout = timeout
|
||||
# ``None`` falls back to HINDSIGHT_API_LLM_TIMEOUT, then DEFAULT_LLM_TIMEOUT — never None,
|
||||
# so the hard ``asyncio.wait_for`` backstop in ``call`` is always bounded.
|
||||
self.timeout = timeout if timeout is not None else float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT)))
|
||||
self._litellm: Any = None
|
||||
# User-configured extra params merged as top-level kwargs into every
|
||||
# completion call so LiteLLM normalizes them per-provider (e.g. maps
|
||||
@@ -209,7 +215,10 @@ class LiteLLMLLM(LLMInterface):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self._stage_label}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._acompletion(**call_kwargs)
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
content = response.choices[0].message.content or ""
|
||||
finish_reason = response.choices[0].finish_reason
|
||||
@@ -304,6 +313,25 @@ class LiteLLMLLM(LLMInterface):
|
||||
logger.error(f"LiteLLM returned invalid JSON after {max_retries + 1} attempts")
|
||||
raise
|
||||
|
||||
except (TimeoutError, asyncio.TimeoutError, LiteLLMTimeout) as e:
|
||||
# litellm/httpx don't always honor their own ``timeout=`` (e.g. a connection held
|
||||
# open with no token progress), so ``wait_for`` is the hard cap that cancels a hung
|
||||
# call regardless — otherwise one straggler pins a worker slot and stalls its gather.
|
||||
last_exception = e
|
||||
exc_name = type(e).__name__
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"LiteLLM call exceeded timeout={self.timeout}s ({exc_name}, scope={scope}), retrying..."
|
||||
)
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
logger.error(
|
||||
f"LiteLLM call timed out after {self.timeout}s on {attempt + 1} attempts "
|
||||
f"({exc_name}, scope={scope})"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
error_str = str(e).lower()
|
||||
# Fast fail on auth errors
|
||||
@@ -354,7 +382,10 @@ class LiteLLMLLM(LLMInterface):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self._stage_label}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._acompletion(**call_kwargs)
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
message = response.choices[0].message
|
||||
content = message.content
|
||||
@@ -424,6 +455,23 @@ class LiteLLMLLM(LLMInterface):
|
||||
output_tokens=output_tokens,
|
||||
)
|
||||
|
||||
except (TimeoutError, asyncio.TimeoutError, LiteLLMTimeout) as e:
|
||||
# See ``call`` — hard cap so a hung completion cannot block
|
||||
# forever and pin a worker slot / concurrency permit.
|
||||
last_exception = e
|
||||
exc_name = type(e).__name__
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"LiteLLM tool call exceeded timeout={self.timeout}s ({exc_name}, scope={scope}), retrying..."
|
||||
)
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
logger.error(
|
||||
f"LiteLLM tool call timed out after {self.timeout}s on {attempt + 1} attempts "
|
||||
f"({exc_name}, scope={scope})"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
error_str = str(e).lower()
|
||||
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
|
||||
|
||||
@@ -67,7 +67,7 @@ class LiteLLMRouterLLM(LiteLLMLLM):
|
||||
model: str,
|
||||
config: dict[str, Any],
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float = 300.0,
|
||||
timeout: float | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(
|
||||
|
||||
@@ -26,6 +26,8 @@ import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from email.utils import parsedate_to_datetime
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qs, urlparse, urlunparse
|
||||
|
||||
@@ -34,7 +36,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
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError, ProviderRateLimitResetError
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
@@ -83,6 +85,49 @@ def _strip_code_fences(content: str) -> str:
|
||||
return content
|
||||
|
||||
|
||||
# Reasoning/thinking tags emitted by extended-thinking models. Some providers
|
||||
# (e.g. MiniMax-M3) leak the chain-of-thought wrapped in these tags into the
|
||||
# response body instead of a separate reasoning_content field. Each entry is
|
||||
# (open_tag, close_tag); the open tag also matches when the close tag is missing
|
||||
# (truncated output) so a dangling block is removed to end-of-string.
|
||||
_REASONING_TAG_PAIRS: tuple[tuple[str, str], ...] = (
|
||||
("<think>", "</think>"),
|
||||
("<thinking>", "</thinking>"),
|
||||
("<thought>", "</thought>"),
|
||||
("<reasoning>", "</reasoning>"),
|
||||
("|startthink|", "|endthink|"),
|
||||
)
|
||||
|
||||
|
||||
def _strip_reasoning_tags(text: str) -> str:
|
||||
"""Strip extended-thinking/reasoning blocks from an LLM response.
|
||||
|
||||
Removes the full set of tag styles emitted by reasoning models:
|
||||
``<think>``, ``<thinking>``, ``<thought>``, ``<reasoning>`` and the
|
||||
``|startthink|...|endthink|`` markers. Both the structured (JSON) path and
|
||||
the free-form path must call this — otherwise a non-structured response
|
||||
(e.g. a mental-model markdown blob from MiniMax-M3) leaks the raw
|
||||
``<think>...</think>`` verbatim into stored memories.
|
||||
|
||||
Handles two cases:
|
||||
1. Closed blocks: ``<think>...</think>`` removed wherever they appear.
|
||||
2. Unclosed blocks: a dangling ``<think>`` with no closing tag (model output
|
||||
truncated mid-thought) is removed from the open tag to end-of-string.
|
||||
|
||||
Returns the input unchanged (modulo surrounding whitespace) when no tags are
|
||||
present.
|
||||
"""
|
||||
if not text:
|
||||
return text
|
||||
for open_tag, close_tag in _REASONING_TAG_PAIRS:
|
||||
open_re = re.escape(open_tag)
|
||||
close_re = re.escape(close_tag)
|
||||
# Closed blocks first, then any remaining unclosed (truncated) block.
|
||||
text = re.sub(rf"{open_re}.*?{close_re}", "", text, flags=re.DOTALL)
|
||||
text = re.sub(rf"{open_re}.*", "", text, flags=re.DOTALL)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def _response_get(response: Any, key: str, default: Any = None) -> Any:
|
||||
if isinstance(response, dict):
|
||||
return response.get(key, default)
|
||||
@@ -234,6 +279,122 @@ def _summarize_status_error(e: APIStatusError, body_max: int = 400) -> str:
|
||||
return f"HTTP {e.status_code}: {body_str or '<no body>'}"
|
||||
|
||||
|
||||
_RATE_LIMIT_RESET_AT_RE = re.compile(
|
||||
r"\breset at\s+"
|
||||
r"(?P<reset_at>\d{4}-\d{2}-\d{2}[ T]\d{2}:\d{2}:\d{2}(?:\s*(?:Z|[+-]\d{2}:?\d{2}))?)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_RATE_LIMIT_WINDOW_RE = re.compile(
|
||||
r"\b(?:for|in)\s+(?P<amount>\d+)\s*(?P<unit>second|minute|hour|day)s?\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def _status_error_body_text(e: APIStatusError) -> str:
|
||||
body: Any = getattr(e, "body", None)
|
||||
if body is None:
|
||||
try:
|
||||
body = e.response.text
|
||||
except Exception:
|
||||
body = None
|
||||
if isinstance(body, (dict, list)):
|
||||
try:
|
||||
return json.dumps(body, default=str, ensure_ascii=False)
|
||||
except Exception:
|
||||
return str(body)
|
||||
return str(body or "").strip()
|
||||
|
||||
|
||||
def _parse_retry_after_header(value: str | None, now: datetime) -> datetime | None:
|
||||
if not value:
|
||||
return None
|
||||
raw = value.strip()
|
||||
try:
|
||||
seconds = float(raw)
|
||||
except ValueError:
|
||||
seconds = -1.0
|
||||
if seconds >= 0:
|
||||
return now + timedelta(seconds=seconds)
|
||||
|
||||
try:
|
||||
parsed = parsedate_to_datetime(raw)
|
||||
except (TypeError, ValueError, IndexError, OverflowError):
|
||||
return None
|
||||
if parsed.tzinfo is None:
|
||||
parsed = parsed.replace(tzinfo=UTC)
|
||||
return parsed.astimezone(UTC)
|
||||
|
||||
|
||||
def _parse_reset_at_datetime(value: str) -> datetime | None:
|
||||
raw = value.strip().replace(" ", "T")
|
||||
if raw.endswith("Z"):
|
||||
raw = f"{raw[:-1]}+00:00"
|
||||
elif re.search(r"[+-]\d{4}$", raw):
|
||||
raw = f"{raw[:-2]}:{raw[-2:]}"
|
||||
try:
|
||||
parsed = datetime.fromisoformat(raw)
|
||||
except ValueError:
|
||||
return None
|
||||
if parsed.tzinfo is None:
|
||||
# Some providers (z.ai included) return a wall-clock reset timestamp
|
||||
# without a zone. Interpret it in the host's local zone so logs, status
|
||||
# pages, and the queued next_retry_at describe the same operator-facing
|
||||
# clock instead of silently shifting by UTC offset.
|
||||
parsed = parsed.astimezone()
|
||||
return parsed.astimezone(UTC)
|
||||
|
||||
|
||||
def _rate_limit_retry_at(e: APIStatusError) -> datetime | None:
|
||||
now = datetime.now(UTC)
|
||||
response = getattr(e, "response", None)
|
||||
headers = getattr(response, "headers", None)
|
||||
if headers is not None:
|
||||
retry_at = _parse_retry_after_header(headers.get("retry-after") or headers.get("Retry-After"), now)
|
||||
if retry_at is not None and retry_at > now:
|
||||
return retry_at
|
||||
|
||||
body_text = _status_error_body_text(e)
|
||||
reset_match = _RATE_LIMIT_RESET_AT_RE.search(body_text)
|
||||
if reset_match:
|
||||
retry_at = _parse_reset_at_datetime(reset_match.group("reset_at"))
|
||||
if retry_at is not None and retry_at > now:
|
||||
return retry_at
|
||||
|
||||
window_match = _RATE_LIMIT_WINDOW_RE.search(body_text)
|
||||
if not window_match:
|
||||
return None
|
||||
amount = int(window_match.group("amount"))
|
||||
unit = window_match.group("unit").lower()
|
||||
if unit == "second":
|
||||
seconds = amount
|
||||
elif unit == "minute":
|
||||
seconds = amount * 60
|
||||
elif unit == "hour":
|
||||
seconds = amount * 3600
|
||||
else:
|
||||
seconds = amount * 86400
|
||||
return now + timedelta(seconds=seconds)
|
||||
|
||||
|
||||
def _raise_provider_quota_defer(
|
||||
e: APIStatusError, *, provider: str, model: str, scope: str, max_backoff: float
|
||||
) -> None:
|
||||
if e.status_code != 429:
|
||||
return
|
||||
retry_at = _rate_limit_retry_at(e)
|
||||
if retry_at is None:
|
||||
return
|
||||
if (retry_at - datetime.now(UTC)).total_seconds() <= max_backoff:
|
||||
return
|
||||
summary = _summarize_status_error(e)
|
||||
raise ProviderRateLimitResetError(
|
||||
retry_at=retry_at,
|
||||
message=(
|
||||
f"Provider quota exhausted ({provider}/{model}, scope={scope}); retry at {retry_at.isoformat()}: {summary}"
|
||||
),
|
||||
) from e
|
||||
|
||||
|
||||
class OpenAICompatibleLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider for OpenAI-compatible APIs.
|
||||
@@ -269,7 +430,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
base_url: Base URL for the API (uses defaults for groq/ollama/lmstudio if empty).
|
||||
model: Model name.
|
||||
reasoning_effort: Reasoning effort level for supported models ("low", "medium", "high").
|
||||
timeout: Request timeout in seconds (uses env var or 300s default).
|
||||
timeout: Request timeout in seconds (uses env var or 120s default).
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto").
|
||||
extra_body: Extra body params merged into every API call.
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
@@ -290,6 +451,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"openrouter",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
"fireworks",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
@@ -315,6 +477,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
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.
|
||||
@@ -335,6 +499,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"openrouter",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"atlas",
|
||||
"ollama-cloud",
|
||||
)
|
||||
and not self.api_key
|
||||
@@ -617,15 +782,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
scope=scope,
|
||||
)
|
||||
|
||||
# Strip reasoning model thinking tags
|
||||
# Strip reasoning model thinking tags (closed and unclosed).
|
||||
# Supports: <think>, <thinking>, <thought>, <reasoning>, |startthink|/|endthink|
|
||||
original_len = len(content)
|
||||
content = re.sub(r"<think>.*?</think>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"<thinking>.*?</thinking>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"<thought>.*?</thought>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"<reasoning>.*?</reasoning>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"\|startthink\|.*?\|endthink\|", "", content, flags=re.DOTALL)
|
||||
content = content.strip()
|
||||
content = _strip_reasoning_tags(content)
|
||||
if len(content) < original_len:
|
||||
logger.debug(f"Stripped {original_len - len(content)} chars of reasoning tokens")
|
||||
|
||||
@@ -674,6 +834,13 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
scope=scope,
|
||||
)
|
||||
|
||||
# Free-form (non-structured) output also leaks reasoning tags:
|
||||
# reasoning models like MiniMax-M3 wrap their chain-of-thought
|
||||
# in <think>...</think> in the response body. Without this strip
|
||||
# a mental-model markdown blob is stored verbatim with the raw
|
||||
# thinking tags. Mirrors the structured-output path above.
|
||||
result = _strip_reasoning_tags(result)
|
||||
|
||||
# Record token usage metrics
|
||||
duration = time.time() - start_time
|
||||
usage = response.usage
|
||||
@@ -761,6 +928,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
logger.error(f"Auth error (HTTP {e.status_code}), not retrying: {str(e)}")
|
||||
raise
|
||||
|
||||
_raise_provider_quota_defer(
|
||||
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
|
||||
)
|
||||
|
||||
# Handle tool_use_failed error - model outputted in tool call format
|
||||
if e.status_code == 400 and response_format is not None:
|
||||
try:
|
||||
@@ -814,7 +985,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
f"scope={scope}): {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except ProviderResponseError as e:
|
||||
last_exception = e
|
||||
if e.retryable and attempt < max_retries:
|
||||
@@ -1047,6 +1217,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
f"not retrying: {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
_raise_provider_quota_defer(
|
||||
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
|
||||
)
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
@@ -1060,7 +1234,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
f"({self.provider}/{self.model}, scope={scope}): {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@ import time
|
||||
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||
|
||||
from ...config import get_config
|
||||
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, TokenUsageSummary, ToolCall
|
||||
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, StructuredOutputResult, TokenUsageSummary, ToolCall
|
||||
from .prompts import (
|
||||
_extract_directive_rules,
|
||||
build_final_prompt,
|
||||
@@ -141,7 +141,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 +151,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 +186,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)
|
||||
|
||||
@@ -259,11 +259,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 +441,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 +471,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 +539,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 +554,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 +604,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 +618,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 +680,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 +730,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 +745,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 +807,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 +822,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 +862,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 +877,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 +1177,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,
|
||||
)
|
||||
|
||||
|
||||
@@ -307,7 +329,8 @@ class ReflectResult(BaseModel):
|
||||
],
|
||||
"experience": [],
|
||||
"opinion": [],
|
||||
"mental_models": [],
|
||||
"observation": [],
|
||||
"mental-models": [],
|
||||
"directives": [
|
||||
{
|
||||
"id": "directive-123",
|
||||
@@ -324,7 +347,7 @@ class ReflectResult(BaseModel):
|
||||
|
||||
text: str = Field(description="The formulated answer text")
|
||||
based_on: dict[str, Any] = Field(
|
||||
description="Facts used to formulate the answer, organized by type (world, experience, mental_models, directives)"
|
||||
description="Facts used to formulate the answer, organized by type (world, experience, observation, mental-models, directives)"
|
||||
)
|
||||
structured_output: dict[str, Any] | None = Field(
|
||||
default=None,
|
||||
|
||||
@@ -14,6 +14,7 @@ from typing import Any, Literal, cast
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, create_model, field_validator
|
||||
|
||||
from ..llm_interface import ProviderRateLimitResetError
|
||||
from ..llm_wrapper import LLMConfig, OutputTooLongError, sanitize_llm_output
|
||||
from ..operation_metadata import RetainExtractionErrors
|
||||
from ..response_models import TokenUsage
|
||||
@@ -451,6 +452,11 @@ def chunk_text(text: str, max_chars: int, structured_chunk_size: int | None = No
|
||||
``structured_chunk_size``. When unset, that limit defaults to ``max_chars``.
|
||||
For plain text, uses sentence-aware splitting.
|
||||
|
||||
The result is idempotent: re-chunking any chunk this returns yields that chunk
|
||||
unchanged. The streaming retain pipeline pre-chunks each document once and then
|
||||
re-chunks every piece during extraction; if a piece re-split, its sub-chunks
|
||||
would inherit one chunk_index and collide on ``chunk_id`` (issue #2301).
|
||||
|
||||
Args:
|
||||
text: Input text to chunk (plain text, JSON conversation, or JSONL)
|
||||
max_chars: Target maximum characters per chunk
|
||||
@@ -469,11 +475,23 @@ def chunk_text(text: str, max_chars: int, structured_chunk_size: int | None = No
|
||||
# Try to parse as JSON conversation array
|
||||
try:
|
||||
parsed = json.loads(text)
|
||||
if isinstance(parsed, list) and all(isinstance(turn, dict) for turn in parsed):
|
||||
# This looks like a conversation - chunk at turn boundaries
|
||||
return _chunk_conversation(parsed, max_chars, structured_limit)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
parsed = None
|
||||
|
||||
if isinstance(parsed, list) and all(isinstance(turn, dict) for turn in parsed):
|
||||
# This looks like a conversation - chunk at turn boundaries
|
||||
return _chunk_conversation(parsed, max_chars, structured_limit)
|
||||
|
||||
if isinstance(parsed, dict):
|
||||
# A single JSON object — e.g. one JSONL line handed back to the extractor
|
||||
# after the producer already pre-chunked it. It is one structured unit:
|
||||
# keep it whole up to the structured limit, else split it as text within
|
||||
# the chunk budget. Without this, a lone object (one line, so _chunk_jsonl
|
||||
# declines) would fall through to plain-text splitting and re-split a chunk
|
||||
# the producer deliberately kept whole — breaking idempotency (issue #2301).
|
||||
if len(text) <= structured_limit:
|
||||
return [text]
|
||||
return _split_oversized_unit(text, max_chars)
|
||||
|
||||
# Try to parse as JSONL (newline-delimited JSON objects, e.g. session logs)
|
||||
jsonl_chunks = _chunk_jsonl(text, max_chars, structured_limit)
|
||||
@@ -515,10 +533,12 @@ def _chunk_conversation(turns: list[dict], max_chars: int, structured_limit: int
|
||||
turn_size = turn_unit_size + 1 # +1 for comma
|
||||
|
||||
# A turn too large to keep whole even alone: flush, then split it as
|
||||
# text so no chunk runs far over budget (the extractor won't re-chunk).
|
||||
# text. Fragment within min(structured_limit, max_chars) so no fragment
|
||||
# exceeds the chunk budget — otherwise a downstream re-chunk would split
|
||||
# it again and collide on chunk_id (issue #2301).
|
||||
if turn_unit_size > structured_limit:
|
||||
_flush()
|
||||
chunks.extend(_split_oversized_unit(turn_json, structured_limit))
|
||||
chunks.extend(_split_oversized_unit(turn_json, min(structured_limit, max_chars)))
|
||||
continue
|
||||
|
||||
# If adding this turn would exceed limit and we have turns, save current chunk
|
||||
@@ -581,10 +601,12 @@ def _chunk_jsonl(text: str, max_chars: int, structured_limit: int) -> list[str]
|
||||
line_size = len(line) + 1 # +1 for the joining newline
|
||||
|
||||
# A line too large to keep whole even alone: flush, then split it as
|
||||
# text so no chunk runs far over budget (the extractor won't re-chunk).
|
||||
# text. Fragment within min(structured_limit, max_chars) so no fragment
|
||||
# exceeds the chunk budget — otherwise a downstream re-chunk would split
|
||||
# it again and collide on chunk_id (issue #2301).
|
||||
if line_unit_size > structured_limit:
|
||||
_flush()
|
||||
chunks.extend(_split_oversized_unit(line, structured_limit))
|
||||
chunks.extend(_split_oversized_unit(line, min(structured_limit, max_chars)))
|
||||
continue
|
||||
|
||||
# If adding this line would exceed the limit and we have lines, flush.
|
||||
@@ -1792,10 +1814,21 @@ 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])
|
||||
quota_errors = [err for _, err in failed_chunks if isinstance(err, ProviderRateLimitResetError)]
|
||||
if quota_errors and len(quota_errors) == len(failed_chunks):
|
||||
retry_at = max(err.retry_at for err in quota_errors)
|
||||
raise ProviderRateLimitResetError(
|
||||
retry_at=retry_at,
|
||||
message=(
|
||||
f"Fact extraction deferred by provider quota: {len(failed_chunks)}/{len(chunks)} chunks failed. "
|
||||
f"First failures: {failed_summary}. Provider detail: {quota_errors[0]}"
|
||||
),
|
||||
) from quota_errors[0]
|
||||
|
||||
# Fail the entire retain — partial extraction is not acceptable.
|
||||
# All successfully extracted facts are discarded because the transaction
|
||||
# hasn't committed yet. The worker poller will retry the entire task.
|
||||
failed_summary = ", ".join(f"chunk {idx}: {type(err).__name__}" for idx, err in failed_chunks[:5])
|
||||
raise RuntimeError(
|
||||
f"Fact extraction failed: {len(failed_chunks)}/{len(chunks)} chunks failed. "
|
||||
f"First failures: {failed_summary}"
|
||||
|
||||
@@ -1615,8 +1615,19 @@ async def _streaming_retain_batch(
|
||||
# Check if facts are already committed (recovery from previous crash).
|
||||
# If so, skip extraction+writes and jump straight to final ANN pass.
|
||||
# ---------------------------------------------------------------------------
|
||||
# Only the call that starts a document at chunk 0 may take the whole-document
|
||||
# skip. When an oversized single item is split into several sequential
|
||||
# sub-batches that SHARE one document_id AND one operation_id (see
|
||||
# _split_contents_into_sub_batches), the first sub-batch commits its chunks
|
||||
# and stamps effective_doc_id into result_metadata.facts_committed_document_ids.
|
||||
# Without the offset gate, every later sub-batch (chunk_index_offset > 0) would
|
||||
# then see its own document already "committed" and skip extraction, dropping
|
||||
# all chunks past the first slice. A non-zero offset inherently means this call
|
||||
# continues a document another sub-batch already started, so it must always do
|
||||
# its work — crash-safety for those chunks still comes from the per-chunk hash
|
||||
# recovery (existing_chunk_hashes) below.
|
||||
facts_already_committed = False
|
||||
if operation_id:
|
||||
if operation_id and chunk_index_offset == 0:
|
||||
try:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
row = await conn.fetchrow(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -97,6 +97,10 @@ class PrecheckContext:
|
||||
- ``bank_id``: parsed from the URL path.
|
||||
- ``request_context``: the authenticated :class:`RequestContext` (tenant
|
||||
already resolved by the tenant extension).
|
||||
- ``content_length``: value of the ``Content-Length`` request header as an
|
||||
int, or ``None`` when the header is absent or unparseable (e.g. chunked
|
||||
transfer encoding). Lets a precheck make size-aware decisions — such as
|
||||
an upper-bound cost estimate — without reading or deserialising the body.
|
||||
|
||||
Implementations should keep precheck cheap and side-effect-free. The
|
||||
full per-request validators (``validate_retain`` / ``validate_recall``
|
||||
@@ -107,6 +111,7 @@ class PrecheckContext:
|
||||
operation: str
|
||||
bank_id: str
|
||||
request_context: "RequestContext"
|
||||
content_length: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -203,6 +208,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
|
||||
|
||||
@@ -12,6 +12,7 @@ from datetime import datetime, timezone
|
||||
from typing import Any, Callable
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from mcp.types import ToolAnnotations
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
@@ -199,6 +200,47 @@ def build_content_dict(
|
||||
return content_dict, None
|
||||
|
||||
|
||||
# MCP tool annotations. Hindsight is a closed memory store (no open-world / internet
|
||||
# access), so openWorldHint=False throughout. readOnlyHint lets clients group and
|
||||
# auto-approve safe reads; destructiveHint flags tools that delete or clear memory.
|
||||
_READ_ONLY_TOOLS = {
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_banks",
|
||||
"get_bank",
|
||||
"get_bank_stats",
|
||||
"list_mental_models",
|
||||
"get_mental_model",
|
||||
"list_directives",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"list_operations",
|
||||
"get_operation",
|
||||
"list_tags",
|
||||
}
|
||||
_DESTRUCTIVE_TOOLS = {
|
||||
"delete_bank",
|
||||
"clear_memories",
|
||||
"clear_mental_model",
|
||||
"delete_mental_model",
|
||||
"delete_directive",
|
||||
"delete_document",
|
||||
"invalidate_memory",
|
||||
}
|
||||
|
||||
|
||||
def _tool_annotations(name: str) -> ToolAnnotations:
|
||||
if name in _READ_ONLY_TOOLS:
|
||||
return ToolAnnotations(readOnlyHint=True, openWorldHint=False)
|
||||
if name in _DESTRUCTIVE_TOOLS:
|
||||
return ToolAnnotations(readOnlyHint=False, destructiveHint=True, openWorldHint=False)
|
||||
# Everything else writes but does not destructively delete/clear memory
|
||||
# (retain, create_*, update_*, refresh_mental_model, cancel_operation).
|
||||
return ToolAnnotations(readOnlyHint=False, destructiveHint=False, openWorldHint=False)
|
||||
|
||||
|
||||
def register_mcp_tools(
|
||||
mcp: FastMCP,
|
||||
memory: MemoryEngine,
|
||||
@@ -552,7 +594,7 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
@mcp.tool(description=description, annotations=_tool_annotations("retain"))
|
||||
async def retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
@@ -608,7 +650,7 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
@mcp.tool(description=description, annotations=_tool_annotations("retain"))
|
||||
async def retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
@@ -666,7 +708,7 @@ def _register_sync_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("sync_retain"))
|
||||
async def sync_retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
@@ -724,7 +766,7 @@ def _register_sync_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("sync_retain"))
|
||||
async def sync_retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
@@ -785,12 +827,13 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
@mcp.tool(description=description, annotations=_tool_annotations("recall"))
|
||||
async def recall(
|
||||
query: str,
|
||||
max_tokens: int = 4096,
|
||||
budget: str = "high",
|
||||
types: list[str] | None = None,
|
||||
prefer_observations: bool = False,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list[dict] | None = None,
|
||||
@@ -803,6 +846,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
|
||||
@@ -831,6 +878,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),
|
||||
@@ -857,12 +905,13 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
@mcp.tool(description=description, annotations=_tool_annotations("recall"))
|
||||
async def recall(
|
||||
query: str,
|
||||
max_tokens: int = 4096,
|
||||
budget: str = "high",
|
||||
types: list[str] | None = None,
|
||||
prefer_observations: bool = False,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list[dict] | None = None,
|
||||
@@ -874,6 +923,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
|
||||
@@ -901,6 +954,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),
|
||||
@@ -931,7 +985,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("reflect"))
|
||||
async def reflect(
|
||||
query: str,
|
||||
context: str | None = None,
|
||||
@@ -941,6 +995,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
include_based_on: bool = False,
|
||||
include_trace: bool = False,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -971,6 +1026,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
tags: Optional tags to filter memories by (e.g., ['project:alpha'])
|
||||
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
|
||||
include_based_on: Include source facts used for synthesis. Defaults to false because broad reflections can exceed MCP client result limits.
|
||||
include_trace: Include the reflection's internal trace fields (tool_trace/llm_trace and directives_applied). Defaults to false because the trace can be tens of KB and overflow MCP client context; enable only for debugging.
|
||||
bank_id: Optional bank to reflect in (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
@@ -1000,6 +1056,15 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
result_data = json.loads(reflect_result.model_dump_json(indent=2))
|
||||
if not include_based_on:
|
||||
result_data.pop("based_on", None)
|
||||
if not include_trace:
|
||||
# The agentic reflect loop's trace fields can be tens of KB (full
|
||||
# mental-model text) and silently overflow MCP client context; the
|
||||
# REST API omits them by default too. directives_applied is built by
|
||||
# the engine "for the trace" and carries full directive content, so it
|
||||
# belongs with tool_trace/llm_trace here. Opt in via include_trace.
|
||||
result_data.pop("tool_trace", None)
|
||||
result_data.pop("llm_trace", None)
|
||||
result_data.pop("directives_applied", None)
|
||||
if response_schema is not None and hasattr(reflect_result, "structured_output"):
|
||||
result_data["structured_output"] = reflect_result.structured_output
|
||||
return json.dumps(result_data, indent=2)
|
||||
@@ -1012,7 +1077,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("reflect"))
|
||||
async def reflect(
|
||||
query: str,
|
||||
context: str | None = None,
|
||||
@@ -1022,6 +1087,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
include_based_on: bool = False,
|
||||
include_trace: bool = False,
|
||||
) -> dict:
|
||||
"""
|
||||
Generate thoughtful analysis by synthesizing stored memories with the bank's personality.
|
||||
@@ -1051,6 +1117,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
tags: Optional tags to filter memories by (e.g., ['project:alpha'])
|
||||
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
|
||||
include_based_on: Include source facts used for synthesis. Defaults to false because broad reflections can exceed MCP client result limits.
|
||||
include_trace: Include the reflection's internal trace fields (tool_trace/llm_trace and directives_applied). Defaults to false because the trace can be tens of KB and overflow MCP client context; enable only for debugging.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
@@ -1079,6 +1146,15 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
result_data = reflect_result.model_dump()
|
||||
if not include_based_on:
|
||||
result_data.pop("based_on", None)
|
||||
if not include_trace:
|
||||
# The agentic reflect loop's trace fields can be tens of KB (full
|
||||
# mental-model text) and silently overflow MCP client context; the
|
||||
# REST API omits them by default too. directives_applied is built by
|
||||
# the engine "for the trace" and carries full directive content, so it
|
||||
# belongs with tool_trace/llm_trace here. Opt in via include_trace.
|
||||
result_data.pop("tool_trace", None)
|
||||
result_data.pop("llm_trace", None)
|
||||
result_data.pop("directives_applied", None)
|
||||
if response_schema is not None and hasattr(reflect_result, "structured_output"):
|
||||
result_data["structured_output"] = reflect_result.structured_output
|
||||
return result_data
|
||||
@@ -1093,7 +1169,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
def _register_list_banks(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the list_banks tool."""
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_banks"))
|
||||
async def list_banks() -> str:
|
||||
"""
|
||||
List all available memory banks.
|
||||
@@ -1118,7 +1194,7 @@ def _register_list_banks(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
def _register_create_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the create_bank tool."""
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_bank"))
|
||||
async def create_bank(bank_id: str, name: str | None = None, mission: str | None = None) -> str:
|
||||
"""
|
||||
Create a new memory bank or get an existing one.
|
||||
@@ -1182,7 +1258,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_mental_models"))
|
||||
async def list_mental_models(
|
||||
tags: list[str] | None = None,
|
||||
detail: str = "full",
|
||||
@@ -1221,7 +1297,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_mental_models"))
|
||||
async def list_mental_models(
|
||||
tags: list[str] | None = None,
|
||||
detail: str = "full",
|
||||
@@ -1262,7 +1338,7 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_mental_model"))
|
||||
async def get_mental_model(
|
||||
mental_model_id: str,
|
||||
detail: str = "full",
|
||||
@@ -1302,7 +1378,7 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_mental_model"))
|
||||
async def get_mental_model(
|
||||
mental_model_id: str,
|
||||
detail: str = "full",
|
||||
@@ -1344,7 +1420,7 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_mental_model"))
|
||||
async def create_mental_model(
|
||||
name: str,
|
||||
source_query: str,
|
||||
@@ -1428,7 +1504,7 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_mental_model"))
|
||||
async def create_mental_model(
|
||||
name: str,
|
||||
source_query: str,
|
||||
@@ -1510,7 +1586,7 @@ def _register_update_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("update_mental_model"))
|
||||
async def update_mental_model(
|
||||
mental_model_id: str,
|
||||
name: str | None = None,
|
||||
@@ -1571,7 +1647,7 @@ def _register_update_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("update_mental_model"))
|
||||
async def update_mental_model(
|
||||
mental_model_id: str,
|
||||
name: str | None = None,
|
||||
@@ -1634,7 +1710,7 @@ def _register_delete_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_mental_model"))
|
||||
async def delete_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -1670,7 +1746,7 @@ def _register_delete_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_mental_model"))
|
||||
async def delete_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
@@ -1708,7 +1784,7 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("refresh_mental_model"))
|
||||
async def refresh_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -1752,7 +1828,7 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("refresh_mental_model"))
|
||||
async def refresh_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
@@ -1796,7 +1872,7 @@ def _register_clear_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("clear_mental_model"))
|
||||
async def clear_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -1842,7 +1918,7 @@ def _register_clear_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("clear_mental_model"))
|
||||
async def clear_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
@@ -1893,7 +1969,7 @@ def _register_list_directives(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_directives"))
|
||||
async def list_directives(
|
||||
tags: list[str] | None = None,
|
||||
active_only: bool = True,
|
||||
@@ -1931,7 +2007,7 @@ def _register_list_directives(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_directives"))
|
||||
async def list_directives(
|
||||
tags: list[str] | None = None,
|
||||
active_only: bool = True,
|
||||
@@ -1971,7 +2047,7 @@ def _register_create_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_directive"))
|
||||
async def create_directive(
|
||||
name: str,
|
||||
content: str,
|
||||
@@ -2017,7 +2093,7 @@ def _register_create_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_directive"))
|
||||
async def create_directive(
|
||||
name: str,
|
||||
content: str,
|
||||
@@ -2065,7 +2141,7 @@ def _register_delete_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_directive"))
|
||||
async def delete_directive(
|
||||
directive_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2101,7 +2177,7 @@ def _register_delete_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_directive"))
|
||||
async def delete_directive(
|
||||
directive_id: str,
|
||||
) -> dict:
|
||||
@@ -2144,7 +2220,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_memories"))
|
||||
async def list_memories(
|
||||
type: str | None = None,
|
||||
q: str | None = None,
|
||||
@@ -2159,7 +2235,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
browse/search without relevance ranking.
|
||||
|
||||
Args:
|
||||
type: Filter by fact type: 'world', 'experience', or 'opinion'
|
||||
type: Filter by fact type: 'world', 'experience', or 'observation'
|
||||
q: Optional text search query to filter memories
|
||||
limit: Maximum number of results (default: 100)
|
||||
offset: Pagination offset (default: 0)
|
||||
@@ -2188,7 +2264,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_memories"))
|
||||
async def list_memories(
|
||||
type: str | None = None,
|
||||
q: str | None = None,
|
||||
@@ -2202,7 +2278,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
browse/search without relevance ranking.
|
||||
|
||||
Args:
|
||||
type: Filter by fact type: 'world', 'experience', or 'opinion'
|
||||
type: Filter by fact type: 'world', 'experience', or 'observation'
|
||||
q: Optional text search query to filter memories
|
||||
limit: Maximum number of results (default: 100)
|
||||
offset: Pagination offset (default: 0)
|
||||
@@ -2234,7 +2310,7 @@ def _register_get_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_memory"))
|
||||
async def get_memory(
|
||||
memory_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2270,7 +2346,7 @@ def _register_get_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_memory"))
|
||||
async def get_memory(
|
||||
memory_id: str,
|
||||
) -> dict:
|
||||
@@ -2321,7 +2397,7 @@ def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(description=_EDIT_DOC)
|
||||
@mcp.tool(description=_EDIT_DOC, annotations=_tool_annotations("update_memory"))
|
||||
async def update_memory(
|
||||
memory_id: str,
|
||||
text: str | None = None,
|
||||
@@ -2367,7 +2443,7 @@ def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(description=_EDIT_DOC)
|
||||
@mcp.tool(description=_EDIT_DOC, annotations=_tool_annotations("update_memory"))
|
||||
async def update_memory(
|
||||
memory_id: str,
|
||||
text: str | None = None,
|
||||
@@ -2426,7 +2502,7 @@ def _register_invalidate_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPT
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(description=_INVALIDATE_DOC)
|
||||
@mcp.tool(description=_INVALIDATE_DOC, annotations=_tool_annotations("invalidate_memory"))
|
||||
async def invalidate_memory(
|
||||
memory_id: str,
|
||||
reason: str | None = None,
|
||||
@@ -2466,7 +2542,7 @@ def _register_invalidate_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPT
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(description=_INVALIDATE_DOC)
|
||||
@mcp.tool(description=_INVALIDATE_DOC, annotations=_tool_annotations("invalidate_memory"))
|
||||
async def invalidate_memory(
|
||||
memory_id: str,
|
||||
reason: str | None = None,
|
||||
@@ -2513,7 +2589,7 @@ def _register_list_documents(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_documents"))
|
||||
async def list_documents(
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
@@ -2551,7 +2627,7 @@ def _register_list_documents(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_documents"))
|
||||
async def list_documents(
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
@@ -2591,7 +2667,7 @@ def _register_get_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsC
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_document"))
|
||||
async def get_document(
|
||||
document_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2627,7 +2703,7 @@ def _register_get_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsC
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_document"))
|
||||
async def get_document(
|
||||
document_id: str,
|
||||
) -> dict:
|
||||
@@ -2665,7 +2741,7 @@ def _register_delete_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_document"))
|
||||
async def delete_document(
|
||||
document_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2699,7 +2775,7 @@ def _register_delete_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_document"))
|
||||
async def delete_document(
|
||||
document_id: str,
|
||||
) -> dict:
|
||||
@@ -2740,7 +2816,7 @@ def _register_list_operations(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_operations"))
|
||||
async def list_operations(
|
||||
status: str | None = None,
|
||||
limit: int = 20,
|
||||
@@ -2777,7 +2853,7 @@ def _register_list_operations(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_operations"))
|
||||
async def list_operations(
|
||||
status: str | None = None,
|
||||
limit: int = 20,
|
||||
@@ -2816,7 +2892,7 @@ def _register_get_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_operation"))
|
||||
async def get_operation(
|
||||
operation_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2850,7 +2926,7 @@ def _register_get_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_operation"))
|
||||
async def get_operation(
|
||||
operation_id: str,
|
||||
) -> dict:
|
||||
@@ -2886,7 +2962,7 @@ def _register_cancel_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("cancel_operation"))
|
||||
async def cancel_operation(
|
||||
operation_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2918,7 +2994,7 @@ def _register_cancel_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("cancel_operation"))
|
||||
async def cancel_operation(
|
||||
operation_id: str,
|
||||
) -> dict:
|
||||
@@ -2957,7 +3033,7 @@ def _register_list_tags(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConf
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_tags"))
|
||||
async def list_tags(
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
@@ -2994,7 +3070,7 @@ def _register_list_tags(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConf
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_tags"))
|
||||
async def list_tags(
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
@@ -3033,7 +3109,7 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_bank"))
|
||||
async def get_bank(
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
@@ -3066,7 +3142,7 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_bank"))
|
||||
async def get_bank() -> dict:
|
||||
"""
|
||||
Get the profile of this memory bank.
|
||||
@@ -3096,7 +3172,7 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
|
||||
def _register_get_bank_stats(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the get_bank_stats tool (multi-bank only)."""
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_bank_stats"))
|
||||
async def get_bank_stats(
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
@@ -3169,7 +3245,7 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("update_bank"))
|
||||
async def update_bank(
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
@@ -3230,7 +3306,7 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("update_bank"))
|
||||
async def update_bank(
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
@@ -3293,7 +3369,7 @@ def _register_delete_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_bank"))
|
||||
async def delete_bank(
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
@@ -3325,7 +3401,7 @@ def _register_delete_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_bank"))
|
||||
async def delete_bank() -> dict:
|
||||
"""
|
||||
Delete this memory bank and all its data.
|
||||
@@ -3356,7 +3432,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("clear_memories"))
|
||||
async def clear_memories(
|
||||
type: str | None = None,
|
||||
bank_id: str | None = None,
|
||||
@@ -3367,7 +3443,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
Optionally filter by fact type to only clear specific kinds of memories.
|
||||
|
||||
Args:
|
||||
type: Optional fact type filter: 'world', 'experience', or 'opinion'. If not specified, clears all.
|
||||
type: Optional fact type filter: 'world', 'experience', or 'observation'. If not specified, clears all.
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
@@ -3391,7 +3467,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("clear_memories"))
|
||||
async def clear_memories(
|
||||
type: str | None = None,
|
||||
) -> dict:
|
||||
@@ -3401,7 +3477,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
Optionally filter by fact type to only clear specific kinds of memories.
|
||||
|
||||
Args:
|
||||
type: Optional fact type filter: 'world', 'experience', or 'opinion'. If not specified, clears all.
|
||||
type: Optional fact type filter: 'world', 'experience', or 'observation'. If not specified, clears all.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
|
||||
@@ -11,15 +11,17 @@ This module provides metrics for:
|
||||
- Database connection pool metrics
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import importlib
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
|
||||
_resource_mod = importlib.import_module("resource") if importlib.util.find_spec("resource") else None
|
||||
import threading
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from typing import TYPE_CHECKING, Callable
|
||||
from typing import TYPE_CHECKING, Callable, NamedTuple
|
||||
|
||||
from opentelemetry import metrics
|
||||
from opentelemetry.exporter.prometheus import PrometheusMetricReader
|
||||
@@ -75,6 +77,28 @@ LLM_DURATION_BUCKETS = (0.1, 0.25, 0.5, 1.0, 2.0, 3.0, 5.0, 10.0, 15.0, 30.0, 60
|
||||
# HTTP request duration buckets (millisecond-level for fast endpoints)
|
||||
HTTP_DURATION_BUCKETS = (0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0)
|
||||
|
||||
# How often the backlog / queue-depth gauge caches are refreshed (seconds).
|
||||
# The counts are aggregate COUNT queries, so a background task refreshes a
|
||||
# cache and the observable gauges read from it — keeping the /metrics scrape
|
||||
# path synchronous (the same reason the db-pool gauges read cached state).
|
||||
BACKLOG_METRICS_REFRESH_SECONDS = 30
|
||||
|
||||
|
||||
class _AsyncOpKey(NamedTuple):
|
||||
"""Cache / label key for the async-operation queue gauge."""
|
||||
|
||||
tenant: str
|
||||
operation_type: str
|
||||
status: str
|
||||
bank_id: str | None
|
||||
|
||||
|
||||
class _BacklogKey(NamedTuple):
|
||||
"""Cache / label key for the consolidation backlog and failed gauges."""
|
||||
|
||||
tenant: str
|
||||
bank_id: str | None
|
||||
|
||||
|
||||
def get_token_bucket(token_count: int) -> str:
|
||||
"""
|
||||
@@ -113,6 +137,27 @@ def get_token_bucket(token_count: int) -> str:
|
||||
return "50k+"
|
||||
|
||||
|
||||
# Template unbounded id segments before a path is used as the low-cardinality
|
||||
# "endpoint" metric label. A raw per-bank path segment (e.g. user-123) would
|
||||
# otherwise create one never-evicted OTel series per bank.
|
||||
_METRIC_BANK_SEGMENT_RE = re.compile(r"(/banks/)[^/]+")
|
||||
_METRIC_UUID_RE = re.compile(r"/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}")
|
||||
_METRIC_NUMERIC_ID_RE = re.compile(r"/\d+(?=/|$)")
|
||||
|
||||
|
||||
def normalize_http_endpoint(path: str) -> str:
|
||||
"""Template high-cardinality id segments in an HTTP path for safe metric labeling.
|
||||
|
||||
Collapses the "/banks/<id>" segment (any bank id, including non-numeric ones like
|
||||
"user-123"), UUIDs, and numeric ids to placeholders so the "endpoint" metric label
|
||||
has bounded cardinality. Analogous to get_token_bucket for token counts.
|
||||
"""
|
||||
path = _METRIC_BANK_SEGMENT_RE.sub(r"\g<1>{bank_id}", path)
|
||||
path = _METRIC_UUID_RE.sub("/{id}", path)
|
||||
path = _METRIC_NUMERIC_ID_RE.sub("/{id}", path)
|
||||
return path
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Global meter instance
|
||||
@@ -201,6 +246,19 @@ class MetricsCollectorBase:
|
||||
"""Context manager to record operation duration and status."""
|
||||
raise NotImplementedError
|
||||
|
||||
def record_operation_result(
|
||||
self,
|
||||
operation: str,
|
||||
bank_id: str,
|
||||
success: bool,
|
||||
duration: float,
|
||||
source: str = "api",
|
||||
budget: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
):
|
||||
"""Record a single completed operation with an explicit success label."""
|
||||
raise NotImplementedError
|
||||
|
||||
def record_llm_call(
|
||||
self,
|
||||
provider: str,
|
||||
@@ -254,6 +312,19 @@ class NoOpMetricsCollector(MetricsCollectorBase):
|
||||
"""No-op context manager."""
|
||||
yield
|
||||
|
||||
def record_operation_result(
|
||||
self,
|
||||
operation: str,
|
||||
bank_id: str,
|
||||
success: bool,
|
||||
duration: float,
|
||||
source: str = "api",
|
||||
budget: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
):
|
||||
"""No-op operation result recording."""
|
||||
pass
|
||||
|
||||
def record_llm_call(
|
||||
self,
|
||||
provider: str,
|
||||
@@ -361,6 +432,13 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
# DB pool metrics holder (set via set_db_pool)
|
||||
self._db_pool: "asyncpg.Pool | None" = None
|
||||
|
||||
# Backlog / queue-depth gauge caches, refreshed by a background task
|
||||
# (see _setup_backlog_metrics) so the scrape path stays synchronous.
|
||||
self._async_ops_counts: dict[_AsyncOpKey, int] = {}
|
||||
self._consolidation_backlog: dict[_BacklogKey, int] = {}
|
||||
self._consolidation_failed: dict[_BacklogKey, int] = {}
|
||||
self._backlog_task: "asyncio.Task | None" = None
|
||||
|
||||
@contextmanager
|
||||
def record_operation(
|
||||
self,
|
||||
@@ -386,18 +464,6 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
max_tokens: Optional max tokens for the operation
|
||||
"""
|
||||
start_time = time.time()
|
||||
attributes = {
|
||||
"operation": operation,
|
||||
"source": source,
|
||||
"tenant": _get_tenant(),
|
||||
}
|
||||
if self._include_bank_id:
|
||||
attributes["bank_id"] = bank_id
|
||||
if budget:
|
||||
attributes["budget"] = budget
|
||||
if max_tokens:
|
||||
attributes["max_tokens"] = str(max_tokens)
|
||||
|
||||
success = True
|
||||
cancelled = False
|
||||
try:
|
||||
@@ -416,14 +482,51 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
raise
|
||||
finally:
|
||||
if not cancelled:
|
||||
duration = time.time() - start_time
|
||||
attributes["success"] = str(success).lower()
|
||||
self.record_operation_result(
|
||||
operation,
|
||||
bank_id,
|
||||
success=success,
|
||||
duration=time.time() - start_time,
|
||||
source=source,
|
||||
budget=budget,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
|
||||
# Record duration
|
||||
self.operation_duration.record(duration, attributes)
|
||||
def record_operation_result(
|
||||
self,
|
||||
operation: str,
|
||||
bank_id: str,
|
||||
success: bool,
|
||||
duration: float,
|
||||
source: str = "api",
|
||||
budget: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
):
|
||||
"""Record a single completed operation (duration + count) with a success label.
|
||||
|
||||
# Record operation count
|
||||
self.operation_total.add(1, attributes)
|
||||
Direct (non-context-manager) recording for code paths that need explicit
|
||||
success control rather than the exception-based ``record_operation`` — e.g.
|
||||
the async worker, where deferrals/retries are not terminal outcomes and must
|
||||
not be counted as completions.
|
||||
"""
|
||||
attributes = {
|
||||
"operation": operation,
|
||||
"source": source,
|
||||
"tenant": _get_tenant(),
|
||||
}
|
||||
if self._include_bank_id:
|
||||
attributes["bank_id"] = bank_id
|
||||
if budget:
|
||||
attributes["budget"] = budget
|
||||
if max_tokens:
|
||||
attributes["max_tokens"] = str(max_tokens)
|
||||
attributes["success"] = str(success).lower()
|
||||
|
||||
# Record duration
|
||||
self.operation_duration.record(duration, attributes)
|
||||
|
||||
# Record operation count
|
||||
self.operation_total.add(1, attributes)
|
||||
|
||||
def record_llm_call(
|
||||
self,
|
||||
@@ -628,6 +731,10 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
"""
|
||||
self._db_pool = pool
|
||||
self._setup_db_pool_metrics()
|
||||
from .config import get_config
|
||||
|
||||
if get_config().metrics_backlog_enabled:
|
||||
self._setup_backlog_metrics()
|
||||
|
||||
def _setup_db_pool_metrics(self):
|
||||
"""Set up observable gauges for database pool metrics."""
|
||||
@@ -693,6 +800,192 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
unit="{connections}",
|
||||
)
|
||||
|
||||
def _setup_backlog_metrics(self):
|
||||
"""Observable gauges for the async-operation queue and the
|
||||
consolidation backlog.
|
||||
|
||||
These mirror fields the bank-stats endpoint already computes
|
||||
(``operations_by_status``, ``pending_consolidation``,
|
||||
``failed_consolidation``) but expose them as scrapable gauges, so
|
||||
queue depth and backlog can be trended and alerted on instead of only
|
||||
polled per-bank over HTTP. The two motivating questions both come for
|
||||
free here: "is the worker keeping up?" (async-op queue) and "is the
|
||||
knowledge base caught up?" (consolidation backlog) — including the
|
||||
``processing`` state, which is the only signal that surfaces a hung
|
||||
operation stuck holding a worker slot.
|
||||
|
||||
Counts are aggregate ``COUNT`` queries, so a background task refreshes
|
||||
a cache every ``BACKLOG_METRICS_REFRESH_SECONDS`` and these callbacks
|
||||
read it — keeping the scrape path synchronous, the same approach as
|
||||
the db-pool gauges above.
|
||||
"""
|
||||
if self._backlog_task is not None:
|
||||
return # already started for this collector
|
||||
|
||||
def get_async_operations(_options):
|
||||
for key, value in list(self._async_ops_counts.items()):
|
||||
attrs = {"tenant": key.tenant, "operation_type": key.operation_type, "status": key.status}
|
||||
if key.bank_id is not None:
|
||||
attrs["bank_id"] = key.bank_id
|
||||
yield metrics.Observation(value, attrs)
|
||||
|
||||
def get_consolidation_backlog(_options):
|
||||
for key, value in list(self._consolidation_backlog.items()):
|
||||
attrs = {"tenant": key.tenant}
|
||||
if key.bank_id is not None:
|
||||
attrs["bank_id"] = key.bank_id
|
||||
yield metrics.Observation(value, attrs)
|
||||
|
||||
def get_consolidation_failed(_options):
|
||||
for key, value in list(self._consolidation_failed.items()):
|
||||
attrs = {"tenant": key.tenant}
|
||||
if key.bank_id is not None:
|
||||
attrs["bank_id"] = key.bank_id
|
||||
yield metrics.Observation(value, attrs)
|
||||
|
||||
self.meter.create_observable_gauge(
|
||||
name="hindsight.async_operations",
|
||||
callbacks=[get_async_operations],
|
||||
description="Async operations in a non-terminal state, by operation_type and status "
|
||||
"(pending=queued backlog, processing=in-flight, failed=stranded)",
|
||||
unit="{operations}",
|
||||
)
|
||||
self.meter.create_observable_gauge(
|
||||
name="hindsight.consolidation.backlog",
|
||||
callbacks=[get_consolidation_backlog],
|
||||
description="Source memories (experience/world) not yet consolidated into observations",
|
||||
unit="{memories}",
|
||||
)
|
||||
self.meter.create_observable_gauge(
|
||||
name="hindsight.consolidation.failed",
|
||||
callbacks=[get_consolidation_failed],
|
||||
description="Source memories whose consolidation permanently failed "
|
||||
"(recoverable via the consolidation recovery endpoint)",
|
||||
unit="{memories}",
|
||||
)
|
||||
|
||||
# Drive the caches from a background task on the running loop.
|
||||
# set_db_pool runs during async startup, so a loop is normally present;
|
||||
# if not, the gauges simply stay empty rather than crashing collection.
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
logger.warning("No running event loop; backlog metrics disabled")
|
||||
return
|
||||
# Process-lifetime task: there is no collector teardown hook to cancel it
|
||||
# on, so it's torn down with the event loop at process shutdown. If a
|
||||
# shutdown path is ever added, cancel self._backlog_task there.
|
||||
self._backlog_task = loop.create_task(self._backlog_refresh_loop())
|
||||
|
||||
async def _backlog_refresh_loop(self):
|
||||
"""Periodically refresh the backlog / queue-depth caches."""
|
||||
while True:
|
||||
try:
|
||||
await self._refresh_backlog()
|
||||
except Exception:
|
||||
logger.debug("Backlog metrics refresh failed", exc_info=True)
|
||||
await asyncio.sleep(BACKLOG_METRICS_REFRESH_SECONDS)
|
||||
|
||||
async def _refresh_backlog(self):
|
||||
"""Recount the async-operation queue and consolidation backlog across
|
||||
every provisioned Hindsight schema.
|
||||
|
||||
Per-bank labels are gated behind ``metrics_include_bank_id`` (off by
|
||||
default) to keep cardinality bounded; when off, counts are aggregated
|
||||
per tenant/schema. All SQL here is PostgreSQL-specific (``FILTER``,
|
||||
``information_schema``), which is consistent with this collector
|
||||
already being bound to an asyncpg pool.
|
||||
"""
|
||||
if self._db_pool is None:
|
||||
return
|
||||
|
||||
async_ops: dict[_AsyncOpKey, int] = {}
|
||||
backlog: dict[_BacklogKey, int] = {}
|
||||
failed: dict[_BacklogKey, int] = {}
|
||||
per_bank = self._include_bank_id
|
||||
bank_sel = "bank_id, " if per_bank else ""
|
||||
bank_grp = " GROUP BY bank_id" if per_bank else ""
|
||||
|
||||
async with self._db_pool.acquire() as conn:
|
||||
# memory_units is the central per-tenant table; its presence marks a
|
||||
# provisioned Hindsight schema.
|
||||
schema_rows = await conn.fetch(
|
||||
"SELECT table_schema FROM information_schema.tables WHERE table_name = 'memory_units'"
|
||||
)
|
||||
for schema_row in schema_rows:
|
||||
schema = schema_row["table_schema"]
|
||||
|
||||
# Worker queue depth — mirrors operations_by_status, split by
|
||||
# operation_type. Terminal states (completed/cancelled) are
|
||||
# excluded on purpose: a gauge of finished work grows without
|
||||
# bound and says nothing about current load.
|
||||
# Index: idx_async_operations_status.
|
||||
ops_grp = "operation_type, status" + (", bank_id" if per_bank else "")
|
||||
try:
|
||||
rows = await conn.fetch(
|
||||
f"SELECT operation_type, status, {bank_sel}COUNT(*) AS count "
|
||||
f'FROM "{schema}".async_operations '
|
||||
"WHERE status IN ('pending', 'processing', 'failed') "
|
||||
f"GROUP BY {ops_grp}"
|
||||
)
|
||||
for row in rows:
|
||||
bank = row["bank_id"] if per_bank else None
|
||||
key = _AsyncOpKey(schema, row["operation_type"] or "unknown", row["status"], bank)
|
||||
async_ops[key] = async_ops.get(key, 0) + int(row["count"])
|
||||
except Exception:
|
||||
logger.debug("Async-ops queue query failed for schema %s", schema, exc_info=True)
|
||||
|
||||
# Consolidation backlog + stranded counts. Two separate COUNT(*)
|
||||
# queries rather than one with two FILTERs — each WHERE matches a
|
||||
# partial-index predicate exactly:
|
||||
# idx_memory_units_unconsolidated WHERE consolidated_at IS NULL ...
|
||||
# idx_memory_units_consolidation_failed WHERE consolidation_failed_at IS NOT NULL ...
|
||||
# GROUP BY bank_id still composes — bank_id is each index's lead column.
|
||||
#
|
||||
# The backlog count runs with seqscan disabled in a scoped
|
||||
# transaction. The partial index matches its predicate, but
|
||||
# `consolidated_at IS NULL` is true for a large fraction of the
|
||||
# table (every observation has a null consolidated_at), so the
|
||||
# planner misjudges selectivity and otherwise seq-scans the whole
|
||||
# (largest) table on every refresh — verified on a 114k-row table
|
||||
# via EXPLAIN: seq scan ~92 ms vs index scan ~0.1 ms. SET LOCAL
|
||||
# forces the index path and resets at transaction end. The failed
|
||||
# count below needs no such nudge: `consolidation_failed_at IS NOT
|
||||
# NULL` is rare, so its index is chosen on cost.
|
||||
try:
|
||||
async with conn.transaction():
|
||||
await conn.execute("SET LOCAL enable_seqscan = off")
|
||||
rows = await conn.fetch(
|
||||
f"SELECT {bank_sel}COUNT(*) AS count "
|
||||
f'FROM "{schema}".memory_units '
|
||||
"WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')"
|
||||
f"{bank_grp}"
|
||||
)
|
||||
for row in rows:
|
||||
bank = row["bank_id"] if per_bank else None
|
||||
key = _BacklogKey(schema, bank)
|
||||
backlog[key] = backlog.get(key, 0) + int(row["count"])
|
||||
except Exception:
|
||||
logger.debug("Consolidation backlog query failed for schema %s", schema, exc_info=True)
|
||||
|
||||
try:
|
||||
rows = await conn.fetch(
|
||||
f"SELECT {bank_sel}COUNT(*) AS count "
|
||||
f'FROM "{schema}".memory_units '
|
||||
"WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')"
|
||||
f"{bank_grp}"
|
||||
)
|
||||
for row in rows:
|
||||
bank = row["bank_id"] if per_bank else None
|
||||
key = _BacklogKey(schema, bank)
|
||||
failed[key] = failed.get(key, 0) + int(row["count"])
|
||||
except Exception:
|
||||
logger.debug("Consolidation failed query failed for schema %s", schema, exc_info=True)
|
||||
|
||||
self._async_ops_counts = async_ops
|
||||
self._consolidation_backlog = backlog
|
||||
self._consolidation_failed = failed
|
||||
|
||||
|
||||
# Global metrics collector instance (defaults to no-op)
|
||||
_metrics_collector: MetricsCollectorBase = NoOpMetricsCollector()
|
||||
|
||||
@@ -686,24 +686,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
|
||||
|
||||
@@ -20,9 +20,23 @@ from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..engine.schema import fq_table_explicit as fq_table
|
||||
from ..metrics import get_metrics_collector
|
||||
from .exceptions import DeferOperation, RetryTaskAt
|
||||
from .stage import StageHolder, bind_holder
|
||||
|
||||
# Map DB operation_type -> metric `operation` label, collapsing the retain
|
||||
# variants onto "retain" so async worker completions land on the same
|
||||
# operation="retain" series the synchronous API path emits. Unknown types
|
||||
# pass through unchanged.
|
||||
_RETAIN_OP_TYPES = {"retain", "batch_retain", "file_convert_retain"}
|
||||
|
||||
|
||||
def _metric_operation_label(operation_type: str | None) -> str:
|
||||
if operation_type in _RETAIN_OP_TYPES:
|
||||
return "retain"
|
||||
return operation_type or "unknown"
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.db.base import DatabaseBackend, DatabaseConnection
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
@@ -701,6 +715,24 @@ class WorkerPoller:
|
||||
"""
|
||||
task_type = task.task_dict.get("type", "unknown")
|
||||
bank_id = task.task_dict.get("bank_id", "unknown")
|
||||
# Operation metric (source="worker"): record on terminal outcomes only, so
|
||||
# async worker throughput and latency (retain, consolidation and the other
|
||||
# worker task types) are visible in Prometheus. Prefer the DB-authoritative
|
||||
# operation_type.
|
||||
#
|
||||
# success semantics are deliberately narrow: success=false means the task
|
||||
# raised out to the poller (an unexpected error, or retry-exhausted). It does
|
||||
# NOT capture deterministic failures that the executor handles itself and
|
||||
# returns from normally (file_convert_retain, non-retryable errors via
|
||||
# memory_engine.execute_task) — those record success=true here. Treat this as
|
||||
# a completion-throughput signal, not a failure-rate one: for authoritative
|
||||
# failure visibility use the hindsight_async_operations{status="failed"} gauge,
|
||||
# which reads each operation's final DB status.
|
||||
op_label = _metric_operation_label(task.task_dict.get("operation_type") or task_type)
|
||||
op_start = time.time()
|
||||
metrics = get_metrics_collector()
|
||||
# None = not a terminal outcome (deferred/retried) → no metric.
|
||||
terminal_success: bool | None = None
|
||||
|
||||
# Bind the stage holder in this task's own contextvar scope so engine
|
||||
# code running under us can update it via stage.set_stage(). If holder
|
||||
@@ -717,14 +749,28 @@ class WorkerPoller:
|
||||
task.task_dict["_schema"] = task.schema
|
||||
await self._executor(task.task_dict)
|
||||
logger.debug(f"Task {task.operation_id} execution finished")
|
||||
terminal_success = True
|
||||
except DeferOperation as e:
|
||||
# Deferral is not a terminal outcome — do not record a completion.
|
||||
await self._defer_operation(task.operation_id, e.exec_date, e.reason, task.schema)
|
||||
except RetryTaskAt as e:
|
||||
# Retry is not a terminal outcome — do not record a completion.
|
||||
await self._schedule_retry(task.operation_id, e.retry_at, str(e), task.schema)
|
||||
except Exception as e:
|
||||
logger.error(f"Task {task.operation_id} failed: {e}")
|
||||
traceback.print_exc()
|
||||
await self._mark_failed(task.operation_id, str(e), task.schema)
|
||||
terminal_success = False
|
||||
|
||||
# Record the metric outside the executor's exception scope so a metrics
|
||||
# reporting failure can never be mistaken for a task failure and flip terminal state.
|
||||
if terminal_success is not None:
|
||||
try:
|
||||
metrics.record_operation_result(
|
||||
op_label, bank_id, success=terminal_success, duration=time.time() - op_start, source="worker"
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(f"Failed to record worker operation metric for {task.operation_id}", exc_info=True)
|
||||
|
||||
async def recover_own_tasks(self) -> int:
|
||||
"""
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api-slim"
|
||||
version = "0.8.2"
|
||||
version = "0.8.3"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -60,10 +60,10 @@ dependencies = [
|
||||
"pyasn1>=0.6.3", # DoS vulnerability fix
|
||||
"urllib3>=2.7.0", # Decompression-bomb safeguards bypass + sensitive header forwarding fixes
|
||||
"langchain-core>=1.2.22", # Path traversal in legacy load_prompt functions fix
|
||||
"langsmith>=0.6.3", # SSRF via tracing header injection fix
|
||||
"langsmith>=0.8.18", # GHSA-f4xh-w4cj-qxq8: arbitrary server-side file read in TracingMiddleware fix (supersedes >=0.6.3 SSRF tracing-header-injection floor)
|
||||
"protobuf>=6.33.5", # JSON recursion depth bypass fix
|
||||
"pillow>=12.1.1", # Out-of-bounds write in PSD image loading fix
|
||||
"cryptography>=46.0.6,<47", # Incomplete DNS name constraint enforcement fix; cap <47 — 47.0.0 SIGILLs on some ARM64 Linux VMs (Docker/Podman on Apple Silicon), pyca/cryptography#14733
|
||||
"cryptography>=48.0.1", # GHSA-537c-gmf6-5ccf: bundled-OpenSSL OOB read fix needs >=48.0.1. Prior <47 cap (47.0.0 SIGILL on ARM64 Docker/Podman, pyca/cryptography#14733) lifted — 47/48/49 verified importing + RSA sign/verify cleanly on linux/arm64 (Docker on Apple Silicon) and native arm64 macOS; upstream issue closed unconfirmed.
|
||||
"filelock>=3.20.1", # TOCTOU race condition fix
|
||||
"authlib>=1.6.9", # Account takeover/JWS header injection vulnerability fix
|
||||
"pyjwt>=2.12.0", # Accepts unknown crit header extensions fix
|
||||
|
||||
@@ -547,3 +547,36 @@ async def test_run_migration_with_schema_only_runs_requested_schema(monkeypatch)
|
||||
assert calls["run_migrations"] == [("resolved::postgresql://test", "tenant_demo")]
|
||||
assert calls["ensure_vector_extension"] == [("resolved::postgresql://test", "pgvector", "tenant_demo")]
|
||||
assert calls["ensure_text_search_extension"] == [("resolved::postgresql://test", "native", "", "tenant_demo")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("ensure_extensions", "expected"),
|
||||
[(True, True), (False, False)],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_migration_threads_ensure_extensions_flag(monkeypatch, ensure_extensions, expected):
|
||||
"""The --skip-extension-reconcile flag (ensure_extensions=False) must reach run_migrations_for_schemas.
|
||||
|
||||
The post-migration vector/text-search reconcile only does work on a backend change, so operators
|
||||
can skip it on a no-change re-migration over many tenant schemas. Verify the flag is threaded through
|
||||
rather than silently dropped.
|
||||
"""
|
||||
monkeypatch.setenv("HINDSIGHT_API_DATABASE_URL", "postgresql://test")
|
||||
captured: dict = {}
|
||||
|
||||
async def fake_resolve_database_url(db_url: str) -> str:
|
||||
return f"resolved::{db_url}"
|
||||
|
||||
def fake_run_migrations_for_schemas(database_url, schemas, **kwargs):
|
||||
captured["ensure_extensions"] = kwargs.get("ensure_extensions")
|
||||
|
||||
monkeypatch.setattr(admin_cli, "load_extension", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(admin_cli, "resolve_database_url", fake_resolve_database_url)
|
||||
|
||||
from hindsight_api import migrations as migrations_module
|
||||
|
||||
monkeypatch.setattr(migrations_module, "run_migrations_for_schemas", fake_run_migrations_for_schemas)
|
||||
|
||||
await admin_cli._run_migration("postgresql://test", schema="tenant_demo", ensure_extensions=ensure_extensions)
|
||||
|
||||
assert captured["ensure_extensions"] is expected
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
"""Regression tests for issue #1002 — Anthropic structured output via forced tool_use.
|
||||
|
||||
When strict_schema=True, AnthropicLLM.call() must request the schema through a single
|
||||
forced tool_use tool (tool_choice={"type":"tool",...}) and read the validated args from
|
||||
the tool_use block, NOT inject the schema as text and json.loads() the reply (which caused
|
||||
a ~1:1 invalid-JSON retry storm / OOM in production).
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class _Decision(BaseModel):
|
||||
action: str
|
||||
reason: str
|
||||
|
||||
|
||||
def _make_anthropic_provider():
|
||||
with patch("anthropic.AsyncAnthropic") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
from hindsight_api.engine.providers.anthropic_llm import AnthropicLLM
|
||||
|
||||
provider = AnthropicLLM(
|
||||
provider="anthropic",
|
||||
api_key="fake-key",
|
||||
base_url="",
|
||||
model="claude-sonnet-4-20250514",
|
||||
)
|
||||
provider._client = MagicMock()
|
||||
return provider
|
||||
|
||||
|
||||
def _tool_use_response(args: dict):
|
||||
block = MagicMock()
|
||||
block.type = "tool_use"
|
||||
block.name = "structured_response"
|
||||
block.input = args
|
||||
resp = MagicMock()
|
||||
resp.content = [block]
|
||||
resp.usage = MagicMock(input_tokens=5, output_tokens=2, cache_read_input_tokens=0)
|
||||
resp.stop_reason = "tool_use"
|
||||
return resp
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strict_schema_uses_forced_tool_choice():
|
||||
"""strict_schema=True ⇒ a single tool is defined and tool_choice forces it (no schema text-injection)."""
|
||||
provider = _make_anthropic_provider()
|
||||
provider._client.messages.create = AsyncMock(return_value=_tool_use_response({"action": "skip", "reason": "dup"}))
|
||||
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
|
||||
result = await provider.call(
|
||||
messages=[{"role": "user", "content": "decide"}],
|
||||
response_format=_Decision,
|
||||
strict_schema=True,
|
||||
scope="test",
|
||||
max_retries=0,
|
||||
)
|
||||
kwargs = provider._client.messages.create.call_args.kwargs
|
||||
# forced tool_use requested
|
||||
assert "tools" in kwargs and len(kwargs["tools"]) == 1
|
||||
assert kwargs["tool_choice"] == {"type": "tool", "name": "structured_response"}
|
||||
# schema NOT injected as text into the system prompt
|
||||
assert "valid JSON matching this schema" not in (kwargs.get("system") or "")
|
||||
# validated model returned straight from tool_use.input
|
||||
assert isinstance(result, _Decision)
|
||||
assert result.action == "skip"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strict_schema_tool_use_never_hits_json_retry_loop():
|
||||
"""A tool_use response is structurally valid → no second messages.create call (no retry storm)."""
|
||||
provider = _make_anthropic_provider()
|
||||
create = AsyncMock(return_value=_tool_use_response({"action": "keep", "reason": "novel"}))
|
||||
provider._client.messages.create = create
|
||||
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
|
||||
await provider.call(
|
||||
messages=[{"role": "user", "content": "x"}],
|
||||
response_format=_Decision,
|
||||
strict_schema=True,
|
||||
scope="test",
|
||||
max_retries=10, # would allow 11 attempts on the old text-parse path
|
||||
)
|
||||
assert create.await_count == 1 # exactly one call — the bug was N retries on malformed text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_strict_keeps_text_injection_fallback():
|
||||
"""strict_schema=False (default) preserves the legacy schema-in-prompt behavior."""
|
||||
provider = _make_anthropic_provider()
|
||||
block = MagicMock()
|
||||
block.type = "text"
|
||||
block.text = '{"action":"skip","reason":"d"}'
|
||||
resp = MagicMock()
|
||||
resp.content = [block]
|
||||
resp.usage = MagicMock(input_tokens=5, output_tokens=2, cache_read_input_tokens=0)
|
||||
resp.stop_reason = "end_turn"
|
||||
provider._client.messages.create = AsyncMock(return_value=resp)
|
||||
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
|
||||
result = await provider.call(
|
||||
messages=[{"role": "user", "content": "decide"}],
|
||||
response_format=_Decision,
|
||||
strict_schema=False,
|
||||
scope="test",
|
||||
max_retries=0,
|
||||
)
|
||||
kwargs = provider._client.messages.create.call_args.kwargs
|
||||
assert "tools" not in kwargs # no forced tool when not strict
|
||||
assert "valid JSON matching this schema" in (kwargs.get("system") or "")
|
||||
assert isinstance(result, _Decision)
|
||||
@@ -0,0 +1,90 @@
|
||||
"""Regression test: submitting an async op for a bank that doesn't exist must
|
||||
raise a clean validation error, not a raw asyncpg `ForeignKeyViolationError`.
|
||||
|
||||
`_submit_async_operation` inserts into `async_operations`, which has an FK to
|
||||
`banks.bank_id`. If a caller submits for a missing bank (typo, race against a
|
||||
deletion, integration that derives bank IDs before the bank is created), the
|
||||
INSERT raises `asyncpg.exceptions.ForeignKeyViolationError`. The FastAPI
|
||||
endpoint's broad `except Exception` then surfaces it as a 500 — but this is
|
||||
a client error, not a server error, and should be a 404.
|
||||
|
||||
This test exercises the call directly via `MemoryEngine.submit_async_*` so
|
||||
the failure mode is observable without spinning up the HTTP layer.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.extensions.operation_validator import OperationValidationError
|
||||
|
||||
pytestmark = pytest.mark.xdist_group("async_submit_bank_not_found_tests")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_inline_execution(memory):
|
||||
"""Prevent SyncTaskBackend from running the submitted op inline so we
|
||||
only test the submit-path failure, not downstream execution."""
|
||||
|
||||
async def _noop(_payload):
|
||||
return None
|
||||
|
||||
original = memory._task_backend.submit_task
|
||||
memory._task_backend.submit_task = _noop
|
||||
yield
|
||||
memory._task_backend.submit_task = original
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consolidation_submit_on_missing_bank_raises_validation_error(
|
||||
memory, request_context, no_inline_execution
|
||||
):
|
||||
"""A `/consolidate` submit against a bank that doesn't exist must raise
|
||||
OperationValidationError(404), not a raw asyncpg FK violation that bubbles
|
||||
out as a 500 from the API."""
|
||||
missing_bank = f"does-not-exist-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
with pytest.raises(OperationValidationError) as exc_info:
|
||||
await memory.submit_async_consolidation(
|
||||
bank_id=missing_bank,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert missing_bank in exc_info.value.reason
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_consolidation_submit_on_missing_bank_raises_validation_error(
|
||||
memory, request_context, no_inline_execution
|
||||
):
|
||||
"""Scoped consolidates (with `observation_scopes`) take the
|
||||
`dedupe_by_bank=False` branch, which historically skipped the bank lock
|
||||
entirely and went straight to the FK-violating INSERT. Same 404 contract."""
|
||||
missing_bank = f"does-not-exist-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
with pytest.raises(OperationValidationError) as exc_info:
|
||||
await memory.submit_async_consolidation(
|
||||
bank_id=missing_bank,
|
||||
request_context=request_context,
|
||||
observation_scopes=[{"tag": "anything"}],
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert missing_bank in exc_info.value.reason
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_maintenance_on_missing_bank_short_circuits(memory, request_context, no_inline_execution):
|
||||
"""`submit_async_graph_maintenance` has its own short-circuit that checks
|
||||
the per-bank queue before calling `_submit_async_operation`. A missing
|
||||
bank means an empty queue, so it returns `no_work=True` without reaching
|
||||
the FK-violating INSERT. This test pins that behaviour."""
|
||||
missing_bank = f"does-not-exist-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
result = await memory.submit_async_graph_maintenance(
|
||||
bank_id=missing_bank,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result == {"operation_id": None, "no_work": True}
|
||||
@@ -0,0 +1,244 @@
|
||||
"""
|
||||
Tests for the async-operation queue and consolidation backlog gauges
|
||||
(``_setup_backlog_metrics`` / ``_refresh_backlog`` in metrics.py).
|
||||
|
||||
These gauges expose, as scrapable time-series, the same counts the bank-stats
|
||||
endpoint already returns per bank (``operations_by_status``,
|
||||
``pending_consolidation``, ``failed_consolidation``):
|
||||
|
||||
- ``hindsight_async_operations{operation_type,status}`` — worker queue depth
|
||||
(pending=backlog, processing=in-flight, failed=stranded)
|
||||
- ``hindsight_consolidation_backlog`` — source memories not yet consolidated
|
||||
- ``hindsight_consolidation_failed`` — source memories permanently failed
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.metrics import MetricsCollector, _AsyncOpKey, _BacklogKey
|
||||
|
||||
|
||||
class _FakeTxn:
|
||||
async def __aenter__(self):
|
||||
return None
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
return False
|
||||
|
||||
|
||||
class _FakeConn:
|
||||
"""asyncpg-like connection whose fetch() is dispatched by SQL substring."""
|
||||
|
||||
def __init__(self, fetch_fn):
|
||||
self._fetch_fn = fetch_fn
|
||||
self.executed = []
|
||||
|
||||
async def fetch(self, sql, *args):
|
||||
return self._fetch_fn(sql, *args)
|
||||
|
||||
async def execute(self, sql, *args):
|
||||
self.executed.append(sql)
|
||||
|
||||
def transaction(self):
|
||||
return _FakeTxn()
|
||||
|
||||
|
||||
class _FakeAcquire:
|
||||
def __init__(self, conn):
|
||||
self._conn = conn
|
||||
|
||||
async def __aenter__(self):
|
||||
return self._conn
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
return False
|
||||
|
||||
|
||||
class _FakePool:
|
||||
def __init__(self, fetch_fn):
|
||||
self._conn = _FakeConn(fetch_fn)
|
||||
|
||||
def acquire(self):
|
||||
return _FakeAcquire(self._conn)
|
||||
|
||||
|
||||
def _collector(include_bank_id=False):
|
||||
mock_config = MagicMock()
|
||||
mock_config.metrics_include_bank_id = include_bank_id
|
||||
with (
|
||||
patch("hindsight_api.metrics.get_meter", return_value=MagicMock()),
|
||||
patch("hindsight_api.config.get_config", return_value=mock_config),
|
||||
):
|
||||
return MetricsCollector()
|
||||
|
||||
|
||||
def _set_db_pool_with_backlog_enabled(collector, pool):
|
||||
"""Call set_db_pool with the backlog flag forced on (it's off by default)."""
|
||||
mock_config = MagicMock()
|
||||
mock_config.metrics_backlog_enabled = True
|
||||
with patch("hindsight_api.config.get_config", return_value=mock_config):
|
||||
collector.set_db_pool(pool)
|
||||
|
||||
|
||||
def _rows_for(sql):
|
||||
"""Canned results, keyed off distinctive substrings of each query."""
|
||||
if "information_schema.tables" in sql:
|
||||
return [{"table_schema": "public"}]
|
||||
if "async_operations" in sql:
|
||||
return [
|
||||
{"operation_type": "retain", "status": "pending", "count": 5},
|
||||
{"operation_type": "consolidation", "status": "pending", "count": 12},
|
||||
{"operation_type": "consolidation", "status": "processing", "count": 1},
|
||||
{"operation_type": "consolidation", "status": "failed", "count": 2},
|
||||
]
|
||||
if "memory_units" in sql and "consolidated_at IS NULL" in sql:
|
||||
return [{"count": 42}]
|
||||
if "memory_units" in sql and "consolidation_failed_at IS NOT NULL" in sql:
|
||||
return [{"count": 3}]
|
||||
return []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_backlog_aggregates_queue_and_consolidation():
|
||||
collector = _collector(include_bank_id=False)
|
||||
collector._db_pool = _FakePool(lambda sql, *a: _rows_for(sql))
|
||||
|
||||
await collector._refresh_backlog()
|
||||
|
||||
# Worker queue depth keyed by (schema, operation_type, status, bank=None)
|
||||
assert collector._async_ops_counts[("public", "retain", "pending", None)] == 5
|
||||
assert collector._async_ops_counts[("public", "consolidation", "pending", None)] == 12
|
||||
assert collector._async_ops_counts[("public", "consolidation", "processing", None)] == 1
|
||||
assert collector._async_ops_counts[("public", "consolidation", "failed", None)] == 2
|
||||
# Consolidation backlog (source memories), keyed by (schema, bank=None)
|
||||
assert collector._consolidation_backlog[("public", None)] == 42
|
||||
assert collector._consolidation_failed[("public", None)] == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_backlog_uses_index_matched_predicates_not_filter_scan():
|
||||
"""Backlog/failed must be two separate COUNT(*) queries whose WHERE matches
|
||||
a partial-index predicate exactly (no FILTER over a full-table scan), and
|
||||
the queue query must exclude terminal statuses."""
|
||||
captured = []
|
||||
collector = _collector()
|
||||
collector._db_pool = _FakePool(lambda sql, *a: (captured.append(sql), _rows_for(sql))[1])
|
||||
await collector._refresh_backlog()
|
||||
|
||||
mem_queries = [s for s in captured if "memory_units" in s and "COUNT(*)" in s]
|
||||
assert len(mem_queries) == 2 # split, not a single two-FILTER aggregate
|
||||
assert all("FILTER" not in s for s in mem_queries)
|
||||
assert any("consolidated_at IS NULL AND fact_type IN ('experience', 'world')" in s for s in mem_queries)
|
||||
assert any("consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')" in s for s in mem_queries)
|
||||
|
||||
ops_sql = next(s for s in captured if "async_operations" in s and "GROUP BY" in s)
|
||||
assert "status IN ('pending', 'processing', 'failed')" in ops_sql
|
||||
assert "completed" not in ops_sql and "cancelled" not in ops_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_backlog_count_runs_with_seqscan_disabled():
|
||||
"""`consolidated_at IS NULL` is true for a large fraction of the table, so
|
||||
the planner misjudges selectivity and won't use the partial index without a
|
||||
nudge — the backlog count must issue SET LOCAL enable_seqscan=off."""
|
||||
collector = _collector()
|
||||
pool = _FakePool(lambda sql, *a: _rows_for(sql))
|
||||
collector._db_pool = pool
|
||||
await collector._refresh_backlog()
|
||||
|
||||
assert any("enable_seqscan" in s.lower() and "off" in s.lower() for s in pool._conn.executed)
|
||||
# the result is still correct under the nudge
|
||||
assert collector._consolidation_backlog[("public", None)] == 42
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_backlog_per_bank_labels_and_group_by_when_enabled():
|
||||
"""With metrics_include_bank_id on, bank_id enters the cache key and the
|
||||
SQL switches to GROUP BY bank_id."""
|
||||
captured = []
|
||||
|
||||
def fetch(sql, *a):
|
||||
captured.append(sql)
|
||||
if "information_schema.tables" in sql:
|
||||
return [{"table_schema": "public"}]
|
||||
if "async_operations" in sql:
|
||||
return [{"operation_type": "retain", "status": "pending", "bank_id": "bankA", "count": 4}]
|
||||
if "memory_units" in sql and "consolidated_at IS NULL" in sql:
|
||||
return [{"bank_id": "bankA", "count": 11}]
|
||||
if "memory_units" in sql and "consolidation_failed_at IS NOT NULL" in sql:
|
||||
return [{"bank_id": "bankA", "count": 2}]
|
||||
return []
|
||||
|
||||
collector = _collector(include_bank_id=True)
|
||||
collector._db_pool = _FakePool(fetch)
|
||||
await collector._refresh_backlog()
|
||||
|
||||
assert collector._async_ops_counts[("public", "retain", "pending", "bankA")] == 4
|
||||
assert collector._consolidation_backlog[("public", "bankA")] == 11
|
||||
assert collector._consolidation_failed[("public", "bankA")] == 2
|
||||
# bank_id must be grouped in every per-bank count query
|
||||
assert all("GROUP BY bank_id" in s for s in captured if "memory_units" in s and "COUNT(*)" in s)
|
||||
|
||||
|
||||
def test_gauges_register_and_emit_cached_values_without_bank_id():
|
||||
collector = _collector(include_bank_id=False)
|
||||
# Sync call: no running loop, so gauges register but no background task spawns.
|
||||
_set_db_pool_with_backlog_enabled(collector, MagicMock())
|
||||
|
||||
gauges = {
|
||||
c.kwargs["name"]: c.kwargs["callbacks"][0]
|
||||
for c in collector.meter.create_observable_gauge.call_args_list
|
||||
if "callbacks" in c.kwargs
|
||||
}
|
||||
assert "hindsight.async_operations" in gauges
|
||||
assert "hindsight.consolidation.backlog" in gauges
|
||||
assert "hindsight.consolidation.failed" in gauges
|
||||
|
||||
collector._async_ops_counts = {
|
||||
_AsyncOpKey("public", "retain", "pending", None): 7,
|
||||
_AsyncOpKey("public", "consolidation", "processing", None): 1,
|
||||
}
|
||||
collector._consolidation_backlog = {_BacklogKey("public", None): 9}
|
||||
|
||||
obs = list(gauges["hindsight.async_operations"](None))
|
||||
by_label = {(o.attributes["operation_type"], o.attributes["status"]): o.value for o in obs}
|
||||
assert by_label[("retain", "pending")] == 7
|
||||
assert by_label[("consolidation", "processing")] == 1
|
||||
assert all("bank_id" not in o.attributes for o in obs) # cardinality guard
|
||||
|
||||
backlog_obs = list(gauges["hindsight.consolidation.backlog"](None))
|
||||
assert backlog_obs[0].value == 9
|
||||
assert backlog_obs[0].attributes["tenant"] == "public"
|
||||
|
||||
|
||||
def test_gauge_emits_bank_id_attribute_when_present():
|
||||
collector = _collector(include_bank_id=True)
|
||||
_set_db_pool_with_backlog_enabled(collector, MagicMock())
|
||||
gauges = {
|
||||
c.kwargs["name"]: c.kwargs["callbacks"][0]
|
||||
for c in collector.meter.create_observable_gauge.call_args_list
|
||||
if "callbacks" in c.kwargs
|
||||
}
|
||||
collector._consolidation_backlog = {_BacklogKey("public", "bankA"): 4}
|
||||
obs = list(gauges["hindsight.consolidation.backlog"](None))
|
||||
assert obs[0].value == 4
|
||||
assert obs[0].attributes["bank_id"] == "bankA"
|
||||
|
||||
|
||||
def test_backlog_gauges_not_registered_when_flag_disabled():
|
||||
"""Backlog metrics are off by default: set_db_pool must not register the
|
||||
gauges unless metrics_backlog_enabled is set."""
|
||||
collector = _collector()
|
||||
mock_config = MagicMock()
|
||||
mock_config.metrics_backlog_enabled = False
|
||||
with patch("hindsight_api.config.get_config", return_value=mock_config):
|
||||
collector.set_db_pool(MagicMock())
|
||||
|
||||
names = [
|
||||
c.kwargs.get("name") for c in collector.meter.create_observable_gauge.call_args_list if "callbacks" in c.kwargs
|
||||
]
|
||||
assert "hindsight.async_operations" not in names
|
||||
assert "hindsight.consolidation.backlog" not in names
|
||||
assert "hindsight.consolidation.failed" not in names
|
||||
assert collector._backlog_task is None
|
||||
@@ -186,6 +186,43 @@ async def test_invalidate_drops_entry() -> None:
|
||||
assert calls[0] == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalidate_detaches_in_flight_loader() -> None:
|
||||
cache = BankStatsCache(ttl_seconds=60, max_entries=100)
|
||||
stale_started = asyncio.Event()
|
||||
release_stale = asyncio.Event()
|
||||
fresh_started = asyncio.Event()
|
||||
|
||||
async def stale_loader() -> dict[str, Any]:
|
||||
stale_started.set()
|
||||
await release_stale.wait()
|
||||
return {"v": "stale"}
|
||||
|
||||
async def fresh_loader() -> dict[str, Any]:
|
||||
fresh_started.set()
|
||||
return {"v": "fresh"}
|
||||
|
||||
stale_task = asyncio.create_task(cache.get_or_load("schema", "bank", stale_loader))
|
||||
await stale_started.wait()
|
||||
await cache.invalidate("schema", "bank")
|
||||
|
||||
# A request after invalidation must start a new load instead of joining the
|
||||
# pre-invalidation query, which may contain data from before a bank write.
|
||||
fresh_result = await asyncio.wait_for(cache.get_or_load("schema", "bank", fresh_loader), timeout=1)
|
||||
assert fresh_started.is_set()
|
||||
assert fresh_result == {"v": "fresh"}
|
||||
|
||||
release_stale.set()
|
||||
assert await stale_task == {"v": "stale"}
|
||||
|
||||
# The stale loader completed last, but must not overwrite the fresh value.
|
||||
async def should_not_run() -> dict[str, Any]:
|
||||
raise AssertionError("fresh value was not cached")
|
||||
|
||||
cached = await cache.get_or_load("schema", "bank", should_not_run)
|
||||
assert cached == {"v": "fresh"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_drops_all_entries() -> None:
|
||||
cache = BankStatsCache(ttl_seconds=60, max_entries=100)
|
||||
@@ -199,3 +236,27 @@ async def test_clear_drops_all_entries() -> None:
|
||||
await cache.get_or_load("s", "a", loader)
|
||||
await cache.get_or_load("s", "b", loader)
|
||||
assert calls[0] == 4
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_detaches_in_flight_loaders() -> None:
|
||||
cache = BankStatsCache(ttl_seconds=60, max_entries=100)
|
||||
stale_started = asyncio.Event()
|
||||
release_stale = asyncio.Event()
|
||||
|
||||
async def stale_loader() -> dict[str, Any]:
|
||||
stale_started.set()
|
||||
await release_stale.wait()
|
||||
return {"v": "stale"}
|
||||
|
||||
async def fresh_loader() -> dict[str, Any]:
|
||||
return {"v": "fresh"}
|
||||
|
||||
stale_task = asyncio.create_task(cache.get_or_load("schema", "bank", stale_loader))
|
||||
await stale_started.wait()
|
||||
await cache.clear()
|
||||
assert await cache.get_or_load("schema", "bank", fresh_loader) == {"v": "fresh"}
|
||||
|
||||
release_stale.set()
|
||||
assert await stale_task == {"v": "stale"}
|
||||
assert await cache.get_or_load("schema", "bank", fresh_loader) == {"v": "fresh"}
|
||||
|
||||
@@ -0,0 +1,164 @@
|
||||
"""Regression tests: get_bank_stats cache must be invalidated by mutations.
|
||||
|
||||
`get_bank_stats` is served from a short-TTL per-process cache (`BankStatsCache`).
|
||||
`delete_bank` already invalidates that cache after it mutates counts, but the
|
||||
other operations that change the same counts — `delete_memory_unit`,
|
||||
`delete_document`, `clear_observations`, and `update_document` (when a tag
|
||||
change deletes observations) — did not, so a client polling stats right after a
|
||||
deletion would see pre-mutation counts until the TTL expired (up to a minute).
|
||||
|
||||
Each test pins a long TTL on the engine's stats cache so that, *without* the
|
||||
invalidation fix, the second `get_bank_stats` call would be served the stale
|
||||
cached value and the assertion would fail.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.bank_stats_cache import BankStatsCache
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
# A TTL long enough that, absent invalidation, the warmed cache would still be
|
||||
# served on the post-mutation read within the same test.
|
||||
_PINNED_TTL_SECONDS = 300.0
|
||||
|
||||
|
||||
async def _insert_memory(conn, bank_id: str, text: str, fact_type: str = "experience") -> uuid.UUID:
|
||||
"""Insert a memory unit directly, bypassing the LLM retain pipeline."""
|
||||
mem_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, created_at, updated_at, consolidated_at)
|
||||
VALUES ($1, $2, $3, $4, NOW(), NOW(), NOW(), NOW())
|
||||
""",
|
||||
mem_id,
|
||||
bank_id,
|
||||
text,
|
||||
fact_type,
|
||||
)
|
||||
return mem_id
|
||||
|
||||
|
||||
async def _insert_observation(conn, bank_id: str, text: str, source_memory_ids: list[uuid.UUID]) -> uuid.UUID:
|
||||
"""Insert an observation unit directly."""
|
||||
obs_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_units (
|
||||
id, bank_id, text, fact_type, event_date, source_memory_ids, proof_count, created_at, updated_at
|
||||
) VALUES ($1, $2, $3, 'observation', NOW(), $4, $5, NOW(), NOW())
|
||||
""",
|
||||
obs_id,
|
||||
bank_id,
|
||||
text,
|
||||
source_memory_ids,
|
||||
len(source_memory_ids),
|
||||
)
|
||||
return obs_id
|
||||
|
||||
|
||||
async def _insert_document(conn, bank_id: str, doc_id: str) -> None:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO documents (id, bank_id, original_text, content_hash)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
""",
|
||||
doc_id,
|
||||
bank_id,
|
||||
f"text-for-{doc_id}",
|
||||
doc_id,
|
||||
)
|
||||
|
||||
|
||||
async def _attach_unit_to_doc(conn, unit_id: uuid.UUID, doc_id: str) -> None:
|
||||
await conn.execute("UPDATE memory_units SET document_id = $1 WHERE id = $2", doc_id, unit_id)
|
||||
|
||||
|
||||
async def _ensure_bank(memory: MemoryEngine, bank_id: str, request_context: RequestContext) -> None:
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
|
||||
def _pin_cache(memory: MemoryEngine) -> None:
|
||||
"""Replace the stats cache with one that has a deterministic long TTL."""
|
||||
memory._bank_stats_cache = BankStatsCache(ttl_seconds=_PINNED_TTL_SECONDS, max_entries=128)
|
||||
|
||||
|
||||
class TestBankStatsCacheInvalidation:
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_memory_unit_invalidates_stats_cache(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
bank_id = f"test-stats-cache-delunit-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
|
||||
await _insert_memory(conn, bank_id, "Bob enjoys cycling.")
|
||||
|
||||
_pin_cache(memory)
|
||||
try:
|
||||
before = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert before["node_counts"].get("experience") == 2
|
||||
|
||||
await memory.delete_memory_unit(str(m1), request_context=request_context)
|
||||
|
||||
after = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
# Without invalidation the long-TTL cache would still report 2.
|
||||
assert after["node_counts"].get("experience") == 1
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_document_invalidates_stats_cache(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
bank_id = f"test-stats-cache-deldoc-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
document_id = f"doc-{uuid.uuid4().hex[:8]}"
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_document(conn, bank_id, document_id)
|
||||
unit_id = await _insert_memory(conn, bank_id, "Alice works at Acme.")
|
||||
await _attach_unit_to_doc(conn, unit_id, document_id)
|
||||
|
||||
_pin_cache(memory)
|
||||
try:
|
||||
before = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert before["total_documents"] == 1
|
||||
assert before["node_counts"].get("experience") == 1
|
||||
|
||||
await memory.delete_document(document_id, bank_id, request_context=request_context)
|
||||
|
||||
after = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
# Without invalidation the long-TTL cache would still report 1 document.
|
||||
assert after["total_documents"] == 0
|
||||
assert after["node_counts"].get("experience", 0) == 0
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_observations_invalidates_stats_cache(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
bank_id = f"test-stats-cache-clearobs-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
|
||||
await _insert_observation(conn, bank_id, "Alice enjoys hiking regularly.", [m1])
|
||||
|
||||
_pin_cache(memory)
|
||||
try:
|
||||
before = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert before["total_observations"] == 1
|
||||
|
||||
await memory.clear_observations(bank_id, request_context=request_context)
|
||||
|
||||
after = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
# Without invalidation the long-TTL cache would still report 1 observation.
|
||||
assert after["total_observations"] == 0
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -322,3 +322,76 @@ def test_plain_text_lines_not_treated_as_jsonl():
|
||||
# Sanity: these are not JSON objects (so the JSONL path correctly declined).
|
||||
with pytest.raises(json.JSONDecodeError):
|
||||
json.loads(chunks[0])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Idempotency — re-chunking a produced chunk must be a no-op (issue #2301)
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# The streaming retain pipeline pre-chunks each document once (producer) and then
|
||||
# re-chunks every piece during extraction (consumer), stamping all sub-chunks of
|
||||
# one piece with that piece's single chunk_index. If a piece re-split, its
|
||||
# sub-chunks would derive the same chunk_id = {bank}_{doc}_{index} and the
|
||||
# ON CONFLICT upsert would fail with CardinalityViolationError. This can only
|
||||
# happen when structured_chunk_size > max_chars (a chunk legitimately exceeds the
|
||||
# re-chunk budget); the defaults (structured == max_chars) never trip it.
|
||||
|
||||
|
||||
def _assert_idempotent(text: str, *, max_chars: int, structured_chunk_size: int) -> list[str]:
|
||||
chunks = chunk_text(text, max_chars=max_chars, structured_chunk_size=structured_chunk_size)
|
||||
for chunk in chunks:
|
||||
rechunked = chunk_text(chunk, max_chars=max_chars, structured_chunk_size=structured_chunk_size)
|
||||
assert rechunked == [chunk], (
|
||||
f"re-chunking a produced chunk split it again ({len(chunk)} chars -> "
|
||||
f"{len(rechunked)} pieces) — not idempotent (issue #2301)"
|
||||
)
|
||||
return chunks
|
||||
|
||||
|
||||
def test_conversation_turn_over_chunk_size_is_rechunk_stable():
|
||||
"""A conversation turn larger than max_chars but kept whole by the larger
|
||||
structured cap must survive a re-chunk unchanged (issue #2301)."""
|
||||
content = json.dumps([{"role": "assistant", "content": "x" * 6000}])
|
||||
|
||||
_assert_idempotent(content, max_chars=3000, structured_chunk_size=5000)
|
||||
|
||||
|
||||
def test_jsonl_line_over_chunk_size_is_rechunk_stable():
|
||||
"""A single oversized JSONL line, kept whole within the structured cap, must
|
||||
not be re-split when handed back through chunk_text (issue #2301)."""
|
||||
text = "\n".join([json.dumps({"event": "x" * 3800}), json.dumps({"event": "small"})])
|
||||
|
||||
_assert_idempotent(text, max_chars=3000, structured_chunk_size=4500)
|
||||
|
||||
|
||||
def test_oversized_unit_fragments_stay_within_chunk_budget():
|
||||
"""A unit past even the structured cap is fragmented as text; no fragment may
|
||||
exceed max_chars, so a re-chunk leaves the fragments intact (issue #2301)."""
|
||||
text = "\n".join([json.dumps({"event": "z" * 9000}), json.dumps({"e": "s"})])
|
||||
|
||||
chunks = _assert_idempotent(text, max_chars=3000, structured_chunk_size=4500)
|
||||
|
||||
assert all(len(c) <= 3000 for c in chunks)
|
||||
|
||||
|
||||
def test_single_json_object_kept_whole_within_structured_cap():
|
||||
"""A lone JSON object over max_chars but within the structured cap is returned
|
||||
whole rather than plain-text-split (the basis of re-chunk stability)."""
|
||||
obj = json.dumps({"role": "assistant", "content": "x" * 4000})
|
||||
|
||||
assert chunk_text(obj, max_chars=3000, structured_chunk_size=5000) == [obj]
|
||||
|
||||
|
||||
def test_rechunk_preserves_one_chunk_id_per_pre_chunk():
|
||||
"""End-to-end of the producer/consumer chunk_id derivation: each pre-chunk
|
||||
(one global index) must re-chunk to exactly one piece, so the derived
|
||||
chunk_ids stay unique within an upsert batch (issue #2301)."""
|
||||
content = json.dumps([{"role": "assistant", "content": "x" * 6000}])
|
||||
|
||||
pre_chunks = chunk_text(content, max_chars=3000, structured_chunk_size=5000)
|
||||
chunk_ids = []
|
||||
for global_idx, pre in enumerate(pre_chunks):
|
||||
for _ in chunk_text(pre, max_chars=3000, structured_chunk_size=5000):
|
||||
chunk_ids.append(f"bank_doc_{global_idx}")
|
||||
|
||||
assert len(chunk_ids) == len(set(chunk_ids)), f"duplicate chunk_ids in one batch: {chunk_ids}"
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Tests for ``CODEX_HOME`` resolution of the Codex ``auth.json`` location.
|
||||
|
||||
Codex stores its OAuth credentials under a configurable home directory. The
|
||||
canonical ``@openai/codex`` CLI honors the ``CODEX_HOME`` environment variable
|
||||
and falls back to ``~/.codex``. Hindsight's Codex auth/LLM/embeddings paths
|
||||
must resolve the same way so that a user who relocates ``CODEX_HOME`` is still
|
||||
authenticated.
|
||||
"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from hindsight_api.engine.providers.codex_auth import (
|
||||
CodexAuthManager,
|
||||
default_codex_auth_file,
|
||||
)
|
||||
from hindsight_api.engine.providers.codex_llm import CodexLLM
|
||||
|
||||
|
||||
def _write_auth(auth_dir: Path, access_token: str = "at-test") -> Path:
|
||||
"""Write a minimal chatgpt-mode auth.json under ``auth_dir``."""
|
||||
auth_dir.mkdir(parents=True, exist_ok=True)
|
||||
auth_file = auth_dir / "auth.json"
|
||||
auth_file.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"auth_mode": "chatgpt",
|
||||
"tokens": {
|
||||
"access_token": access_token,
|
||||
"refresh_token": "rt-test",
|
||||
"account_id": "acct-test",
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
return auth_file
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# default_codex_auth_file()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_default_auth_file_falls_back_to_home_codex_when_unset(tmp_path, monkeypatch):
|
||||
monkeypatch.delenv("CODEX_HOME", raising=False)
|
||||
monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path))
|
||||
|
||||
assert default_codex_auth_file() == tmp_path / ".codex" / "auth.json"
|
||||
|
||||
|
||||
def test_default_auth_file_honors_codex_home_when_set(tmp_path, monkeypatch):
|
||||
codex_home = tmp_path / "custom-codex"
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
|
||||
assert default_codex_auth_file() == codex_home / "auth.json"
|
||||
|
||||
|
||||
def test_default_auth_file_empty_codex_home_falls_back(tmp_path, monkeypatch):
|
||||
"""An empty ``CODEX_HOME`` is treated as unset (matches shell semantics)."""
|
||||
monkeypatch.setenv("CODEX_HOME", "")
|
||||
monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path))
|
||||
|
||||
assert default_codex_auth_file() == tmp_path / ".codex" / "auth.json"
|
||||
|
||||
|
||||
def test_default_auth_file_resolved_lazily(tmp_path, monkeypatch):
|
||||
"""The env var is read on each call, not cached at import time."""
|
||||
monkeypatch.setenv("CODEX_HOME", str(tmp_path / "a"))
|
||||
assert default_codex_auth_file() == tmp_path / "a" / "auth.json"
|
||||
|
||||
monkeypatch.setenv("CODEX_HOME", str(tmp_path / "b"))
|
||||
assert default_codex_auth_file() == tmp_path / "b" / "auth.json"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CodexAuthManager.from_file() — honors CODEX_HOME by default
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_auth_manager_from_file_uses_codex_home(tmp_path, monkeypatch):
|
||||
codex_home = tmp_path / "custom-codex"
|
||||
_write_auth(codex_home, access_token="at-from-codex-home")
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
|
||||
mgr = CodexAuthManager.from_file()
|
||||
|
||||
assert mgr.access_token == "at-from-codex-home"
|
||||
assert mgr._auth_file == codex_home / "auth.json"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CodexLLM — loads credentials from CODEX_HOME
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_codex_llm_loads_from_codex_home(tmp_path, monkeypatch):
|
||||
codex_home = tmp_path / "custom-codex"
|
||||
_write_auth(codex_home, access_token="at-llm")
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
|
||||
llm = CodexLLM(
|
||||
provider="codex",
|
||||
api_key="ignored",
|
||||
base_url="",
|
||||
model="gpt-5-codex",
|
||||
)
|
||||
|
||||
assert llm.access_token == "at-llm"
|
||||
assert llm._auth_file == codex_home / "auth.json"
|
||||
@@ -8,9 +8,12 @@ relevance score, independent of the cross-encoder model's score calibration.
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.search.reranking import apply_combined_scoring, _RECENCY_ALPHA, _TEMPORAL_ALPHA
|
||||
from hindsight_api.engine.search.reranking import (
|
||||
_RECENCY_ALPHA,
|
||||
_TEMPORAL_ALPHA,
|
||||
apply_combined_scoring,
|
||||
compute_recency_decay,
|
||||
)
|
||||
from hindsight_api.engine.search.types import MergedCandidate, RetrievalResult, ScoredResult
|
||||
|
||||
UTC = timezone.utc
|
||||
@@ -200,3 +203,52 @@ class TestBoostFormula:
|
||||
|
||||
def test_empty_list_is_noop(self):
|
||||
apply_combined_scoring([], now=NOW) # must not raise
|
||||
|
||||
|
||||
class TestRecencyDecayFunction:
|
||||
"""The configurable age→freshness curve (compute_recency_decay)."""
|
||||
|
||||
def test_linear_is_default_and_unchanged(self):
|
||||
"""Default function reproduces the historical linear decay over 365 days."""
|
||||
assert compute_recency_decay(0) == 1.0
|
||||
assert abs(compute_recency_decay(182.5) - 0.5) < 1e-6 # neutral at half the window
|
||||
assert compute_recency_decay(400) == 0.1 # floored past the window
|
||||
|
||||
def test_linear_window_is_configurable(self):
|
||||
"""A custom window moves the neutral crossing; 730d window → neutral at 365d."""
|
||||
assert abs(compute_recency_decay(365, "linear", linear_window_days=730) - 0.5) < 1e-6
|
||||
|
||||
def test_exponential_neutral_at_halflife(self):
|
||||
"""Exponential decay is exactly neutral (0.5) at the configured half-life."""
|
||||
assert compute_recency_decay(0, "exponential", halflife_days=90) == 1.0
|
||||
assert abs(compute_recency_decay(90, "exponential", halflife_days=90) - 0.5) < 1e-9
|
||||
assert abs(compute_recency_decay(180, "exponential", halflife_days=90) - 0.25) < 1e-9
|
||||
|
||||
def test_exponential_penalises_old_less_harshly_than_linear(self):
|
||||
"""A 1-year-old memory keeps more freshness under a 90d-halflife exponential
|
||||
than under the linear floor — the curve never hard-cuts to 0.1."""
|
||||
lin = compute_recency_decay(365, "linear")
|
||||
exp = compute_recency_decay(365, "exponential", halflife_days=180)
|
||||
assert exp > lin
|
||||
|
||||
def test_none_is_always_neutral(self):
|
||||
"""'none' disables the recency signal — always neutral, no boost."""
|
||||
assert compute_recency_decay(0, "none") == 0.5
|
||||
assert compute_recency_decay(10_000, "none") == 0.5
|
||||
|
||||
def test_future_dates_clamp_to_max(self):
|
||||
"""Negative ages (future-dated memories) never exceed full freshness."""
|
||||
assert compute_recency_decay(-100, "linear") == 1.0
|
||||
assert compute_recency_decay(-100, "exponential", halflife_days=90) == 1.0
|
||||
|
||||
def test_nonpositive_halflife_falls_back_to_neutral(self):
|
||||
"""A misconfigured (<=0) half-life degrades to neutral rather than dividing by zero."""
|
||||
assert compute_recency_decay(30, "exponential", halflife_days=0) == 0.5
|
||||
|
||||
def test_function_threads_through_apply_combined_scoring(self):
|
||||
"""The decay function chosen at the call site is what scores sr.recency."""
|
||||
old = NOW - timedelta(days=180)
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=old)
|
||||
apply_combined_scoring([sr], now=NOW, recency_decay_function="none")
|
||||
assert sr.recency == 0.5
|
||||
assert abs(sr.weight - 0.5) < 1e-9 # neutral → no recency boost
|
||||
|
||||
@@ -25,6 +25,7 @@ def setup_test_env():
|
||||
"HINDSIGHT_API_LLM_MODEL",
|
||||
"HINDSIGHT_API_LLM_REASONING_EFFORT",
|
||||
"HINDSIGHT_API_LLM_BEDROCK_SERVICE_TIER",
|
||||
"HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER",
|
||||
"HINDSIGHT_API_SEMANTIC_MIN_SIMILARITY",
|
||||
"HINDSIGHT_API_DATABASE_URL",
|
||||
"HINDSIGHT_API_MIGRATION_DATABASE_URL",
|
||||
@@ -452,6 +453,53 @@ def test_llm_output_language_empty_string_is_unset(monkeypatch):
|
||||
assert config.llm_output_language is None
|
||||
|
||||
|
||||
def test_markitdown_ocr_defaults_disabled(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.file_parser_markitdown_ocr_enabled is False
|
||||
|
||||
|
||||
def test_markitdown_ocr_does_not_fall_back_to_main_llm_config(monkeypatch):
|
||||
from hindsight_api.config import DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT, HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED", "true")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "anthropic")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "main-key")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_BASE_URL", "https://main.example/v1")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_MODEL", "main-vision-model")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.file_parser_markitdown_ocr_enabled is True
|
||||
assert config.file_parser_markitdown_ocr_api_key is None
|
||||
assert config.file_parser_markitdown_ocr_base_url is None
|
||||
assert config.file_parser_markitdown_ocr_model is None
|
||||
assert config.file_parser_markitdown_ocr_prompt == DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
|
||||
|
||||
|
||||
def test_markitdown_ocr_uses_explicit_config(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED", "true")
|
||||
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY", "parser-key")
|
||||
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL", "https://parser.example/v1")
|
||||
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL", "parser-vision-model")
|
||||
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT", "Extract this document exactly.")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "main-key")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_BASE_URL", "https://main.example/v1")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_MODEL", "main-vision-model")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.file_parser_markitdown_ocr_enabled is True
|
||||
assert config.file_parser_markitdown_ocr_api_key == "parser-key"
|
||||
assert config.file_parser_markitdown_ocr_base_url == "https://parser.example/v1"
|
||||
assert config.file_parser_markitdown_ocr_model == "parser-vision-model"
|
||||
assert config.file_parser_markitdown_ocr_prompt == "Extract this document exactly."
|
||||
|
||||
|
||||
def test_llm_reasoning_effort_defaults_to_low(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
@@ -580,3 +628,81 @@ def test_bedrock_service_tier_rejects_invalid_value(monkeypatch):
|
||||
assert "HINDSIGHT_API_LLM_BEDROCK_SERVICE_TIER" in error_message
|
||||
assert "standard" in error_message
|
||||
assert "'standard' is not a valid Bedrock service tier" in error_message
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Gemini service tier (HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_gemini_service_tier_defaults_to_none(monkeypatch):
|
||||
"""Gemini service tier defaults to None (standard tier) when unset."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.delenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", raising=False)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_service_tier is None
|
||||
|
||||
|
||||
def test_gemini_service_tier_flex(monkeypatch):
|
||||
"""Flex tier is accepted for Gemini."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "flex")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "gemini")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "fake-key")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_service_tier == "flex"
|
||||
|
||||
|
||||
def test_gemini_service_tier_accepts_mixed_case_provider(monkeypatch):
|
||||
"""Gemini tier parsing follows provider's case-insensitive handling."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "flex")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "Gemini")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "fake-key")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_service_tier == "flex"
|
||||
|
||||
|
||||
def test_gemini_service_tier_rejects_invalid_value(monkeypatch):
|
||||
"""Unknown Gemini service tiers are rejected early."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "standard")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "gemini")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "fake-key")
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
HindsightConfig.from_env()
|
||||
|
||||
error_message = str(exc_info.value)
|
||||
assert "HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER" in error_message
|
||||
assert "standard" in error_message
|
||||
|
||||
|
||||
def test_gemini_service_tier_ignored_for_non_gemini_provider(monkeypatch):
|
||||
"""Invalid Gemini-only tiers do not break unrelated providers."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "standard")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_service_tier is None
|
||||
|
||||
|
||||
def test_gemini_service_tier_empty_env_is_unset(monkeypatch):
|
||||
"""Empty env values are treated as unset for templated deployments."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_service_tier is None
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
"""Tests for write-side validation of bank disposition config overrides.
|
||||
|
||||
Disposition traits (skepticism / literalism / empathy) are integers on a 1-5
|
||||
scale. The ``PATCH /v1/{tenant}/banks/{id}/config`` write path must reject
|
||||
out-of-contract values (floats, 0-1 scales, ints outside 1-5) at write time;
|
||||
otherwise a single malformed bank 500s the entire bank list because the read
|
||||
overlay injects the stored value verbatim into a strict
|
||||
``DispositionTraits(int, ge=1, le=5)``. See issue #2348.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config_resolver import _validate_disposition_updates
|
||||
|
||||
|
||||
_DISPOSITION_FIELD_NAMES = (
|
||||
"disposition_skepticism",
|
||||
"disposition_literalism",
|
||||
"disposition_empathy",
|
||||
)
|
||||
|
||||
|
||||
class TestValidateDispositionUpdates:
|
||||
def test_no_op_passes(self):
|
||||
_validate_disposition_updates({})
|
||||
_validate_disposition_updates({"unrelated_field": 123})
|
||||
|
||||
def test_valid_in_range_integers_pass(self):
|
||||
for key in _DISPOSITION_FIELD_NAMES:
|
||||
for value in (1, 2, 3, 4, 5):
|
||||
_validate_disposition_updates({key: value})
|
||||
|
||||
def test_none_clears_override(self):
|
||||
# None is the "unset this per-bank override" sentinel (field is int | None).
|
||||
for key in _DISPOSITION_FIELD_NAMES:
|
||||
_validate_disposition_updates({key: None})
|
||||
|
||||
def test_out_of_range_integer_raises(self):
|
||||
for key in _DISPOSITION_FIELD_NAMES:
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: 0})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: 6})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: -1})
|
||||
|
||||
def test_float_raises(self):
|
||||
# The reported v0.8.3 case: a 0-1 scale used by mistake.
|
||||
for key in _DISPOSITION_FIELD_NAMES:
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: 0.7})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: 3.0}) # float, even if in 1-5 range
|
||||
|
||||
def test_bool_raises(self):
|
||||
# bool is an int subclass and would sneak past a naive isinstance(int) check.
|
||||
for key in _DISPOSITION_FIELD_NAMES:
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: True})
|
||||
|
||||
def test_string_raises(self):
|
||||
for key in _DISPOSITION_FIELD_NAMES:
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_disposition_updates({key: "3"})
|
||||
@@ -124,6 +124,10 @@ def test_openai_codex_provider_uses_codex_oauth_token_and_configured_batch_size(
|
||||
)
|
||||
|
||||
monkeypatch.setenv("HOME", str(tmp_path))
|
||||
# Codex auth resolves via CODEX_HOME first (falling back to ~/.codex), so a
|
||||
# CODEX_HOME leaking in from the runner's environment would point auth.json
|
||||
# away from the tmp_path fixture. Pin resolution to the patched HOME.
|
||||
monkeypatch.delenv("CODEX_HOME", raising=False)
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_PROVIDER"] = "openai-codex"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL"] = "text-embedding-3-small"
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
"""Regression test: `enqueue_graph_maintenance` must insert unit_ids in a
|
||||
deterministic sorted order so concurrent transactions can't deadlock on the
|
||||
graph_maintenance_queue unique-key check.
|
||||
|
||||
Symptom (production): under load, concurrent `PATCH /memories/{id}` requests
|
||||
on the same bank generate overlapping `victim_ids` sets (the surviving units
|
||||
whose outgoing links pointed at the updated unit). Each transaction inserts
|
||||
those victims into `graph_maintenance_queue` with
|
||||
`ON CONFLICT (bank_id, unit_id) DO NOTHING`. The conflict check takes a
|
||||
short-lived row-level lock per (bank_id, unit_id) being inserted, and when
|
||||
two transactions insert overlapping sets in different orders Postgres
|
||||
detects a deadlock and aborts one of them — surfacing as
|
||||
`asyncpg.exceptions.DeadlockDetectedError` from the API, which becomes a 500.
|
||||
|
||||
The fix sorts the input list inside both `ops_postgresql` and `ops_oracle`
|
||||
before passing it to the INSERT, so every transaction acquires the per-row
|
||||
locks in the same global (sorted-UUID) order. With a total order over the
|
||||
lock set, deadlock is mathematically impossible — Postgres still serializes
|
||||
the conflicting inserts but they queue cleanly instead of cycling.
|
||||
|
||||
This test pins that post-condition by capturing the array passed to the
|
||||
underlying `conn.execute` (PG path) / `conn.executemany` (Oracle path) and
|
||||
asserting it's sorted.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.db.ops_oracle import OracleOps
|
||||
from hindsight_api.engine.db.ops_postgresql import PostgreSQLOps
|
||||
|
||||
|
||||
def _shuffled_uuids(n: int) -> list[uuid.UUID]:
|
||||
"""Generate n UUIDs in a deliberately non-monotonic order. Hex literals
|
||||
avoid `uuid.uuid4()` because uuid4 is random and we want determinism."""
|
||||
raw = [
|
||||
"ffffffff-ffff-4fff-8fff-ffffffffffff",
|
||||
"00000000-0000-4000-8000-000000000001",
|
||||
"88888888-8888-4888-8888-888888888888",
|
||||
"11111111-1111-4111-8111-111111111111",
|
||||
"ccccccccc-cccc-4ccc-8ccc-cccccccccccc"[:36],
|
||||
"44444444-4444-4444-8444-444444444444",
|
||||
]
|
||||
return [uuid.UUID(s) for s in raw[:n]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pg_enqueue_graph_maintenance_inserts_in_sorted_order():
|
||||
"""The PostgreSQL ops impl must pass the unit_ids to the INSERT in
|
||||
sorted order, regardless of how the caller ordered them."""
|
||||
ops = PostgreSQLOps()
|
||||
conn = AsyncMock()
|
||||
|
||||
unit_ids = _shuffled_uuids(6)
|
||||
assert unit_ids != sorted(unit_ids), "test inputs must be unsorted"
|
||||
|
||||
await ops.enqueue_graph_maintenance(
|
||||
conn=conn,
|
||||
table="graph_maintenance_queue",
|
||||
bank_id="test-bank",
|
||||
unit_ids=unit_ids,
|
||||
)
|
||||
|
||||
assert conn.execute.await_count == 1
|
||||
_sql, bank_id_arg, ids_arg = conn.execute.await_args.args
|
||||
assert bank_id_arg == "test-bank"
|
||||
assert ids_arg == sorted(unit_ids), f"expected sorted unit_ids for deadlock-free concurrent inserts, got {ids_arg}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oracle_enqueue_graph_maintenance_inserts_in_sorted_order():
|
||||
"""The Oracle ops impl applies the same sort. `executemany` receives a
|
||||
list of (bank_id, unit_id) tuples; the unit_id projection must be
|
||||
sorted."""
|
||||
ops = OracleOps()
|
||||
conn = AsyncMock()
|
||||
|
||||
unit_ids = _shuffled_uuids(6)
|
||||
assert unit_ids != sorted(unit_ids), "test inputs must be unsorted"
|
||||
|
||||
await ops.enqueue_graph_maintenance(
|
||||
conn=conn,
|
||||
table="graph_maintenance_queue",
|
||||
bank_id="test-bank",
|
||||
unit_ids=unit_ids,
|
||||
)
|
||||
|
||||
assert conn.executemany.await_count == 1
|
||||
_sql, rows = conn.executemany.await_args.args
|
||||
assert [r[0] for r in rows] == ["test-bank"] * len(unit_ids)
|
||||
assert [r[1] for r in rows] == sorted(unit_ids), (
|
||||
f"expected sorted unit_ids for deadlock-free concurrent inserts, got {[r[1] for r in rows]}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pg_empty_unit_ids_short_circuits():
|
||||
"""Empty input must remain a no-op — the early return predates this fix
|
||||
and must continue to skip the INSERT entirely."""
|
||||
ops = PostgreSQLOps()
|
||||
conn = AsyncMock()
|
||||
await ops.enqueue_graph_maintenance(conn, "graph_maintenance_queue", "b", [])
|
||||
conn.execute.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oracle_empty_unit_ids_short_circuits():
|
||||
ops = OracleOps()
|
||||
conn = AsyncMock()
|
||||
await ops.enqueue_graph_maintenance(conn, "graph_maintenance_queue", "b", [])
|
||||
conn.executemany.assert_not_awaited()
|
||||
@@ -0,0 +1,70 @@
|
||||
"""ensure_vector_extension must not create the (unused) global memory_units index.
|
||||
|
||||
For per-bank backends (pgvector / pgvectorscale / vchord) every vector search is
|
||||
bank + fact_type scoped and served by the per-(bank, fact_type) partial indexes
|
||||
created at bank-creation time. The global `idx_memory_units_embedding` is never
|
||||
chosen by the planner (migration d5e6f7a8b9c0 drops it for exactly this reason),
|
||||
so the post-migration reconcile must not recreate it on a fresh schema.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import create_engine, text
|
||||
|
||||
from hindsight_api._vector_index import uses_per_bank_vector_indexes
|
||||
from hindsight_api.config import HindsightConfig
|
||||
from hindsight_api.migrations import ensure_vector_extension, run_migrations
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def vec_db_url():
|
||||
"""A dedicated pg0 instance so the test owns its schema/index state."""
|
||||
from hindsight_api.pg0 import EmbeddedPostgres
|
||||
|
||||
pg0 = EmbeddedPostgres(name="hindsight-vecidx-test", port=5570)
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
return loop.run_until_complete(pg0.ensure_running())
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
|
||||
def test_per_bank_backend_does_not_create_global_memory_units_index(vec_db_url):
|
||||
config = HindsightConfig.from_env()
|
||||
vec = config.vector_extension
|
||||
if not uses_per_bank_vector_indexes(vec):
|
||||
pytest.skip(f"backend {vec!r} uses a global vector index by design (no per-bank indexes)")
|
||||
|
||||
schema = "vecidx_fresh"
|
||||
|
||||
engine = create_engine(vec_db_url)
|
||||
try:
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
conn.commit()
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
run_migrations(vec_db_url, schema=schema)
|
||||
# Fresh, empty schema (no banks yet) → the reconcile must be a no-op for the
|
||||
# global index, not recreate it.
|
||||
ensure_vector_extension(vec_db_url, vector_extension=vec, schema=schema)
|
||||
|
||||
engine = create_engine(vec_db_url)
|
||||
try:
|
||||
with engine.connect() as conn:
|
||||
global_index_count = conn.execute(
|
||||
text(
|
||||
"SELECT COUNT(*) FROM pg_indexes "
|
||||
"WHERE schemaname = :schema AND tablename = 'memory_units' "
|
||||
"AND indexname = 'idx_memory_units_embedding'"
|
||||
),
|
||||
{"schema": schema},
|
||||
).scalar()
|
||||
conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
conn.commit()
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
assert global_index_count == 0
|
||||
@@ -911,7 +911,7 @@ class TestPrecheckHttpWiring:
|
||||
def _build_app(validator):
|
||||
"""Mirror the precheck wiring from ``hindsight_api.api.http`` in a
|
||||
standalone FastAPI app."""
|
||||
from fastapi import Depends, FastAPI, HTTPException
|
||||
from fastapi import Depends, FastAPI, HTTPException, Request
|
||||
from pydantic import BaseModel, model_validator
|
||||
|
||||
from hindsight_api.extensions import PrecheckContext
|
||||
@@ -952,12 +952,23 @@ class TestPrecheckHttpWiring:
|
||||
def _precheck_for(operation: str):
|
||||
async def _dep(
|
||||
bank_id: str,
|
||||
request: Request,
|
||||
request_context: RequestContext = Depends(_request_context),
|
||||
) -> None:
|
||||
cl_header = request.headers.get("content-length")
|
||||
content_length: int | None = None
|
||||
if cl_header is not None:
|
||||
try:
|
||||
parsed = int(cl_header)
|
||||
except ValueError:
|
||||
parsed = -1
|
||||
if parsed >= 0:
|
||||
content_length = parsed
|
||||
ctx = PrecheckContext(
|
||||
operation=operation,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
content_length=content_length,
|
||||
)
|
||||
result = await validator.precheck(ctx)
|
||||
if not result.allowed:
|
||||
@@ -1081,3 +1092,86 @@ class TestPrecheckHttpWiring:
|
||||
resp = client.get("/v1/default/banks/precheck-bank/memories/list")
|
||||
assert resp.status_code == 200
|
||||
assert len(validator.precheck_calls) == 0
|
||||
|
||||
def test_precheck_context_carries_content_length(self):
|
||||
"""Content-Length header is exposed to the precheck so a validator
|
||||
can make size-aware decisions (e.g. upper-bound cost estimate)
|
||||
before the body is deserialised."""
|
||||
validator = RecordingPrecheckValidator(reject=False)
|
||||
app, _ = self._build_app(validator)
|
||||
client = TestClient(app)
|
||||
|
||||
# Body must contain at least 500 'x' bytes; check the surfaced
|
||||
# Content-Length is within a tight band around that floor (allows
|
||||
# for JSON envelope + httpx's serialisation choices without
|
||||
# depending on exact byte counts).
|
||||
payload = {"items": [{"content": "x" * 500}]}
|
||||
resp = client.post(
|
||||
"/v1/default/banks/precheck-bank/memories",
|
||||
json=payload,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert len(validator.precheck_calls) == 1
|
||||
ctx = validator.precheck_calls[0]
|
||||
assert ctx.content_length is not None
|
||||
assert 500 <= ctx.content_length <= 600
|
||||
|
||||
def test_precheck_context_content_length_zero_is_not_none(self):
|
||||
"""An empty POST body has Content-Length: 0. That should surface
|
||||
as the int 0, not None — None means 'unknown', 0 means 'known to
|
||||
be empty'."""
|
||||
validator = RecordingPrecheckValidator(reject=False)
|
||||
app, _ = self._build_app(validator)
|
||||
client = TestClient(app)
|
||||
|
||||
# Empty body fails Pydantic parse (422), but precheck runs first
|
||||
# and records the Content-Length.
|
||||
client.post(
|
||||
"/v1/default/banks/precheck-bank/memories",
|
||||
content=b"",
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
assert len(validator.precheck_calls) >= 1
|
||||
ctx = validator.precheck_calls[-1]
|
||||
assert ctx.content_length == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_precheck_context_content_length_none_when_header_missing(self):
|
||||
"""When the Content-Length header isn't set (e.g. chunked transfer
|
||||
encoding) the validator sees None, not a crash and not a default 0."""
|
||||
from starlette.requests import Request as _StarletteRequest
|
||||
|
||||
from hindsight_api.extensions import PrecheckContext
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
validator = RecordingPrecheckValidator(reject=False)
|
||||
|
||||
# Replicate the wiring's parse step inline so the test exercises
|
||||
# the same code-path semantics introduced in
|
||||
# ``hindsight_api.api.http._precheck_dep``.
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/v1/default/banks/bank-x/memories",
|
||||
"headers": [], # no content-length
|
||||
"query_string": b"",
|
||||
}
|
||||
req = _StarletteRequest(scope)
|
||||
cl_header = req.headers.get("content-length")
|
||||
content_length: int | None = None
|
||||
if cl_header is not None:
|
||||
try:
|
||||
parsed = int(cl_header)
|
||||
except ValueError:
|
||||
parsed = -1
|
||||
if parsed >= 0:
|
||||
content_length = parsed
|
||||
|
||||
ctx = PrecheckContext(
|
||||
operation="retain",
|
||||
bank_id="bank-x",
|
||||
request_context=RequestContext(),
|
||||
content_length=content_length,
|
||||
)
|
||||
await validator.precheck(ctx)
|
||||
assert validator.precheck_calls[-1].content_length is None
|
||||
|
||||
@@ -77,6 +77,27 @@ async def test_dry_run_extracts_without_persisting(api_client, memory):
|
||||
assert after["total"] == before["total"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dry_run_rejects_empty_content(api_client, memory):
|
||||
"""Empty/whitespace-only content is rejected by request validation (422) before the
|
||||
billable LLM extraction call runs — matching retain (RetainItem.content) and recall
|
||||
(RecallRequest.query), which already reject empty input."""
|
||||
bank_id = f"dryrun-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=RequestContext())
|
||||
|
||||
before = await memory.list_memory_units(bank_id=bank_id, request_context=RequestContext())
|
||||
for content in ("", " ", "\n\t "):
|
||||
resp = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/memories/dry-run-extract",
|
||||
json={"content": content},
|
||||
)
|
||||
assert resp.status_code == 422, resp.text
|
||||
|
||||
# Rejected before extraction: nothing was persisted.
|
||||
after = await memory.list_memory_units(bank_id=bank_id, request_context=RequestContext())
|
||||
assert after["total"] == before["total"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dry_run_disabled_returns_404(api_client, memory):
|
||||
"""With HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT=false the endpoint is removed (returns 404)."""
|
||||
|
||||
@@ -346,6 +346,149 @@ async def test_markitdown_converter():
|
||||
assert "test document" in result.lower() or "multiple lines" in result.lower()
|
||||
|
||||
|
||||
def test_markitdown_converter_does_not_enable_ocr_by_default(monkeypatch):
|
||||
"""Markitdown should keep its local/default behavior unless OCR is explicitly enabled."""
|
||||
import markitdown
|
||||
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
calls = []
|
||||
|
||||
class FakeMarkItDown:
|
||||
def __init__(self, **kwargs):
|
||||
calls.append(kwargs)
|
||||
|
||||
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
|
||||
|
||||
MarkitdownParser()
|
||||
|
||||
assert calls == [{}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_markitdown_image_without_ocr_has_actionable_error(monkeypatch):
|
||||
"""Image uploads should explain that MarkItDown OCR is disabled instead of surfacing a low-level error."""
|
||||
import markitdown
|
||||
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
class FakeMarkItDown:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def convert(self, path):
|
||||
raise AssertionError("MarkItDown should not be called when image OCR is disabled")
|
||||
|
||||
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
|
||||
|
||||
parser = MarkitdownParser()
|
||||
with pytest.raises(RuntimeError, match="Image OCR is not enabled for the markitdown parser"):
|
||||
await parser.convert(b"\x89PNG\r\n\x1a\n", "screenshot.png")
|
||||
|
||||
|
||||
def test_markitdown_converter_can_enable_ocr(monkeypatch):
|
||||
"""When enabled, Markitdown receives an OpenAI-compatible client, model, and OCR prompt."""
|
||||
import markitdown
|
||||
import openai
|
||||
|
||||
from hindsight_api.config import DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
markitdown_calls = []
|
||||
openai_calls = []
|
||||
|
||||
class FakeMarkItDown:
|
||||
def __init__(self, **kwargs):
|
||||
markitdown_calls.append(kwargs)
|
||||
|
||||
class FakeOpenAI:
|
||||
def __init__(self, **kwargs):
|
||||
openai_calls.append(kwargs)
|
||||
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
|
||||
monkeypatch.setattr(openai, "OpenAI", FakeOpenAI)
|
||||
|
||||
MarkitdownParser(
|
||||
ocr_enabled=True,
|
||||
ocr_api_key="parser-key",
|
||||
ocr_base_url="https://vision.example/v1",
|
||||
ocr_model="vision-model",
|
||||
)
|
||||
|
||||
assert openai_calls == [
|
||||
{
|
||||
"api_key": "parser-key",
|
||||
"base_url": "https://vision.example/v1",
|
||||
}
|
||||
]
|
||||
assert markitdown_calls[0]["llm_client"].__class__ is FakeOpenAI
|
||||
assert markitdown_calls[0]["llm_model"] == "vision-model"
|
||||
assert markitdown_calls[0]["llm_prompt"] == DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
|
||||
|
||||
|
||||
def test_markitdown_converter_requires_model_when_ocr_enabled(monkeypatch):
|
||||
"""OCR should fail fast when enabled without a model."""
|
||||
import markitdown
|
||||
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
class FakeMarkItDown:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
|
||||
|
||||
with pytest.raises(ValueError, match="no model"):
|
||||
MarkitdownParser(ocr_enabled=True, ocr_api_key="parser-key")
|
||||
|
||||
|
||||
def test_markitdown_converter_requires_base_url_when_ocr_enabled(monkeypatch):
|
||||
"""OCR should fail fast when enabled without a dedicated OpenAI-compatible endpoint."""
|
||||
import markitdown
|
||||
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
class FakeMarkItDown:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
|
||||
|
||||
with pytest.raises(ValueError, match="no base URL"):
|
||||
MarkitdownParser(ocr_enabled=True, ocr_api_key="parser-key", ocr_model="vision-model")
|
||||
|
||||
|
||||
def test_markitdown_converter_reports_missing_openai_when_ocr_enabled(monkeypatch):
|
||||
"""Missing OpenAI SDK should not be reported as missing MarkItDown."""
|
||||
import builtins
|
||||
import markitdown
|
||||
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
real_import = builtins.__import__
|
||||
|
||||
class FakeMarkItDown:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def fake_import(name, globals=None, locals=None, fromlist=(), level=0):
|
||||
if name == "openai":
|
||||
raise ImportError("no openai")
|
||||
return real_import(name, globals, locals, fromlist, level)
|
||||
|
||||
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
|
||||
monkeypatch.setattr(builtins, "__import__", fake_import)
|
||||
|
||||
with pytest.raises(RuntimeError, match="openai package is required"):
|
||||
MarkitdownParser(
|
||||
ocr_enabled=True,
|
||||
ocr_api_key="parser-key",
|
||||
ocr_base_url="https://vision.example/v1",
|
||||
ocr_model="vision-model",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_converter_registry():
|
||||
"""Test file parser registry."""
|
||||
|
||||
@@ -12,6 +12,8 @@ import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
@@ -21,6 +23,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
from testcontainers.core.container import DockerContainer
|
||||
from testcontainers.core.docker_client import DockerClient as _DockerClient
|
||||
|
||||
_has_testcontainers = True
|
||||
except ImportError:
|
||||
@@ -38,6 +41,8 @@ SEAWEEDFS_S3_PORT = 8333
|
||||
TEST_BUCKET = "hindsight-test"
|
||||
ACCESS_KEY = "test_access_key"
|
||||
SECRET_KEY = "test_secret_key"
|
||||
_PORT_MAPPING_RETRY_TIMEOUT_SECONDS = 10.0
|
||||
_PORT_MAPPING_RETRY_INTERVAL_SECONDS = 0.1
|
||||
|
||||
# SeaweedFS S3 IAM config granting full access to our test credentials
|
||||
_S3_CONFIG = {
|
||||
@@ -64,6 +69,33 @@ def _docker_available() -> bool:
|
||||
return False
|
||||
|
||||
|
||||
if _has_testcontainers:
|
||||
|
||||
@contextmanager
|
||||
def _retry_testcontainers_port_mapping() -> Iterator[None]:
|
||||
original_port = _DockerClient.port
|
||||
|
||||
def port_with_retry(self: _DockerClient, container_id: str, port: int) -> str:
|
||||
deadline = time.monotonic() + _PORT_MAPPING_RETRY_TIMEOUT_SECONDS
|
||||
while True:
|
||||
try:
|
||||
return original_port(self, container_id, port)
|
||||
except ConnectionError:
|
||||
# Docker Desktop can report a container as running before its
|
||||
# published port appears in NetworkSettings.Ports. This affects
|
||||
# both Ryuk's 8080 lookup inside testcontainers and the
|
||||
# SeaweedFS S3 port lookup below.
|
||||
if time.monotonic() >= deadline:
|
||||
raise
|
||||
time.sleep(_PORT_MAPPING_RETRY_INTERVAL_SECONDS)
|
||||
|
||||
_DockerClient.port = port_with_retry
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_DockerClient.port = original_port
|
||||
|
||||
|
||||
def _wait_for_seaweedfs(endpoint: str, timeout: int = 30) -> None:
|
||||
"""Poll SeaweedFS S3 endpoint until ready."""
|
||||
deadline = time.time() + timeout
|
||||
@@ -101,11 +133,11 @@ def seaweedfs_container():
|
||||
.with_command(f"server -s3 -s3.port={SEAWEEDFS_S3_PORT} -s3.config=/etc/seaweedfs/s3.json -ip.bind=0.0.0.0")
|
||||
)
|
||||
|
||||
container.start()
|
||||
|
||||
try:
|
||||
host = container.get_container_host_ip()
|
||||
port = container.get_exposed_port(SEAWEEDFS_S3_PORT)
|
||||
with _retry_testcontainers_port_mapping():
|
||||
container.start()
|
||||
host = container.get_container_host_ip()
|
||||
port = container.get_exposed_port(SEAWEEDFS_S3_PORT)
|
||||
endpoint = f"http://{host}:{port}"
|
||||
|
||||
_wait_for_seaweedfs(endpoint, timeout=240)
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
"""Plumbing tests for the Gemini service tier flag."""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.llm_wrapper import LLMConfig
|
||||
|
||||
|
||||
def test_llm_config_threads_gemini_service_tier_to_provider_impl():
|
||||
"""End-to-end: LLMConfig -> create_llm_provider -> GeminiLLM carries the tier."""
|
||||
pytest.importorskip("google.genai")
|
||||
with patch("google.genai.Client", return_value=MagicMock()):
|
||||
llm = LLMConfig(
|
||||
provider="gemini",
|
||||
api_key="fake-key",
|
||||
base_url="",
|
||||
model="gemini-2.5-flash",
|
||||
gemini_service_tier="flex",
|
||||
)
|
||||
|
||||
assert llm._provider_impl._service_tier == "flex"
|
||||
|
||||
|
||||
def test_llm_provider_from_env_validates_gemini_service_tier(monkeypatch):
|
||||
"""Direct env construction rejects the same invalid tiers as HindsightConfig."""
|
||||
from hindsight_api.config import clear_config_cache
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "gemini")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "fake-key")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "standard")
|
||||
clear_config_cache()
|
||||
|
||||
with pytest.raises(ValueError, match="HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER"):
|
||||
LLMProvider.from_env()
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_llm_provider_from_env_ignores_gemini_tier_for_non_gemini(monkeypatch):
|
||||
"""Invalid Gemini-only tier env values do not break other providers."""
|
||||
from hindsight_api.config import clear_config_cache
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "standard")
|
||||
clear_config_cache()
|
||||
|
||||
provider = LLMProvider.from_env()
|
||||
|
||||
assert provider.gemini_service_tier is None
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_llm_provider_from_env_keeps_lightweight_loader(monkeypatch):
|
||||
"""Reading the Gemini tier must not construct the full application config."""
|
||||
from hindsight_api.config import clear_config_cache
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "gemini")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "fake-key")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "flex")
|
||||
monkeypatch.setenv("HINDSIGHT_API_RETAIN_MAX_COMPLETION_TOKENS", "1000")
|
||||
monkeypatch.setenv("HINDSIGHT_API_RETAIN_CHUNK_SIZE", "2000")
|
||||
clear_config_cache()
|
||||
|
||||
with patch("google.genai.Client", return_value=MagicMock()):
|
||||
provider = LLMProvider.from_env()
|
||||
|
||||
assert provider.gemini_service_tier == "flex"
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_llm_provider_constructor_validates_gemini_service_tier():
|
||||
"""Direct Gemini construction rejects invalid tiers before API calls."""
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
with pytest.raises(ValueError, match="HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER"):
|
||||
LLMProvider(
|
||||
provider="gemini",
|
||||
api_key="fake-key",
|
||||
base_url="",
|
||||
model="gemini-2.5-flash",
|
||||
gemini_service_tier="standard",
|
||||
)
|
||||
|
||||
|
||||
def test_vertexai_ignores_gemini_service_tier():
|
||||
"""The Gemini-only tier flag is not forwarded to Vertex AI providers."""
|
||||
from hindsight_api.engine.llm_wrapper import create_llm_provider
|
||||
|
||||
with patch("hindsight_api.engine.providers.GeminiLLM") as mock_gemini:
|
||||
create_llm_provider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="gemini-2.5-flash",
|
||||
reasoning_effort="low",
|
||||
gemini_service_tier="flex",
|
||||
)
|
||||
|
||||
assert mock_gemini.call_args.kwargs["gemini_service_tier"] is None
|
||||
@@ -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
|
||||
@@ -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,71 @@
|
||||
"""Regression test: list_banks must apply the same disposition + mission
|
||||
config overlay that get_bank_profile applies.
|
||||
|
||||
Bug (reproduced live against 0.8.1): for a bank whose disposition and
|
||||
mission were evolved/overridden via bank *config* (the banks.config JSONB:
|
||||
reflect_mission, disposition_skepticism/literalism/empathy), the single-bank
|
||||
get path returns the real values while the list path returns the stale legacy
|
||||
DB-column defaults ({skepticism:3, literalism:3, empathy:3} and "").
|
||||
|
||||
Root cause: MemoryEngine.get_bank_profile overlays the resolved bank config
|
||||
on top of the legacy banks.disposition/banks.mission columns, but
|
||||
MemoryEngine.list_banks returned bank_utils.list_banks rows straight from
|
||||
those columns with no overlay. The two endpoints disagreed for the same bank.
|
||||
|
||||
This test sets disposition + mission through the config path (so the legacy
|
||||
columns keep their defaults) and asserts list_banks agrees with
|
||||
get_bank_profile for that bank.
|
||||
|
||||
Runs via: uv run pytest tests/test_list_banks_config_overlay.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_banks_overlays_config_disposition_and_mission(memory):
|
||||
bank_id = "list_banks_config_overlay_bank"
|
||||
request_context = RequestContext(api_key=None, api_key_id=None, tenant_id=None, internal=False)
|
||||
|
||||
# Values that differ from the 3/3/3 defaults on every trait, and a
|
||||
# clearly non-empty mission, so a stale-default regression is unmissable.
|
||||
overrides = {
|
||||
"reflect_mission": "I am the shared long-term memory for this regression test.",
|
||||
"disposition_skepticism": 4,
|
||||
"disposition_literalism": 5,
|
||||
"disposition_empathy": 2,
|
||||
}
|
||||
|
||||
try:
|
||||
# Create the bank. Its legacy banks.disposition/banks.mission columns
|
||||
# keep their defaults (3/3/3 and "") — the real values live in config.
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Set disposition + mission via the *config* path (banks.config JSONB),
|
||||
# exactly the path that triggered the live bug.
|
||||
await memory._config_resolver.update_bank_config(bank_id, overrides, request_context)
|
||||
|
||||
# Source of truth: the single-bank get path already overlays config.
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
assert profile["mission"] == overrides["reflect_mission"]
|
||||
assert profile["disposition"] == {"skepticism": 4, "literalism": 5, "empathy": 2}
|
||||
|
||||
# The list path must agree with the get path for this bank.
|
||||
banks = await memory.list_banks(request_context=request_context)
|
||||
entry = next((b for b in banks if b["bank_id"] == bank_id), None)
|
||||
assert entry is not None, f"bank {bank_id!r} not present in list_banks output"
|
||||
|
||||
assert entry["mission"] == profile["mission"], (
|
||||
"list_banks returned a different mission than get_bank_profile: "
|
||||
f"list={entry['mission']!r} get={profile['mission']!r}"
|
||||
)
|
||||
assert entry["disposition"] == profile["disposition"], (
|
||||
"list_banks returned a different disposition than get_bank_profile: "
|
||||
f"list={entry['disposition']!r} get={profile['disposition']!r}"
|
||||
)
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -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
|
||||
@@ -0,0 +1,82 @@
|
||||
"""
|
||||
Regression test for the hard timeout on the LiteLLM provider.
|
||||
|
||||
A completion that never returns — a connection held open with no token
|
||||
progress, or one straggler inside a concurrent ``asyncio.gather`` fan-out —
|
||||
must not block forever. ``call`` / ``call_with_tools`` wrap the request in
|
||||
``asyncio.wait_for`` so it is cancelled after ``timeout`` seconds and surfaced
|
||||
as a retryable ``TimeoutError`` instead of pinning a worker slot and a
|
||||
concurrency permit indefinitely.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
|
||||
from hindsight_api.engine.providers.litellm_llm import LiteLLMLLM
|
||||
|
||||
|
||||
def _make_provider(timeout: float | None) -> LiteLLMLLM:
|
||||
return LiteLLMLLM(
|
||||
provider="litellm",
|
||||
api_key="unused",
|
||||
base_url="http://localhost:0/v1",
|
||||
model="litellm_proxy/test-model",
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
async def test_call_cancels_hung_completion(monkeypatch):
|
||||
"""A hung ``_acompletion`` is cancelled per attempt and raises TimeoutError."""
|
||||
provider = _make_provider(timeout=0.1)
|
||||
calls = 0
|
||||
|
||||
async def _hang(**kwargs):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
await asyncio.Event().wait() # never resolves
|
||||
|
||||
monkeypatch.setattr(provider, "_acompletion", _hang)
|
||||
|
||||
started = time.monotonic()
|
||||
with pytest.raises((TimeoutError, asyncio.TimeoutError)):
|
||||
await provider.call(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
max_retries=1,
|
||||
initial_backoff=0.01,
|
||||
max_backoff=0.01,
|
||||
)
|
||||
elapsed = time.monotonic() - started
|
||||
|
||||
# max_retries=1 -> attempts 0 and 1, each bounded by the timeout.
|
||||
assert calls == 2
|
||||
# Bounded by ~2 * timeout + backoff — nowhere near hanging forever.
|
||||
assert elapsed < 2.0
|
||||
|
||||
|
||||
async def test_call_with_tools_cancels_hung_completion(monkeypatch):
|
||||
provider = _make_provider(timeout=0.1)
|
||||
|
||||
async def _hang(**kwargs):
|
||||
await asyncio.Event().wait()
|
||||
|
||||
monkeypatch.setattr(provider, "_acompletion", _hang)
|
||||
|
||||
with pytest.raises((TimeoutError, asyncio.TimeoutError)):
|
||||
await provider.call_with_tools(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
tools=[],
|
||||
max_retries=0,
|
||||
initial_backoff=0.01,
|
||||
max_backoff=0.01,
|
||||
)
|
||||
|
||||
|
||||
async def test_unset_timeout_falls_back_to_default(monkeypatch):
|
||||
"""``None`` must resolve to a finite default — never ``None``, which would
|
||||
make ``asyncio.wait_for`` wait forever and reintroduce the hang."""
|
||||
monkeypatch.delenv(ENV_LLM_TIMEOUT, raising=False)
|
||||
provider = _make_provider(timeout=None)
|
||||
assert provider.timeout == DEFAULT_LLM_TIMEOUT
|
||||
@@ -120,7 +120,7 @@ async def test_anthropic_no_extra_body_omits_key():
|
||||
# ─── Gemini ───────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _make_gemini_provider(extra_body=None):
|
||||
def _make_gemini_provider(extra_body=None, gemini_service_tier=None):
|
||||
pytest.importorskip("google.genai")
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
@@ -132,6 +132,7 @@ def _make_gemini_provider(extra_body=None):
|
||||
base_url="",
|
||||
model="gemini-2.5-flash",
|
||||
extra_body=extra_body,
|
||||
gemini_service_tier=gemini_service_tier,
|
||||
)
|
||||
provider._client = MagicMock()
|
||||
return provider
|
||||
@@ -176,6 +177,174 @@ async def test_gemini_explicit_temperature_overrides_extra_body():
|
||||
assert config_arg.temperature == 0.9
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_service_tier_applies_to_http_options_extra_body():
|
||||
"""The native Gemini service tier flag reaches GenerateContentConfig."""
|
||||
provider = _make_gemini_provider(gemini_service_tier="flex")
|
||||
provider._client.aio.models.generate_content = AsyncMock(return_value=_fake_gemini_response())
|
||||
|
||||
await provider.call(messages=[{"role": "user", "content": "hi"}], scope="test")
|
||||
|
||||
config_arg = provider._client.aio.models.generate_content.call_args.kwargs.get("config")
|
||||
assert config_arg.http_options.extra_body["service_tier"] == "flex"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_extra_body_service_tier_takes_precedence():
|
||||
"""The explicit extra_body escape hatch wins over the native flag."""
|
||||
provider = _make_gemini_provider(
|
||||
extra_body={"http_options": {"extra_body": {"service_tier": "standard"}}},
|
||||
gemini_service_tier="flex",
|
||||
)
|
||||
provider._client.aio.models.generate_content = AsyncMock(return_value=_fake_gemini_response())
|
||||
|
||||
await provider.call(messages=[{"role": "user", "content": "hi"}], scope="test")
|
||||
|
||||
config_arg = provider._client.aio.models.generate_content.call_args.kwargs.get("config")
|
||||
assert config_arg.http_options.extra_body["service_tier"] == "standard"
|
||||
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 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
||||
@@ -39,11 +39,20 @@ async def hundred_tenant_schemas(memory: MemoryEngine):
|
||||
"""Create N_TENANTS isolated schemas cloning the loop's tables; drop them after."""
|
||||
prefix = f"mt{uuid.uuid4().hex[:8]}"
|
||||
schemas = [f"{prefix}_{i:03d}" for i in range(N_TENANTS)]
|
||||
# Create all schemas + their tables in ONE transaction so the schemas become
|
||||
# visible to other connections only once fully built. Without this, each DDL
|
||||
# autocommits, leaving a window where a schema exists with only some of its
|
||||
# tables. The global maintenance routines (schemas_with_expired_rows /
|
||||
# banks_needing_consolidation) discover schemas by table presence and are run
|
||||
# concurrently by test_maintenance_routines on another xdist worker against
|
||||
# the shared test DB; they would query a not-yet-created table in a half-built
|
||||
# schema and fail with `relation "<schema>.<table>" does not exist`.
|
||||
async with memory._pool.acquire() as conn:
|
||||
for s in schemas:
|
||||
await conn.execute(f'CREATE SCHEMA "{s}"')
|
||||
for table in _CLONED_TABLES:
|
||||
await conn.execute(f'CREATE TABLE "{s}".{table} (LIKE public.{table} INCLUDING DEFAULTS)')
|
||||
async with conn.transaction():
|
||||
for s in schemas:
|
||||
await conn.execute(f'CREATE SCHEMA "{s}"')
|
||||
for table in _CLONED_TABLES:
|
||||
await conn.execute(f'CREATE TABLE "{s}".{table} (LIKE public.{table} INCLUDING DEFAULTS)')
|
||||
try:
|
||||
yield prefix, schemas
|
||||
finally:
|
||||
|
||||
@@ -134,6 +134,33 @@ async def test_banks_needing_consolidation_includes_in_flight_after_completion(m
|
||||
assert bank in {r["bank_id"] for r in rows}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_banks_needing_consolidation_skips_schema_with_vanished_table(memory: MemoryEngine):
|
||||
"""A schema discovered via its ``memory_units`` table but missing the
|
||||
``banks`` table the routine joins must be skipped, not abort the scan.
|
||||
|
||||
This reproduces the time-of-check/time-of-use race deterministically: the
|
||||
routine snapshots schemas owning ``memory_units`` from ``pg_class`` and then
|
||||
joins each schema's ``banks`` table. A tenant being dropped or migrated (and,
|
||||
in the test suite, the concurrent multi-tenant maintenance test) can leave a
|
||||
schema whose ``banks`` table is gone. Before the fix the dynamic query raised
|
||||
``undefined_table`` and aborted the whole routine (migration c7e9f1a3b5d2)."""
|
||||
schema = f"mtvanish{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(f'CREATE SCHEMA "{schema}"')
|
||||
# Discovered by the FOR loop (has memory_units) but the JOIN target
|
||||
# `banks` is absent — exactly a half-built / vanishing schema.
|
||||
await conn.execute(f'CREATE TABLE "{schema}".memory_units (LIKE public.memory_units INCLUDING DEFAULTS)')
|
||||
|
||||
# Must not raise; the bad schema is simply skipped.
|
||||
rows = await conn.fetch("SELECT schema_name, bank_id FROM public.banks_needing_consolidation()")
|
||||
assert schema not in {r["schema_name"] for r in rows}
|
||||
finally:
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_schemas_with_expired_rows(memory: MemoryEngine):
|
||||
"""Returns schemas holding a row older than p_days; respects the p_days<=0 guard."""
|
||||
|
||||
@@ -423,6 +423,28 @@ def test_global_mcp_enabled_tools_intersects_with_single_bank_mode(mock_memory):
|
||||
assert "list_banks" not in tools # single-bank mode excludes it regardless
|
||||
|
||||
|
||||
def test_mcp_instructions_append_to_retain_and_recall_descriptions(mock_memory):
|
||||
"""HINDSIGHT_API_MCP_INSTRUCTIONS customizes retain/recall tool descriptions."""
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from hindsight_api.api.mcp import create_mcp_server
|
||||
|
||||
custom_instructions = "Also store every action you take."
|
||||
mock_cfg = MagicMock()
|
||||
mock_cfg.mcp_enabled_tools = ["retain", "recall", "reflect"]
|
||||
mock_cfg.mcp_instructions = custom_instructions
|
||||
|
||||
with patch("hindsight_api.api.mcp._get_raw_config", return_value=mock_cfg):
|
||||
mcp_server = create_mcp_server(mock_memory, multi_bank=True)
|
||||
|
||||
tools = _tools(mcp_server)
|
||||
expected_suffix = f"Additional instructions: {custom_instructions}"
|
||||
|
||||
assert expected_suffix in tools["retain"].description
|
||||
assert expected_suffix in tools["recall"].description
|
||||
assert expected_suffix not in tools["reflect"].description
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_logic_from_url_path():
|
||||
"""Test that routing correctly selects server based on URL structure.
|
||||
@@ -430,14 +452,6 @@ async def test_routing_logic_from_url_path():
|
||||
Simulates the path parsing logic from MCPMiddleware.__call__ after the
|
||||
prefix has been stripped. Any first path segment is treated as a bank_id.
|
||||
"""
|
||||
from hindsight_api.api.mcp import MCPMiddleware
|
||||
|
||||
# Mock memory
|
||||
mock_memory = MagicMock()
|
||||
|
||||
# Create middleware
|
||||
middleware = MCPMiddleware(None, mock_memory)
|
||||
|
||||
# Simulate different URL patterns and verify routing
|
||||
# Path is what remains after stripping the /mcp prefix
|
||||
test_cases = [
|
||||
|
||||
@@ -1911,3 +1911,107 @@ class TestBankToolFiltering:
|
||||
# Filter bypassed — config resolver was never consulted, all tools visible
|
||||
assert "recall" in visible
|
||||
mock_memory_with_resolver._config_resolver.get_bank_config.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestToolAnnotations:
|
||||
"""Every MCP tool must carry read-only / destructive hints (openWorldHint=False)."""
|
||||
|
||||
async def test_read_only_tool(self, mock_memory):
|
||||
ann = _tools(_make_mcp_server(mock_memory, {"recall"}))["recall"].annotations
|
||||
assert ann is not None
|
||||
assert ann.readOnlyHint is True
|
||||
assert ann.openWorldHint is False
|
||||
|
||||
async def test_reflect_is_read_only(self, mock_memory):
|
||||
# reflect synthesizes an answer and persists nothing (memory_engine.reflect_async),
|
||||
# so it carries readOnlyHint=True like recall.
|
||||
ann = _tools(_make_mcp_server(mock_memory, {"reflect"}))["reflect"].annotations
|
||||
assert ann is not None
|
||||
assert ann.readOnlyHint is True
|
||||
assert ann.openWorldHint is False
|
||||
|
||||
async def test_destructive_tool(self, mock_memory):
|
||||
ann = _tools(_make_mcp_server(mock_memory, {"delete_bank"}))["delete_bank"].annotations
|
||||
assert ann is not None
|
||||
assert ann.readOnlyHint is False
|
||||
assert ann.destructiveHint is True
|
||||
|
||||
async def test_write_tool_is_not_destructive(self, mock_memory):
|
||||
ann = _tools(_make_mcp_server(mock_memory, {"retain"}))["retain"].annotations
|
||||
assert ann is not None
|
||||
assert ann.readOnlyHint is False
|
||||
assert ann.destructiveHint is False
|
||||
|
||||
async def test_annotations_apply_in_single_bank_mode(self, mock_memory):
|
||||
ann = _tools(_make_mcp_server(mock_memory, {"recall"}, include_bank_id=False))["recall"].annotations
|
||||
assert ann is not None
|
||||
assert ann.readOnlyHint is True
|
||||
|
||||
|
||||
def _reflect_mcp_with_trace(include_bank_id_param: bool):
|
||||
"""An MCP server whose reflect returns a result carrying tool_trace/llm_trace."""
|
||||
from fastmcp import FastMCP
|
||||
|
||||
# Mirrors ReflectResult: the agentic loop's trace fields are large and present.
|
||||
reflect_payload = {
|
||||
"text": "answer",
|
||||
"based_on": {"world": []},
|
||||
"tool_trace": [{"tool": "recall", "output": "x" * 1000}],
|
||||
"llm_trace": [{"model": "test", "output": "y" * 1000}],
|
||||
"directives_applied": [{"id": "d1", "name": "Tone", "content": "z" * 1000}],
|
||||
}
|
||||
memory = MagicMock()
|
||||
memory.reflect_async = AsyncMock(
|
||||
return_value=MagicMock(
|
||||
model_dump_json=lambda indent=None: json.dumps(reflect_payload),
|
||||
model_dump=lambda: dict(reflect_payload),
|
||||
structured_output=None,
|
||||
)
|
||||
)
|
||||
mcp = FastMCP("test")
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=lambda: "test-bank",
|
||||
include_bank_id_param=include_bank_id_param,
|
||||
tools={"reflect"},
|
||||
)
|
||||
register_mcp_tools(mcp, memory, config)
|
||||
return mcp
|
||||
|
||||
|
||||
def _reflect_result_data(result) -> dict:
|
||||
"""The multi-bank reflect returns a JSON string; single-bank returns a dict."""
|
||||
return json.loads(result) if isinstance(result, str) else result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestReflectTraceOmission:
|
||||
"""reflect must not leak the agentic tool_trace/llm_trace/directives_applied into MCP responses by default."""
|
||||
|
||||
@pytest.mark.parametrize("multi_bank", [True, False])
|
||||
async def test_trace_omitted_by_default(self, multi_bank):
|
||||
mcp = _reflect_mcp_with_trace(multi_bank)
|
||||
data = _reflect_result_data(await _tools(mcp)["reflect"].fn(query="q"))
|
||||
assert data["text"] == "answer"
|
||||
assert "tool_trace" not in data
|
||||
assert "llm_trace" not in data
|
||||
# directives_applied is built "for the trace" and carries full directive content,
|
||||
# so it must be omitted by default like the other trace fields.
|
||||
assert "directives_applied" not in data
|
||||
|
||||
@pytest.mark.parametrize("multi_bank", [True, False])
|
||||
async def test_trace_included_when_requested(self, multi_bank):
|
||||
mcp = _reflect_mcp_with_trace(multi_bank)
|
||||
data = _reflect_result_data(await _tools(mcp)["reflect"].fn(query="q", include_trace=True))
|
||||
assert "tool_trace" in data
|
||||
assert "llm_trace" in data
|
||||
assert "directives_applied" in data
|
||||
|
||||
@pytest.mark.parametrize("multi_bank", [True, False])
|
||||
async def test_based_on_flag_is_independent_of_trace(self, multi_bank):
|
||||
# include_based_on keeps based_on but must not pull the trace back in.
|
||||
mcp = _reflect_mcp_with_trace(multi_bank)
|
||||
data = _reflect_result_data(await _tools(mcp)["reflect"].fn(query="q", include_based_on=True))
|
||||
assert "based_on" in data
|
||||
assert "tool_trace" not in data
|
||||
assert "directives_applied" not in data
|
||||
|
||||
@@ -334,6 +334,9 @@ def _make_minimal_engine():
|
||||
mock_embeddings = MagicMock()
|
||||
mock_embeddings.dimension = 384
|
||||
|
||||
from hindsight_api.config import clear_config_cache
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
@@ -343,11 +346,17 @@ def _make_minimal_engine():
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
from hindsight_api.config import clear_config_cache
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
clear_config_cache()
|
||||
return MemoryEngine(db_url="postgresql://localhost/hindsight_test", embeddings=mock_embeddings)
|
||||
engine = MemoryEngine(db_url="postgresql://localhost/hindsight_test", embeddings=mock_embeddings)
|
||||
|
||||
# Constructing the engine above repopulated the process-global config cache
|
||||
# from the patched env (provider="none" forces retain_extraction_mode="chunks").
|
||||
# Now that the patched env is gone, drop that cache so the leaked "none"/chunks
|
||||
# config does not bleed into other tests on this xdist worker — their retains
|
||||
# would silently skip entity extraction (0 unit_entities) and fail unrelated
|
||||
# assertions. The next get_config() rebuilds from the real env.
|
||||
clear_config_cache()
|
||||
return engine
|
||||
|
||||
|
||||
def test_engine_memory_defense_shares_ext_ctx() -> None:
|
||||
|
||||
@@ -11,6 +11,7 @@ from hindsight_api.metrics import (
|
||||
get_token_bucket,
|
||||
create_metrics_collector,
|
||||
initialize_metrics,
|
||||
normalize_http_endpoint,
|
||||
)
|
||||
|
||||
|
||||
@@ -221,6 +222,17 @@ class TestMetricsCollector:
|
||||
assert reflect_attrs["operation"] == "reflect"
|
||||
assert reflect_attrs["source"] == "api"
|
||||
|
||||
def test_record_operation_result_records_with_explicit_success(self, collector):
|
||||
"""Direct recording path used by the worker (source=worker, explicit success)."""
|
||||
collector.record_operation_result("retain", bank_id="test_bank", success=False, duration=1.5, source="worker")
|
||||
|
||||
duration, attributes = collector.operation_duration.record.call_args[0]
|
||||
assert duration == 1.5
|
||||
assert attributes["operation"] == "retain"
|
||||
assert attributes["source"] == "worker"
|
||||
assert attributes["success"] == "false"
|
||||
collector.operation_total.add.assert_called_once_with(1, attributes)
|
||||
|
||||
def test_record_operation_includes_bank_id_when_enabled(self):
|
||||
"""Test that bank_id is included in attributes when metrics_include_bank_id is enabled."""
|
||||
mock_config = MagicMock()
|
||||
@@ -324,6 +336,24 @@ class TestGetTokenBucket:
|
||||
assert get_token_bucket(1000000) == "50k+"
|
||||
|
||||
|
||||
class TestNormalizeHttpEndpoint:
|
||||
"""Tests for normalize_http_endpoint (low-cardinality HTTP metric labels)."""
|
||||
|
||||
def test_templates_high_cardinality_segments(self):
|
||||
"""Bank ids (incl. non-numeric), UUIDs, and numeric ids collapse to placeholders."""
|
||||
cases = [
|
||||
("/v1/default/banks/user-1680/memories/recall", "/v1/default/banks/{bank_id}/memories/recall"),
|
||||
("/v1/default/banks/tenant-acme/memories", "/v1/default/banks/{bank_id}/memories"),
|
||||
("/v1/default/banks/user-1680", "/v1/default/banks/{bank_id}"),
|
||||
("/v1/default/banks/3f8c1e2a-1111-2222-3333-444455556666/config", "/v1/default/banks/{bank_id}/config"),
|
||||
("/v1/default/banks/42/config", "/v1/default/banks/{bank_id}/config"),
|
||||
("/v1/default/banks", "/v1/default/banks"),
|
||||
("/health", "/health"),
|
||||
]
|
||||
for raw, expected in cases:
|
||||
assert normalize_http_endpoint(raw) == expected, raw
|
||||
|
||||
|
||||
class TestLLMMetrics:
|
||||
"""Tests for LLM-specific metrics recording."""
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ from pathlib import Path
|
||||
import pytest
|
||||
from alembic import command
|
||||
from alembic.config import Config
|
||||
from sqlalchemy import create_engine, text, inspect
|
||||
from sqlalchemy import create_engine, text
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
@@ -33,8 +33,17 @@ def _upgrade(db_url: str, revision: str) -> None:
|
||||
command.upgrade(_alembic_cfg(db_url), revision)
|
||||
|
||||
|
||||
def _downgrade(db_url: str, revision: str) -> None:
|
||||
command.downgrade(_alembic_cfg(db_url), revision)
|
||||
def _reset_public_schema(db_url: str) -> None:
|
||||
engine = create_engine(db_url, isolation_level="AUTOCOMMIT")
|
||||
try:
|
||||
with engine.connect() as conn:
|
||||
# This test rewinds/replays migration history against a persistent
|
||||
# pg0 instance. Rebuild only its dedicated public schema so a
|
||||
# previous run cannot leave alembic_version ahead of the real DDL.
|
||||
conn.execute(text("DROP SCHEMA IF EXISTS public CASCADE"))
|
||||
conn.execute(text("CREATE SCHEMA public"))
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -43,16 +52,16 @@ def _downgrade(db_url: str, revision: str) -> None:
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def pre_backsweep_db_url():
|
||||
def pre_backsweep_db_url() -> str:
|
||||
"""
|
||||
Spin up a dedicated pg0 instance and ensure schema is at the revision
|
||||
just before the backsweep so each test can seed orphan data and then
|
||||
apply the backsweep itself.
|
||||
|
||||
Because pg0 data directories persist across test runs, the DB may
|
||||
already be at head. We upgrade to head first (to ensure all tables
|
||||
exist), then stamp the revision back to pre-backsweep so Alembic
|
||||
treats the backsweep as not-yet-applied.
|
||||
Because pg0 data directories persist across test runs, the DB may already
|
||||
have schema from a previous test run. Reset this test's dedicated schema
|
||||
first, then migrate to the real pre-backsweep revision instead of stamping
|
||||
a head schema backward.
|
||||
"""
|
||||
from hindsight_api.pg0 import EmbeddedPostgres
|
||||
|
||||
@@ -63,10 +72,8 @@ def pre_backsweep_db_url():
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
# Ensure all tables exist (upgrade to head), then stamp back to
|
||||
# pre-backsweep so the backsweep migration will actually run.
|
||||
_upgrade(url, "heads")
|
||||
command.stamp(_alembic_cfg(url), "f6g7h8i9j0k1")
|
||||
_reset_public_schema(url)
|
||||
_upgrade(url, "f6g7h8i9j0k1")
|
||||
return url
|
||||
|
||||
|
||||
@@ -75,7 +82,7 @@ def pre_backsweep_db_url():
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_backsweep_removes_orphans_and_preserves_legit_rows(pre_backsweep_db_url):
|
||||
def test_backsweep_removes_orphans_and_preserves_legit_rows(pre_backsweep_db_url: str) -> None:
|
||||
"""
|
||||
Seed four kinds of rows then apply the backsweep migration and verify:
|
||||
|
||||
@@ -114,15 +121,20 @@ def test_backsweep_removes_orphans_and_preserves_legit_rows(pre_backsweep_db_url
|
||||
conn.execute(text("INSERT INTO banks (bank_id) VALUES (:b)"), {"b": alive_bank})
|
||||
|
||||
# --- seed memory_units ---
|
||||
def insert_mu(uid, bank, fact_type, sources=None):
|
||||
def insert_mu(
|
||||
uid: uuid.UUID,
|
||||
bank: str,
|
||||
fact_type: str,
|
||||
sources: list[uuid.UUID] | None = None,
|
||||
) -> None:
|
||||
src_arr = "{" + ",".join(str(s) for s in (sources or [])) + "}"
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO memory_units
|
||||
(id, bank_id, text, fact_type, source_memory_ids)
|
||||
(id, bank_id, text, event_date, fact_type, source_memory_ids)
|
||||
VALUES
|
||||
(:id, :bank, :text, :ft, CAST(:src AS uuid[]))
|
||||
(:id, :bank, :text, now(), :ft, CAST(:src AS uuid[]))
|
||||
"""
|
||||
),
|
||||
{"id": uid, "bank": bank, "text": "test", "ft": fact_type, "src": src_arr},
|
||||
@@ -150,7 +162,7 @@ def test_backsweep_removes_orphans_and_preserves_legit_rows(pre_backsweep_db_url
|
||||
# --- verify ---
|
||||
with engine.connect() as conn:
|
||||
|
||||
def exists(uid):
|
||||
def exists(uid: uuid.UUID) -> bool:
|
||||
return conn.execute(text("SELECT 1 FROM memory_units WHERE id = :id"), {"id": uid}).fetchone() is not None
|
||||
|
||||
# Must be gone
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
"""Provider quota reset windows defer worker retries instead of failing retains."""
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from openai import APIStatusError
|
||||
|
||||
from hindsight_api.engine.llm_interface import ProviderRateLimitResetError
|
||||
from hindsight_api.engine.providers.openai_compatible_llm import OpenAICompatibleLLM
|
||||
|
||||
|
||||
def _llm() -> OpenAICompatibleLLM:
|
||||
return OpenAICompatibleLLM(
|
||||
provider="zai",
|
||||
model="glm-5-turbo",
|
||||
api_key="test",
|
||||
base_url="https://example.com/v1",
|
||||
)
|
||||
|
||||
|
||||
def _usage_limit_error(reset_at: str) -> APIStatusError:
|
||||
body = {
|
||||
"code": "1308",
|
||||
"message": f"Usage limit reached for 5 hour. Your limit will reset at {reset_at}",
|
||||
}
|
||||
response = MagicMock()
|
||||
response.status_code = 429
|
||||
response.text = '{"code": "1308", "message": "usage limit"}'
|
||||
response.headers = {}
|
||||
return APIStatusError("rate limited", response=response, body=body)
|
||||
|
||||
|
||||
def _short_retry_after_error() -> APIStatusError:
|
||||
response = MagicMock()
|
||||
response.status_code = 429
|
||||
response.text = '{"code": "rate_limit", "message": "retry shortly"}'
|
||||
response.headers = {"retry-after": "1"}
|
||||
return APIStatusError("rate limited", response=response, body={"message": "retry shortly"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_limit_429_with_reset_defers_without_inner_retry() -> None:
|
||||
llm = _llm()
|
||||
reset_at = (datetime.now(UTC) + timedelta(hours=5)).replace(microsecond=0)
|
||||
create = AsyncMock(side_effect=_usage_limit_error(reset_at.isoformat().replace("+00:00", "Z")))
|
||||
llm._client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create)))
|
||||
|
||||
with patch(
|
||||
"hindsight_api.engine.providers.openai_compatible_llm.asyncio.sleep",
|
||||
new_callable=AsyncMock,
|
||||
) as sleep:
|
||||
with pytest.raises(ProviderRateLimitResetError) as exc_info:
|
||||
await llm.call(
|
||||
messages=[{"role": "user", "content": "x"}],
|
||||
scope="retain_extract_facts",
|
||||
max_retries=2,
|
||||
)
|
||||
|
||||
assert create.await_count == 1
|
||||
sleep.assert_not_awaited()
|
||||
assert abs((exc_info.value.retry_at - reset_at).total_seconds()) < 1
|
||||
assert "Provider quota exhausted" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_short_retry_after_429_uses_normal_retry_loop() -> None:
|
||||
llm = _llm()
|
||||
create = AsyncMock(side_effect=_short_retry_after_error())
|
||||
llm._client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create)))
|
||||
|
||||
with patch(
|
||||
"hindsight_api.engine.providers.openai_compatible_llm.asyncio.sleep",
|
||||
new_callable=AsyncMock,
|
||||
) as sleep:
|
||||
with pytest.raises(APIStatusError):
|
||||
await llm.call(
|
||||
messages=[{"role": "user", "content": "x"}],
|
||||
scope="retain_extract_facts",
|
||||
max_retries=2,
|
||||
max_backoff=60,
|
||||
)
|
||||
|
||||
assert create.await_count == 3
|
||||
assert sleep.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extract_facts_from_text_preserves_provider_quota_reset(monkeypatch) -> None:
|
||||
from hindsight_api.engine.retain import fact_extraction
|
||||
|
||||
retry_at = (datetime.now(UTC) + timedelta(hours=2)).replace(microsecond=0)
|
||||
|
||||
async def quota_limited_chunk(**_: object) -> None:
|
||||
raise ProviderRateLimitResetError(retry_at=retry_at, message="quota resets later")
|
||||
|
||||
monkeypatch.setattr(fact_extraction, "_extract_facts_with_auto_split", quota_limited_chunk)
|
||||
|
||||
with pytest.raises(ProviderRateLimitResetError) as exc_info:
|
||||
await fact_extraction.extract_facts_from_text(
|
||||
text="Alice moved to Berlin.",
|
||||
event_date=None,
|
||||
llm_config=object(),
|
||||
agent_name="TestAgent",
|
||||
config=SimpleNamespace(retain_chunk_size=1000, retain_structured_chunk_size=None),
|
||||
)
|
||||
|
||||
assert exc_info.value.retry_at == retry_at
|
||||
assert "Fact extraction deferred by provider quota" in str(exc_info.value)
|
||||
@@ -0,0 +1,184 @@
|
||||
"""Tests for the recall `prefer_observations` deduplication flag.
|
||||
|
||||
When the caller recalls raw facts ('world'/'experience') together with
|
||||
'observation' and sets prefer_observations=True, any raw fact that a returned
|
||||
observation was consolidated from (tracked via memory_units.source_memory_ids)
|
||||
is dropped so the observation supersedes it — no duplicate content.
|
||||
|
||||
Dedup is provenance-based, not semantic: a raw fact that is semantically
|
||||
similar to an observation but NOT listed in its source_memory_ids must survive.
|
||||
|
||||
No LLM required — inserts memory_units directly via SQL with real embeddings.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api import MemoryEngine, RequestContext
|
||||
from hindsight_api.engine.retain import embedding_utils
|
||||
|
||||
RC = RequestContext(tenant_id="default")
|
||||
|
||||
QUERY = "Alice mountain hiking"
|
||||
|
||||
# Two raw facts the observation is consolidated from (must be dropped when the
|
||||
# flag is on), one raw fact that is semantically similar but NOT a source (must
|
||||
# survive), and the observation itself.
|
||||
SRC1_TEXT = "Alice loves hiking in the mountains"
|
||||
SRC2_TEXT = "Alice hikes the Alps every summer"
|
||||
NON_SRC_TEXT = "Alice enjoys exploring mountain hiking trails"
|
||||
OBS_TEXT = "Alice is an avid mountain hiker"
|
||||
|
||||
|
||||
async def _insert_unit(
|
||||
conn,
|
||||
*,
|
||||
unit_id: str,
|
||||
text: str,
|
||||
bank_id: str,
|
||||
embedding_str: str,
|
||||
fact_type: str = "world",
|
||||
source_memory_ids: list[uuid.UUID] | None = None,
|
||||
) -> None:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_units (id, bank_id, text, fact_type, embedding, source_memory_ids)
|
||||
VALUES ($1, $2, $3, $4, $5::vector, $6::uuid[])
|
||||
""",
|
||||
unit_id,
|
||||
bank_id,
|
||||
text,
|
||||
fact_type,
|
||||
embedding_str,
|
||||
source_memory_ids,
|
||||
)
|
||||
|
||||
|
||||
def _to_str(emb: list[float]) -> str:
|
||||
return "[" + ",".join(str(v) for v in emb) + "]"
|
||||
|
||||
|
||||
def _result_ids(result) -> set[str]:
|
||||
return {str(r.id) for r in result.results}
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def seeded_obs_memory(memory_no_llm_verify: MemoryEngine):
|
||||
"""Seed two source facts, one non-source fact, and an observation over the two sources."""
|
||||
engine = memory_no_llm_verify
|
||||
bank_id = f"test-prefer-obs-{uuid.uuid4().hex[:8]}"
|
||||
await engine.get_bank_profile(bank_id, request_context=RC)
|
||||
|
||||
src1_id = str(uuid.uuid4())
|
||||
src2_id = str(uuid.uuid4())
|
||||
non_src_id = str(uuid.uuid4())
|
||||
obs_id = str(uuid.uuid4())
|
||||
|
||||
embeddings = await embedding_utils.generate_embeddings_batch(
|
||||
engine.embeddings,
|
||||
[SRC1_TEXT, SRC2_TEXT, NON_SRC_TEXT, OBS_TEXT],
|
||||
)
|
||||
|
||||
pool = await engine._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_unit(conn, unit_id=src1_id, text=SRC1_TEXT, bank_id=bank_id, embedding_str=_to_str(embeddings[0]))
|
||||
await _insert_unit(conn, unit_id=src2_id, text=SRC2_TEXT, bank_id=bank_id, embedding_str=_to_str(embeddings[1]))
|
||||
await _insert_unit(
|
||||
conn, unit_id=non_src_id, text=NON_SRC_TEXT, bank_id=bank_id, embedding_str=_to_str(embeddings[2])
|
||||
)
|
||||
await _insert_unit(
|
||||
conn,
|
||||
unit_id=obs_id,
|
||||
text=OBS_TEXT,
|
||||
bank_id=bank_id,
|
||||
embedding_str=_to_str(embeddings[3]),
|
||||
fact_type="observation",
|
||||
source_memory_ids=[uuid.UUID(src1_id), uuid.UUID(src2_id)],
|
||||
)
|
||||
|
||||
ids = {"src1": src1_id, "src2": src2_id, "non_src": non_src_id, "obs": obs_id}
|
||||
yield engine, bank_id, ids
|
||||
|
||||
await engine.delete_bank(bank_id, request_context=RC)
|
||||
|
||||
|
||||
class TestPreferObservations:
|
||||
async def test_disabled_returns_sources_and_observation(self, seeded_obs_memory):
|
||||
"""Without the flag, the source facts AND the observation are all returned."""
|
||||
engine, bank_id, ids = seeded_obs_memory
|
||||
result = await engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=QUERY,
|
||||
request_context=RC,
|
||||
fact_type=["world", "experience", "observation"],
|
||||
prefer_observations=False,
|
||||
max_tokens=10000,
|
||||
)
|
||||
found = _result_ids(result)
|
||||
assert ids["src1"] in found
|
||||
assert ids["src2"] in found
|
||||
assert ids["obs"] in found
|
||||
|
||||
async def test_enabled_drops_source_facts_keeps_observation(self, seeded_obs_memory):
|
||||
"""With the flag, the observation supersedes the facts it was consolidated from."""
|
||||
engine, bank_id, ids = seeded_obs_memory
|
||||
result = await engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=QUERY,
|
||||
request_context=RC,
|
||||
fact_type=["world", "experience", "observation"],
|
||||
prefer_observations=True,
|
||||
max_tokens=10000,
|
||||
)
|
||||
found = _result_ids(result)
|
||||
assert ids["obs"] in found, "the observation must remain"
|
||||
assert ids["src1"] not in found, "source fact 1 is superseded by the observation"
|
||||
assert ids["src2"] not in found, "source fact 2 is superseded by the observation"
|
||||
|
||||
async def test_enabled_keeps_non_source_fact(self, seeded_obs_memory):
|
||||
"""Dedup is provenance-based: a similar fact NOT in source_memory_ids survives."""
|
||||
engine, bank_id, ids = seeded_obs_memory
|
||||
result = await engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=QUERY,
|
||||
request_context=RC,
|
||||
fact_type=["world", "experience", "observation"],
|
||||
prefer_observations=True,
|
||||
max_tokens=10000,
|
||||
)
|
||||
found = _result_ids(result)
|
||||
assert ids["non_src"] in found, "a non-source fact must not be dropped, even if semantically similar"
|
||||
|
||||
async def test_noop_without_observation_type(self, seeded_obs_memory):
|
||||
"""The flag is a no-op when 'observation' is not among the requested types."""
|
||||
engine, bank_id, ids = seeded_obs_memory
|
||||
result = await engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=QUERY,
|
||||
request_context=RC,
|
||||
fact_type=["world", "experience"],
|
||||
prefer_observations=True,
|
||||
max_tokens=10000,
|
||||
)
|
||||
found = _result_ids(result)
|
||||
assert ids["src1"] in found
|
||||
assert ids["src2"] in found
|
||||
|
||||
|
||||
def test_flag_is_opt_in_by_default():
|
||||
"""prefer_observations is opt-in: off at the API surface and the engine method.
|
||||
|
||||
The engine default in particular must stay False so internal callers — notably
|
||||
consolidation, which needs the raw facts it folds into observations — are never
|
||||
silently deduped.
|
||||
"""
|
||||
import inspect
|
||||
|
||||
from hindsight_api.api.http import RecallRequest
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
assert RecallRequest(query="anything").prefer_observations is False
|
||||
engine_default = inspect.signature(MemoryEngine.recall_async).parameters["prefer_observations"].default
|
||||
assert engine_default is False
|
||||
@@ -100,7 +100,11 @@ async def test_recall_async_passes_question_date_to_combined_scoring(monkeypatch
|
||||
)
|
||||
|
||||
def apply_combined_scoring(
|
||||
scored_results: list[ScoredResult], *, now: datetime, is_passthrough_reranker: bool
|
||||
scored_results: list[ScoredResult],
|
||||
*,
|
||||
now: datetime,
|
||||
is_passthrough_reranker: bool,
|
||||
**_kwargs: object,
|
||||
) -> None:
|
||||
nonlocal captured_now
|
||||
assert is_passthrough_reranker is False
|
||||
|
||||
@@ -0,0 +1,78 @@
|
||||
"""Tests for _strip_reasoning_tags helper in OpenAI-compatible LLM provider."""
|
||||
|
||||
from hindsight_api.engine.providers.openai_compatible_llm import _strip_reasoning_tags
|
||||
|
||||
|
||||
class TestStripReasoningTags:
|
||||
"""Test reasoning/thinking tag stripping from LLM responses."""
|
||||
|
||||
def test_plain_text_unchanged(self):
|
||||
"""Text without reasoning tags passes through (modulo edge whitespace)."""
|
||||
content = "User prefers functional programming patterns."
|
||||
assert _strip_reasoning_tags(content) == content
|
||||
|
||||
def test_empty_string(self):
|
||||
"""Empty string passes through."""
|
||||
assert _strip_reasoning_tags("") == ""
|
||||
|
||||
def test_closed_think_stripped(self):
|
||||
"""A closed <think>...</think> block is removed."""
|
||||
content = "<think>let me reason</think>The answer is 42."
|
||||
assert _strip_reasoning_tags(content) == "The answer is 42."
|
||||
|
||||
def test_closed_thinking_stripped(self):
|
||||
assert _strip_reasoning_tags("<thinking>reasoning</thinking>Result") == "Result"
|
||||
|
||||
def test_closed_thought_stripped(self):
|
||||
assert _strip_reasoning_tags("<thought>hmm</thought>Result") == "Result"
|
||||
|
||||
def test_closed_reasoning_stripped(self):
|
||||
assert _strip_reasoning_tags("<reasoning>step by step</reasoning>Result") == "Result"
|
||||
|
||||
def test_startthink_endthink_stripped(self):
|
||||
"""The |startthink|...|endthink| marker style is removed."""
|
||||
content = "|startthink|internal monologue|endthink|Final output"
|
||||
assert _strip_reasoning_tags(content) == "Final output"
|
||||
|
||||
def test_multiline_think_stripped(self):
|
||||
"""DOTALL: a multi-line thinking block is fully removed."""
|
||||
content = "<think>\nline one\nline two\n</think>\nThe real content."
|
||||
assert _strip_reasoning_tags(content) == "The real content."
|
||||
|
||||
def test_unclosed_think_stripped_to_end(self):
|
||||
"""An unclosed <think> (truncated output) is removed to end-of-string."""
|
||||
content = "Partial answer.\n<think>I started thinking but got cut off"
|
||||
assert _strip_reasoning_tags(content) == "Partial answer."
|
||||
|
||||
def test_unclosed_thinking_stripped_to_end(self):
|
||||
content = "result text\n<thinking>dangling reasoning with no close"
|
||||
assert _strip_reasoning_tags(content) == "result text"
|
||||
|
||||
def test_only_unclosed_think_becomes_empty(self):
|
||||
"""Content that is entirely an unclosed thinking block collapses to empty."""
|
||||
content = "<think>everything is reasoning and it never closed"
|
||||
assert _strip_reasoning_tags(content) == ""
|
||||
|
||||
def test_multiple_blocks_stripped(self):
|
||||
"""Multiple closed blocks are all removed."""
|
||||
content = "<think>a</think>Hello <think>b</think>World"
|
||||
assert _strip_reasoning_tags(content) == "Hello World"
|
||||
|
||||
def test_mental_model_markdown_contamination(self):
|
||||
"""Real-world MiniMax-M3 free-form leak: <think> wrapping a markdown mental model."""
|
||||
content = (
|
||||
"<think>\n"
|
||||
"The user keeps asking about FP. I should consolidate this.\n"
|
||||
"</think>\n"
|
||||
"# Mental Model: Coding Preferences\n\n"
|
||||
"The user prefers functional programming patterns and immutable data."
|
||||
)
|
||||
result = _strip_reasoning_tags(content)
|
||||
assert "<think>" not in result
|
||||
assert "</think>" not in result
|
||||
assert result.startswith("# Mental Model: Coding Preferences")
|
||||
|
||||
def test_unclosed_think_after_json_payload(self):
|
||||
"""Truncated <think> trailing valid JSON is stripped (closing tag absent)."""
|
||||
content = '{"facts": [{"what": "test"}]}\n<think>oops truncated'
|
||||
assert _strip_reasoning_tags(content) == '{"facts": [{"what": "test"}]}'
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Regression test for https://github.com/vectorize-io/hindsight/issues/2301
|
||||
|
||||
Raising ``HINDSIGHT_API_RETAIN_STRUCTURED_CHUNK_SIZE`` above
|
||||
``HINDSIGHT_API_RETAIN_CHUNK_SIZE`` and retaining a JSONL/conversation document
|
||||
whose line/turn overflows the chunk size used to crash with::
|
||||
|
||||
asyncpg.exceptions.CardinalityViolationError:
|
||||
ON CONFLICT DO UPDATE command cannot affect row a second time
|
||||
|
||||
The streaming retain pipeline pre-chunks each document once (one ``chunk_index``
|
||||
per piece) and then re-chunks every piece during extraction. With the structured
|
||||
cap above the chunk size, a pre-chunk could legitimately exceed the re-chunk
|
||||
budget, so it re-split into several sub-chunks that all inherited the one
|
||||
``chunk_index`` — colliding on ``chunk_id = {bank}_{doc}_{index}`` in a single
|
||||
upsert batch. ``chunk_text`` is now idempotent, so the re-chunk is a no-op.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
|
||||
def _ts() -> float:
|
||||
return datetime.now(timezone.utc).timestamp()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _structured_chunk_env(monkeypatch):
|
||||
# Structured cap ABOVE the chunk size — the configuration that triggers #2301.
|
||||
# Default-scale sizes (matching the issue) so the document yields only a
|
||||
# handful of chunks: a smaller chunk size explodes the embedding/link work and
|
||||
# destabilises the shared session fixture under xdist.
|
||||
monkeypatch.setenv("HINDSIGHT_API_RETAIN_CHUNK_SIZE", "3000")
|
||||
monkeypatch.setenv("HINDSIGHT_API_RETAIN_STRUCTURED_CHUNK_SIZE", "4500")
|
||||
monkeypatch.setenv("HINDSIGHT_API_ENABLE_AUTO_CONSOLIDATION", "false")
|
||||
monkeypatch.setenv("HINDSIGHT_API_ENABLE_OBSERVATIONS", "false")
|
||||
clear_config_cache()
|
||||
yield
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_jsonl_line_over_chunk_size_retains_without_collision(memory, request_context):
|
||||
"""A JSONL document with a line longer than the chunk size retains cleanly
|
||||
when the structured cap is raised above it (issue #2301)."""
|
||||
import json
|
||||
|
||||
bank_id = f"test_2301_{_ts()}"
|
||||
document_id = "doc-2301"
|
||||
|
||||
try:
|
||||
body = "\n".join(
|
||||
[
|
||||
json.dumps({"role": "user", "content": "short opening line"}),
|
||||
# Between chunk size (3000) and structured cap (4500): kept whole by
|
||||
# the producer, would re-split on re-chunk without the fix.
|
||||
json.dumps({"role": "assistant", "content": "k" * 3800}),
|
||||
# Past even the structured cap: fragmented as text.
|
||||
json.dumps({"role": "assistant", "content": "m" * 9000}),
|
||||
json.dumps({"role": "user", "content": "short closing line"}),
|
||||
]
|
||||
)
|
||||
|
||||
# The bug raised CardinalityViolationError here.
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=body,
|
||||
context="jsonl with oversized line",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
chunks = await memory.list_document_chunks(bank_id, document_id, limit=10000, request_context=request_context)
|
||||
indices = sorted(c["chunk_index"] for c in chunks["items"])
|
||||
# No collisions: chunk_index values are unique.
|
||||
assert len(indices) == len(set(indices)), f"duplicate chunk_index values: {indices}"
|
||||
assert indices == list(range(len(indices))), f"chunk_index sequence not contiguous: {indices}"
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -0,0 +1,168 @@
|
||||
"""Regression: sub-batch slices that each span MULTIPLE extraction chunks must
|
||||
keep full chunk coverage on BOTH the sync (inline) and async (submitted) retain
|
||||
paths.
|
||||
|
||||
Two distinct bugs hid behind the same symptom — ingesting a large plain-text
|
||||
document dropped most of its body (and any fact past the first slice). Both only
|
||||
trigger when an oversized single item is split into sequential sub-batches whose
|
||||
*slices each re-chunk into several extraction chunks* (the default config: batch
|
||||
tokens 10k → ~30k-char slices, re-chunked at 3k → ~10 chunks/slice):
|
||||
|
||||
1. chunk_index offset (sync + async). retain_batch_async advanced the per-document
|
||||
chunk_index cursor by re-chunking ``item["content"]`` AFTER the orchestrator
|
||||
had consumed (popped) it — ``chunk_text("")`` returns ``[""]`` (count 1), so
|
||||
the cursor moved by 1 per sub-batch instead of by the real chunk count. Later
|
||||
slices restarted ~1 slot in, colliding ``chunk_id = {bank}_{doc}_{index}`` and
|
||||
overwriting earlier chunks via upsert.
|
||||
|
||||
2. whole-document recovery skip (async only). All sub-batches of one submitted
|
||||
operation share one ``operation_id``; the first slice stamps the document into
|
||||
``result_metadata.facts_committed_document_ids``. The crash-recovery fast-path
|
||||
then saw every later slice's document already "committed" and skipped
|
||||
extraction entirely, so only the first slice survived.
|
||||
|
||||
The existing #1888 coverage tests use ``RETAIN_BATCH_TOKENS=100`` (a ~300-char
|
||||
budget, under the chunk size) so every slice collapses to ONE chunk — which masks
|
||||
both bugs (offset-by-1 happens to equal the real count, and a 1-chunk doc isn't
|
||||
re-sliced). These tests size the body so each slice fans out to ~6 chunks, with
|
||||
globally-unique tokens so no chunk-hash dedup hides a dropped slice, and assert
|
||||
full coverage + contiguous indices + a needle planted in a late slice.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
# The async test submits via submit_async_retain, which inserts parent/child rows
|
||||
# into async_operations. test_worker.py drives its own WorkerPoller.claim_batch()
|
||||
# against the same pool, so running the two files on different xdist workers lets
|
||||
# them steal each other's pending rows. Share the "worker_tests" group so they
|
||||
# serialize on the same xdist process (matches test_async_batch_retain.py).
|
||||
pytestmark = pytest.mark.xdist_group("worker_tests")
|
||||
|
||||
# Planted in a late paragraph so it lands in a late sub-batch slice — the first
|
||||
# thing either bug drops (mirrors the field-reported "165 commits" fact that
|
||||
# vanished on the async path). A single no-space token so it can't straddle a
|
||||
# chunk boundary (a multi-word phrase can split across two chunks at this test's
|
||||
# small 500-char chunk size and read as "dropped" when it wasn't).
|
||||
NEEDLE = "NEEDLE_165_COMMITS_MERGED_INTO_THE_MAIN_BRANCH"
|
||||
|
||||
|
||||
def _ts() -> float:
|
||||
return datetime.now(timezone.utc).timestamp()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _multichunk_split_env(monkeypatch):
|
||||
# Small extraction chunks (500 chars) with a batch-token budget whose char
|
||||
# budget (700 * 3 = 2100) spans several chunks, so each oversized sub-batch
|
||||
# slice fans out to ~6 extraction chunks. Skip consolidation/observations to
|
||||
# keep the test fast and deterministic.
|
||||
monkeypatch.setenv("HINDSIGHT_API_RETAIN_CHUNK_SIZE", "500")
|
||||
monkeypatch.setenv("HINDSIGHT_API_RETAIN_BATCH_TOKENS", "700")
|
||||
monkeypatch.setenv("HINDSIGHT_API_ENABLE_AUTO_CONSOLIDATION", "false")
|
||||
monkeypatch.setenv("HINDSIGHT_API_ENABLE_OBSERVATIONS", "false")
|
||||
clear_config_cache()
|
||||
yield
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def _make_body(paragraphs: int = 24, needle_at: int = 20) -> str:
|
||||
"""Plain-text transcript whose every token is unique across the whole body,
|
||||
so no two extraction chunks can hash-collide (a real content-hash collision
|
||||
would legitimately dedup and mask a dropped slice). The needle sits in a late
|
||||
paragraph."""
|
||||
lines = []
|
||||
for i in range(paragraphs):
|
||||
toks = " ".join(f"w{i:03d}t{j:03d}" for j in range(60))
|
||||
if i == needle_at:
|
||||
lines.append(f"[Turn {i}] Assistant: {NEEDLE} fact {toks}")
|
||||
else:
|
||||
lines.append(f"[Turn {i}] Assistant: progress {i}: {toks}")
|
||||
return "\n\n".join(lines)
|
||||
|
||||
|
||||
async def _chunk_coverage(memory, bank_id, document_id, request_context):
|
||||
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc is not None
|
||||
original_len = len(doc["original_text"])
|
||||
chunks = await memory.list_document_chunks(bank_id, document_id, limit=10000, request_context=request_context)
|
||||
items = chunks["items"]
|
||||
sum_chunk_text = sum(len(c["chunk_text"]) for c in items)
|
||||
indices = sorted(c["chunk_index"] for c in items)
|
||||
needle_present = any(NEEDLE in c["chunk_text"] for c in items)
|
||||
return original_len, sum_chunk_text, indices, needle_present
|
||||
|
||||
|
||||
def _assert_full_coverage(label, original_len, sum_chunk_text, indices, needle_present):
|
||||
# Sanity: the body must actually fan out to many chunks across several
|
||||
# multi-chunk slices, or the test wouldn't exercise the bug at all.
|
||||
assert len(indices) >= 16, f"{label}: only {len(indices)} chunks — body too small to exercise multi-chunk slices"
|
||||
assert sum_chunk_text >= original_len * 0.9, (
|
||||
f"{label}: chunks cover only {sum_chunk_text}/{original_len} chars "
|
||||
f"(~{100 * sum_chunk_text // original_len}%) — a sub-batch slice was overwritten or skipped"
|
||||
)
|
||||
assert indices == list(range(len(indices))), (
|
||||
f"{label}: chunk_index sequence is not contiguous: {indices} — sub-batch slices collided on chunk_id"
|
||||
)
|
||||
assert needle_present, f"{label}: the late-slice needle fact was dropped (offset collision or recovery skip)"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_inline_multichunk_subbatch_coverage(memory, request_context):
|
||||
"""Sync inline path (retain_batch_async): an oversized doc whose slices each
|
||||
span several extraction chunks must keep full coverage (offset bug)."""
|
||||
bank_id = f"test_multichunk_sync_{_ts()}"
|
||||
document_id = "doc-multichunk-sync"
|
||||
try:
|
||||
body = _make_body()
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": body, "context": "big doc", "document_id": document_id}],
|
||||
request_context=request_context,
|
||||
)
|
||||
cov = await _chunk_coverage(memory, bank_id, document_id, request_context)
|
||||
_assert_full_coverage("sync", *cov)
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(600)
|
||||
async def test_async_submit_multichunk_subbatch_coverage(memory, request_context):
|
||||
"""Async submit path (submit_async_retain → child op → worker): the same
|
||||
oversized doc must keep full coverage too. Exercises both the offset bug and
|
||||
the shared-operation_id whole-document recovery skip."""
|
||||
import asyncio
|
||||
|
||||
bank_id = f"test_multichunk_async_{_ts()}"
|
||||
document_id = "doc-multichunk-async"
|
||||
try:
|
||||
body = _make_body()
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": body, "context": "big doc", "document_id": document_id}],
|
||||
request_context=request_context,
|
||||
)
|
||||
operation_id = result["operation_id"]
|
||||
|
||||
# SyncTaskBackend (test backend) drains children inline; wait for the
|
||||
# parent to reach a terminal state before reading chunks.
|
||||
status = None
|
||||
for _ in range(600):
|
||||
status = await memory.get_operation_status(
|
||||
bank_id=bank_id, operation_id=operation_id, request_context=request_context
|
||||
)
|
||||
if status["status"] in ("completed", "failed"):
|
||||
break
|
||||
await asyncio.sleep(0.1)
|
||||
assert status is not None and status["status"] == "completed", (
|
||||
f"async retain did not complete: {status['status'] if status else 'no status'}"
|
||||
)
|
||||
|
||||
cov = await _chunk_coverage(memory, bank_id, document_id, request_context)
|
||||
_assert_full_coverage("async", *cov)
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -23,6 +23,7 @@ from hindsight_api.engine.search.tags import (
|
||||
TagGroupNot,
|
||||
TagGroupOr,
|
||||
build_tag_groups_where_clause,
|
||||
build_tags_where_clause,
|
||||
build_tags_where_clause_simple,
|
||||
filter_results_by_tag_groups,
|
||||
filter_results_by_tags,
|
||||
@@ -136,6 +137,46 @@ class TestTagsWhereClauseBuilder:
|
||||
assert "@>" in result
|
||||
assert "<@" in result
|
||||
|
||||
# ---- Test "exact" mode with the empty scope ([]) = untagged/global only ----
|
||||
|
||||
def test_tags_match_exact_empty_list_matches_untagged_only(self):
|
||||
"""match='exact' with [] filters to untagged rows only (no bind param)."""
|
||||
result = build_tags_where_clause_simple([], 5, match="exact")
|
||||
assert "IS NULL" in result
|
||||
assert "= '{}'" in result
|
||||
# Untagged-only is param-free: callers append no tags param for an empty list.
|
||||
assert "$5" not in result
|
||||
# Must not use set-equality operators (which would need a bound scope).
|
||||
assert "@>" not in result
|
||||
assert "<@" not in result
|
||||
|
||||
def test_tags_match_exact_empty_list_with_table_alias(self):
|
||||
"""Empty-scope exact clause respects the table alias."""
|
||||
result = build_tags_where_clause_simple([], 5, table_alias="mu.", match="exact")
|
||||
assert "mu.tags IS NULL" in result
|
||||
assert "mu.tags = '{}'" in result
|
||||
|
||||
def test_tags_match_exact_none_matches_untagged_only(self):
|
||||
"""match='exact' with None (no tags) selects the global scope, like the graph endpoint."""
|
||||
result = build_tags_where_clause_simple(None, 5, match="exact")
|
||||
assert "IS NULL" in result
|
||||
assert "= '{}'" in result
|
||||
assert "$5" not in result
|
||||
|
||||
def test_tags_match_any_empty_list_still_no_filter(self):
|
||||
"""Empty list only filters under 'exact'; other modes treat [] as no filter."""
|
||||
assert build_tags_where_clause_simple([], 5, match="any") == ""
|
||||
assert build_tags_where_clause_simple([], 5, match="any_strict") == ""
|
||||
|
||||
@pytest.mark.parametrize("tags", [None, []])
|
||||
def test_tags_where_clause_exact_empty_scope_keeps_param_offset(self, tags):
|
||||
"""The parameterized builder must not consume a bind index for the empty scope,
|
||||
so following clauses stay aligned with their params."""
|
||||
clause, params, next_offset = build_tags_where_clause(tags, param_offset=4, match="exact")
|
||||
assert clause == "AND (tags IS NULL OR tags = '{}')"
|
||||
assert params == []
|
||||
assert next_offset == 4
|
||||
|
||||
# ---- Test table alias with all modes ----
|
||||
|
||||
def test_tags_match_any_with_table_alias(self):
|
||||
@@ -255,6 +296,20 @@ class TestFilterResultsByTags:
|
||||
assert len(filtered) == 1
|
||||
assert filtered[0].tags == ["a"]
|
||||
|
||||
def test_exact_mode_empty_scope_matches_untagged_only(self):
|
||||
"""'exact' mode with [] should keep only untagged results (NULL or empty)."""
|
||||
results = [MockResult(["a"]), MockResult(["a", "b"]), MockResult(None), MockResult([])]
|
||||
filtered = filter_results_by_tags(results, [], match="exact")
|
||||
assert len(filtered) == 2
|
||||
assert all(not r.tags for r in filtered)
|
||||
|
||||
def test_exact_mode_none_matches_untagged_only(self):
|
||||
"""'exact' mode with None (no tags) selects the global scope (untagged only)."""
|
||||
results = [MockResult(["a"]), MockResult(None), MockResult([])]
|
||||
filtered = filter_results_by_tags(results, None, match="exact")
|
||||
assert len(filtered) == 2
|
||||
assert all(not r.tags for r in filtered)
|
||||
|
||||
def test_all_mode_includes_untagged(self):
|
||||
"""'all' mode should include untagged results."""
|
||||
results = [MockResult(["a", "b"]), MockResult(None), MockResult([])]
|
||||
@@ -502,6 +557,16 @@ class TestBuildTagGroupsWhereClause:
|
||||
assert len(params) == 2
|
||||
assert next_offset == 3
|
||||
|
||||
def test_exact_leaf_empty_scope_matches_untagged_only(self):
|
||||
"""An exact leaf with [] becomes an untagged-only clause with no bind param."""
|
||||
groups = [TagGroupLeaf(tags=[], match="exact")]
|
||||
clause, params, next_offset = build_tag_groups_where_clause(groups, 5)
|
||||
assert "IS NULL" in clause
|
||||
assert "= '{}'" in clause
|
||||
assert "$5" not in clause # param-free
|
||||
assert params == []
|
||||
assert next_offset == 5 # offset unchanged — no param consumed
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Unit Tests for filter_results_by_tag_groups (Python-side)
|
||||
@@ -531,6 +596,14 @@ class TestFilterResultsByTagGroups:
|
||||
assert len(filtered) == 1
|
||||
assert filtered[0].tags == ["step:5"]
|
||||
|
||||
def test_exact_leaf_empty_scope_matches_untagged_only(self):
|
||||
"""An exact leaf with [] keeps only untagged results (matches SQL builder)."""
|
||||
groups = [TagGroupLeaf(tags=[], match="exact")]
|
||||
results = [MockResult(["a"]), MockResult(["a", "b"]), MockResult(None), MockResult([])]
|
||||
filtered = filter_results_by_tag_groups(results, groups)
|
||||
assert len(filtered) == 2
|
||||
assert all(not r.tags for r in filtered)
|
||||
|
||||
def test_single_leaf_all_strict_matches_superset(self):
|
||||
"""Single all_strict leaf matches results that contain all tags."""
|
||||
groups = [TagGroupLeaf(tags=["user:alice", "step:5"], match="all_strict")]
|
||||
@@ -904,6 +977,37 @@ async def test_recall_with_empty_tags_returns_all(api_client, test_bank_id):
|
||||
assert any("Rachel" in t for t in texts), "Should find Rachel"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_empty_tags_exact_returns_untagged_only(api_client, test_bank_id):
|
||||
"""tags=[] with tags_match='exact' returns only untagged/global memories."""
|
||||
# One untagged (global) memory and one tagged memory.
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{"content": "Sam studies astronomy."}, # no tags -> global scope
|
||||
{"content": "Tina studies geology.", "tags": ["user_tina"]},
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# exact match on the empty scope -> only the untagged memory.
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories/recall",
|
||||
json={"query": "Who studies what?", "budget": "low", "tags": [], "tags_match": "exact"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
results = response.json()["results"]
|
||||
|
||||
texts = [r["text"] for r in results]
|
||||
assert any("Sam" in t for t in texts), "Should find the untagged memory"
|
||||
assert not any("Tina" in t for t in texts), "Should NOT find the tagged memory"
|
||||
# Every returned memory must be untagged.
|
||||
for r in results:
|
||||
assert not r.get("tags"), f"Expected untagged result, got tags={r.get('tags')}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multi_user_agent_visibility(api_client):
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
"""Tests for cached / thoughts token propagation through TokenUsage,
|
||||
LLMToolCallResult, TokenUsageSummary, and RetainResult.
|
||||
|
||||
The Gemini 2.5+ family (and any future provider with prompt caching +
|
||||
reasoning tokens) reports four distinct token counts on every response:
|
||||
prompt, candidates (visible output), cached_content, and thoughts. The
|
||||
last two are billed separately by the provider but were previously not
|
||||
threaded through to downstream return contexts, so application-layer
|
||||
metering had no way to attribute prompt-cache hit rate or reasoning cost
|
||||
per operation.
|
||||
|
||||
These tests pin the propagation: when a provider populates cached or
|
||||
thoughts on the way out, every accumulator and aggregate type carries
|
||||
the value through unchanged.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.reflect.agent import _generate_structured_output
|
||||
from hindsight_api.engine.reflect.models import StructuredOutputResult, TokenUsageSummary
|
||||
from hindsight_api.engine.response_models import LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.extensions.operation_validator import RetainResult
|
||||
|
||||
|
||||
def test_token_usage_carries_cached_and_thoughts():
|
||||
"""TokenUsage defaults both new fields to 0 and accepts non-zero values."""
|
||||
u = TokenUsage(input_tokens=1500, output_tokens=500, total_tokens=2000)
|
||||
assert u.cached_tokens == 0
|
||||
assert u.thoughts_tokens == 0
|
||||
|
||||
u = TokenUsage(
|
||||
input_tokens=1500,
|
||||
output_tokens=500,
|
||||
total_tokens=2000,
|
||||
cached_tokens=200,
|
||||
thoughts_tokens=80,
|
||||
)
|
||||
assert u.cached_tokens == 200
|
||||
assert u.thoughts_tokens == 80
|
||||
|
||||
|
||||
def test_token_usage_aggregates_thoughts_tokens():
|
||||
"""TokenUsage.__add__ sums thoughts_tokens alongside the existing fields.
|
||||
|
||||
Multi-iteration agentic loops accumulate per-call usage via ``+``. If
|
||||
thoughts_tokens isn't summed, the per-op total undercounts reasoning
|
||||
spend by a factor of N (the number of LLM sub-calls).
|
||||
"""
|
||||
a = TokenUsage(input_tokens=10, output_tokens=5, total_tokens=15, cached_tokens=2, thoughts_tokens=7)
|
||||
b = TokenUsage(input_tokens=20, output_tokens=8, total_tokens=28, cached_tokens=3, thoughts_tokens=11)
|
||||
c = a + b
|
||||
assert c.input_tokens == 30
|
||||
assert c.output_tokens == 13
|
||||
assert c.total_tokens == 43
|
||||
assert c.cached_tokens == 5
|
||||
assert c.thoughts_tokens == 18
|
||||
|
||||
|
||||
def test_llm_tool_call_result_carries_cached_and_thoughts():
|
||||
"""call_with_tools returns LLMToolCallResult — both new fields default to 0
|
||||
and accept non-zero values from the provider."""
|
||||
r = LLMToolCallResult(content="ok", input_tokens=1234, output_tokens=56)
|
||||
assert r.cached_tokens == 0
|
||||
assert r.thoughts_tokens == 0
|
||||
|
||||
r = LLMToolCallResult(
|
||||
content="ok",
|
||||
input_tokens=1234,
|
||||
output_tokens=56,
|
||||
cached_tokens=200,
|
||||
thoughts_tokens=78,
|
||||
)
|
||||
assert r.cached_tokens == 200
|
||||
assert r.thoughts_tokens == 78
|
||||
|
||||
|
||||
def test_token_usage_summary_carries_cached_and_thoughts():
|
||||
"""TokenUsageSummary is what reflect agent returns to its caller — needs
|
||||
to propagate the aggregate so per-op cost attribution works."""
|
||||
s = TokenUsageSummary(
|
||||
input_tokens=10000,
|
||||
output_tokens=200,
|
||||
total_tokens=10200,
|
||||
cached_tokens=3000,
|
||||
thoughts_tokens=150,
|
||||
)
|
||||
assert s.cached_tokens == 3000
|
||||
assert s.thoughts_tokens == 150
|
||||
|
||||
|
||||
def test_token_usage_summary_defaults_cached_and_thoughts_to_zero():
|
||||
"""Defaults preserve backward compatibility for callers built before the
|
||||
fields existed."""
|
||||
s = TokenUsageSummary(input_tokens=100, output_tokens=50, total_tokens=150)
|
||||
assert s.cached_tokens == 0
|
||||
assert s.thoughts_tokens == 0
|
||||
|
||||
|
||||
def test_retain_result_carries_cached_input_and_thoughts():
|
||||
"""RetainResult is the contract between the engine and any metering
|
||||
extension. The two new fields are optional (None) so older extensions
|
||||
that don't read them are unaffected; engines that DO populate them get
|
||||
end-to-end attribution into the metering hook."""
|
||||
|
||||
class _Ctx:
|
||||
pass
|
||||
|
||||
r = RetainResult(
|
||||
bank_id="b",
|
||||
contents=[],
|
||||
request_context=_Ctx(),
|
||||
document_id=None,
|
||||
fact_type_override=None,
|
||||
unit_ids=[],
|
||||
llm_input_tokens=1000,
|
||||
llm_output_tokens=50,
|
||||
llm_total_tokens=1050,
|
||||
llm_cached_input_tokens=300,
|
||||
llm_thoughts_tokens=25,
|
||||
)
|
||||
assert r.llm_cached_input_tokens == 300
|
||||
assert r.llm_thoughts_tokens == 25
|
||||
|
||||
# Defaults stay None for engines that don't surface the data, so
|
||||
# downstream extensions can use ``or 0`` without breaking on a
|
||||
# core-only build.
|
||||
r2 = RetainResult(
|
||||
bank_id="b",
|
||||
contents=[],
|
||||
request_context=_Ctx(),
|
||||
document_id=None,
|
||||
fact_type_override=None,
|
||||
unit_ids=[],
|
||||
)
|
||||
assert r2.llm_cached_input_tokens is None
|
||||
assert r2.llm_thoughts_tokens is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_structured_output_returns_dataclass_on_no_fields():
|
||||
"""_generate_structured_output returns a StructuredOutputResult, not a tuple.
|
||||
|
||||
Regression guard: the function and all six call sites must agree on a single
|
||||
return type. A previous tuple-based contract drifted out of sync (the failure
|
||||
branch returned 3 values while callers unpacked 5), which would crash reflect
|
||||
with a ValueError on any structured-output failure. An empty schema exercises
|
||||
the no-LLM-call branch deterministically.
|
||||
"""
|
||||
result = await _generate_structured_output(
|
||||
answer="anything",
|
||||
response_schema={},
|
||||
llm_config=None,
|
||||
reflect_id="test",
|
||||
)
|
||||
assert isinstance(result, StructuredOutputResult)
|
||||
assert result.structured_output is None
|
||||
assert result.input_tokens == 0
|
||||
assert result.output_tokens == 0
|
||||
assert result.cached_tokens == 0
|
||||
assert result.thoughts_tokens == 0
|
||||
@@ -13,6 +13,7 @@ Tests cover:
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
@@ -70,6 +71,99 @@ async def clean_operations(pool):
|
||||
)
|
||||
|
||||
|
||||
def test_metric_operation_label_normalises_retain_variants():
|
||||
"""Worker completion metrics collapse retain variants onto operation="retain"
|
||||
so they share the API path's series; other types pass through unchanged."""
|
||||
from hindsight_api.worker.poller import _metric_operation_label
|
||||
|
||||
assert _metric_operation_label("retain") == "retain"
|
||||
assert _metric_operation_label("batch_retain") == "retain"
|
||||
assert _metric_operation_label("file_convert_retain") == "retain"
|
||||
assert _metric_operation_label("consolidation") == "consolidation"
|
||||
assert _metric_operation_label("reflect") == "reflect"
|
||||
assert _metric_operation_label(None) == "unknown"
|
||||
|
||||
|
||||
class TestWorkerOperationMetrics:
|
||||
"""_execute_task_inner emits operation metrics on terminal outcomes only (no DB)."""
|
||||
|
||||
def _make_poller(self, executor):
|
||||
from hindsight_api.worker import WorkerPoller
|
||||
|
||||
poller = WorkerPoller(backend=MagicMock(), worker_id="w-test", executor=executor)
|
||||
# Stub terminal-state handlers so _execute_task_inner never touches the DB.
|
||||
poller._mark_failed = AsyncMock()
|
||||
poller._defer_operation = AsyncMock()
|
||||
poller._schedule_retry = AsyncMock()
|
||||
return poller
|
||||
|
||||
async def _run(self, executor, task_type="batch_retain"):
|
||||
from hindsight_api.worker.poller import ClaimedTask
|
||||
|
||||
poller = self._make_poller(executor)
|
||||
task = ClaimedTask(
|
||||
operation_id=str(uuid.uuid4()),
|
||||
task_dict={"type": task_type, "operation_type": task_type, "bank_id": "bank-1"},
|
||||
schema=None,
|
||||
)
|
||||
collector = MagicMock()
|
||||
with patch("hindsight_api.worker.poller.get_metrics_collector", return_value=collector):
|
||||
await poller._execute_task_inner(task)
|
||||
return collector
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_executor_returning_normally_records_success(self):
|
||||
"""Success is inferred from the executor returning without raising to the
|
||||
poller. This deliberately includes deterministic failures that
|
||||
memory_engine.execute_task handles itself and returns from normally
|
||||
(file_convert_retain, non-retryable errors) — at the poller boundary they
|
||||
are indistinguishable from a clean completion, so they also record
|
||||
success=true. The worker counter is therefore a completion-throughput
|
||||
signal; authoritative failure visibility comes from the
|
||||
hindsight_async_operations{status="failed"} gauge, which reads each
|
||||
operation's final DB status.
|
||||
"""
|
||||
collector = await self._run(AsyncMock()) # executor returns normally
|
||||
collector.record_operation_result.assert_called_once()
|
||||
call = collector.record_operation_result.call_args
|
||||
assert call.args[0] == "retain" # batch_retain normalised
|
||||
assert call.kwargs["success"] is True
|
||||
assert call.kwargs["source"] == "worker"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_records_failure(self):
|
||||
async def boom(_):
|
||||
raise RuntimeError("kaboom")
|
||||
|
||||
collector = await self._run(boom)
|
||||
collector.record_operation_result.assert_called_once()
|
||||
assert collector.record_operation_result.call_args.kwargs["success"] is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deferral_not_counted(self):
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from hindsight_api.worker.exceptions import DeferOperation
|
||||
|
||||
async def defer(_):
|
||||
raise DeferOperation(exec_date=datetime.now(timezone.utc), reason="later")
|
||||
|
||||
collector = await self._run(defer)
|
||||
collector.record_operation_result.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_not_counted(self):
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from hindsight_api.worker.exceptions import RetryTaskAt
|
||||
|
||||
async def retry(_):
|
||||
raise RetryTaskAt(retry_at=datetime.now(timezone.utc), message="transient")
|
||||
|
||||
collector = await self._run(retry)
|
||||
collector.record_operation_result.assert_not_called()
|
||||
|
||||
|
||||
def test_all_operation_types_have_slot_reservation_config():
|
||||
"""Every operation_type used in memory_engine must be listed in
|
||||
WORKER_SLOT_RESERVATION_TYPES so it can be reserved via env var.
|
||||
@@ -832,6 +926,34 @@ class TestWorkerPoller:
|
||||
# Defensive: confirm it wasn't a RetryTaskAt masquerading as Defer.
|
||||
assert not isinstance(exc_info.value, RetryTaskAt)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memory_engine_provider_quota_reset_becomes_defer_operation(self, memory, monkeypatch):
|
||||
"""Provider quota windows should park worker tasks until the reset time."""
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
from hindsight_api.engine.llm_interface import ProviderRateLimitResetError
|
||||
from hindsight_api.worker.exceptions import DeferOperation, RetryTaskAt
|
||||
|
||||
retry_at = (datetime.now(UTC) + timedelta(hours=5)).replace(microsecond=0)
|
||||
|
||||
async def quota_limited_retain(_task_dict: object) -> None:
|
||||
raise ProviderRateLimitResetError(retry_at=retry_at, message="quota resets later")
|
||||
|
||||
monkeypatch.setattr(memory, "_handle_batch_retain", quota_limited_retain)
|
||||
|
||||
with pytest.raises(DeferOperation) as exc_info:
|
||||
await memory.execute_task(
|
||||
{
|
||||
"type": "batch_retain",
|
||||
"bank_id": "test-provider-quota-defer",
|
||||
"contents": [{"content": "x"}],
|
||||
}
|
||||
)
|
||||
|
||||
assert exc_info.value.exec_date == retry_at
|
||||
assert exc_info.value.reason == "quota resets later"
|
||||
assert not isinstance(exc_info.value, RetryTaskAt)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claim_batch_skips_consolidation_when_same_bank_processing(self, pool, backend, clean_operations):
|
||||
"""Test that pending consolidation is skipped if same bank has one processing."""
|
||||
@@ -947,9 +1069,11 @@ class TestWorkerPoller:
|
||||
|
||||
claimed = await poller.claim_batch()
|
||||
|
||||
# Should claim the retain task (non-consolidation tasks are unaffected)
|
||||
assert len(claimed) == 1
|
||||
assert claimed[0].operation_id == str(retain_op_id)
|
||||
# Should claim the retain task (non-consolidation tasks are unaffected).
|
||||
# Filter to our bank — parallel tests may contribute other claims.
|
||||
my_claims = [c for c in claimed if c.task_dict.get("bank_id") == bank_id]
|
||||
assert len(my_claims) == 1, f"Expected 1 claim for our bank, got {len(my_claims)}"
|
||||
assert my_claims[0].operation_id == str(retain_op_id)
|
||||
|
||||
|
||||
class TestWorkerRecovery:
|
||||
@@ -1586,7 +1710,9 @@ class TestDynamicTenantDiscovery:
|
||||
# First claim_batch should call list_tenants
|
||||
claimed1 = await poller.claim_batch()
|
||||
assert mock_extension.list_tenants_calls == 1
|
||||
assert len(claimed1) == 2
|
||||
# Filter to our bank — parallel tests may contribute other claims.
|
||||
my_claims1 = [c for c in claimed1 if c.task_dict.get("bank_id") == bank_id]
|
||||
assert len(my_claims1) == 2, f"Expected 2 claims for our bank, got {len(my_claims1)}"
|
||||
|
||||
# Add more tasks
|
||||
for i in range(2):
|
||||
@@ -1605,7 +1731,8 @@ class TestDynamicTenantDiscovery:
|
||||
# Second claim_batch should call list_tenants again
|
||||
claimed2 = await poller.claim_batch()
|
||||
assert mock_extension.list_tenants_calls == 2
|
||||
assert len(claimed2) == 2
|
||||
my_claims2 = [c for c in claimed2 if c.task_dict.get("bank_id") == bank_id]
|
||||
assert len(my_claims2) == 2, f"Expected 2 claims for our bank, got {len(my_claims2)}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poller_picks_up_new_tenants_without_restart(self, pool, backend, clean_operations):
|
||||
@@ -1650,10 +1777,12 @@ class TestDynamicTenantDiscovery:
|
||||
tenant_extension=dynamic_extension,
|
||||
)
|
||||
|
||||
# First poll - only public schema
|
||||
# First poll - only public schema. Filter to our bank — parallel tests
|
||||
# may contribute other claims.
|
||||
claimed1 = await poller.claim_batch()
|
||||
assert len(claimed1) == 1
|
||||
assert claimed1[0].schema is None # public is represented as None
|
||||
my_claims1 = [c for c in claimed1 if c.task_dict.get("bank_id") == bank_id]
|
||||
assert len(my_claims1) == 1, f"Expected 1 claim for our bank, got {len(my_claims1)}"
|
||||
assert my_claims1[0].schema is None # public is represented as None
|
||||
assert dynamic_extension.list_tenants_calls == 1
|
||||
|
||||
# Simulate tenant list changing (but we won't add a non-existent schema)
|
||||
@@ -1675,12 +1804,14 @@ class TestDynamicTenantDiscovery:
|
||||
|
||||
# Second poll - list_tenants should be called again
|
||||
claimed2 = await poller.claim_batch()
|
||||
assert len(claimed2) == 1
|
||||
my_claims2 = [c for c in claimed2 if c.task_dict.get("bank_id") == bank_id]
|
||||
assert len(my_claims2) == 1, f"Expected 1 claim for our bank, got {len(my_claims2)}"
|
||||
assert dynamic_extension.list_tenants_calls == 2 # Called again on second poll
|
||||
|
||||
# Third poll with no tasks - still calls list_tenants
|
||||
claimed3 = await poller.claim_batch()
|
||||
assert len(claimed3) == 0
|
||||
my_claims3 = [c for c in claimed3 if c.task_dict.get("bank_id") == bank_id]
|
||||
assert len(my_claims3) == 0, f"Expected 0 claims for our bank, got {len(my_claims3)}"
|
||||
assert dynamic_extension.list_tenants_calls == 3 # Called again even with no tasks
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api"
|
||||
version = "0.8.2"
|
||||
version = "0.8.3"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim[all]==0.8.2",
|
||||
"hindsight-api-slim[all]==0.8.3",
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "hindsight-cli"
|
||||
version = "0.8.2"
|
||||
version = "0.8.3"
|
||||
edition = "2021"
|
||||
authors = ["Hindsight Team"]
|
||||
description = "A beautiful CLI for Hindsight - semantic memory system"
|
||||
|
||||
@@ -425,8 +425,8 @@ impl ApiClient {
|
||||
.client
|
||||
.list_documents(
|
||||
agent_id,
|
||||
limit.map(|l| l as i64),
|
||||
offset.map(|o| o as i64),
|
||||
limit.map(|l| l as u64),
|
||||
offset.map(|o| o as u64),
|
||||
q,
|
||||
None,
|
||||
None,
|
||||
@@ -524,8 +524,8 @@ impl ApiClient {
|
||||
bank_id,
|
||||
None, // consolidation_state
|
||||
None, // document_id
|
||||
limit,
|
||||
offset,
|
||||
limit.map(|l| l as u64),
|
||||
offset.map(|o| o as u64),
|
||||
q,
|
||||
None, // state
|
||||
type_filter,
|
||||
@@ -546,7 +546,12 @@ impl ApiClient {
|
||||
self.runtime.block_on(async {
|
||||
let response = self
|
||||
.client
|
||||
.list_entities(bank_id, limit, offset, None)
|
||||
.list_entities(
|
||||
bank_id,
|
||||
limit.map(|l| l as u64),
|
||||
offset.map(|o| o as u64),
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
Ok(response.into_inner())
|
||||
})
|
||||
@@ -664,7 +669,17 @@ impl ApiClient {
|
||||
self.runtime.block_on(async {
|
||||
let response = self
|
||||
.client
|
||||
.get_graph(bank_id, None, None, limit, None, None, None, type_filter, None)
|
||||
.get_graph(
|
||||
bank_id,
|
||||
None,
|
||||
None,
|
||||
limit.map(|l| l as u64),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
type_filter,
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
Ok(response.into_inner())
|
||||
})
|
||||
@@ -726,7 +741,14 @@ impl ApiClient {
|
||||
self.runtime.block_on(async {
|
||||
let response = self
|
||||
.client
|
||||
.list_tags(bank_id, limit, offset, q, None, None)
|
||||
.list_tags(
|
||||
bank_id,
|
||||
limit.map(|l| l as u64),
|
||||
offset.map(|o| o as u64),
|
||||
q,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
Ok(response.into_inner())
|
||||
})
|
||||
|
||||
@@ -340,6 +340,7 @@ impl App {
|
||||
max_tokens: query_max_tokens,
|
||||
trace: false,
|
||||
query_timestamp: None,
|
||||
prefer_observations: false,
|
||||
include: None,
|
||||
tags: None,
|
||||
tags_match: TagsMatch::Any,
|
||||
|
||||
@@ -277,6 +277,7 @@ pub fn recall(
|
||||
tags: Vec<String>,
|
||||
tags_match: Option<String>,
|
||||
query_timestamp: Option<String>,
|
||||
prefer_observations: bool,
|
||||
verbose: bool,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
@@ -310,6 +311,7 @@ pub fn recall(
|
||||
max_tokens,
|
||||
trace,
|
||||
query_timestamp,
|
||||
prefer_observations,
|
||||
include,
|
||||
tags: if tags.is_empty() { None } else { Some(tags) },
|
||||
tags_match: parse_tags_match(&tags_match),
|
||||
|
||||
@@ -532,6 +532,10 @@ enum MemoryCommands {
|
||||
/// Reference timestamp for recall (ISO 8601, e.g. 2023-05-30T23:40:00)
|
||||
#[arg(long)]
|
||||
query_timestamp: Option<String>,
|
||||
|
||||
/// Prefer observations: drop raw facts a returned observation was consolidated from (no effect unless observation + a raw type are both recalled)
|
||||
#[arg(long)]
|
||||
prefer_observations: bool,
|
||||
},
|
||||
|
||||
/// Generate answers using bank identity (reflect/reasoning)
|
||||
@@ -1401,6 +1405,7 @@ fn run() -> Result<()> {
|
||||
tags,
|
||||
tags_match,
|
||||
query_timestamp,
|
||||
prefer_observations,
|
||||
} => commands::memory::recall(
|
||||
&client,
|
||||
&bank_id,
|
||||
@@ -1414,6 +1419,7 @@ fn run() -> Result<()> {
|
||||
tags,
|
||||
tags_match,
|
||||
query_timestamp,
|
||||
prefer_observations,
|
||||
verbose,
|
||||
output_format,
|
||||
),
|
||||
|
||||
@@ -7,7 +7,7 @@ info:
|
||||
name: Apache 2.0
|
||||
url: https://www.apache.org/licenses/LICENSE-2.0.html
|
||||
title: Hindsight HTTP API
|
||||
version: 0.8.2
|
||||
version: 0.8.3
|
||||
servers:
|
||||
- url: /
|
||||
paths:
|
||||
@@ -80,6 +80,7 @@ paths:
|
||||
required: false
|
||||
schema:
|
||||
default: 1000
|
||||
minimum: 0
|
||||
title: Limit
|
||||
type: integer
|
||||
style: form
|
||||
@@ -212,6 +213,7 @@ paths:
|
||||
required: false
|
||||
schema:
|
||||
default: 100
|
||||
minimum: 0
|
||||
title: Limit
|
||||
type: integer
|
||||
style: form
|
||||
@@ -221,6 +223,7 @@ paths:
|
||||
required: false
|
||||
schema:
|
||||
default: 0
|
||||
minimum: 0
|
||||
title: Offset
|
||||
type: integer
|
||||
style: form
|
||||
@@ -732,6 +735,7 @@ paths:
|
||||
schema:
|
||||
default: 100
|
||||
description: Maximum number of entities to return
|
||||
minimum: 0
|
||||
title: Limit
|
||||
type: integer
|
||||
style: form
|
||||
@@ -743,6 +747,7 @@ paths:
|
||||
schema:
|
||||
default: 0
|
||||
description: Offset for pagination
|
||||
minimum: 0
|
||||
title: Offset
|
||||
type: integer
|
||||
style: form
|
||||
@@ -792,6 +797,7 @@ paths:
|
||||
schema:
|
||||
default: 1000
|
||||
description: Maximum number of co-occurrence edges to return
|
||||
minimum: 0
|
||||
title: Limit
|
||||
type: integer
|
||||
style: form
|
||||
@@ -1690,6 +1696,7 @@ paths:
|
||||
required: false
|
||||
schema:
|
||||
default: 100
|
||||
minimum: 0
|
||||
title: Limit
|
||||
type: integer
|
||||
style: form
|
||||
@@ -1699,6 +1706,7 @@ paths:
|
||||
required: false
|
||||
schema:
|
||||
default: 0
|
||||
minimum: 0
|
||||
title: Offset
|
||||
type: integer
|
||||
style: form
|
||||
@@ -2046,6 +2054,7 @@ paths:
|
||||
schema:
|
||||
default: 100
|
||||
description: Maximum number of tags to return
|
||||
minimum: 0
|
||||
title: Limit
|
||||
type: integer
|
||||
style: form
|
||||
@@ -2057,6 +2066,7 @@ paths:
|
||||
schema:
|
||||
default: 0
|
||||
description: Offset for pagination
|
||||
minimum: 0
|
||||
title: Offset
|
||||
type: integer
|
||||
style: form
|
||||
@@ -3541,7 +3551,7 @@ paths:
|
||||
This endpoint handles file upload, conversion, and memory creation in a single operation.
|
||||
|
||||
**Features:**
|
||||
- Supports PDF, DOCX, PPTX, XLSX, images (with OCR), audio (with transcription)
|
||||
- Supports PDF, DOCX, PPTX, XLSX, images (parser-dependent OCR), audio (with transcription)
|
||||
- Automatic file-to-markdown conversion using pluggable parsers
|
||||
- Files stored in object storage (PostgreSQL by default, S3 for production)
|
||||
- Each file becomes a separate document with optional metadata/tags
|
||||
@@ -7113,6 +7123,17 @@ components:
|
||||
type: string
|
||||
nullable: true
|
||||
type: array
|
||||
prefer_observations:
|
||||
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."
|
||||
title: Prefer Observations
|
||||
type: boolean
|
||||
budget:
|
||||
$ref: '#/components/schemas/Budget'
|
||||
max_tokens:
|
||||
@@ -7137,7 +7158,9 @@ components:
|
||||
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)."
|
||||
\ (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."
|
||||
enum:
|
||||
- any
|
||||
- all
|
||||
@@ -7843,12 +7866,13 @@ components:
|
||||
type: integer
|
||||
output_tokens:
|
||||
default: 0
|
||||
description: Number of output/completion tokens generated
|
||||
description: Number of visible output/completion tokens generated (excludes
|
||||
reasoning/thoughts)
|
||||
title: Output Tokens
|
||||
type: integer
|
||||
total_tokens:
|
||||
default: 0
|
||||
description: Total tokens (input + output)
|
||||
description: "Total tokens (input + output, excludes thoughts)"
|
||||
title: Total Tokens
|
||||
type: integer
|
||||
cached_tokens:
|
||||
@@ -7856,6 +7880,13 @@ components:
|
||||
description: "Cached/cache-read prompt tokens, when reported by the provider"
|
||||
title: Cached Tokens
|
||||
type: integer
|
||||
thoughts_tokens:
|
||||
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.
|
||||
title: Thoughts Tokens
|
||||
type: integer
|
||||
title: TokenUsage
|
||||
ToolCallsIncludeOptions:
|
||||
description: Options for including tool calls in reflect results.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user