Compare commits

..
1863 changed files with 37164 additions and 296492 deletions
-15
View File
@@ -1,15 +0,0 @@
{
"$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
"name": "hindsight",
"description": "Official Hindsight integrations for Claude Code",
"owner": {
"name": "vectorize-io"
},
"plugins": [
{
"name": "hindsight-memory",
"description": "Automatic long-term memory for Claude Code via Hindsight",
"source": "./hindsight-integrations/claude-code"
}
]
}
-205
View File
@@ -1,205 +0,0 @@
---
name: code-review
description: Review changed code against project standards. Checks for missing tests, dead code, type safety, lint issues, and coding conventions. Run after completing any implementation work.
user_invocable: true
---
# Code Review
Review all changed code against the project's quality standards and coding conventions.
## Code Standards
Read and internalize these standards before writing code. The review steps below verify compliance.
### Python Style
- Python 3.11+, type hints required
- Async throughout (asyncpg, async FastAPI)
- Pydantic models for request/response
- Ruff for linting (line-length 120)
- No Python files at project root - maintain clean directory structure
- **Never use multi-item tuple return values** — not even for internal/private functions. Always use a dataclass or Pydantic model. No exceptions, no "it's just two values" shortcuts. If a function returns more than one value, define a named type for it.
### Type Safety with Pydantic Models
**NEVER use raw `dict` types for structured data** — this applies to all code, including internal helpers and private functions. If the dict has known keys, it must be a dataclass or Pydantic model:
- Use Pydantic `BaseModel` for all data structures passed between functions
- Use `@dataclass` for lightweight internal data containers when Pydantic validation isn't needed
- Add `@field_validator` for type coercion (e.g., ensuring datetimes are timezone-aware)
- Avoid `dict.get()` patterns - use typed model attributes instead
- Parse external data (JSON, API responses) into Pydantic models at the boundary
- This catches type errors at parse time, not deep in business logic
- The only acceptable `dict` usage is for truly dynamic/unknown keys (e.g., arbitrary metadata, JSON blobs with no fixed schema)
```python
# BAD - error-prone dict access
def process(data: dict) -> str:
return data.get("name", "") # No validation, silent failures
# GOOD - typed and validated
class UserData(BaseModel):
name: str
created_at: datetime
@field_validator("created_at", mode="before")
@classmethod
def ensure_tz_aware(cls, v):
if isinstance(v, str):
v = datetime.fromisoformat(v.replace("Z", "+00:00"))
if v.tzinfo is None:
return v.replace(tzinfo=timezone.utc)
return v
def process(data: UserData) -> str:
return data.name # Type-safe, validated at construction
```
### TypeScript Style
- Next.js App Router for control plane
- Tailwind CSS with shadcn/ui components
### Code Comments
- **Always comment non-trivial technical decisions** with the reasoning behind the choice. If someone would ask "why is it done this way?", there should be a comment.
- **Keep comments up to date with history** — when changing an approach, update the comment to explain what was tried before and why it was changed. Comments serve as a tracker of previous implementations that likely had problems.
- Don't comment obvious code — only where the "why" isn't self-evident from the code itself.
```python
# BAD - no context for future readers
results = await asyncio.gather(*tasks, return_exceptions=True)
# GOOD - explains the non-obvious choice
# Use return_exceptions=True to avoid cancelling sibling tasks on failure.
# Previously we used TaskGroup but it cancelled all tasks when one failed,
# causing partial writes that left orphaned entity links (see #412).
results = await asyncio.gather(*tasks, return_exceptions=True)
```
### Branch Hygiene
- **Always start new feature branches from `origin/main`** — rebase to ensure a clean base.
- **Only include commits relevant to the PR/branch/feature** — no unrelated changes. If the branch contains commits that don't belong, they must be removed before merging.
### General Principles
- Don't add features, refactor code, or make "improvements" beyond what was asked
- Don't add unnecessary error handling for impossible scenarios
- Don't create helpers or abstractions for one-time operations
- No backwards-compatibility hacks (unused vars, re-exports, "removed" comments)
- Three similar lines of code is better than a premature abstraction
## Review Steps
### 1. Check branch hygiene
- Run `git log --oneline main..HEAD` to list all commits on the branch.
- Verify every commit is relevant to the feature/PR. Flag any unrelated commits.
- Check the branch is based on a recent `origin/main` (no stale base).
### 2. Identify changed files
Run `git diff --name-only HEAD` (unstaged) and `git diff --cached --name-only` (staged) to get all changed files. If there are no local changes, diff against the base branch using `git diff main...HEAD --name-only` and `git diff main...HEAD` to review all commits on the current branch.
### 3. Run linters
```bash
./scripts/hooks/lint.sh
```
Report any failures. Do NOT fix them yourself — just report.
### 4. Check for dead code
For each changed Python file, check for:
- Unused imports (Ruff should catch these, but verify)
- Functions/methods/classes that were added but are never called from anywhere
- Variables assigned but never read
- Commented-out code blocks that should be removed
For each changed TypeScript file, check for:
- Unused imports
- Unused variables or functions
- Commented-out code
### 5. Check type safety (Python)
For each changed Python file, check for violations:
- **No raw `dict` for structured data** — must use Pydantic model or dataclass, even for internal/private functions (only exception: truly dynamic/unknown keys)
- **No multi-item tuple returns** — must use dataclass or Pydantic model, even for internal/private functions (no exceptions)
- **Missing type hints** on function parameters and return types
- **Missing `@field_validator`** for datetime fields that should be timezone-aware
### 6. Check for missing tests
For each new or significantly changed function/endpoint/class:
- Check if there is a corresponding test addition or update
- New API endpoints MUST have integration tests
- New utility functions MUST have unit tests
- Bug fixes SHOULD have a regression test
Flag any new logic that lacks test coverage.
### 7. Check API consistency
If any files in `hindsight-api-slim/hindsight_api/api/` were changed:
- Were the OpenAPI specs regenerated? (`./scripts/generate-openapi.sh`)
- Were the client SDKs regenerated? (`./scripts/generate-clients.sh`)
- Were the control plane proxy routes updated? (`hindsight-control-plane/src/app/api/`)
### 8. Check code comments
For each non-trivial change:
- **New non-obvious logic** — is there a comment explaining the reasoning?
- **Changed approach** — does the comment include what was done before and why it changed?
- **Stale comments** — do existing comments near the changed code still accurately describe the behavior?
### 9. Check integration completeness
If any files in `hindsight-integrations/` were added or changed, verify:
- **Tests exist** — the integration must have tests that simulate/exercise the external framework (not just pure unit tests of helpers). Check for a `tests/` directory with meaningful test files.
- **CI job exists** — check `.github/workflows/test.yml` for a corresponding `test-<name>-integration` job. If missing, flag it.
- **Release process** — check that the integration name is in the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh`. If missing, flag it.
- **Code standards** — the integration code must follow all Python style rules (type hints, no raw dicts, no tuple returns, etc.).
### 10. Check MCP tool registration completeness
If any new MCP tools were added or existing tools renamed in `hindsight-api-slim/hindsight_api/mcp_tools.py`:
- **`_ALL_TOOLS` set** in `mcp_tools.py` — must include the new tool name
- **`tools_to_register` default set** in `register_mcp_tools()` in `mcp_tools.py` — must include the new tool name
- **`_SINGLE_BANK_TOOLS` set** in `hindsight-api-slim/hindsight_api/api/mcp.py` — must include the new tool if it is bank-scoped (not a bank-management tool like `list_banks`/`create_bank`)
- **`MCP_TOOL_GROUPS`** in `hindsight-control-plane/src/components/bank-config-view.tsx` — must include the new tool in the appropriate group for the UI tool selector
- **Tool count assertions** in tests (e.g., `test_mcp_tools.py`) — must be updated to reflect the new count
### 11. Review against other coding standards
Check the diff for violations of the standards listed above:
- Python files at project root (not allowed)
- Missing async patterns (should be async throughout)
- Pydantic models for request/response
- Line length > 120 chars
- New features/code beyond what was asked (over-engineering)
- Unnecessary error handling for impossible scenarios
- Premature abstractions or speculative helpers
- Backwards-compatibility hacks (unused vars, re-exports, "removed" comments)
### 12. Report findings
Present a clear summary organized by severity:
**Must fix** — issues that will break CI or violate hard project rules:
- Unrelated commits on the branch
- Lint failures
- Missing type hints on public functions
- Raw dict usage for structured data (including internal code)
- Multi-item tuple returns (including internal code)
- Missing tests for new endpoints
- New integration missing tests, CI job, or release-integration.sh entry
**Should fix** — issues that hurt code quality:
- Dead code / unused imports missed by linter
- Missing tests for non-trivial utility functions
- Over-engineering beyond the task scope
**Note** — observations that may or may not need action:
- API changes that might need client regeneration
- Patterns that deviate from nearby code style
For each finding, include the file path, line number, and a brief explanation.
Do NOT auto-fix any issues. Report all findings and let the user decide what to address. If there are no findings, confirm the code looks good.
+1 -12
View File
@@ -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, volcano
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai
HINDSIGHT_API_LLM_PROVIDER=openai
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
@@ -20,16 +20,6 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
# HINDSIGHT_API_LLM_VERTEXAI_REGION=us-central1
# HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY=/path/to/service-account-key.json # Optional, uses ADC if not set
# Example: MiniMax configuration (1M context window)
# HINDSIGHT_API_LLM_PROVIDER=minimax
# HINDSIGHT_API_LLM_API_KEY=your-minimax-api-key
# HINDSIGHT_API_LLM_MODEL=MiniMax-M2.7
# Example: DeepSeek configuration (https://api.deepseek.com)
# HINDSIGHT_API_LLM_PROVIDER=deepseek
# HINDSIGHT_API_LLM_API_KEY=your-deepseek-api-key
# HINDSIGHT_API_LLM_MODEL=deepseek-v4-flash # or deepseek-v4-pro / deepseek-chat / deepseek-reasoner
# Example: LM Studio local configuration (Qwen 2.5 32B recommended)
# HINDSIGHT_API_LLM_PROVIDER=lmstudio
# HINDSIGHT_API_LLM_API_KEY=lmstudio
@@ -49,7 +39,6 @@ HINDSIGHT_API_LOG_LEVEL=info
# Database (Optional - uses embedded pg0 by default)
# HINDSIGHT_API_DATABASE_URL=postgresql://user:pass@host:5432/db
# HINDSIGHT_API_MIGRATION_DATABASE_URL= # Direct PostgreSQL URL for migrations (bypasses PgBouncer). Falls back to DATABASE_URL.
# HINDSIGHT_API_DATABASE_SCHEMA=public # PostgreSQL schema name (default: public)
# Vector Extension (Optional - uses pgvector by default)
-6
View File
@@ -1,6 +0,0 @@
version: 2
updates:
- package-ecosystem: "github-actions"
directory: "/"
schedule:
interval: "weekly"
+5 -5
View File
@@ -21,20 +21,20 @@ jobs:
build:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- uses: actions/setup-node@v6
- uses: actions/checkout@v4
- uses: actions/setup-node@v4
with:
node-version: 20
cache: npm
cache-dependency-path: package-lock.json
- uses: astral-sh/setup-uv@v7
- uses: astral-sh/setup-uv@v4
- run: npm ci --workspace=hindsight-docs
- run: uv run generate-llms-full
- run: npm run build --workspace=hindsight-docs
env:
UMAMI_URL: https://analytics.hindsight.vectorize.io
UMAMI_WEBSITE_ID: ${{ secrets.UMAMI_WEBSITE_ID }}
- uses: actions/upload-pages-artifact@v5
- uses: actions/upload-pages-artifact@v3
with:
path: hindsight-docs/build
deploy:
@@ -44,5 +44,5 @@ jobs:
runs-on: ubuntu-latest
needs: build
steps:
- uses: actions/deploy-pages@v5
- uses: actions/deploy-pages@v4
id: deployment
-197
View File
@@ -1,197 +0,0 @@
name: Performance Tests
on:
schedule:
# Run daily at 06:00 UTC
- cron: "0 6 * * *"
workflow_dispatch:
inputs:
scale:
description: "Test scale (perf-test)"
type: choice
options:
- tiny
- small
- medium
- large
default: large
suite:
description: "Perf-test suite to run (blank = all)"
type: choice
options:
- ""
- retain
- recall
- recall-with-observations
- consolidation
default: ""
locomo_max_conversations:
description: "LoComo max conversations (0 = skip, blank = all)"
type: number
default: 0
locomo_skip:
description: "Skip LoComo job"
type: boolean
default: false
ref:
description: "Git ref to test (branch, tag, or SHA). Defaults to main."
type: string
default: ""
concurrency:
group: perf-test
cancel-in-progress: true
jobs:
perf-test:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
with:
ref: ${{ inputs.ref || github.ref }}
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
prune-cache: false
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version-file: ".python-version"
- name: Cache HuggingFace models
uses: actions/cache@v5
with:
path: ~/.cache/huggingface
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
restore-keys: |
${{ runner.os }}-huggingface-
- name: Pre-download models
working-directory: ./hindsight-api-slim
run: |
uv run --frozen --all-extras --index-strategy unsafe-best-match python -c "
from sentence_transformers import SentenceTransformer
print('Downloading embedding model...')
SentenceTransformer('BAAI/bge-small-en-v1.5')
print('Model downloaded successfully')
"
- name: Install hindsight-dev dependencies
run: |
cd hindsight-dev && uv sync --frozen --all-extras --index-strategy unsafe-best-match
- name: "Suite: retain"
if: inputs.suite == '' || inputs.suite == 'retain'
run: |
./scripts/benchmarks/run-perf-test.sh \
--scale ${{ inputs.scale || 'large' }} \
--suite retain \
--output perf-results-retain.json
- name: "Suite: recall"
if: inputs.suite == '' || inputs.suite == 'recall'
run: |
./scripts/benchmarks/run-perf-test.sh \
--scale ${{ inputs.scale || 'large' }} \
--suite recall \
--output perf-results-recall.json
- name: "Suite: recall-with-observations"
if: inputs.suite == '' || inputs.suite == 'recall-with-observations'
run: |
./scripts/benchmarks/run-perf-test.sh \
--scale ${{ inputs.scale || 'large' }} \
--suite recall-with-observations \
--output perf-results-recall-with-observations.json
- name: "Suite: consolidation"
if: inputs.suite == '' || inputs.suite == 'consolidation'
run: |
./scripts/benchmarks/run-perf-test.sh \
--scale ${{ inputs.scale || 'large' }} \
--suite consolidation \
--output perf-results-consolidation.json
- name: Upload perf results
if: always()
uses: actions/upload-artifact@v7
with:
name: perf-results-${{ github.sha }}
path: hindsight-dev/perf-results-*.json
retention-days: 90
locomo:
if: inputs.locomo_skip != true
runs-on: ubuntu-latest
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
HINDSIGHT_API_JUDGE_LLM_PROVIDER: vertexai
HINDSIGHT_API_JUDGE_LLM_MODEL: google/gemini-2.5-flash-lite
HINDSIGHT_API_ANSWER_LLM_PROVIDER: vertexai
HINDSIGHT_API_ANSWER_LLM_MODEL: google/gemini-2.5-flash
steps:
- uses: actions/checkout@v6
with:
ref: ${{ inputs.ref || github.ref }}
- name: Setup GCP credentials
run: |
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
prune-cache: false
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version-file: ".python-version"
- name: Cache HuggingFace models
uses: actions/cache@v5
with:
path: ~/.cache/huggingface
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
restore-keys: |
${{ runner.os }}-huggingface-
- name: Pre-download models
working-directory: ./hindsight-api-slim
run: |
uv run --frozen --all-extras --index-strategy unsafe-best-match python -c "
from sentence_transformers import SentenceTransformer
print('Downloading embedding model...')
SentenceTransformer('BAAI/bge-small-en-v1.5')
print('Model downloaded successfully')
"
- name: Install hindsight-dev dependencies
run: |
cd hindsight-dev && uv sync --frozen --all-extras --index-strategy unsafe-best-match
- name: Run LoComo benchmark
run: |
MAX_CONV_ARG=""
if [ "${{ inputs.locomo_max_conversations }}" != "0" ] && [ -n "${{ inputs.locomo_max_conversations }}" ]; then
MAX_CONV_ARG="--max-conversations ${{ inputs.locomo_max_conversations }}"
fi
uv run python hindsight-dev/benchmarks/locomo/locomo_benchmark.py \
--wait-consolidation \
$MAX_CONV_ARG
- name: Upload LoComo results
if: always()
uses: actions/upload-artifact@v7
with:
name: locomo-results-${{ github.sha }}
path: hindsight-dev/benchmarks/locomo/results/
retention-days: 90
-131
View File
@@ -1,131 +0,0 @@
name: Release Integration
on:
push:
tags:
- 'integrations/**'
jobs:
publish:
runs-on: ubuntu-latest
permissions:
id-token: write # for PyPI trusted publishing
steps:
- uses: actions/checkout@v6
- name: Extract integration info
id: info
run: |
# refs/tags/integrations/litellm/v0.1.0 → integration=litellm, version=0.1.0
TAG="${GITHUB_REF#refs/tags/}"
INTEGRATION=$(echo "$TAG" | cut -d'/' -f2)
VERSION=$(echo "$TAG" | cut -d'/' -f3 | sed 's/^v//')
echo "integration=$INTEGRATION" >> $GITHUB_OUTPUT
echo "version=$VERSION" >> $GITHUB_OUTPUT
echo "tag=$TAG" >> $GITHUB_OUTPUT
echo "Integration: $INTEGRATION, Version: $VERSION"
- name: Detect integration type
id: type
run: |
if [ -f "hindsight-integrations/${{ steps.info.outputs.integration }}/pyproject.toml" ]; then
echo "type=python" >> $GITHUB_OUTPUT
elif [ -f "hindsight-integrations/${{ steps.info.outputs.integration }}/package.json" ]; then
echo "type=typescript" >> $GITHUB_OUTPUT
else
echo "type=plugin" >> $GITHUB_OUTPUT
fi
# ── Python integrations (litellm, pydantic-ai, crewai) ──────────────────
- name: Install uv
if: steps.type.outputs.type == 'python'
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
- name: Set up Python
if: steps.type.outputs.type == 'python'
uses: actions/setup-python@v6
with:
python-version-file: ".python-version"
- name: Build Python package
if: steps.type.outputs.type == 'python'
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
run: uv build --out-dir dist
- name: Publish Python package to PyPI
if: steps.type.outputs.type == 'python'
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: ./hindsight-integrations/${{ steps.info.outputs.integration }}/dist
skip-existing: true
# ── TypeScript integrations (ai-sdk, chat, openclaw) ────────────────────
# ── Plugin integrations (claude-code) — no package to publish ───────────
- name: Plugin release
if: steps.type.outputs.type == 'plugin'
run: |
echo "Plugin integration ${{ steps.info.outputs.integration }} v${{ steps.info.outputs.version }} — no package to publish."
echo "Users install via: claude plugin marketplace add vectorize-io/hindsight --sparse hindsight-integrations"
# ── TypeScript integrations (ai-sdk, chat, openclaw) ────────────────────
- name: Set up Node.js
if: steps.type.outputs.type == 'typescript'
uses: actions/setup-node@v6
with:
node-version: '22'
registry-url: 'https://registry.npmjs.org'
cache: 'npm'
cache-dependency-path: package-lock.json
- name: Check integration lockfile
if: steps.type.outputs.type == 'typescript'
run: ./scripts/check-integration-lockfiles.sh
# Some integrations depend on workspace packages (hindsight-client,
# hindsight-all, hindsight-agent-sdk) via file: refs. Install from root
# so npm resolves them, then build the workspace deps before the integration.
- name: Install root workspace dependencies
if: steps.type.outputs.type == 'typescript'
run: npm ci
- name: Build workspace deps (hindsight-client, hindsight-all, hindsight-agent-sdk)
if: steps.type.outputs.type == 'typescript'
run: |
npm run build --workspace=hindsight-clients/typescript
npm run build --workspace=hindsight-all-npm
npm run build --workspace=hindsight-tools/hindsight-agent-sdk
- name: Install integration dependencies
if: steps.type.outputs.type == 'typescript'
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
run: npm ci
- name: Build TypeScript package
if: steps.type.outputs.type == 'typescript'
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
run: npm run build
- name: Publish TypeScript package to npm
if: steps.type.outputs.type == 'typescript'
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
run: |
set +e
OUTPUT=$(npm publish --access public 2>&1)
EXIT_CODE=$?
echo "$OUTPUT"
if [ $EXIT_CODE -ne 0 ]; then
if echo "$OUTPUT" | grep -q "cannot publish over"; then
echo "Package version already published, skipping..."
exit 0
fi
exit $EXIT_CODE
fi
env:
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
-65
View File
@@ -1,65 +0,0 @@
name: Release Tool
on:
push:
tags:
- 'tools/**'
jobs:
publish:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- name: Extract tool info
id: info
run: |
# refs/tags/tools/self-driving-agents/v0.0.1 → tool=self-driving-agents, version=0.0.1
TAG="${GITHUB_REF#refs/tags/}"
TOOL=$(echo "$TAG" | cut -d'/' -f2)
VERSION=$(echo "$TAG" | cut -d'/' -f3 | sed 's/^v//')
echo "tool=$TOOL" >> $GITHUB_OUTPUT
echo "version=$VERSION" >> $GITHUB_OUTPUT
echo "tag=$TAG" >> $GITHUB_OUTPUT
echo "Tool: $TOOL, Version: $VERSION"
- name: Set up Node.js
uses: actions/setup-node@v6
with:
node-version: '22'
registry-url: 'https://registry.npmjs.org'
cache: 'npm'
cache-dependency-path: package-lock.json
# Tools live under hindsight-tools/ and may depend on workspace packages
# (e.g. @vectorize-io/hindsight-client). Install from root so npm resolves
# workspace deps, then build any required workspace packages first.
- name: Install root workspace dependencies
run: npm ci
- name: Build hindsight-client (workspace dep)
run: npm run build --workspace=hindsight-clients/typescript
- name: Build hindsight-agent-sdk (workspace dep)
run: npm run build --workspace=hindsight-tools/hindsight-agent-sdk
- name: Build tool
run: npm run build --workspace=hindsight-tools/${{ steps.info.outputs.tool }}
- name: Publish to npm
working-directory: ./hindsight-tools/${{ steps.info.outputs.tool }}
run: |
set +e
OUTPUT=$(npm publish --access public 2>&1)
EXIT_CODE=$?
echo "$OUTPUT"
if [ $EXIT_CODE -ne 0 ]; then
if echo "$OUTPUT" | grep -q "cannot publish over"; then
echo "Package version already published, skipping..."
exit 0
fi
exit $EXIT_CODE
fi
env:
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
+200 -86
View File
@@ -13,15 +13,15 @@ jobs:
id-token: write
steps:
- uses: actions/checkout@v6
- uses: actions/checkout@v4
- name: Install uv
uses: astral-sh/setup-uv@v7
uses: astral-sh/setup-uv@v5
with:
enable-cache: true
- name: Set up Python
uses: actions/setup-python@v6
uses: actions/setup-python@v5
with:
python-version-file: ".python-version"
@@ -30,39 +30,37 @@ jobs:
working-directory: ./hindsight-clients/python
run: uv build --out-dir dist
- name: Build hindsight-api-slim
working-directory: ./hindsight-api-slim
run: uv build --out-dir dist
- name: Build hindsight-api
working-directory: ./hindsight-api
run: uv build --out-dir dist
- name: Build hindsight-all
working-directory: ./hindsight-all
working-directory: ./hindsight
run: uv build --out-dir dist
- name: Build hindsight-all-slim
working-directory: ./hindsight-all-slim
- name: Build hindsight-litellm
working-directory: ./hindsight-integrations/litellm
run: uv build --out-dir dist
- name: Build hindsight-embed
working-directory: ./hindsight-embed
run: uv build --out-dir dist
# Publish in order (client and api-slim first, then api/all wrappers which depend on them)
- name: Build hindsight-crewai
working-directory: ./hindsight-integrations/crewai
run: uv build --out-dir dist
- name: Build hindsight-pydantic-ai
working-directory: ./hindsight-integrations/pydantic-ai
run: uv build --out-dir dist
# Publish in order (client and api first, then hindsight-all which depends on them)
- name: Publish hindsight-client to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: ./hindsight-clients/python/dist
skip-existing: true
- name: Publish hindsight-api-slim to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: ./hindsight-api-slim/dist
skip-existing: true
- name: Publish hindsight-api to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
@@ -72,13 +70,13 @@ jobs:
- name: Publish hindsight-all to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: ./hindsight-all/dist
packages-dir: ./hindsight/dist
skip-existing: true
- name: Publish hindsight-all-slim to PyPI
- name: Publish hindsight-litellm to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: ./hindsight-all-slim/dist
packages-dir: ./hindsight-integrations/litellm/dist
skip-existing: true
- name: Publish hindsight-embed to PyPI
@@ -87,18 +85,31 @@ jobs:
packages-dir: ./hindsight-embed/dist
skip-existing: true
- name: Publish hindsight-crewai to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: ./hindsight-integrations/crewai/dist
skip-existing: true
- name: Publish hindsight-pydantic-ai to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: ./hindsight-integrations/pydantic-ai/dist
skip-existing: true
# Upload artifacts for GitHub release
- name: Upload artifacts
uses: actions/upload-artifact@v7
uses: actions/upload-artifact@v4
with:
name: python-packages
path: |
hindsight-clients/python/dist/*
hindsight-api-slim/dist/*
hindsight-api/dist/*
hindsight-all/dist/*
hindsight-all-slim/dist/*
hindsight/dist/*
hindsight-integrations/litellm/dist/*
hindsight-embed/dist/*
hindsight-integrations/crewai/dist/*
hindsight-integrations/pydantic-ai/dist/*
retention-days: 1
release-typescript-client:
@@ -106,10 +117,10 @@ jobs:
environment: npm
steps:
- uses: actions/checkout@v6
- uses: actions/checkout@v4
- name: Set up Node.js
uses: actions/setup-node@v6
uses: actions/setup-node@v4
with:
node-version: '20'
registry-url: 'https://registry.npmjs.org'
@@ -144,35 +155,35 @@ jobs:
run: npm pack
- name: Upload artifacts
uses: actions/upload-artifact@v7
uses: actions/upload-artifact@v4
with:
name: typescript-client
path: hindsight-clients/typescript/*.tgz
retention-days: 1
release-hindsight-all-npm:
release-openclaw-integration:
runs-on: ubuntu-latest
environment: npm
steps:
- uses: actions/checkout@v6
- uses: actions/checkout@v4
- name: Set up Node.js
uses: actions/setup-node@v6
uses: actions/setup-node@v4
with:
node-version: '22'
registry-url: 'https://registry.npmjs.org'
cache: 'npm'
cache-dependency-path: package-lock.json
- name: Install dependencies
run: npm ci --workspace=hindsight-all-npm
working-directory: ./hindsight-integrations/openclaw
run: npm ci
- name: Build
run: npm run build --workspace=hindsight-all-npm
working-directory: ./hindsight-integrations/openclaw
run: npm run build
- name: Publish to npm
working-directory: ./hindsight-all-npm
working-directory: ./hindsight-integrations/openclaw
run: |
set +e
OUTPUT=$(npm publish --access public 2>&1)
@@ -189,14 +200,112 @@ jobs:
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
- name: Pack for GitHub release
working-directory: ./hindsight-all-npm
working-directory: ./hindsight-integrations/openclaw
run: npm pack
- name: Upload artifacts
uses: actions/upload-artifact@v7
uses: actions/upload-artifact@v4
with:
name: hindsight-all-npm
path: hindsight-all-npm/*.tgz
name: openclaw-integration
path: hindsight-integrations/openclaw/*.tgz
retention-days: 1
release-ai-sdk-integration:
runs-on: ubuntu-latest
environment: npm
steps:
- uses: actions/checkout@v4
- name: Set up Node.js
uses: actions/setup-node@v4
with:
node-version: '22'
registry-url: 'https://registry.npmjs.org'
- name: Install dependencies
working-directory: ./hindsight-integrations/ai-sdk
run: npm ci
- name: Build
working-directory: ./hindsight-integrations/ai-sdk
run: npm run build
- name: Publish to npm
working-directory: ./hindsight-integrations/ai-sdk
run: |
set +e
OUTPUT=$(npm publish --access public 2>&1)
EXIT_CODE=$?
echo "$OUTPUT"
if [ $EXIT_CODE -ne 0 ]; then
if echo "$OUTPUT" | grep -q "cannot publish over"; then
echo "Package version already published, skipping..."
exit 0
fi
exit $EXIT_CODE
fi
env:
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
- name: Pack for GitHub release
working-directory: ./hindsight-integrations/ai-sdk
run: npm pack
- name: Upload artifacts
uses: actions/upload-artifact@v4
with:
name: ai-sdk-integration
path: hindsight-integrations/ai-sdk/*.tgz
retention-days: 1
release-chat-integration:
runs-on: ubuntu-latest
environment: npm
steps:
- uses: actions/checkout@v4
- name: Set up Node.js
uses: actions/setup-node@v4
with:
node-version: '22'
registry-url: 'https://registry.npmjs.org'
- name: Install dependencies
working-directory: ./hindsight-integrations/chat
run: npm ci
- name: Build
working-directory: ./hindsight-integrations/chat
run: npm run build
- name: Publish to npm
working-directory: ./hindsight-integrations/chat
run: |
set +e
OUTPUT=$(npm publish --access public 2>&1)
EXIT_CODE=$?
echo "$OUTPUT"
if [ $EXIT_CODE -ne 0 ]; then
if echo "$OUTPUT" | grep -q "cannot publish over"; then
echo "Package version already published, skipping..."
exit 0
fi
exit $EXIT_CODE
fi
env:
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
- name: Pack for GitHub release
working-directory: ./hindsight-integrations/chat
run: npm pack
- name: Upload artifacts
uses: actions/upload-artifact@v4
with:
name: chat-integration
path: hindsight-integrations/chat/*.tgz
retention-days: 1
release-control-plane:
@@ -204,10 +313,10 @@ jobs:
environment: npm
steps:
- uses: actions/checkout@v6
- uses: actions/checkout@v4
- name: Set up Node.js
uses: actions/setup-node@v6
uses: actions/setup-node@v4
with:
node-version: '20'
registry-url: 'https://registry.npmjs.org'
@@ -255,7 +364,7 @@ jobs:
run: npm pack
- name: Upload artifacts
uses: actions/upload-artifact@v7
uses: actions/upload-artifact@v4
with:
name: control-plane
path: hindsight-control-plane/*.tgz
@@ -278,13 +387,9 @@ jobs:
target: aarch64-apple-darwin
artifact_name: hindsight
asset_name: hindsight-darwin-arm64
- os: ubuntu-24.04-arm
target: aarch64-unknown-linux-gnu
artifact_name: hindsight
asset_name: hindsight-linux-arm64
steps:
- uses: actions/checkout@v6
- uses: actions/checkout@v4
- name: Install Rust
uses: dtolnay/rust-toolchain@stable
@@ -302,7 +407,7 @@ jobs:
chmod +x artifacts/${{ matrix.asset_name }}
- name: Upload artifacts
uses: actions/upload-artifact@v7
uses: actions/upload-artifact@v4
with:
name: rust-cli-${{ matrix.asset_name }}
path: artifacts/${{ matrix.asset_name }}
@@ -343,7 +448,7 @@ jobs:
PRELOAD_ML_MODELS=false
steps:
- uses: actions/checkout@v6
- uses: actions/checkout@v4
- name: Free Disk Space
uses: jlumbroso/free-disk-space@main
@@ -357,13 +462,13 @@ jobs:
swap-storage: true
- name: Set up QEMU
uses: docker/setup-qemu-action@v4
uses: docker/setup-qemu-action@v3
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v4
uses: docker/setup-buildx-action@v3
- name: Log in to GitHub Container Registry
uses: docker/login-action@v4
uses: docker/login-action@v3
with:
registry: ghcr.io
username: ${{ github.actor }}
@@ -375,7 +480,7 @@ jobs:
- name: Extract metadata for release tags
id: meta
uses: docker/metadata-action@v6
uses: docker/metadata-action@v5
with:
images: ghcr.io/${{ github.repository_owner }}/${{ matrix.image_name }}
flavor: |
@@ -391,7 +496,7 @@ jobs:
# # Step 1: Build for local testing (single platform, no push)
# # This creates an identical image to what will be released, just for one platform
# - name: Build image for testing
# uses: docker/build-push-action@v7
# uses: docker/build-push-action@v6
# with:
# context: .
# file: docker/standalone/Dockerfile
@@ -410,7 +515,7 @@ jobs:
# Build multi-platform and push to release tags
- name: Build and push release images
uses: docker/build-push-action@v7
uses: docker/build-push-action@v6
with:
context: .
file: docker/standalone/Dockerfile
@@ -428,10 +533,10 @@ jobs:
packages: write
steps:
- uses: actions/checkout@v6
- uses: actions/checkout@v4
- name: Install Helm
uses: azure/setup-helm@v5
uses: azure/setup-helm@v4
with:
version: 'latest'
@@ -448,7 +553,7 @@ jobs:
run: helm push helm-packages/*.tgz oci://ghcr.io/${{ github.repository_owner }}/charts
- name: Upload artifacts
uses: actions/upload-artifact@v7
uses: actions/upload-artifact@v4
with:
name: helm-chart
path: helm-packages/*.tgz
@@ -456,67 +561,73 @@ jobs:
create-github-release:
runs-on: ubuntu-latest
needs: [release-python-packages, release-typescript-client, release-hindsight-all-npm, release-control-plane, release-rust-cli, release-docker-images, release-helm-chart]
needs: [release-python-packages, release-typescript-client, release-openclaw-integration, release-ai-sdk-integration, release-chat-integration, release-control-plane, release-rust-cli, release-docker-images, release-helm-chart]
permissions:
contents: write
steps:
- uses: actions/checkout@v6
- uses: actions/checkout@v4
- name: Extract version from tag
id: get_version
run: echo "VERSION=${GITHUB_REF#refs/tags/v}" >> $GITHUB_OUTPUT
- name: Download Python packages
uses: actions/download-artifact@v8
uses: actions/download-artifact@v4
with:
name: python-packages
path: ./artifacts/python-packages
- name: Download TypeScript client
uses: actions/download-artifact@v8
uses: actions/download-artifact@v4
with:
name: typescript-client
path: ./artifacts/typescript-client
- name: Download OpenClaw Integration
uses: actions/download-artifact@v4
with:
name: openclaw-integration
path: ./artifacts/openclaw-integration
- name: Download AI SDK Integration
uses: actions/download-artifact@v4
with:
name: ai-sdk-integration
path: ./artifacts/ai-sdk-integration
- name: Download Chat Integration
uses: actions/download-artifact@v4
with:
name: chat-integration
path: ./artifacts/chat-integration
- name: Download Control Plane
uses: actions/download-artifact@v8
uses: actions/download-artifact@v4
with:
name: control-plane
path: ./artifacts/control-plane
- name: Download hindsight-embed npm wrapper
uses: actions/download-artifact@v8
with:
name: hindsight-all-npm
path: ./artifacts/hindsight-all-npm
- name: Download Rust CLI (Linux)
uses: actions/download-artifact@v8
uses: actions/download-artifact@v4
with:
name: rust-cli-hindsight-linux-amd64
path: ./artifacts/rust-cli-linux
- name: Download Rust CLI (Linux ARM)
uses: actions/download-artifact@v8
with:
name: rust-cli-hindsight-linux-arm64
path: ./artifacts/rust-cli-linux-arm64
- name: Download Rust CLI (macOS Intel)
uses: actions/download-artifact@v8
uses: actions/download-artifact@v4
with:
name: rust-cli-hindsight-darwin-amd64
path: ./artifacts/rust-cli-darwin-amd64
- name: Download Rust CLI (macOS ARM)
uses: actions/download-artifact@v8
uses: actions/download-artifact@v4
with:
name: rust-cli-hindsight-darwin-arm64
path: ./artifacts/rust-cli-darwin-arm64
- name: Download Helm chart
uses: actions/download-artifact@v8
uses: actions/download-artifact@v4
with:
name: helm-chart
path: ./artifacts/helm-chart
@@ -526,20 +637,23 @@ jobs:
mkdir -p release-assets
# Python packages
cp artifacts/python-packages/hindsight-clients/python/dist/* release-assets/ || true
cp artifacts/python-packages/hindsight-api-slim/dist/* release-assets/ || true
cp artifacts/python-packages/hindsight-api/dist/* release-assets/ || true
cp artifacts/python-packages/hindsight-all/dist/* release-assets/ || true
cp artifacts/python-packages/hindsight-all-slim/dist/* release-assets/ || true
cp artifacts/python-packages/hindsight/dist/* release-assets/ || true
cp artifacts/python-packages/hindsight-integrations/litellm/dist/* release-assets/ || true
cp artifacts/python-packages/hindsight-integrations/pydantic-ai/dist/* release-assets/ || true
cp artifacts/python-packages/hindsight-embed/dist/* release-assets/ || true
# TypeScript client
cp artifacts/typescript-client/*.tgz release-assets/ || true
# hindsight-embed npm wrapper
cp artifacts/hindsight-all-npm/*.tgz release-assets/ || true
# OpenClaw Integration
cp artifacts/openclaw-integration/*.tgz release-assets/ || true
# AI SDK Integration
cp artifacts/ai-sdk-integration/*.tgz release-assets/ || true
# Chat Integration
cp artifacts/chat-integration/*.tgz release-assets/ || true
# Control Plane
cp artifacts/control-plane/*.tgz release-assets/ || true
# Rust CLI binaries
cp artifacts/rust-cli-linux/hindsight-linux-amd64 release-assets/ || true
cp artifacts/rust-cli-linux-arm64/hindsight-linux-arm64 release-assets/ || true
cp artifacts/rust-cli-darwin-amd64/hindsight-darwin-amd64 release-assets/ || true
cp artifacts/rust-cli-darwin-arm64/hindsight-darwin-arm64 release-assets/ || true
# Helm chart
@@ -547,7 +661,7 @@ jobs:
ls -la release-assets/
- name: Create GitHub Release
uses: softprops/action-gh-release@v3
uses: softprops/action-gh-release@v2
with:
files: release-assets/*
generate_release_notes: true
+187 -2322
View File
File diff suppressed because it is too large Load Diff
+1 -2
View File
@@ -50,8 +50,7 @@ hindsight-dev/benchmarks/perf/results/
benchmarks/results/
hindsight-cli/target
hindsight-clients/rust/target
.claude/*
!.claude/skills/
.claude
whats-next.md
TASK.md
# Changelog is now tracked in hindsight-docs/src/pages/changelog.md
-7
View File
@@ -1,7 +0,0 @@
{
"semi": true,
"singleQuote": false,
"tabWidth": 2,
"trailingComma": "es5",
"printWidth": 100
}
+75 -97
View File
@@ -11,32 +11,26 @@ Hindsight is an agent memory system that provides long-term memory for AI agents
## Development Commands
### Local Development (API + UI)
```bash
# Start both API server and control plane UI
./scripts/dev/start.sh
```
### API Server (Python/FastAPI)
```bash
# Start API server only (loads .env automatically)
# Start API server (loads .env automatically)
./scripts/dev/start-api.sh
# Run all tests (parallelized with pytest-xdist)
cd hindsight-api-slim && uv run pytest tests/
cd hindsight-api && uv run pytest tests/
# Run specific test file
cd hindsight-api-slim && uv run pytest tests/test_http_api_integration.py -v
cd hindsight-api && uv run pytest tests/test_http_api_integration.py -v
# Run single test function
cd hindsight-api-slim && uv run pytest tests/test_retain.py::test_retain_simple -v
cd hindsight-api && uv run pytest tests/test_retain.py::test_retain_simple -v
# Lint and format
cd hindsight-api-slim && uv run ruff check .
cd hindsight-api-slim && uv run ruff format .
cd hindsight-api && uv run ruff check .
cd hindsight-api && uv run ruff format .
# Type checking (uses ty - extremely fast type checker from Astral)
cd hindsight-api-slim && uv run ty check hindsight_api/
cd hindsight-api && uv run ty check hindsight_api/
```
### Control Plane (Next.js)
@@ -68,9 +62,8 @@ cd hindsight-control-plane && npm run dev
./scripts/benchmarks/run-locomo.sh
# Performance benchmarks
./scripts/benchmarks/run-perf-test.sh # System perf (mock LLM + pg0)
./scripts/benchmarks/run-perf-test.sh --scale tiny # Quick smoke test
./scripts/benchmarks/run-consolidation.sh
./scripts/benchmarks/run-retain-perf.sh --document <path> # Requires API server running
# Results viewer
./scripts/benchmarks/start-visualizer.sh # View results at localhost:8001
@@ -79,17 +72,18 @@ cd hindsight-control-plane && npm run dev
## Architecture
### Monorepo Structure
- **hindsight-api-slim/**: Core FastAPI server with memory engine (Python, uv)
- **hindsight-api/**: Core FastAPI server with memory engine (Python, uv)
- **hindsight/**: Embedded Python bundle (hindsight-all package)
- **hindsight-control-plane/**: Admin UI (Next.js, npm)
- **hindsight-cli/**: CLI tool (Rust, cargo, uses progenitor for API client)
- **hindsight-clients/**: Generated SDK clients (Python, TypeScript, Rust)
- **hindsight-docs/**: Docusaurus documentation site
- **hindsight-integrations/**: Framework integrations (LiteLLM, CrewAI, LangGraph, Pydantic AI, AG2, Claude Code, etc.)
- **hindsight-integrations/**: Framework integrations (LiteLLM, OpenAI)
- **hindsight-dev/**: Development tools and benchmarks
### Core Engine (hindsight-api-slim/hindsight_api/engine/)
- `memory_engine.py`: Main orchestrator for retain/recall/reflect operations
- `llm_wrapper.py`: LLM abstraction supporting OpenAI, Anthropic, Gemini, VertexAI, Groq, MiniMax, Ollama, LM Studio, LiteLLM, Claude Code
### Core Engine (hindsight-api/hindsight_api/engine/)
- `memory_engine.py`: Main orchestrator (~170KB) for retain/recall/reflect operations
- `llm_wrapper.py`: LLM abstraction supporting OpenAI, Anthropic, Gemini, Groq, Ollama, LM Studio
- `embeddings.py`: Embedding generation (local sentence-transformers or TEI)
- `cross_encoder.py`: Reranking (local or TEI)
- `entity_resolver.py`: Entity extraction and normalization
@@ -102,13 +96,13 @@ cd hindsight-control-plane && npm run dev
**search/**: Multi-strategy retrieval
- `retrieval.py`: Main retrieval orchestrator
- `graph_retrieval.py`: Graph retrieval abstract base class
- `link_expansion_retrieval.py`: Link expansion graph retrieval
- `graph_retrieval.py`: Entity/relationship graph traversal
- `mpfp_retrieval.py`: Multi-Path Fact Propagation retrieval
- `fusion.py`: Reciprocal rank fusion for combining results
- `reranking.py`: Cross-encoder reranking
### API Layer (hindsight-api-slim/hindsight_api/api/)
- `http.py`: FastAPI HTTP routers for all REST endpoints
### API Layer (hindsight-api/hindsight_api/api/)
- `http.py`: FastAPI HTTP routers (~80KB) for all REST endpoints
- `mcp.py`: Model Context Protocol server implementation
Main operations:
@@ -117,23 +111,18 @@ Main operations:
- **Reflect**: Disposition-aware reasoning using memories and mental models.
### Database
PostgreSQL with pgvector. Schema managed via Alembic migrations in `hindsight-api-slim/hindsight_api/alembic/`. Migrations run automatically on API startup.
PostgreSQL with pgvector. Schema managed via Alembic migrations in `hindsight-api/hindsight_api/alembic/`. Migrations run automatically on API startup.
Key tables: `banks`, `memory_units`, `documents`, `entities`, `entity_links`
### Adding Database Migrations
Hindsight runs the same Alembic tree against PostgreSQL and Oracle 23ai. Each
migration file dispatches through `run_for_dialect`, which calls either
`_pg_upgrade` or `_oracle_upgrade` based on the live connection. A pytest lint
(`tests/test_migration_shape.py`) fails CI if a migration omits the dispatcher.
1. **Create a new migration file** in `hindsight-api-slim/hindsight_api/alembic/versions/`:
1. **Create a new migration file** in `hindsight-api/hindsight_api/alembic/versions/`:
- File name format: `<revision_id>_<description>.py` (e.g., `f1a2b3c4d5e6_add_new_index.py`)
- Use a unique hex revision ID (12 chars)
- Set `down_revision` to the previous migration's revision ID
2. **Migration template** (the `script.py.mako` template scaffolds this; fill in the bodies):
2. **Migration template**:
```python
"""Description of the migration
@@ -144,61 +133,28 @@ migration file dispatches through `run_for_dialect`, which calls either
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "f1a2b3c4d5e6"
down_revision: str | Sequence[str] | None = "<previous_revision_id>"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _pg_schema_prefix() -> str:
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _pg_schema_prefix()
def upgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"CREATE INDEX ... ON {schema}table_name(...)")
def _pg_downgrade() -> None:
schema = _pg_schema_prefix()
op.execute(f"DROP INDEX IF EXISTS {schema}index_name")
def _oracle_upgrade() -> None:
# Oracle 23ai equivalent. Use op.get_bind().exec_driver_sql for forms
# that Alembic core does not model (vector/text indexes, partitions).
op.execute("CREATE INDEX ... ON table_name(...)")
def _oracle_downgrade() -> None:
op.execute("DROP INDEX IF EXISTS index_name")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
schema = _get_schema_prefix()
op.execute(f"DROP INDEX IF EXISTS {schema}index_name")
```
**Dialect-only migrations.** If a change genuinely doesn't apply to one
dialect (e.g. enabling `pg_trgm` is PG-only), omit the unused slot:
```python
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade) # oracle slot intentionally absent → no-op
```
Make the asymmetry deliberate. Don't leave an Oracle slot empty just because
you didn't think about it — copy-pasting a PG migration without the Oracle
half is exactly how schemas drift.
3. **Run migrations locally**:
```bash
# Set database URL and run migrations for the base schema plus all tenants
# Set database URL and run migrations
uv run hindsight-admin run-db-migration
# Run on a specific tenant schema
@@ -208,17 +164,11 @@ migration file dispatches through `run_for_dialect`, which calls either
## Key Conventions
### Code Quality
**Before writing code, read `.claude/skills/code-review/SKILL.md`** for the full coding standards (Python style, type safety, TypeScript style, general principles).
**Always run the lint script after making Python or TypeScript/Node changes:**
```bash
./scripts/hooks/lint.sh
```
**After completing any implementation work, run `/code-review`** to verify your changes against project standards (missing tests, dead code, type safety, etc.). Fix any "must fix" issues before considering the task done.
**MANDATORY: Run `/code-review` before pushing code or creating a pull request.** Do not push or create a PR until all "must fix" issues are resolved.
This runs the same checks as the pre-commit hook (Ruff for Python, ESLint/Prettier for TypeScript).
### Memory Banks
- Each bank is an isolated memory store (like a "brain" for one user/agent)
@@ -250,20 +200,48 @@ When adding or modifying parameters in the dataplane API (hindsight-api), you mu
- Update the client type definition in `lib/api.ts`
- Update any UI components that need to use the new parameter
### Adding New Integrations
### Python Style
- Python 3.11+, type hints required
- Async throughout (asyncpg, async FastAPI)
- Pydantic models for request/response
- Ruff for linting (line-length 120)
- No Python files at project root - maintain clean directory structure
- **Never use multi-item tuple return values** - prefer dataclass or Pydantic model for structured returns
Every new integration in `hindsight-integrations/` must satisfy all of the following before it can be merged:
### Type Safety with Pydantic Models
**NEVER use raw `dict` types for structured data.** Always use Pydantic models:
- Use Pydantic `BaseModel` for all data structures passed between functions
- Add `@field_validator` for type coercion (e.g., ensuring datetimes are timezone-aware)
- Avoid `dict.get()` patterns - use typed model attributes instead
- Parse external data (JSON, API responses) into Pydantic models at the boundary
- This catches type errors at parse time, not deep in business logic
1. **Tests are required** — tests must simulate or exercise the external system (mock the framework's interfaces and verify the integration actually calls Hindsight correctly). Pure unit tests of helper functions are not sufficient.
2. **CI job** — add a test job in `.github/workflows/test.yml` following the existing pattern (e.g., `test-crewai-integration`). The job must build, install deps, and run `uv run pytest tests -v`. Also add the integration to `detect-changes` outputs so it only runs when its files change.
3. **Release process** — add the integration name to the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh` so it can be released via the standard release workflow.
4. **Follow project code standards** — Python style, type safety, no raw dicts for structured data, no multi-item tuple returns (see `.claude/skills/code-review/SKILL.md`).
```python
# BAD - error-prone dict access
def process(data: dict) -> str:
return data.get("name", "") # No validation, silent failures
If any of these are missing, the integration is incomplete and must not be pushed or merged.
# GOOD - typed and validated
class UserData(BaseModel):
name: str
created_at: datetime
### Changelogs
@field_validator("created_at", mode="before")
@classmethod
def ensure_tz_aware(cls, v):
if isinstance(v, str):
v = datetime.fromisoformat(v.replace("Z", "+00:00"))
if v.tzinfo is None:
return v.replace(tzinfo=timezone.utc)
return v
Never add "Unreleased" entries to changelogs (e.g. `hindsight-docs/src/pages/changelog/**`). Changelog entries are written by the release script (`./scripts/release-integration.sh`) when a version is actually cut. If a bug fix or feature needs documenting before release, describe it in the PR/commit — the release tooling will surface it in the published changelog section.
def process(data: UserData) -> str:
return data.name # Type-safe, validated at construction
```
### TypeScript Style
- Next.js App Router for control plane
- Tailwind CSS with shadcn/ui components
### Adding New API Configuration Flags
@@ -273,24 +251,24 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
#### Adding a New Configuration Field
1. **config.py** (`hindsight-api-slim/hindsight_api/config.py`):
1. **config.py** (`hindsight-api/hindsight_api/config.py`):
- Add `ENV_*` constant for the environment variable name (e.g., `ENV_MY_SETTING = "HINDSIGHT_API_MY_SETTING"`)
- Add `DEFAULT_*` constant for the default value
- Add field to `HindsightConfig` dataclass with type annotation
- **Mark as configurable** by adding to `_CONFIGURABLE_FIELDS` set if the field should be overridable per-tenant/bank via API
- **Mark as hierarchical or static** by adding to `_HIERARCHICAL_FIELDS` set (hierarchical) or leaving it out (static)
- Add initialization in `from_env()` method
```python
# Configurable field (can be overridden per-tenant/bank via API)
_CONFIGURABLE_FIELDS = {
# Hierarchical field (can be overridden per-bank)
_HIERARCHICAL_FIELDS = {
...,
"my_setting", # Add here for configurable
"my_setting", # Add here for hierarchical
}
# Static field - just don't add to _CONFIGURABLE_FIELDS
# Static field - just don't add to _HIERARCHICAL_FIELDS
```
2. **main.py** (`hindsight-api-slim/hindsight_api/main.py`):
2. **main.py** (`hindsight-api/hindsight_api/main.py`):
- Add field to the manual `HindsightConfig()` constructor call (search for "CLI override")
3. **Use hierarchical config in MemoryEngine**:
@@ -330,14 +308,14 @@ cp .env.example .env
# Edit .env with LLM API key
# Python deps
uv sync --directory hindsight-api-slim/
uv sync --directory hindsight-api/
# Node deps (uses npm workspaces)
npm install
```
Required env vars:
- `HINDSIGHT_API_LLM_PROVIDER`: openai, anthropic, gemini, groq, minimax, ollama, lmstudio
- `HINDSIGHT_API_LLM_PROVIDER`: openai, anthropic, gemini, groq, ollama, lmstudio
- `HINDSIGHT_API_LLM_API_KEY`: Your API key
- `HINDSIGHT_API_LLM_MODEL`: Model name (e.g., gpt-4o-mini, claude-sonnet-4-20250514)
+2 -6
View File
@@ -7,12 +7,10 @@
[![CI](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml/badge.svg)](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml)
[![Slack Community](https://img.shields.io/badge/Slack-Join%20Community-4A154B?logo=slack)](https://join.slack.com/t/hindsight-space/shared_invite/zt-3nhbm4w29-LeSJ5Ixi6j8PdiYOCPlOgg)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![gitcgr](https://gitcgr.com/badge/vectorize-io/hindsight.svg)](https://gitcgr.com/vectorize-io/hindsight)
![PyPI - Downloads](https://img.shields.io/pypi/dm/hindsight-api?label=PyPI)
![NPM Downloads](https://img.shields.io/npm/dm/%40vectorize-io%2Fhindsight-client?logoColor=orange&label=NPM&color=blue&link=https%3A%2F%2Fwww.npmjs.com%2Fpackage%2F%40vectorize-io%2Fhindsight-client)
<br/>
<a href="https://trendshift.io/repositories/15603" target="_blank"><img src="https://trendshift.io/api/badge/repositories/15603" alt="vectorize-io%2Fhindsight | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
</div>
---
@@ -71,7 +69,7 @@ docker run --rm -it --pull always -p 8888:8888 -p 9999:9999 \
>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`, and `lmstudio`. The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
@@ -84,8 +82,6 @@ cd docker/docker-compose
docker compose up
```
> Oracle AI Database is also supported for enterprise deployments with full feature parity. See the [storage documentation](https://hindsight.vectorize.io/developer/storage) for details.
>API: http://localhost:8888
>UI: http://localhost:9999
Generated
-171
View File
@@ -1,171 +0,0 @@
{
"version": "5",
"specifiers": {
"jsr:@std/assert@^1.0.17": "1.0.19",
"jsr:@std/assert@^1.0.19": "1.0.19",
"jsr:@std/expect@*": "1.0.18",
"jsr:@std/internal@^1.0.12": "1.0.12",
"jsr:@std/path@^1.1.4": "1.1.4",
"jsr:@std/testing@*": "1.0.17"
},
"jsr": {
"@std/[email protected]": {
"integrity": "eaada96ee120cb980bc47e040f82814d786fe8162ecc53c91d8df60b8755991e",
"dependencies": [
"jsr:@std/internal"
]
},
"@std/[email protected]": {
"integrity": "8566eab35200466f8609eb7e7aed062ed0db314e9a258d5d201b1b8997ce801a",
"dependencies": [
"jsr:@std/assert@^1.0.19",
"jsr:@std/internal",
"jsr:@std/path"
]
},
"@std/[email protected]": {
"integrity": "972a634fd5bc34b242024402972cd5143eac68d8dffaca5eaa4dba30ce17b027"
},
"@std/[email protected]": {
"integrity": "1d2d43f39efb1b42f0b1882a25486647cb851481862dc7313390b2bb044314b5",
"dependencies": [
"jsr:@std/internal"
]
},
"@std/[email protected]": {
"integrity": "87bdc2700fa98249d48a17cd72413352d3d3680dcfbdb64947fd0982d6bbf681",
"dependencies": [
"jsr:@std/assert@^1.0.17",
"jsr:@std/internal"
]
}
},
"workspace": {
"members": {
"hindsight-all-npm": {
"packageJson": {
"dependencies": [
"npm:@types/node@22",
"npm:tsup@^8.5.1",
"npm:typescript@^5.7.0",
"npm:vitest@^4.1.2"
]
}
},
"hindsight-clients/typescript": {
"packageJson": {
"dependencies": [
"npm:@hey-api/[email protected]",
"npm:@types/jest@29",
"npm:@types/node@20",
"npm:jest@29",
"npm:ts-jest@29",
"npm:tsup@^8.5.1",
"npm:typescript@5"
]
}
},
"hindsight-control-plane": {
"packageJson": {
"dependencies": [
"npm:@chenglou/pretext@^0.0.3",
"npm:@eslint/eslintrc@^3.3.3",
"npm:@eslint/js@^9.39.2",
"npm:@radix-ui/react-alert-dialog@^1.1.15",
"npm:@radix-ui/react-checkbox@^1.3.3",
"npm:@radix-ui/react-dialog@^1.1.15",
"npm:@radix-ui/react-dropdown-menu@^2.1.16",
"npm:@radix-ui/react-label@^2.1.8",
"npm:@radix-ui/react-popover@^1.1.15",
"npm:@radix-ui/react-radio-group@^1.3.8",
"npm:@radix-ui/react-select@^2.2.6",
"npm:@radix-ui/react-slider@^1.3.6",
"npm:@radix-ui/react-slot@^1.2.4",
"npm:@radix-ui/react-switch@^1.2.6",
"npm:@radix-ui/react-tabs@^1.1.13",
"npm:@radix-ui/react-tooltip@^1.2.8",
"npm:@tailwindcss/postcss@^4.1.17",
"npm:@tailwindcss/typography@~0.5.19",
"npm:@types/cytoscape@^3.21.9",
"npm:@types/node@^24.10.0",
"npm:@types/react-dom@^19.2.2",
"npm:@types/react@^19.2.2",
"npm:autoprefixer@^10.4.21",
"npm:class-variance-authority@~0.7.1",
"npm:clsx@^2.1.1",
"npm:cmdk@^1.1.1",
"npm:cytoscape-fcose@^2.2.0",
"npm:cytoscape@^3.33.1",
"npm:eslint-config-next@^16.0.1",
"npm:eslint-plugin-react-hooks@^7.0.1",
"npm:eslint-plugin-react@^7.37.5",
"npm:eslint@^9.39.1",
"npm:[email protected]",
"npm:next-themes@~0.4.6",
"npm:next@^16.1.7",
"npm:postcss@^8.5.6",
"npm:prettier@^3.7.4",
"npm:react-chrono@^2.9.1",
"npm:react-dom@^19.2.0",
"npm:react-is@^19.2.4",
"npm:react-markdown@^10.1.0",
"npm:react18-json-view@~0.2.9",
"npm:react@^19.2.0",
"npm:recharts@^3.5.1",
"npm:remark-gfm@^4.0.1",
"npm:sonner@^2.0.7",
"npm:tailwind-merge@^3.4.0",
"npm:tailwindcss-animate@^1.0.7",
"npm:tailwindcss@^4.1.17",
"npm:[email protected]",
"npm:typescript-eslint@^8.50.0",
"npm:typescript@^5.9.3"
]
}
},
"hindsight-docs": {
"packageJson": {
"dependencies": [
"npm:@docusaurus/[email protected]",
"npm:@docusaurus/[email protected]",
"npm:@docusaurus/[email protected]",
"npm:@docusaurus/theme-common@^3.9.2",
"npm:@docusaurus/theme-mermaid@^3.9.2",
"npm:@docusaurus/[email protected]",
"npm:@docusaurus/[email protected]",
"npm:@easyops-cn/docusaurus-search-local@~0.52.2",
"npm:@mdx-js/react@3",
"npm:clsx@2",
"npm:prism-react-renderer@^2.3.0",
"npm:raw-loader@^4.0.2",
"npm:react-dom@19",
"npm:react-icons@^5.6.0",
"npm:react@19",
"npm:redocusaurus@^2.5.0",
"npm:typescript@~5.6.2"
]
}
},
"hindsight-tools/hindsight-agent-sdk": {
"packageJson": {
"dependencies": [
"npm:@vectorize-io/hindsight-client@~0.5.6",
"npm:typescript@^5.4.0",
"npm:vitest@^4.1.2"
]
}
},
"hindsight-tools/self-driving-agents": {
"packageJson": {
"dependencies": [
"npm:@clack/prompts@^1.2.0",
"npm:@vectorize-io/hindsight-client@~0.5.6",
"npm:picocolors@^1.1.0",
"npm:typescript@^5.4.0",
"npm:vitest@^4.1.2"
]
}
}
}
}
}
+13 -10
View File
@@ -42,22 +42,25 @@ RUN apt-get update && apt-get install -y \
&& pip install --no-cache-dir uv
# Copy dependency files and README (required by pyproject.toml)
COPY hindsight-api-slim/pyproject.toml ./api/
COPY hindsight-api-slim/README.md ./api/
COPY hindsight-api/pyproject.toml ./api/
COPY hindsight-api/README.md ./api/
WORKDIR /app/api
# Sync dependencies using appropriate extras based on INCLUDE_LOCAL_MODELS
# local-ml: torch, sentence-transformers, transformers, einops, flashrank, mlx (optional)
# embedded-db: pg0-embedded (always included for embedded PostgreSQL support)
RUN if [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
uv sync --extra local-ml --extra embedded-db; \
else \
uv sync --extra embedded-db; \
# Remove local ML model dependencies if INCLUDE_LOCAL_MODELS=false
# This creates a smaller image when using external providers (TEI, OpenAI, Cohere)
RUN if [ "$INCLUDE_LOCAL_MODELS" != "true" ]; then \
echo "Removing local-models dependencies (sentence-transformers, torch, transformers)..." && \
sed -i '/"sentence-transformers/d' pyproject.toml && \
sed -i '/"transformers/d' pyproject.toml && \
sed -i '/"torch/d' pyproject.toml; \
fi
# Sync dependencies (will create lock file if needed)
RUN uv sync
# Copy source code (alembic migrations are inside hindsight_api/)
COPY hindsight-api-slim/hindsight_api ./hindsight_api
COPY hindsight-api/hindsight_api ./hindsight_api
# Install the local package (uv sync only installed dependencies, not the package itself)
RUN uv pip install -e .
+7 -114
View File
@@ -1,28 +1,6 @@
#!/bin/bash
set -e
# =============================================================================
# Embedded pg0 data integrity check (#675)
#
# When using embedded pg0, check if the data directory has existing PostgreSQL
# data before starting. If the directory exists but appears empty/corrupt
# (e.g., missing PG_VERSION file), log a warning. This helps diagnose data
# loss scenarios where a container restart caused the data directory to be
# wiped despite a volume mount being present.
# =============================================================================
PG0_DATA_DIR="${HOME}/.pg0"
if [ -d "$PG0_DATA_DIR" ]; then
# Look for actual PostgreSQL data directories (pg0 creates subdirs per instance)
if compgen -G "$PG0_DATA_DIR"/*/PG_VERSION > /dev/null 2>&1; then
echo "✅ Existing pg0 data directory detected at $PG0_DATA_DIR"
elif [ "$(ls -A "$PG0_DATA_DIR" 2>/dev/null)" ]; then
echo "⚠️ WARNING: pg0 data directory exists at $PG0_DATA_DIR but no PG_VERSION found."
echo " This may indicate data corruption or an incomplete previous shutdown."
echo " If you see all migrations running from scratch after this, your data may have been lost."
echo " See: https://github.com/vectorize-io/hindsight/issues/675"
fi
fi
# Service flags (default to true if not set)
ENABLE_API="${HINDSIGHT_ENABLE_API:-true}"
ENABLE_CP="${HINDSIGHT_ENABLE_CP:-true}"
@@ -93,95 +71,24 @@ if [ "${HINDSIGHT_WAIT_FOR_DEPS:-false}" = "true" ]; then
done
fi
# =============================================================================
# Graceful shutdown handler (#675)
#
# Docker sends SIGTERM on `docker stop`/`docker restart`. Without a trap, child
# processes (hindsight-api + pg0, control-plane) are killed abruptly. For the
# embedded pg0 database this can cause data loss when the data directory is on
# a Docker volume that gets remounted after restart.
#
# The trap forwards SIGTERM to all tracked child PIDs so that:
# - hindsight-api receives the signal and can run its shutdown hooks
# - pg0 gets a clean PostgreSQL shutdown (checkpoint + WAL flush)
# - The control-plane Node.js process exits cleanly
# =============================================================================
# Guard against concurrent cleanup (e.g., child crash + SIGTERM arriving together)
SHUTTING_DOWN=false
cleanup() {
if $SHUTTING_DOWN; then return; fi
SHUTTING_DOWN=true
echo ""
echo "🛑 Received shutdown signal, stopping services gracefully..."
for pid in "${PIDS[@]}"; do
if kill -0 "$pid" 2>/dev/null; then
kill -TERM "$pid" 2>/dev/null
fi
done
# Give processes time to shut down cleanly (pg0 needs to flush WAL).
# NOTE: Docker's default stop_grace_period is 10s. If you use the default,
# either set stop_grace_period: 30s in your compose file / docker stop -t 30,
# or Docker will SIGKILL the container before this timeout expires.
local timeout=30
for ((i=1; i<=timeout; i++)); do
local all_stopped=true
for pid in "${PIDS[@]}"; do
if kill -0 "$pid" 2>/dev/null; then
all_stopped=false
break
fi
done
if $all_stopped; then
echo "✅ All services stopped cleanly"
exit 0
fi
sleep 1
done
# Force kill if still running after timeout
echo "⚠️ Timeout reached, forcing shutdown..."
for pid in "${PIDS[@]}"; do
if kill -0 "$pid" 2>/dev/null; then
kill -9 "$pid" 2>/dev/null
fi
done
exit 1
}
trap cleanup SIGTERM SIGINT
# Track PIDs for wait
PIDS=()
# Start API if enabled
if [ "$ENABLE_API" = "true" ]; then
cd /app/api
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:8888/health}"
API_STARTUP_WAIT_SECONDS="${HINDSIGHT_API_STARTUP_WAIT_SECONDS:-300}"
# Run API directly - Python's PYTHONUNBUFFERED=1 handles output buffering
hindsight-api &
API_PID=$!
PIDS+=($API_PID)
# Wait for API to be ready
api_ready=false
for ((i=1; i<=API_STARTUP_WAIT_SECONDS; i++)); do
if ! kill -0 "$API_PID" 2>/dev/null; then
wait "$API_PID"
exit $?
fi
if curl -sf "$API_HEALTH_URL" &>/dev/null; then
api_ready=true
for i in {1..60}; do
if curl -sf http://localhost:8888/health &>/dev/null; then
break
fi
sleep 1
done
if [ "$api_ready" != "true" ]; then
echo "❌ API did not become healthy within ${API_STARTUP_WAIT_SECONDS}s"
exit 1
fi
else
echo "API disabled (HINDSIGHT_ENABLE_API=false)"
fi
@@ -190,7 +97,6 @@ fi
if [ "$ENABLE_CP" = "true" ]; then
echo "🎛️ Starting Control Plane..."
cd /app/control-plane
export HOSTNAME="${HINDSIGHT_CP_HOSTNAME:-0.0.0.0}"
PORT="${HINDSIGHT_CP_PORT:-9999}" node server.js &
CP_PID=$!
PIDS+=($CP_PID)
@@ -217,21 +123,8 @@ if [ ${#PIDS[@]} -eq 0 ]; then
exit 1
fi
# Wait for any process to exit (use wait -n with trap-safe loop)
while true; do
# wait -n returns when any child exits; it also returns on signal delivery
# (the trap handler will run and exit, so this loop is just for robustness).
# `&& true` prevents `set -e` from killing the script when wait -n returns
# non-zero (child exited with error or no backgrounded children remain).
wait -n && true
# Check if any tracked PID has exited
for pid in "${PIDS[@]}"; do
if ! kill -0 "$pid" 2>/dev/null; then
wait "$pid" 2>/dev/null
exit_code=$?
echo "⚠️ Service (PID $pid) exited with code $exit_code"
# Trigger cleanup for remaining services
cleanup
fi
done
done
# Wait for any process to exit
wait -n
# Exit with status of first exited process
exit $?
-18
View File
@@ -49,9 +49,6 @@
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(dirname "$SCRIPT_DIR")"
# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
@@ -181,21 +178,6 @@ for i in $(seq 1 "$TIMEOUT"); do
echo "=== Health Response ==="
curl -s "http://localhost:${HEALTH_PORT}${HEALTH_PATH}" | python3 -m json.tool 2>/dev/null || curl -s "http://localhost:${HEALTH_PORT}${HEALTH_PATH}"
echo ""
# Run retain/recall smoke test for API targets
if [ "$TARGET" != "cp-only" ]; then
echo ""
echo "=== Retain/Recall Smoke Test ==="
if ! "$REPO_ROOT/scripts/smoke-test-slim.sh" "http://localhost:${HEALTH_PORT}"; then
echo ""
echo "=== Container Logs (last 50 lines) ==="
docker logs "$CONTAINER_NAME" 2>&1 | tail -50
echo ""
echo -e "${RED}Smoke test FAILED${NC}"
exit 1
fi
fi
echo ""
echo "=== Container Logs (last 50 lines) ==="
docker logs "$CONTAINER_NAME" 2>&1 | tail -50
+2 -2
View File
@@ -2,8 +2,8 @@ apiVersion: v2
name: hindsight
description: Hindsight helm chart
type: application
version: 0.6.0
appVersion: "0.6.0"
version: 0.4.15
appVersion: "0.4.15"
keywords:
- ai
- memory
@@ -95,27 +95,6 @@ spec:
{{- toYaml .Values.api.readinessProbe | nindent 10 }}
resources:
{{- toYaml .Values.api.resources | nindent 10 }}
{{- if or .Values.api.persistence.modelCache.enabled .Values.api.extraVolumeMounts }}
volumeMounts:
{{- if .Values.api.persistence.modelCache.enabled }}
- name: model-cache
mountPath: /home/hindsight/.cache
{{- end }}
{{- with .Values.api.extraVolumeMounts }}
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}
{{- if or .Values.api.persistence.modelCache.enabled .Values.api.extraVolumes }}
volumes:
{{- if .Values.api.persistence.modelCache.enabled }}
- name: model-cache
persistentVolumeClaim:
claimName: {{ include "hindsight.fullname" . }}-api-model-cache
{{- end }}
{{- with .Values.api.extraVolumes }}
{{- toYaml . | nindent 6 }}
{{- end }}
{{- end }}
{{- with .Values.nodeSelector }}
nodeSelector:
{{- toYaml . | nindent 8 }}
@@ -1,21 +0,0 @@
{{- if and .Values.api.enabled .Values.api.persistence.modelCache.enabled }}
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
name: {{ include "hindsight.fullname" . }}-api-model-cache
labels:
{{- include "hindsight.api.labels" . | nindent 4 }}
{{- with .Values.api.persistence.modelCache.annotations }}
annotations:
{{- toYaml . | nindent 4 }}
{{- end }}
spec:
accessModes:
{{- toYaml .Values.api.persistence.modelCache.accessModes | nindent 4 }}
{{- if .Values.api.persistence.modelCache.storageClass }}
storageClassName: {{ .Values.api.persistence.modelCache.storageClass }}
{{- end }}
resources:
requests:
storage: {{ .Values.api.persistence.modelCache.size }}
{{- end }}
@@ -95,16 +95,6 @@ spec:
{{- toYaml .Values.worker.readinessProbe | nindent 10 }}
resources:
{{- toYaml .Values.worker.resources | nindent 10 }}
{{- if or .Values.worker.persistence.modelCache.enabled .Values.worker.extraVolumeMounts }}
volumeMounts:
{{- if .Values.worker.persistence.modelCache.enabled }}
- name: model-cache
mountPath: /home/hindsight/.cache
{{- end }}
{{- with .Values.worker.extraVolumeMounts }}
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}
{{- with .Values.nodeSelector }}
nodeSelector:
{{- toYaml . | nindent 8 }}
@@ -117,26 +107,4 @@ spec:
tolerations:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.worker.extraVolumes }}
volumes:
{{- toYaml . | nindent 6 }}
{{- end }}
{{- if .Values.worker.persistence.modelCache.enabled }}
volumeClaimTemplates:
- metadata:
name: model-cache
{{- with .Values.worker.persistence.modelCache.annotations }}
annotations:
{{- toYaml . | nindent 8 }}
{{- end }}
spec:
accessModes:
{{- toYaml .Values.worker.persistence.modelCache.accessModes | nindent 8 }}
{{- if .Values.worker.persistence.modelCache.storageClass }}
storageClassName: {{ .Values.worker.persistence.modelCache.storageClass }}
{{- end }}
resources:
requests:
storage: {{ .Values.worker.persistence.modelCache.size }}
{{- end }}
{{- end }}
-53
View File
@@ -67,33 +67,6 @@ api:
# Pod affinity/anti-affinity (overrides global affinity for this component)
# affinity: {}
# Persistent volume for local model cache (reranker, embeddings)
# Models are downloaded to /home/hindsight/.cache on first use.
# Without persistence, models are re-downloaded on every pod restart.
persistence:
modelCache:
enabled: false
size: 5Gi
storageClass: ""
accessModes:
- ReadWriteOnce
annotations: {}
# Extra volume mounts for the api container
# e.g.
# extraVolumeMounts:
# - name: my-volume
# mountPath: /mnt/my-volume
extraVolumeMounts: []
# Extra volumes for the api pod
# e.g.
# extraVolumes:
# - name: my-volume
# configMap:
# name: my-configmap
extraVolumes: []
# Environment variables
env:
#HINDSIGHT_API_LLM_PROVIDER: "groq"
@@ -167,32 +140,6 @@ worker:
# Pod affinity/anti-affinity (overrides global affinity for this component)
# affinity: {}
# Persistent volume for local model cache (reranker, embeddings)
# Uses volumeClaimTemplates since worker is a StatefulSet.
persistence:
modelCache:
enabled: false
size: 5Gi
storageClass: ""
accessModes:
- ReadWriteOnce
annotations: {}
# Extra volume mounts for the worker container
# e.g.
# extraVolumeMounts:
# - name: my-volume
# mountPath: /mnt/my-volume
extraVolumeMounts: []
# Extra volumes for the worker pod
# e.g.
# extraVolumes:
# - name: my-volume
# configMap:
# name: my-configmap
extraVolumes: []
# Secret environment variables (inherited from api.secrets if not specified)
secrets: {}
-4
View File
@@ -1,4 +0,0 @@
node_modules
dist
*.tgz
.DS_Store
-80
View File
@@ -1,80 +0,0 @@
# @vectorize-io/hindsight-all
Node.js equivalent of the Python [`hindsight-all`](https://pypi.org/project/hindsight-all/) package — programmatic lifecycle manager for a local Hindsight daemon. Use this when you want to embed Hindsight in a Node application without hand-rolling subprocess management.
This package deliberately does **not** ship an HTTP client. Once the daemon is running, talk to it with [`@vectorize-io/hindsight-client`](https://www.npmjs.com/package/@vectorize-io/hindsight-client) against `server.getBaseUrl()`. The two packages compose — one owns the daemon process, the other owns the HTTP API surface.
## Requirements
- **Node.js >= 22** — uses global `fetch` and `AbortSignal.timeout`.
- **`uv` / `uvx`** on `PATH` — used to download and run the underlying `hindsight-embed` daemon on first use. Install via <https://docs.astral.sh/uv/>.
## Install
```bash
npm install @vectorize-io/hindsight-all @vectorize-io/hindsight-client
```
## Example
```ts
import { HindsightServer, consoleLogger } from "@vectorize-io/hindsight-all";
import { HindsightClient } from "@vectorize-io/hindsight-client";
const server = new HindsightServer({
profile: "my-app",
port: 9077,
env: {
HINDSIGHT_API_LLM_PROVIDER: "anthropic",
HINDSIGHT_API_LLM_API_KEY: process.env.ANTHROPIC_API_KEY,
HINDSIGHT_API_LLM_MODEL: "claude-sonnet-4-20250514",
HINDSIGHT_EMBED_DAEMON_IDLE_TIMEOUT: "0",
},
logger: consoleLogger,
});
await server.start();
const client = new HindsightClient({ baseUrl: server.getBaseUrl() });
await client.retain("user-123", "User prefers dark mode and concise answers.", {
documentId: "pref-2026-04-01",
});
const recall = await client.recall("user-123", "what are the user preferences?");
console.log(recall.results);
await server.stop();
```
For a remote Hindsight API, skip `HindsightServer` entirely and just point `HindsightClient` at the remote URL.
## Open config — forward-compatible with new daemon flags
`HindsightServerOptions` is designed so every new environment variable or CLI flag in the underlying Hindsight daemon can be used without waiting for a wrapper release:
- **`env`** accepts an arbitrary `Record<string, string>`. Every entry is exported into the daemon process and written into the profile config via `--env KEY=VALUE`.
- **`extraProfileCreateArgs`** / **`extraDaemonStartArgs`** append raw args to the respective commands.
## Development against a local checkout
If you're hacking on the Python `hindsight-embed` package in the same monorepo, point the server at the local path — it'll use `uv run --directory <path>` instead of `uvx`:
```ts
new HindsightServer({
embedPackagePath: "/path/to/hindsight-embed",
// ...
});
```
## API surface
- `HindsightServer` — daemon lifecycle (`start`, `stop`, `checkHealth`, `getBaseUrl`, `getProfile`).
- `Logger` interface plus `silentLogger` (default) and `consoleLogger` helpers.
- `getEmbedCommand(opts)` — low-level helper that returns the `[cmd, ...args]` tuple used to invoke the underlying Python CLI.
For memory operations (retain, recall, reflect, bank management, stats) use [`@vectorize-io/hindsight-client`](https://www.npmjs.com/package/@vectorize-io/hindsight-client).
## License
MIT
-57
View File
@@ -1,57 +0,0 @@
{
"name": "@vectorize-io/hindsight-all",
"version": "0.6.0",
"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",
"type": "module",
"exports": {
".": {
"types": "./dist/index.d.ts",
"import": "./dist/index.js"
}
},
"keywords": [
"hindsight",
"hindsight-all",
"memory",
"ai",
"agent",
"long-term-memory",
"llm",
"embedded-server"
],
"author": "Vectorize <[email protected]>",
"license": "MIT",
"repository": {
"type": "git",
"url": "https://github.com/vectorize-io/hindsight.git",
"directory": "hindsight-all-npm"
},
"files": [
"dist",
"README.md"
],
"scripts": {
"build": "tsup",
"dev": "tsup --watch",
"clean": "rm -rf dist",
"test": "vitest run src",
"test:watch": "vitest src",
"prepublishOnly": "npm run clean && npm run build"
},
"devDependencies": {
"@types/node": "^22.0.0",
"tsup": "^8.5.1",
"typescript": "^5.7.0",
"vitest": "^4.1.2"
},
"engines": {
"node": ">=22"
},
"overrides": {
"rollup": "^4.59.0",
"picomatch": ">=2.3.2 <3.0.0 || >=4.0.4",
"vite": ">=8.0.5"
}
}
-36
View File
@@ -1,36 +0,0 @@
import { describe, it, expect } from "vitest";
import { getEmbedCommand } from "./command.js";
describe("getEmbedCommand", () => {
it("defaults to uvx hindsight-embed@latest", () => {
expect(getEmbedCommand()).toEqual(["uvx", "hindsight-embed@latest"]);
});
it("honours an explicit version", () => {
expect(getEmbedCommand({ embedVersion: "0.5.0" })).toEqual(["uvx", "[email protected]"]);
});
it("treats an empty version as latest", () => {
expect(getEmbedCommand({ embedVersion: "" })).toEqual(["uvx", "hindsight-embed@latest"]);
});
it("uses uv run --directory when a local path is given", () => {
expect(getEmbedCommand({ embedPackagePath: "/abs/path" })).toEqual([
"uv",
"run",
"--directory",
"/abs/path",
"hindsight-embed",
]);
});
it("local path takes precedence over version", () => {
expect(getEmbedCommand({ embedPackagePath: "/abs/path", embedVersion: "0.5.0" })).toEqual([
"uv",
"run",
"--directory",
"/abs/path",
"hindsight-embed",
]);
});
});
-25
View File
@@ -1,25 +0,0 @@
/**
* Resolve the command that invokes the `hindsight-embed` Python CLI.
*
* - If `embedPackagePath` is set, runs the package from a local checkout via
* `uv run --directory <path> hindsight-embed`. Used for in-repo development.
* - Otherwise runs it via `uvx hindsight-embed@<version>` so no global install
* is required.
*
* Returns the argv as `[command, ...baseArgs]` suitable for `spawn()` /
* `execFile()` (never shell-interpolated).
*/
export interface EmbedCommandOptions {
/** Version spec passed to uvx (e.g. "latest", "0.5.0"). Default: "latest". */
embedVersion?: string;
/** Local checkout path. When set, overrides `embedVersion` and uses `uv run`. */
embedPackagePath?: string;
}
export function getEmbedCommand(opts: EmbedCommandOptions = {}): string[] {
if (opts.embedPackagePath) {
return ["uv", "run", "--directory", opts.embedPackagePath, "hindsight-embed"];
}
const version = opts.embedVersion && opts.embedVersion.length > 0 ? opts.embedVersion : "latest";
return ["uvx", `hindsight-embed@${version}`];
}
-7
View File
@@ -1,7 +0,0 @@
export { HindsightServer } from "./server.js";
export { getEmbedCommand } from "./command.js";
export { silentLogger, consoleLogger } from "./logger.js";
export type { Logger } from "./logger.js";
export type { EmbedCommandOptions } from "./command.js";
export type { HindsightServerOptions } from "./types.js";
-29
View File
@@ -1,29 +0,0 @@
/**
* Pluggable logger interface.
*
* This package does not own any logging infrastructure — consumers inject
* whatever they want (console, pino, openclaw's logger, a no-op). The default
* is silent so embedding this package never adds noise to an unrelated app.
*/
export interface Logger {
debug(msg: string): void;
info(msg: string): void;
warn(msg: string): void;
error(msg: string): void;
}
/** Logger that drops every call. Used when no logger is passed. */
export const silentLogger: Logger = {
debug: () => {},
info: () => {},
warn: () => {},
error: () => {},
};
/** Logger that writes to the standard console. Handy for CLIs and tests. */
export const consoleLogger: Logger = {
debug: (msg) => console.debug(msg),
info: (msg) => console.log(msg),
warn: (msg) => console.warn(msg),
error: (msg) => console.error(msg),
};
-35
View File
@@ -1,35 +0,0 @@
import { describe, it, expect } from "vitest";
import { HindsightServer } from "./server.js";
describe("HindsightServer construction", () => {
it("defaults base URL to http://127.0.0.1:8888", () => {
const server = new HindsightServer();
expect(server.getBaseUrl()).toBe("http://127.0.0.1:8888");
expect(server.getProfile()).toBe("default");
});
it("honours custom profile, port, and host", () => {
const server = new HindsightServer({ profile: "app", port: 9077, host: "0.0.0.0" });
expect(server.getProfile()).toBe("app");
expect(server.getBaseUrl()).toBe("http://0.0.0.0:9077");
});
it("accepts open env pass-through without complaining about unknown keys", () => {
const server = new HindsightServer({
env: {
HINDSIGHT_API_LLM_PROVIDER: "openai",
HINDSIGHT_API_LLM_MODEL: "gpt-4o-mini",
// A field that does not exist today — should still be accepted
HINDSIGHT_FUTURE_FLAG: "enabled",
},
});
expect(server).toBeInstanceOf(HindsightServer);
});
it("exposes checkHealth that returns false when no daemon is running", async () => {
// Random high port that nothing is listening on.
const server = new HindsightServer({ port: 1, readyTimeoutMs: 100 });
const healthy = await server.checkHealth();
expect(healthy).toBe(false);
});
});
-322
View File
@@ -1,322 +0,0 @@
import { spawn } from "child_process";
import { getEmbedCommand } from "./command.js";
import { silentLogger } from "./logger.js";
import type { Logger } from "./logger.js";
import type { HindsightServerOptions } from "./types.js";
const DEFAULT_PORT = 8888;
const DEFAULT_HOST = "127.0.0.1";
const DEFAULT_PROFILE = "default";
const DEFAULT_READY_TIMEOUT_MS = 30_000;
const DEFAULT_READY_POLL_INTERVAL_MS = 1_000;
/**
* Manages the lifecycle of a local Hindsight daemon from a Node.js process.
*
* On {@link start}, this class:
* 1. Resolves the `hindsight-embed` command (via `uvx` or a local `uv run`).
* 2. Runs `profile create <name> --merge --port <port> [--env K=V ...]`
* with every entry in {@link HindsightServerOptions.env} forwarded as
* an `--env` flag.
* 3. Runs `daemon --profile <name> start` and waits for the start command
* to exit.
* 4. Polls `http://host:port/health` until it returns `200` or the
* `readyTimeoutMs` budget is exhausted.
*
* On {@link stop}, it runs `daemon --profile <name> stop` and returns once
* the command exits (or after a short grace period).
*
* This is the Node.js equivalent of the Python `hindsight-all` package's
* `HindsightServer`: a thin programmatic lifecycle wrapper around the
* Hindsight daemon. It does NOT ship an HTTP client — once `start()`
* resolves, use `@vectorize-io/hindsight-client` against `getBaseUrl()` for
* retain / recall / reflect.
*
* The class is deliberately transparent about the daemon: new CLI flags or
* environment variables never require a code change here — callers can pass
* them via `env`, `extraProfileCreateArgs`, or `extraDaemonStartArgs`.
*/
export class HindsightServer {
private readonly profile: string;
private readonly port: number;
private readonly host: string;
private readonly baseUrl: string;
private readonly embedVersion: string | undefined;
private readonly embedPackagePath: string | undefined;
private readonly userEnv: Record<string, string | undefined>;
private readonly extraProfileCreateArgs: string[];
private readonly extraDaemonStartArgs: string[];
private readonly platformCpuWorkaround: boolean;
private readonly readyTimeoutMs: number;
private readonly readyPollIntervalMs: number;
private readonly logger: Logger;
constructor(opts: HindsightServerOptions = {}) {
this.profile = opts.profile ?? DEFAULT_PROFILE;
this.port = opts.port ?? DEFAULT_PORT;
this.host = opts.host ?? DEFAULT_HOST;
this.baseUrl = `http://${this.host}:${this.port}`;
this.embedVersion = opts.embedVersion;
this.embedPackagePath = opts.embedPackagePath;
this.userEnv = opts.env ?? {};
this.extraProfileCreateArgs = opts.extraProfileCreateArgs ?? [];
this.extraDaemonStartArgs = opts.extraDaemonStartArgs ?? [];
this.platformCpuWorkaround = opts.platformCpuWorkaround ?? process.platform === "darwin";
this.readyTimeoutMs = opts.readyTimeoutMs ?? DEFAULT_READY_TIMEOUT_MS;
this.readyPollIntervalMs = opts.readyPollIntervalMs ?? DEFAULT_READY_POLL_INTERVAL_MS;
this.logger = opts.logger ?? silentLogger;
}
/** The base URL the daemon listens on (`http://host:port`). */
getBaseUrl(): string {
return this.baseUrl;
}
/** The profile name this server operates on. */
getProfile(): string {
return this.profile;
}
/**
* Ensure the daemon is configured and running. Idempotent — the underlying
* `profile create --merge` and `daemon start` commands tolerate re-runs.
*/
async start(): Promise<void> {
this.logger.info(`[hindsight] starting daemon for profile "${this.profile}"`);
const env = this.buildEnv();
await this.configureProfile(env);
await this.startDaemon(env);
await this.waitForReady();
this.logger.info(`[hindsight] daemon ready at ${this.baseUrl}`);
}
/** Stop the daemon. Never throws — logs and resolves even on failure. */
async stop(): Promise<void> {
this.logger.info(`[hindsight] stopping daemon for profile "${this.profile}"`);
const [cmd, ...baseArgs] = getEmbedCommand({
embedVersion: this.embedVersion,
embedPackagePath: this.embedPackagePath,
});
const args = [...baseArgs, "daemon", "--profile", this.profile, "stop"];
const child = spawn(cmd, args, { stdio: "pipe" });
this.pipeOutput(child, "daemon.stop");
await new Promise<void>((resolve) => {
const timeout = setTimeout(() => {
this.logger.warn(`[hindsight] daemon stop timed out after 5s`);
resolve();
}, 5_000);
child.on("exit", () => {
clearTimeout(timeout);
this.logger.info(`[hindsight] daemon stopped`);
resolve();
});
child.on("error", (err) => {
clearTimeout(timeout);
this.logger.warn(`[hindsight] error stopping daemon: ${err.message}`);
resolve();
});
});
}
/** Probe `/health` once with a short timeout. */
async checkHealth(): Promise<boolean> {
try {
const res = await fetch(`${this.baseUrl}/health`, {
signal: AbortSignal.timeout(2_000),
});
return res.ok;
} catch {
return false;
}
}
// -------------------------------------------------------------------------
// Internal
// -------------------------------------------------------------------------
/**
* Merge the process env, the caller-supplied `env`, and (on macOS) the
* embeddings CPU workaround. Caller-supplied values always win over the
* workaround; undefined values are dropped.
*/
private buildEnv(): NodeJS.ProcessEnv {
const merged: NodeJS.ProcessEnv = { ...process.env };
if (this.platformCpuWorkaround && process.platform === "darwin") {
merged["HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU"] = "1";
merged["HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU"] = "1";
}
for (const [key, value] of Object.entries(this.userEnv)) {
if (value !== undefined) {
merged[key] = value;
}
}
return merged;
}
/**
* Run `profile create <name> --merge --port <port> [--env K=V ...]`.
* Every entry in the merged env that was passed via {@link userEnv} (or
* auto-applied by the CPU workaround) is forwarded as `--env`.
*/
private async configureProfile(env: NodeJS.ProcessEnv): Promise<void> {
this.logger.info(`[hindsight] configuring profile "${this.profile}"`);
const [cmd, ...baseArgs] = getEmbedCommand({
embedVersion: this.embedVersion,
embedPackagePath: this.embedPackagePath,
});
const createArgs = [
...baseArgs,
"profile",
"create",
this.profile,
"--merge",
"--port",
String(this.port),
];
// Forward every env var that the caller intended for the daemon as --env.
// We only forward keys the caller explicitly set (userEnv) plus the CPU
// workaround values — not the entire process.env, to avoid leaking random
// host state into profile config.
const envForProfile = this.collectProfileEnv(env);
for (const [key, value] of Object.entries(envForProfile)) {
createArgs.push("--env", `${key}=${value}`);
}
createArgs.push(...this.extraProfileCreateArgs);
await this.runCommand(cmd, createArgs, env, "profile.create");
}
/** Collect only the env vars that should be written into the profile file. */
private collectProfileEnv(env: NodeJS.ProcessEnv): Record<string, string> {
const out: Record<string, string> = {};
// 1. User-supplied env — always forwarded.
for (const [key, value] of Object.entries(this.userEnv)) {
if (value !== undefined) {
out[key] = value;
}
}
// 2. CPU workaround — only if auto-applied and not already overridden.
if (this.platformCpuWorkaround && process.platform === "darwin") {
const cpuKeys = [
"HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU",
"HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU",
];
for (const key of cpuKeys) {
if (!(key in out) && env[key] !== undefined) {
out[key] = env[key] as string;
}
}
}
return out;
}
private async startDaemon(env: NodeJS.ProcessEnv): Promise<void> {
const [cmd, ...baseArgs] = getEmbedCommand({
embedVersion: this.embedVersion,
embedPackagePath: this.embedPackagePath,
});
const args = [
...baseArgs,
"daemon",
"--profile",
this.profile,
"start",
...this.extraDaemonStartArgs,
];
await this.runCommand(cmd, args, env, "daemon.start");
}
/**
* Spawn `cmd` with `args`, pipe its output through the logger, and resolve
* once it exits with code 0. Rejects on non-zero exit or spawn error.
*/
private async runCommand(
cmd: string,
args: string[],
env: NodeJS.ProcessEnv,
label: string
): Promise<void> {
const child = spawn(cmd, args, { stdio: "pipe", env });
let output = "";
child.stdout?.on("data", (data: Buffer) => {
const text = data.toString();
output += text;
for (const line of text.trimEnd().split("\n")) {
if (line) this.logger.info(`[hindsight:${label}] ${line}`);
}
});
child.stderr?.on("data", (data: Buffer) => {
const text = data.toString();
output += text;
for (const line of text.trimEnd().split("\n")) {
if (line) this.logger.warn(`[hindsight:${label}] ${line}`);
}
});
await new Promise<void>((resolve, reject) => {
child.on("exit", (code) => {
if (code === 0) {
resolve();
} else {
reject(new Error(`${label} failed with code ${code}: ${output.trim()}`));
}
});
child.on("error", (err) => {
reject(new Error(`${label} failed to spawn: ${err.message}`, { cause: err }));
});
});
}
/** Stream a spawned child's stdout/stderr through the logger without blocking. */
private pipeOutput(child: ReturnType<typeof spawn>, label: string): void {
child.stdout?.on("data", (data: Buffer) => {
for (const line of data.toString().trimEnd().split("\n")) {
if (line) this.logger.info(`[hindsight:${label}] ${line}`);
}
});
child.stderr?.on("data", (data: Buffer) => {
for (const line of data.toString().trimEnd().split("\n")) {
if (line) this.logger.warn(`[hindsight:${label}] ${line}`);
}
});
}
/** Poll `/health` until it succeeds or `readyTimeoutMs` elapses. */
private async waitForReady(): Promise<void> {
const deadline = Date.now() + this.readyTimeoutMs;
let attempt = 0;
while (Date.now() < deadline) {
attempt++;
try {
const res = await fetch(`${this.baseUrl}/health`, {
signal: AbortSignal.timeout(this.readyPollIntervalMs),
});
if (res.ok) {
this.logger.debug(`[hindsight] health check passed (attempt ${attempt})`);
return;
}
} catch {
// expected while the daemon is still booting
}
await new Promise((resolve) => setTimeout(resolve, this.readyPollIntervalMs));
}
throw new Error(
`Hindsight daemon did not become ready within ${this.readyTimeoutMs}ms at ${this.baseUrl}`
);
}
}
-54
View File
@@ -1,54 +0,0 @@
import type { Logger } from "./logger.js";
/**
* Options for {@link HindsightServer}.
*
* The server is intentionally thin and pass-through: anything configurable
* on the daemon side (env vars or CLI flags) can be set here without needing
* a new dedicated option. Use {@link env} for `HINDSIGHT_*` / `OPENAI_API_KEY` /
* custom provider settings, and the two `extra*` arrays to append raw CLI
* args to `profile create` or `daemon start`.
*
* For talking to the daemon after `start()`, use `@vectorize-io/hindsight-client`
* against `server.getBaseUrl()`. This package does not ship its own HTTP
* client.
*/
export interface HindsightServerOptions {
/** Profile name used for `--profile <name>` on every sub-command. Default: `"default"`. */
profile?: string;
/** TCP port the daemon listens on. Default: `8888`. */
port?: number;
/** Hostname the daemon binds to (for health checks). Default: `127.0.0.1`. */
host?: string;
/** Version of the underlying `hindsight-embed` PyPI package to run via `uvx`. Default: `"latest"`. */
embedVersion?: string;
/** Local path to a `hindsight-embed` checkout — takes precedence over `embedVersion`. */
embedPackagePath?: string;
/**
* Environment variables passed to the daemon process AND written into the
* profile via repeated `--env KEY=VALUE` flags. This is the preferred way
* to surface any `HINDSIGHT_API_*` / `HINDSIGHT_EMBED_*` setting — adding a
* new daemon env var never requires a wrapper update.
*
* Values of `undefined` are dropped (so you can spread conditionally).
*/
env?: Record<string, string | undefined>;
/** Extra args appended verbatim to `hindsight-embed profile create <name> --merge ...`. */
extraProfileCreateArgs?: string[];
/** Extra args appended verbatim to `hindsight-embed daemon --profile <name> start ...`. */
extraDaemonStartArgs?: string[];
/**
* On macOS, automatically set
* `HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU=1` and
* `HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU=1` to avoid Metal/MPS crashes in
* daemon mode. Default: `true` on `darwin`, ignored elsewhere. Any value set
* explicitly in {@link env} wins over the auto-applied value.
*/
platformCpuWorkaround?: boolean;
/** Max time (ms) to wait for `/health` to return 200. Default: `30_000`. */
readyTimeoutMs?: number;
/** Polling interval (ms) while waiting for `/health`. Default: `1_000`. */
readyPollIntervalMs?: number;
/** Optional pluggable logger. Default: silent. */
logger?: Logger;
}
-18
View File
@@ -1,18 +0,0 @@
{
"compilerOptions": {
"target": "ES2022",
"module": "ES2022",
"lib": ["ES2022"],
"moduleResolution": "node",
"declaration": true,
"outDir": "./dist",
"rootDir": "./src",
"strict": true,
"esModuleInterop": true,
"skipLibCheck": true,
"forceConsistentCasingInFileNames": true,
"resolveJsonModule": true
},
"include": ["src/**/*"],
"exclude": ["node_modules", "dist", "src/**/*.test.ts"]
}
-11
View File
@@ -1,11 +0,0 @@
import { defineConfig } from "tsup";
export default defineConfig({
entry: ["src/index.ts"],
format: ["esm"],
dts: true,
outDir: "dist",
clean: true,
sourcemap: true,
bundle: true,
});
-8
View File
@@ -1,8 +0,0 @@
import { defineConfig } from "vitest/config";
export default defineConfig({
test: {
include: ["src/**/*.test.ts"],
environment: "node",
},
});
-33
View File
@@ -1,33 +0,0 @@
[build-system]
requires = ["setuptools>=61"]
build-backend = "setuptools.build_meta"
[project]
name = "hindsight-all-slim"
version = "0.6.0"
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.4.17",
"hindsight-client>=0.0.7",
"hindsight-embed>=0.1.0",
]
[tool.uv.sources]
hindsight-api-slim = { workspace = true }
hindsight-client = { workspace = true }
hindsight-embed = { workspace = true }
[project.optional-dependencies]
test = [
"pytest>=7.0.0",
"pytest-asyncio>=0.21.0",
]
[tool.setuptools]
packages = []
[tool.pytest.ini_options]
asyncio_mode = "auto"
asyncio_default_fixture_loop_scope = "function"
-48
View File
@@ -1,48 +0,0 @@
# hindsight-all
All-in-one package for Hindsight - Agent Memory That Works Like Human Memory
## Quick Start
```python
from hindsight import start_server, HindsightClient
# Start server with embedded PostgreSQL
server = start_server(
llm_provider="groq",
llm_api_key="your-api-key",
llm_model="openai/gpt-oss-120b"
)
# Create client
client = HindsightClient(base_url=server.url)
# Store memories
client.put(agent_id="assistant", content="User prefers Python for data analysis")
# Search memories
results = client.search(agent_id="assistant", query="programming preferences")
# Generate contextual response
response = client.think(agent_id="assistant", query="What languages should I recommend?")
# Stop server when done
server.stop()
```
## Using Context Manager
```python
from hindsight import HindsightServer, HindsightClient
with HindsightServer(llm_provider="groq", llm_api_key="...") as server:
client = HindsightClient(base_url=server.url)
# ... use client ...
# Server automatically stops
```
## Installation
```bash
pip install hindsight-all
```
-423
View File
@@ -1,423 +0,0 @@
"""
Wrapper for Hindsight client that adds API namespaces.
Provides organized access to different parts of the Hindsight API through
namespaces like .banks, .mental_models, etc.
"""
from __future__ import annotations
from typing import Any
from hindsight_client import Hindsight
class BanksAPI:
"""Namespace for bank-related operations.
Provides methods to create, delete, and manage memory banks.
"""
def __init__(self, client: Hindsight):
self._client = client
def create(
self,
bank_id: str,
name: str | None = None,
mission: str | None = None,
disposition: dict[str, Any] | None = None,
) -> Any:
"""Create a new bank.
Args:
bank_id: Unique identifier for the bank.
name: Optional display name for the bank.
mission: Optional mission statement for the bank.
disposition: Optional disposition configuration dict.
Returns:
Bank creation response from the API.
"""
return self._client.create_bank(
bank_id=bank_id,
name=name,
mission=mission,
disposition=disposition,
)
def delete(self, bank_id: str) -> Any:
"""Delete a bank.
Args:
bank_id: The ID of the bank to delete.
Returns:
Deletion response from the API.
"""
return self._client.delete_bank(bank_id=bank_id)
def set_mission(self, bank_id: str, mission: str) -> Any:
"""Set or update the mission for a bank.
Args:
bank_id: The ID of the bank.
mission: The mission statement to set.
Returns:
API response confirming the update.
"""
return self._client.set_mission(bank_id=bank_id, mission=mission)
def set_disposition(self, bank_id: str, disposition: dict[str, Any]) -> Any:
"""Set or update the disposition for a bank.
Args:
bank_id: The ID of the bank.
disposition: The disposition configuration dict.
Returns:
API response confirming the update.
"""
return self._client.set_disposition(bank_id=bank_id, disposition=disposition)
def list(self) -> Any:
"""List all banks.
Returns:
List of banks from the API.
"""
from hindsight_client.hindsight_client import _run_async
return _run_async(self._client._banks_api.list_banks())
class MentalModelsAPI:
"""Namespace for mental model operations.
Mental models are reusable knowledge structures that guide agent behavior.
"""
def __init__(self, client: Hindsight):
self._client = client
def create(
self,
bank_id: str,
name: str,
content: str,
tags: list[str] | None = None,
) -> Any:
"""Create a new mental model.
Args:
bank_id: The ID of the bank to add the model to.
name: Name for the mental model.
content: The content/instructions for the mental model.
tags: Optional list of tags for categorization.
Returns:
Creation response from the API.
"""
return self._client.create_mental_model(
bank_id=bank_id,
name=name,
content=content,
tags=tags,
)
def list(self, bank_id: str, tags: list[str] | None = None) -> Any:
"""List all mental models for a bank.
Args:
bank_id: The ID of the bank.
tags: Optional filter by tags.
Returns:
List of mental models.
"""
return self._client.list_mental_models(bank_id=bank_id, tags=tags)
def get(self, bank_id: str, mental_model_id: str) -> Any:
"""Get a specific mental model.
Args:
bank_id: The ID of the bank.
mental_model_id: The ID of the mental model.
Returns:
The mental model details.
"""
return self._client.get_mental_model(bank_id=bank_id, mental_model_id=mental_model_id)
def refresh(self, bank_id: str, mental_model_id: str) -> Any:
"""Refresh a mental model.
Args:
bank_id: The ID of the bank.
mental_model_id: The ID of the mental model to refresh.
Returns:
Refresh response from the API.
"""
return self._client.refresh_mental_model(bank_id=bank_id, mental_model_id=mental_model_id)
def update(
self,
bank_id: str,
mental_model_id: str,
name: str | None = None,
content: str | None = None,
tags: list[str] | None = None,
) -> Any:
"""Update a mental model.
Args:
bank_id: The ID of the bank.
mental_model_id: The ID of the mental model to update.
name: Optional new name.
content: Optional new content.
tags: Optional new tags list.
Returns:
Update response from the API.
"""
return self._client.update_mental_model(
bank_id=bank_id,
mental_model_id=mental_model_id,
name=name,
content=content,
tags=tags,
)
def delete(self, bank_id: str, mental_model_id: str) -> Any:
"""Delete a mental model.
Args:
bank_id: The ID of the bank.
mental_model_id: The ID of the mental model to delete.
Returns:
Deletion response from the API.
"""
return self._client.delete_mental_model(bank_id=bank_id, mental_model_id=mental_model_id)
class DirectivesAPI:
"""Namespace for directive operations.
Directives are explicit instructions that guide agent behavior.
"""
def __init__(self, client: Hindsight):
self._client = client
def create(
self,
bank_id: str,
name: str,
content: str,
tags: list[str] | None = None,
) -> Any:
"""Create a new directive.
Args:
bank_id: The ID of the bank to add the directive to.
name: Name for the directive.
content: The directive content/instructions.
tags: Optional list of tags for categorization.
Returns:
Creation response from the API.
"""
return self._client.create_directive(
bank_id=bank_id,
name=name,
content=content,
tags=tags,
)
def list(self, bank_id: str, tags: list[str] | None = None) -> Any:
"""List all directives for a bank.
Args:
bank_id: The ID of the bank.
tags: Optional filter by tags.
Returns:
List of directives.
"""
return self._client.list_directives(bank_id=bank_id, tags=tags)
def get(self, bank_id: str, directive_id: str) -> Any:
"""Get a specific directive.
Args:
bank_id: The ID of the bank.
directive_id: The ID of the directive.
Returns:
The directive details.
"""
return self._client.get_directive(bank_id=bank_id, directive_id=directive_id)
def update(
self,
bank_id: str,
directive_id: str,
name: str | None = None,
content: str | None = None,
tags: list[str] | None = None,
) -> Any:
"""Update a directive.
Args:
bank_id: The ID of the bank.
directive_id: The ID of the directive to update.
name: Optional new name.
content: Optional new content.
tags: Optional new tags list.
Returns:
Update response from the API.
"""
return self._client.update_directive(
bank_id=bank_id,
directive_id=directive_id,
name=name,
content=content,
tags=tags,
)
def delete(self, bank_id: str, directive_id: str) -> Any:
"""Delete a directive.
Args:
bank_id: The ID of the bank.
directive_id: The ID of the directive to delete.
Returns:
Deletion response from the API.
"""
return self._client.delete_directive(bank_id=bank_id, directive_id=directive_id)
class MemoriesAPI:
"""Namespace for memory operations.
Provides methods to query and retrieve stored memories.
"""
def __init__(self, client: Hindsight):
self._client = client
def list(
self,
bank_id: str,
type: str | None = None,
search_query: str | None = None,
limit: int = 100,
offset: int = 0,
) -> Any:
"""List memories in a bank.
Args:
bank_id: The ID of the bank to query.
type: Optional filter by memory type.
search_query: Optional search query for filtering.
limit: Maximum number of results to return (default: 100).
offset: Number of results to skip for pagination (default: 0).
Returns:
List of memories matching the criteria.
"""
return self._client.list_memories(
bank_id=bank_id,
type=type,
search_query=search_query,
limit=limit,
offset=offset,
)
class HindsightClient(Hindsight):
"""
Enhanced Hindsight client with organized API namespaces.
This wrapper extends the auto-generated Hindsight client with organized
access to different parts of the API through namespaces.
Example:
```python
from hindsight import HindsightClient
client = HindsightClient(base_url="http://localhost:8888")
# Core operations (inherited from Hindsight)
client.retain(bank_id="test", content="Hello")
results = client.recall(bank_id="test", query="Hello")
# Organized API access through namespaces
client.banks.create(bank_id="test", name="Test Bank")
models = client.mental_models.list(bank_id="test")
directives = client.directives.list(bank_id="test")
memories = client.memories.list(bank_id="test")
```
Attributes:
banks: Namespace for bank management operations.
mental_models: Namespace for mental model operations.
directives: Namespace for directive operations.
memories: Namespace for memory listing operations.
"""
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
self._banks_namespace: BanksAPI | None = None
self._mental_models_namespace: MentalModelsAPI | None = None
self._directives_namespace: DirectivesAPI | None = None
self._memories_namespace: MemoriesAPI | None = None
@property
def banks(self) -> BanksAPI:
"""Access bank management operations.
Returns:
BanksAPI instance for bank operations.
"""
if self._banks_namespace is None:
self._banks_namespace = BanksAPI(self)
return self._banks_namespace
@property
def mental_models(self) -> MentalModelsAPI:
"""Access mental model operations.
Returns:
MentalModelsAPI instance for mental model operations.
"""
if self._mental_models_namespace is None:
self._mental_models_namespace = MentalModelsAPI(self)
return self._mental_models_namespace
@property
def directives(self) -> DirectivesAPI:
"""Access directive operations.
Returns:
DirectivesAPI instance for directive operations.
"""
if self._directives_namespace is None:
self._directives_namespace = DirectivesAPI(self)
return self._directives_namespace
@property
def memories(self) -> MemoriesAPI:
"""Access memory listing operations.
Returns:
MemoriesAPI instance for memory operations.
"""
if self._memories_namespace is None:
self._memories_namespace = MemoriesAPI(self)
return self._memories_namespace
View File
@@ -1,56 +0,0 @@
"""
Unit test for _cleanup lock timeout behavior.
Verifies that _cleanup completes even when the lock is held by another thread,
instead of hanging indefinitely (fixes #952).
"""
import threading
import time
from unittest.mock import MagicMock, patch
import pytest
def test_cleanup_completes_when_lock_held():
"""
_cleanup should complete (best-effort) even when self._lock is held
by another thread, e.g. during a long _ensure_started call.
"""
with patch.dict("sys.modules", {
"hindsight_client": MagicMock(),
"hindsight_embed": MagicMock(),
"hindsight.api_namespaces": MagicMock(),
}):
from hindsight.embedded import HindsightEmbedded
client = HindsightEmbedded.__new__(HindsightEmbedded)
client.profile = "test"
client._lock = threading.Lock()
client._closed = False
client._client = None
client._started = False
client._ui = False
# Simulate another thread holding the lock
client._lock.acquire()
cleanup_done = threading.Event()
def run_cleanup():
client._cleanup()
cleanup_done.set()
t = threading.Thread(target=run_cleanup)
t.start()
# Cleanup should complete within the timeout (5s) + margin
assert cleanup_done.wait(timeout=8.0), (
"_cleanup hung instead of timing out on lock acquisition"
)
# Release the lock from the simulating thread
client._lock.release()
t.join(timeout=1.0)
assert client._closed, "Client should be marked as closed after cleanup"
-137
View File
@@ -1,137 +0,0 @@
# Hindsight API
**Memory System for AI Agents** — Temporal + Semantic + Entity Memory Architecture using PostgreSQL with pgvector.
Hindsight gives AI agents persistent memory that works like human memory: it stores facts, tracks entities and relationships, handles temporal reasoning ("what happened last spring?"), and forms opinions based on configurable disposition traits.
## Installation
```bash
pip install hindsight-api
```
## Quick Start
### Run the Server
```bash
# Set your LLM provider
export HINDSIGHT_API_LLM_PROVIDER=openai
export HINDSIGHT_API_LLM_API_KEY=sk-xxxxxxxxxxxx
# Start the server (uses embedded PostgreSQL by default)
hindsight-api
```
The server starts at http://localhost:8888 with:
- REST API for memory operations
- MCP server at `/mcp` for tool-use integration
### Use the Python API
```python
from hindsight_api import MemoryEngine
# Create and initialize the memory engine
memory = MemoryEngine()
await memory.initialize()
# Create a memory bank for your agent
bank = await memory.create_memory_bank(
name="my-assistant",
background="A helpful coding assistant"
)
# Store a memory
await memory.retain(
memory_bank_id=bank.id,
content="The user prefers Python for data science projects"
)
# Recall memories
results = await memory.recall(
memory_bank_id=bank.id,
query="What programming language does the user prefer?"
)
# Reflect with reasoning
response = await memory.reflect(
memory_bank_id=bank.id,
query="Should I recommend Python or R for this ML project?"
)
```
## CLI Options
```bash
hindsight-api --help
# Common options
hindsight-api --port 9000 # Custom port (default: 8888)
hindsight-api --host 127.0.0.1 # Bind to localhost only
hindsight-api --workers 4 # Multiple worker processes
hindsight-api --log-level debug # Verbose logging
```
## Configuration
Configure via environment variables:
| Variable | Description | Default |
|----------|-------------|---------|
| `HINDSIGHT_API_DATABASE_URL` | PostgreSQL connection string | `pg0` (embedded) |
| `HINDSIGHT_API_LLM_PROVIDER` | `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio` | `openai` |
| `HINDSIGHT_API_LLM_API_KEY` | API key for LLM provider | - |
| `HINDSIGHT_API_LLM_MODEL` | Model name | `gpt-4o-mini` |
| `HINDSIGHT_API_HOST` | Server bind address | `0.0.0.0` |
| `HINDSIGHT_API_PORT` | Server port | `8888` |
### Example with External PostgreSQL
```bash
export HINDSIGHT_API_DATABASE_URL=postgresql://user:pass@localhost:5432/hindsight
export HINDSIGHT_API_LLM_PROVIDER=groq
export HINDSIGHT_API_LLM_API_KEY=gsk_xxxxxxxxxxxx
hindsight-api
```
## Docker
```bash
docker run --rm -it -p 8888:8888 \
-e HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY \
-v $HOME/.hindsight-docker:/home/hindsight/.pg0 \
ghcr.io/vectorize-io/hindsight:latest
```
## MCP Server
For local MCP integration without running the full API server:
```bash
hindsight-local-mcp
```
This runs a stdio-based MCP server that can be used directly with MCP-compatible clients.
## Key Features
- **Multi-Strategy Retrieval (TEMPR)** — Semantic, keyword, graph, and temporal search combined with RRF fusion
- **Entity Graph** — Automatic entity extraction and relationship tracking
- **Temporal Reasoning** — Native support for time-based queries
- **Disposition Traits** — Configurable skepticism, literalism, and empathy influence opinion formation
- **Three Memory Types** — World facts, bank actions, and formed opinions with confidence scores
## Documentation
Full documentation: [https://hindsight.vectorize.io](https://hindsight.vectorize.io)
- [Installation Guide](https://hindsight.vectorize.io/developer/installation)
- [Configuration Reference](https://hindsight.vectorize.io/developer/configuration)
- [API Reference](https://hindsight.vectorize.io/api-reference)
- [Python SDK](https://hindsight.vectorize.io/sdks/python)
## License
Apache 2.0
@@ -1,42 +0,0 @@
"""Dialect dispatcher for Alembic migrations.
Each migration file declares a ``_pg_upgrade``/``_oracle_upgrade`` (and matching
downgrades) function and routes ``upgrade()``/``downgrade()`` through
``run_for_dialect``. The helper inspects the live connection's dialect name and
runs the matching function — or no-ops if the migration doesn't apply to the
current backend.
Use ``None`` (or omit the kwarg) when a migration intentionally has no effect
on a dialect; the helper treats it as a no-op.
"""
from __future__ import annotations
from collections.abc import Callable
from alembic import op
DialectFn = Callable[[], None]
_SUPPORTED = ("postgresql", "oracle")
def run_for_dialect(
*,
pg: DialectFn | None = None,
oracle: DialectFn | None = None,
) -> None:
"""Dispatch to the function matching the current bind's dialect.
Args:
pg: Function to run when the active bind is PostgreSQL.
oracle: Function to run when the active bind is Oracle.
Unrecognized dialects raise; an explicit ``None`` for the active dialect
is a no-op (the migration deliberately does nothing here).
"""
name = op.get_bind().dialect.name
if name not in _SUPPORTED:
raise RuntimeError(f"Unsupported dialect for migration dispatch: {name!r}. Expected one of {_SUPPORTED}.")
fn = {"postgresql": pg, "oracle": oracle}[name]
if fn is not None:
fn()
@@ -1,191 +0,0 @@
"""
Alembic environment for Hindsight.
Supports two dialects:
* PostgreSQL (sync psycopg2 driver) — default; uses ``search_path`` for
multi-tenant schema isolation and forces read-write transactions to work
around Supabase's read-only-by-default sessions.
* Oracle 23ai (``oracledb`` driver) — uses ``CURRENT_SCHEMA`` for tenant
isolation; no equivalent of ``search_path`` or read-only session quirks.
Each migration file dispatches its DDL through ``alembic._dialect.run_for_dialect``
so a single revision tree serves both backends.
"""
import logging
import os
from pathlib import Path
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from alembic import context
from dotenv import load_dotenv
from sqlalchemy import Connection, engine_from_config, pool
from sqlalchemy.engine import Engine
from hindsight_api.db_url import is_oracle_url, to_libpq_url
from hindsight_api.models import Base
def load_env() -> None:
"""Load environment variables from .env (skipped if already configured)."""
if os.getenv("HINDSIGHT_API_DATABASE_URL"):
return
root_dir = Path(__file__).parent.parent.parent
env_file = root_dir / ".env"
if env_file.exists():
load_dotenv(env_file)
load_env()
config = context.config
target_metadata = Base.metadata
def _normalize_oracle_url(url: str) -> str:
"""Coerce an Oracle URL into the SQLAlchemy form the oracledb dialect expects.
Two issues to handle:
1. Force the ``oracle+oracledb`` driver — bare ``oracle://`` defaults to
cx_Oracle.
2. Map a path-style service to ``?service_name=...``. SQLAlchemy's oracledb
dialect treats the URL path as a *SID* (legacy), but Oracle Free /
Autonomous DB only register a service name. Without this rewrite we get
``DPY-6003: SID "FREEPDB1" is not registered`` even though the listener
is happy to accept the same name as a service.
"""
parts = urlsplit(url)
if not parts.scheme.startswith("oracle"):
return url
new_scheme = "oracle+oracledb" if parts.scheme == "oracle" else parts.scheme
service = parts.path.lstrip("/")
new_query = parts.query
new_path = parts.path
# Promote /SERVICE to ?service_name=SERVICE unless the caller already
# supplied an explicit ?sid= or ?service_name=.
if service and "service_name=" not in new_query and "sid=" not in new_query:
params = [(k, v) for k, v in parse_qsl(new_query, keep_blank_values=True)]
params.append(("service_name", service))
new_query = urlencode(params)
new_path = ""
return urlunsplit((new_scheme, parts.netloc, new_path, new_query, parts.fragment))
def get_database_url() -> str:
"""Resolve the migration URL from Alembic config or env, normalizing per-dialect."""
database_url = config.get_main_option("sqlalchemy.url")
if not database_url:
database_url = os.getenv("HINDSIGHT_API_DATABASE_URL")
if not database_url:
raise ValueError(
"Database URL not found. "
"Set HINDSIGHT_API_DATABASE_URL environment variable or pass database_url to run_migrations()."
)
if is_oracle_url(database_url):
database_url = _normalize_oracle_url(database_url)
else:
# PG: convert SQLAlchemy-style asyncpg URLs and ?ssl= params to libpq form
# for the sync engine used during migrations.
database_url = to_libpq_url(database_url)
config.set_main_option("sqlalchemy.url", database_url)
return database_url
def run_migrations_offline() -> None:
logging.info("running offline")
database_url = get_database_url()
context.configure(
url=database_url,
target_metadata=target_metadata,
literal_binds=True,
dialect_opts={"paramstyle": "named"},
)
with context.begin_transaction():
context.run_migrations()
def _configure_pg_session(engine: Engine, connection: Connection, target_schema: str | None) -> None:
"""PG-only: ensure the session is RW (Supabase) and bind ``search_path``."""
from sqlalchemy import event, text
@event.listens_for(engine, "connect")
def set_read_write_mode(dbapi_connection, connection_record):
cursor = dbapi_connection.cursor()
cursor.execute("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE")
if target_schema:
cursor.execute(f'CREATE SCHEMA IF NOT EXISTS "{target_schema}"')
cursor.execute(f'SET search_path TO "{target_schema}", public')
cursor.close()
connection.execute(text("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE"))
if target_schema:
connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{target_schema}"'))
connection.execute(text(f'SET search_path TO "{target_schema}", public'))
connection.commit()
def _configure_oracle_session(connection: Connection, target_schema: str | None) -> None:
"""Oracle: switch the session's default schema; tolerate DDL contention."""
from sqlalchemy import text
# Wait up to 30s for DDL locks instead of failing immediately (ORA-00054).
connection.execute(text("ALTER SESSION SET DDL_LOCK_TIMEOUT = 30"))
if target_schema:
connection.execute(text(f'ALTER SESSION SET CURRENT_SCHEMA = "{target_schema}"'))
def run_migrations_online() -> None:
database_url = get_database_url()
target_schema = config.get_main_option("target_schema")
is_oracle = is_oracle_url(database_url)
connectable = engine_from_config(
config.get_section(config.config_ini_section, {}),
prefix="sqlalchemy.",
poolclass=pool.NullPool,
)
with connectable.connect() as connection:
if is_oracle:
_configure_oracle_session(connection, target_schema)
else:
_configure_pg_session(connectable, connection, target_schema)
context_opts = {
"connection": connection,
"target_metadata": target_metadata,
}
if target_schema and not is_oracle:
# Oracle has no equivalent of PG's per-schema version table; the
# ``alembic_version`` table lives in CURRENT_SCHEMA implicitly.
context_opts["version_table_schema"] = target_schema
context.configure(**context_opts)
with context.begin_transaction():
context.run_migrations()
# Always commit. PG needs it for the explicit RW-mode SET to persist;
# Oracle needs it because each DDL auto-commits but the trailing
# ``UPDATE alembic_version`` is plain DML that would otherwise stay in
# an open transaction and roll back when the connection closes —
# producing the "schema is created but the version row is one revision
# behind" failure mode.
connection.commit()
if context.is_offline_mode():
run_migrations_offline()
else:
run_migrations_online()
@@ -1,55 +0,0 @@
"""Recreate entities trigram index on LOWER(canonical_name) for case-insensitive matching
The previous GIN trigram index on canonical_name was case-sensitive, causing
"Alice" and "alice" to have different trigram sets. This recreates it on
LOWER(canonical_name) so the % operator matches case-insensitively.
Revision ID: 2eee35aa3cfc
Revises: d6e7f8a9b0c1
Create Date: 2026-03-31
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "2eee35aa3cfc"
down_revision: str | Sequence[str] | None = "d6e7f8a9b0c1"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
# Drop the old case-sensitive trigram index
op.execute("DROP INDEX IF EXISTS entities_canonical_name_trgm_idx")
# Create case-insensitive trigram index on LOWER(canonical_name)
op.execute(
f"CREATE INDEX IF NOT EXISTS entities_canonical_name_lower_trgm_idx "
f"ON {schema}entities USING GIN (LOWER(canonical_name) gin_trgm_ops)"
)
def _pg_downgrade() -> None:
op.execute("DROP INDEX IF EXISTS entities_canonical_name_lower_trgm_idx")
schema = _get_schema_prefix()
# Restore original case-sensitive index
op.execute(
f"CREATE INDEX IF NOT EXISTS entities_canonical_name_trgm_idx "
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
)
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,50 +0,0 @@
"""Merge divergent migration heads for v0.5.3
v0.5.3 shipped with two migration heads that were never unified:
* ``c4x5y6z7a8b9`` — delta-refresh chain
(``add_last_refreshed_source_query`` ->
``add_structured_content_to_mental_models`` ->
``backsweep_orphan_observations_v2``)
* ``h3i4j5k6l7m8`` — per-bank vector indexes / audit log chain
(the ``merge_heads_and_add_unit_entities_index`` subtree)
Both fork from ``z1u2v3w4x5y6``. Upgrades from v0.5.2 still succeed — the
walker applies the three c4x5 revisions and leaves the database stamped at
both heads — but the result is a split DAG: ``alembic upgrade head``
(singular) is ambiguous, and any future migration has to pick one head as
its parent, orphaning the other.
This revision linearises the DAG into a single head. It has no schema
effect.
Revision ID: 8c6fa6f7230b
Revises: c4x5y6z7a8b9, h3i4j5k6l7m8
Create Date: 2026-04-18
"""
from collections.abc import Sequence
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "8c6fa6f7230b"
down_revision: str | Sequence[str] | None = ("c4x5y6z7a8b9", "h3i4j5k6l7m8")
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _pg_upgrade() -> None:
pass
def _pg_downgrade() -> None:
pass
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,128 +0,0 @@
"""Make memory_links.from_unit_id and memory_links.to_unit_id FKs deferrable.
Revision ID: 9f8e7d6c5b4a
Revises: o1a2b3c4d5e6
Create Date: 2026-05-03
Background
----------
Concurrent retain (which INSERTs into ``memory_links``) and any code path
that DELETEs a row whose deletion cascades into ``memory_links`` (e.g.
delta-retain superseding chunks, which CASCADEs chunks → memory_units →
memory_links) can deadlock under sustained single-tenant write load.
The deadlock cycle:
* Tx A: ``DELETE FROM chunks WHERE chunk_id = ANY(...)``
→ CASCADE acquires row locks on memory_units, then on memory_links rows
where ``to_unit_id`` matches the deleted units.
* Tx B: ``INSERT INTO memory_links (...)`` referencing one of the same
memory_units rows.
→ The immediate FK check takes ``FOR KEY SHARE`` on those memory_units
rows.
The two transactions take row locks on the same memory_units rows in
opposite orders depending on which side started first. PostgreSQL detects
the cycle and aborts one transaction; the loser is killed mid-batch, the
winner continues. Workers then retry, but under sustained write load the
pattern repeats.
Fix
---
Make both ``memory_links → memory_units`` FKs (``from_unit_id`` and
``to_unit_id``) ``DEFERRABLE INITIALLY DEFERRED``. This pushes the FK
check from INSERT time to COMMIT time:
* INSERT no longer takes ``FOR KEY SHARE`` on the memory_units row → no
contention with the cascading DELETE's row lock.
* At COMMIT the engine validates referential integrity in one shot. If a
cascade-DELETE has since removed the referenced unit, the INSERT
transaction commits OR fails with a clean FK violation (sqlstate
23503) instead of a deadlock (sqlstate 40P01).
The ``WHERE EXISTS`` filter already in ``_bulk_insert_links`` continues to
filter out the typical "stale unit_id" case at INSERT time; the deferred
FK is only the backstop for the narrow race window between the EXISTS
probe and COMMIT. ``ON DELETE CASCADE`` semantics are unchanged — only
the *timing* of the constraint check moves.
The ``entity_id`` FK on ``memory_links`` is not changed; entities are not
involved in the observed deadlock cycle and leaving the constraint
immediate keeps the error message specific when an entity row is missing.
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "9f8e7d6c5b4a"
down_revision: str | Sequence[str] | None = "o1a2b3c4d5e6"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
# The two FK constraints installed by the initial schema migration
# (5a366d414dce_initial_schema), mapped to the column they constrain.
# They reference memory_units(id) with ON DELETE CASCADE — that
# semantics is preserved; only the deferral attribute changes.
_FK_COLUMNS: dict[str, str] = {
"fk_memory_links_from_unit_id_memory_units": "from_unit_id",
"fk_memory_links_to_unit_id_memory_units": "to_unit_id",
}
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
# PostgreSQL doesn't allow altering the deferrability of an existing
# constraint with ALTER CONSTRAINT — the constraint must be dropped
# and recreated. DROP IF EXISTS makes the migration safe to re-run
# on schemas where the constraint was already recreated.
for fk_name, column in _FK_COLUMNS.items():
op.execute(f"ALTER TABLE {schema}memory_links DROP CONSTRAINT IF EXISTS {fk_name}")
op.execute(
f"""
ALTER TABLE {schema}memory_links
ADD CONSTRAINT {fk_name}
FOREIGN KEY ({column})
REFERENCES {schema}memory_units (id)
ON DELETE CASCADE
DEFERRABLE INITIALLY DEFERRED
"""
)
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
# Revert to the default (NOT DEFERRABLE) form so a downgrade actually
# restores the prior schema state, even though that re-introduces the
# deadlock window.
for fk_name, column in _FK_COLUMNS.items():
op.execute(f"ALTER TABLE {schema}memory_links DROP CONSTRAINT IF EXISTS {fk_name}")
op.execute(
f"""
ALTER TABLE {schema}memory_links
ADD CONSTRAINT {fk_name}
FOREIGN KEY ({column})
REFERENCES {schema}memory_units (id)
ON DELETE CASCADE
"""
)
def upgrade() -> None:
# PG-only: Oracle's deferrable-FK semantics differ and the deadlock
# cycle was only observed on PostgreSQL. Oracle slot intentionally
# absent → no-op there.
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,64 +0,0 @@
"""Add GIN index on source_memory_ids for observation lookup performance
Without this index, queries using the array overlap operator (&&) or array
containment (@>) on source_memory_ids require a full sequential scan over all
observation memory_units. At ~77k observations this was measured at 45ms per
query, becoming a bottleneck during consolidation recall (57-64s timeouts) and
user recall (18-27s average).
The GIN index reduces these queries to index scans: 45ms → 0.049ms (927x
speedup). Recall dropped from 18-27s to ~6s, and consolidation recall
stabilised from timeout to ~15s.
Created with CONCURRENTLY so the migration does not block reads or writes.
CONCURRENTLY requires running outside a transaction block, so the migration
emits an explicit COMMIT before the statement and uses IF NOT EXISTS for
idempotency.
Revision ID: a2b3c4d5e6f8
Revises: f7g8h9i0j1k2
Create Date: 2026-03-04
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "a2b3c4d5e6f8"
down_revision: str | Sequence[str] | None = "f7g8h9i0j1k2"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
# Commit the current Alembic transaction first.
op.execute("COMMIT")
op.execute(
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
f"ON {schema}memory_units USING GIN (source_memory_ids) "
f"WHERE source_memory_ids IS NOT NULL"
)
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute("COMMIT")
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,48 +0,0 @@
"""Add last_refreshed_source_query column to mental_models
Revision ID: a2v3w4x5y6z7
Revises: z1u2v3w4x5y6
Create Date: 2026-04-15
Tracks the source_query that was used during the most recent refresh.
Used by delta-mode refresh to detect when the query has changed: if it has,
delta mode falls back to a full regeneration because the surgical-edit
assumption (same topic, new facts) no longer holds.
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "a2v3w4x5y6z7"
down_revision: str | Sequence[str] | None = "z1u2v3w4x5y6"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"""
ALTER TABLE {schema}mental_models
ADD COLUMN IF NOT EXISTS last_refreshed_source_query TEXT
""")
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS last_refreshed_source_query")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,62 +0,0 @@
"""Add consolidation_failed_at column to memory_units for tracking persistent LLM failures.
When all LLM retries are exhausted on a single-memory batch, the memory is marked
with consolidation_failed_at instead of consolidated_at, so it is not silently lost
and can be retried later via the API.
Revision ID: a3b4c5d6e7f8
Revises: g7h8i9j0k1l2
Create Date: 2026-03-17
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "a3b4c5d6e7f8"
down_revision: str | Sequence[str] | None = "g7h8i9j0k1l2"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
op.execute(
f"""
ALTER TABLE {schema}memory_units
ADD COLUMN IF NOT EXISTS consolidation_failed_at TIMESTAMPTZ DEFAULT NULL
"""
)
# Index to efficiently query memories that failed consolidation for a given bank
op.execute(
f"""
CREATE INDEX IF NOT EXISTS idx_memory_units_consolidation_failed
ON {schema}memory_units (bank_id, consolidation_failed_at)
WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')
"""
)
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_consolidation_failed")
op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS consolidation_failed_at")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,152 +0,0 @@
"""Fix per-bank vector indexes to match configured extension
Revision ID: a4b5c6d7e8f9
Revises: 2eee35aa3cfc
Create Date: 2026-04-01
Migration d5e6f7a8b9c0 hardcoded HNSW when creating per-bank partial vector
indexes, ignoring HINDSIGHT_API_VECTOR_EXTENSION. Banks that existed when that
migration ran got HNSW indexes even when pgvectorscale (DiskANN) or vchord
was configured.
This migration detects the mismatch and recreates the affected indexes with
the correct type. Skipped entirely when the configured extension is pgvector
(the default), since those indexes are already correct.
"""
import os
from collections.abc import Sequence
from alembic import context, op
from sqlalchemy import text
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "a4b5c6d7e8f9"
down_revision: str | Sequence[str] | None = "2eee35aa3cfc"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_FACT_TYPES: dict[str, str] = {
"world": "worl",
"experience": "expr",
"observation": "obsv",
}
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _target_index_type() -> str | None:
"""Return the target index type, or None if pgvector (no fix needed)."""
ext = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
if ext == "pgvectorscale":
return "diskann"
elif ext == "vchord":
return "vchordrq"
return None
def _vector_index_using_clause() -> str:
"""Return the USING clause based on the configured vector extension."""
ext = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
if ext == "pgvectorscale":
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
elif ext == "vchord":
return "USING vchordrq (embedding vector_l2_ops)"
else:
return "USING hnsw (embedding vector_cosine_ops)"
def _pg_upgrade() -> None:
target = _target_index_type()
if target is None:
# pgvector — indexes are already HNSW, nothing to fix
return
bind = op.get_bind()
schema_name = context.config.get_main_option("target_schema")
schema = _get_schema_prefix()
table_ref = f'"{schema_name}".memory_units' if schema_name else "memory_units"
banks_ref = f'"{schema_name}".banks' if schema_name else "banks"
using_clause = _vector_index_using_clause()
pg_schema = schema_name or "public"
rows = bind.execute(text(f"SELECT bank_id, internal_id FROM {banks_ref}")).fetchall() # noqa: S608
for row in rows:
bank_id = row[0]
internal_id = str(row[1]).replace("-", "")[:16]
escaped_bank_id = bank_id.replace("'", "''")
for ft, ft_short in _FACT_TYPES.items():
idx_name = f"idx_mu_emb_{ft_short}_{internal_id}"
# Check if this index exists and what type it is
idx_info = bind.execute(
text("SELECT indexdef FROM pg_indexes WHERE schemaname = :schema AND indexname = :idx"),
{"schema": pg_schema, "idx": idx_name},
).fetchone()
if idx_info is None:
# Index doesn't exist — create it with the correct type
bind.execute(
text(
f"CREATE INDEX IF NOT EXISTS {idx_name} "
f"ON {table_ref} {using_clause} "
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped_bank_id}'"
)
)
continue
indexdef = idx_info[0].lower()
if target in indexdef:
# Already the correct type
continue
# Wrong type — drop and recreate
bind.execute(text(f"DROP INDEX IF EXISTS {schema}{idx_name}"))
bind.execute(
text(
f"CREATE INDEX IF NOT EXISTS {idx_name} "
f"ON {table_ref} {using_clause} "
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped_bank_id}'"
)
)
def _pg_downgrade() -> None:
# Downgrade recreates indexes as HNSW (the original hardcoded behavior)
target = _target_index_type()
if target is None:
return
bind = op.get_bind()
schema_name = context.config.get_main_option("target_schema")
schema = _get_schema_prefix()
table_ref = f'"{schema_name}".memory_units' if schema_name else "memory_units"
banks_ref = f'"{schema_name}".banks' if schema_name else "banks"
rows = bind.execute(text(f"SELECT bank_id, internal_id FROM {banks_ref}")).fetchall() # noqa: S608
for row in rows:
bank_id = row[0]
internal_id = str(row[1]).replace("-", "")[:16]
escaped_bank_id = bank_id.replace("'", "''")
for ft, ft_short in _FACT_TYPES.items():
idx_name = f"idx_mu_emb_{ft_short}_{internal_id}"
bind.execute(text(f"DROP INDEX IF EXISTS {schema}{idx_name}"))
bind.execute(
text(
f"CREATE INDEX IF NOT EXISTS {idx_name} "
f"ON {table_ref} USING hnsw (embedding vector_cosine_ops) "
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped_bank_id}'"
)
)
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,42 +0,0 @@
"""add content_hash to chunks table for delta retain
Revision ID: b3c4d5e6f7a8
Revises: a3b4c5d6e7f8
Create Date: 2026-03-25
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "b3c4d5e6f7a8"
down_revision: str | Sequence[str] | None = "a3b4c5d6e7f8"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
# Add content_hash column to chunks table for delta comparison
op.execute(f"ALTER TABLE {schema}chunks ADD COLUMN IF NOT EXISTS content_hash TEXT")
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}chunks DROP COLUMN IF EXISTS content_hash")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,54 +0,0 @@
"""Add structured_content JSONB column to mental_models
Revision ID: b3w4x5y6z7a8
Revises: a2v3w4x5y6z7
Create Date: 2026-04-16
Stores the structured representation of a mental model document (sections,
blocks). The plain ``content`` column remains the rendered markdown shown to
users. ``structured_content`` is the source of truth for delta-mode refreshes:
each refresh applies a list of typed operations to the structured doc, then
re-renders to markdown — so unchanged sections come through byte-identical
without an LLM round-trip.
Nullable: existing markdown-only mental models continue to work in full mode;
the column is populated lazily the first time a model is refreshed in delta
mode.
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "b3w4x5y6z7a8"
down_revision: str | Sequence[str] | None = "a2v3w4x5y6z7"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"""
ALTER TABLE {schema}mental_models
ADD COLUMN IF NOT EXISTS structured_content JSONB
""")
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS structured_content")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,91 +0,0 @@
"""Backfill entity_cooccurrences.last_cooccurred from memory_units event time
Revision ID: b5d4e3f2a1c9
Revises: o1a2b3c4d5e6
Create Date: 2026-04-24
The writer path in `entity_resolver.link_units_to_entities_batch` historically
stamped `entity_cooccurrences.last_cooccurred` with `datetime.now(UTC)` at
flush time, ignoring the source memory unit's event date. For normal online
retains that's fine (now ≈ event time), but for any corpus that was
backfilled in a single session — migrating from another memory system, for
example — every co-occurrence collapsed to the import moment, which hid the
underlying knowledge timeline from the dashboard's entity graph recency heat
and from any downstream consumer of the column.
The writer is fixed in the same change set to propagate the unit's event_date;
this migration repairs historical rows by reading the true event time off
`unit_entities × memory_units` (falling back to `created_at` when
`mentioned_at` / `occurred_start` are NULL, so rows never regress).
Oracle slot is intentionally absent: the Oracle baseline (`o1a2b3c4d5e6`)
landed days before this fix, so any Oracle deployment runs the corrected
writer against an effectively empty `entity_cooccurrences` — there is no
historical residue on Oracle to repair. PG-only matches the asymmetry of
the data, not negligence.
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "b5d4e3f2a1c9"
down_revision: str | Sequence[str] | None = "o1a2b3c4d5e6"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
# Recompute last_cooccurred from the true event time per entity pair.
# COALESCE picks the first non-null of mentioned_at / occurred_start /
# created_at so banks without event-time metadata still see a sane value
# (equivalent to the pre-fix behaviour) instead of NULL.
#
# The self-join on `unit_entities` is O(k²) per memory_unit in the number
# of distinct entities mentioned (k). For typical units k is small (single
# digits), but a bank with units containing hundreds of entities and tens
# of millions of co-occurrence rows may want to run this off-hours — the
# whole UPDATE is one statement, so it locks every targeted ec row for
# the duration. The migration is one-time; subsequent online writes
# already carry event time via the writer fix.
op.execute(
f"""
UPDATE {schema}entity_cooccurrences ec
SET last_cooccurred = sub.event_time
FROM (
SELECT
LEAST(ue1.entity_id, ue2.entity_id) AS e1,
GREATEST(ue1.entity_id, ue2.entity_id) AS e2,
MAX(COALESCE(mu.mentioned_at, mu.occurred_start, mu.created_at)) AS event_time
FROM {schema}memory_units mu
JOIN {schema}unit_entities ue1 ON ue1.unit_id = mu.id
JOIN {schema}unit_entities ue2 ON ue2.unit_id = mu.id AND ue1.entity_id <> ue2.entity_id
GROUP BY 1, 2
) sub
WHERE ec.entity_id_1 = sub.e1 AND ec.entity_id_2 = sub.e2
"""
)
def _pg_downgrade() -> None:
# No-op: the previous column value was `now()` at the time of write and
# isn't recoverable. Rolling back the code is sufficient — new writes will
# revert to the old behaviour for subsequent retains.
pass
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade) # oracle slot intentionally absent — see header
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,71 +0,0 @@
"""Add audit_log table for feature usage tracking.
Merge migration that combines the two existing heads (a3b4c5d6e7f8 + c8e5f2a3b4d1).
Stores raw request/response as JSONB for expandability without future migrations.
The metadata JSONB column allows adding arbitrary fields in the future.
Revision ID: c2d3e4f5g6h7
Revises: a3b4c5d6e7f8, c8e5f2a3b4d1
Create Date: 2026-03-26
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "c2d3e4f5g6h7"
down_revision: str | Sequence[str] | None = ("a3b4c5d6e7f8", "c8e5f2a3b4d1")
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
op.execute(
f"""
CREATE TABLE IF NOT EXISTS {schema}audit_log (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
action TEXT NOT NULL,
transport TEXT NOT NULL,
bank_id TEXT,
started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
ended_at TIMESTAMPTZ,
request JSONB,
response JSONB,
metadata JSONB DEFAULT '{{}}'::jsonb
)
"""
)
op.execute(
f"CREATE INDEX IF NOT EXISTS idx_audit_log_action_started ON {schema}audit_log (action, started_at DESC)"
)
op.execute(f"CREATE INDEX IF NOT EXISTS idx_audit_log_bank_started ON {schema}audit_log (bank_id, started_at DESC)")
op.execute(f"CREATE INDEX IF NOT EXISTS idx_audit_log_started ON {schema}audit_log (started_at DESC)")
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"DROP INDEX IF EXISTS {schema}idx_audit_log_started")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_audit_log_bank_started")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_audit_log_action_started")
op.execute(f"DROP TABLE IF EXISTS {schema}audit_log")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,40 +0,0 @@
"""Add history column to mental_models
Revision ID: c3d4e5f6g7h8
Revises: a2b3c4d5e6f7, a2b3c4d5e6f8
Create Date: 2026-03-06
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "c3d4e5f6g7h8"
down_revision: str | Sequence[str] | None = ("a2b3c4d5e6f7", "a2b3c4d5e6f8")
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}mental_models ADD COLUMN IF NOT EXISTS history JSONB DEFAULT '[]'::jsonb")
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS history")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,76 +0,0 @@
"""backsweep_orphan_observations_v2
Re-run of Pass 2 from migration ``g7h8i9j0k1l2_backsweep_orphan_observations``
to sweep observations that became orphaned between then and now.
Why we need it again:
``fact_storage.handle_document_tracking`` (the retain/upsert path) deleted
the existing document via the FK cascade — which removes the source
``memory_units`` — but never invalidated the observations derived from
them. Only the explicit ``MemoryEngine.delete_document`` API called
``_delete_stale_observations_for_memories``. Every document re-ingest
therefore left orphan observations whose ``source_memory_ids`` arrays
pointed at IDs that no longer existed in ``memory_units``.
``handle_document_tracking`` now calls the same cleanup helper before the
cascade, so no new orphans will accumulate going forward. This migration
cleans up the historical residue.
Identical to Pass 2 of g7h8i9j0k1l2. Pass 1 (memory_units whose bank is
gone) is intentionally not re-run; that scenario has no fresh source.
Revision ID: c4x5y6z7a8b9
Revises: b3w4x5y6z7a8
Create Date: 2026-04-16
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "c4x5y6z7a8b9"
down_revision: str | Sequence[str] | None = "b3w4x5y6z7a8"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
mu = f"{schema}memory_units"
# Delete observations whose every source_memory_id refers to a now-deleted
# memory_unit (or the array is empty). Observations with at least one
# surviving source are left alone — the consolidation engine will refresh
# their text on the next pass.
op.execute(
f"""
DELETE FROM {mu} orphan
WHERE orphan.fact_type = 'observation'
AND NOT EXISTS (
SELECT 1
FROM {mu} src
WHERE src.id = ANY(orphan.source_memory_ids)
AND src.bank_id = orphan.bank_id
)
"""
)
def _pg_downgrade() -> None:
# Deleted rows cannot be restored.
pass
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,58 +0,0 @@
"""Add bank_id column to memory_links for direct filtering
The stats endpoint JOINs memory_links to memory_units just to filter by
bank_id. With millions of links this takes 18+ seconds. Adding bank_id
directly to memory_links lets Postgres push the filter down before the JOIN.
Revision ID: c5d6e7f8a9b0
Revises: b3c4d5e6f7a8
Create Date: 2026-03-26
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "c5d6e7f8a9b0"
down_revision: str | Sequence[str] | None = "b3c4d5e6f7a8"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
# 1. Add nullable column
op.execute(f"ALTER TABLE {schema}memory_links ADD COLUMN IF NOT EXISTS bank_id TEXT")
# 2. Backfill from memory_units
op.execute(f"""
UPDATE {schema}memory_links ml
SET bank_id = mu.bank_id
FROM {schema}memory_units mu
WHERE ml.from_unit_id = mu.id
AND ml.bank_id IS NULL
""")
# 3. Set NOT NULL
op.execute(f"ALTER TABLE {schema}memory_links ALTER COLUMN bank_id SET NOT NULL")
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}memory_links DROP COLUMN IF EXISTS bank_id")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,63 +0,0 @@
"""Recreate idx_memory_units_source_memory_ids GIN index with fastupdate=off
GIN indexes use a "fastupdate" pending list by default: small writes are
buffered there and flushed to the main GIN tree in bulk. Flushing requires
AccessExclusiveLock on the index. Under high insert concurrency (e.g. 8
parallel pytest-xdist workers all calling retain_async) two transactions can
each trigger a flush simultaneously and deadlock.
Disabling fastupdate makes every insert write directly to the GIN tree
(slightly slower per insert, but no pending-list lock cycles).
Revision ID: d4e5f6g7h8i9
Revises: d5e6f7a8b9c0
Create Date: 2026-03-11
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "d4e5f6g7h8i9"
down_revision: str | Sequence[str] | None = "d5e6f7a8b9c0"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
# DROP + CREATE CONCURRENTLY must run outside a transaction block.
op.execute("COMMIT")
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
op.execute(
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
f"ON {schema}memory_units USING GIN (source_memory_ids) "
f"WITH (fastupdate=off) "
f"WHERE source_memory_ids IS NOT NULL"
)
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute("COMMIT")
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
op.execute(
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
f"ON {schema}memory_units USING GIN (source_memory_ids) "
f"WHERE source_memory_ids IS NOT NULL"
)
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,149 +0,0 @@
"""Add internal_id to banks and per-(bank, fact_type) partial vector indexes
Revision ID: d5e6f7a8b9c0
Revises: a3b4c5d6e7f8
Create Date: 2026-03-11
This migration:
1. Adds internal_id UUID column to banks (stable identifier for index naming)
2. Drops the global vector index (competes with per-bank partial indexes)
3. Creates per-(bank_id, fact_type) partial vector indexes for all existing banks
using the configured vector extension (HNSW for pgvector, DiskANN for
pgvectorscale, vchordrq for vchord).
(new banks get indexes created at bank-creation time via bank_utils.create_bank_vector_indexes)
Why per-(bank, fact_type) indexes:
- fact_type-only partial indexes are never chosen by the planner when bank_id is in the WHERE
clause, because the idx_memory_units_bank_id B-tree index always wins at planning time.
- Per-(bank, fact_type) partial indexes have both predicates matching → planner selects them.
- The global vector index competes for larger partitions (world, observation) and must be dropped.
"""
import os
from collections.abc import Sequence
from alembic import context, op
from sqlalchemy import text
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "d5e6f7a8b9c0"
down_revision: str | Sequence[str] | None = "c3d4e5f6g7h8"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_FACT_TYPES: dict[str, str] = {
"world": "worl",
"experience": "expr",
"observation": "obsv",
}
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _vector_index_using_clause() -> str:
"""Return the USING clause based on the configured vector extension."""
ext = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
if ext == "pgvectorscale":
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
elif ext == "vchord":
return "USING vchordrq (embedding vector_l2_ops)"
else:
return "USING hnsw (embedding vector_cosine_ops)"
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
# 1. Add internal_id column to banks
op.execute(
f"ALTER TABLE {schema}banks ADD COLUMN IF NOT EXISTS internal_id UUID DEFAULT gen_random_uuid() NOT NULL"
)
op.execute(f"ALTER TABLE {schema}banks ADD CONSTRAINT banks_internal_id_unique UNIQUE (internal_id)")
# 2. Drop any fact_type-only partial indexes that may exist from prior migrations
# (bank_id B-tree always wins over them when bank_id is in the WHERE clause)
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mu_emb_world")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mu_emb_observation")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mu_emb_experience")
# 4. Drop global vector index (competes with per-bank partial indexes)
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_embedding")
# 5. Create per-(bank, fact_type) partial vector indexes for all existing banks
# using the configured extension (HNSW / DiskANN / vchordrq)
bind = op.get_bind()
schema_name = context.config.get_main_option("target_schema")
table_ref = f'"{schema_name}".memory_units' if schema_name else "memory_units"
banks_ref = f'"{schema_name}".banks' if schema_name else "banks"
using_clause = _vector_index_using_clause()
rows = bind.execute(text(f"SELECT bank_id, internal_id FROM {banks_ref}")).fetchall() # noqa: S608
for row in rows:
bank_id = row[0]
internal_id = str(row[1]).replace("-", "")[:16]
escaped_bank_id = bank_id.replace("'", "''")
for ft, ft_short in _FACT_TYPES.items():
idx_name = f"idx_mu_emb_{ft_short}_{internal_id}"
# Index name is schema-unqualified (indexes live in the schema of their table)
bind.execute(
text(
f"CREATE INDEX IF NOT EXISTS {idx_name} "
f"ON {table_ref} {using_clause} "
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped_bank_id}'"
)
)
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
# Drop per-bank HNSW indexes (iterate existing banks)
bind = op.get_bind()
schema_name = context.config.get_main_option("target_schema")
banks_ref = f'"{schema_name}".banks' if schema_name else "banks"
rows = bind.execute(text(f"SELECT internal_id FROM {banks_ref}")).fetchall() # noqa: S608
for row in rows:
internal_id = str(row[0]).replace("-", "")[:16]
for ft_short in _HNSW_FACT_TYPES.values():
idx_name = f"idx_mu_emb_{ft_short}_{internal_id}"
bind.execute(text(f"DROP INDEX IF EXISTS {schema}{idx_name}"))
# Restore the global HNSW index
table_ref = f'"{schema_name}".memory_units' if schema_name else "memory_units"
op.execute(
f"CREATE INDEX IF NOT EXISTS idx_memory_units_embedding ON {table_ref} USING hnsw (embedding vector_cosine_ops)"
)
# Restore old fact_type-only partial indexes
op.execute(
f"CREATE INDEX IF NOT EXISTS idx_mu_emb_world "
f"ON {table_ref} USING hnsw (embedding vector_cosine_ops) "
f"WHERE fact_type = 'world'"
)
op.execute(
f"CREATE INDEX IF NOT EXISTS idx_mu_emb_observation "
f"ON {table_ref} USING hnsw (embedding vector_cosine_ops) "
f"WHERE fact_type = 'observation'"
)
op.execute(
f"CREATE INDEX IF NOT EXISTS idx_mu_emb_experience "
f"ON {table_ref} USING hnsw (embedding vector_cosine_ops) "
f"WHERE fact_type = 'experience'"
)
# Drop internal_id column
op.execute(f"ALTER TABLE {schema}banks DROP CONSTRAINT IF EXISTS banks_internal_id_unique")
op.execute(f"ALTER TABLE {schema}banks DROP COLUMN IF EXISTS internal_id")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,67 +0,0 @@
"""Backfill mental_models.subtype for databases that ran h3c4d5e6f7g8 before the fix
Migration h3c4d5e6f7g8 used CREATE TABLE IF NOT EXISTS to create the
mental_models table with a subtype column. But on databases where the table
already existed (from the reflections -> mental_models rename chain), the
CREATE was a no-op and subtype was never added. A fix was later added to
h3c4d5e6f7g8 (Step 4b), but databases that had already run the migration
never re-execute it. This migration adds the missing columns idempotently.
Revision ID: d5y6z7a8b9c0
Revises: 8c6fa6f7230b
Create Date: 2026-04-18
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "d5y6z7a8b9c0"
down_revision: str | Sequence[str] | None = "8c6fa6f7230b"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
# Add columns that h3c4d5e6f7g8 intended to create but missed when
# the table already existed from the reflections rename chain.
for col_ddl in [
"subtype VARCHAR(32) NOT NULL DEFAULT 'structural'",
"description TEXT NOT NULL DEFAULT ''",
"entity_id UUID",
"observations JSONB DEFAULT '{\"observations\": []}'::jsonb",
"links VARCHAR[]",
"last_updated TIMESTAMP WITH TIME ZONE",
]:
op.execute(f"ALTER TABLE {schema}mental_models ADD COLUMN IF NOT EXISTS {col_ddl}")
# Ensure the CHECK constraint exists
op.execute(f"ALTER TABLE {schema}mental_models DROP CONSTRAINT IF EXISTS ck_mental_models_subtype")
op.execute(f"""
ALTER TABLE {schema}mental_models
ADD CONSTRAINT ck_mental_models_subtype CHECK (subtype IN ('structural', 'emergent', 'pinned', 'learned'))
""")
op.execute(f"CREATE INDEX IF NOT EXISTS idx_mental_models_subtype ON {schema}mental_models(bank_id, subtype)")
def _pg_downgrade() -> None:
# No-op: these columns are part of the intended schema
pass
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,49 +0,0 @@
"""Drop unused metadata column from documents table
Revision ID: d6e7f8a9b0c1
Revises: c2d3e4f5g6h7, c5d6e7f8a9b0
Create Date: 2026-03-30
The metadata column on documents was always stored as an empty dict {}.
Actual document metadata is stored inside retain_params.metadata.
This migration was originally shipped in v0.4.22, then its file was deleted
in v0.5.0 (and its revision ID accidentally reused by 2eee35aa3cfc).
Restoring the file so that databases stamped at this revision can upgrade
cleanly to v0.5.x+.
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "d6e7f8a9b0c1"
down_revision: str | Sequence[str] | None = ("c2d3e4f5g6h7", "c5d6e7f8a9b0")
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}documents DROP COLUMN IF EXISTS metadata")
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}documents ADD COLUMN IF NOT EXISTS metadata jsonb DEFAULT '{{}}'")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,72 +0,0 @@
"""Add webhooks table and next_retry_at to async_operations.
Webhook deliveries are handled as async_operations tasks (operation_type='webhook_delivery')
rather than a dedicated webhook_deliveries table.
Revision ID: e4f5a6b7c8d9
Revises: d2e3f4a5b6c7
Create Date: 2026-03-04
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "e4f5a6b7c8d9"
down_revision: str | Sequence[str] | None = "d2e3f4a5b6c7"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
op.execute(
f"""
CREATE TABLE IF NOT EXISTS {schema}webhooks (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
bank_id TEXT,
url TEXT NOT NULL,
secret TEXT,
event_types TEXT[] NOT NULL DEFAULT '{{}}',
enabled BOOLEAN NOT NULL DEFAULT TRUE,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
)
"""
)
# Index for bank-scoped webhook lookup
op.execute(f"CREATE INDEX IF NOT EXISTS idx_webhooks_bank_id ON {schema}webhooks(bank_id)")
# Add next_retry_at to async_operations for task-owned retry scheduling
op.execute(f"ALTER TABLE {schema}async_operations ADD COLUMN IF NOT EXISTS next_retry_at TIMESTAMPTZ NULL")
# Index for polling: status + next_retry_at
op.execute(
f"CREATE INDEX IF NOT EXISTS idx_async_operations_status_retry "
f"ON {schema}async_operations(status, next_retry_at)"
)
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"DROP INDEX IF EXISTS {schema}idx_async_operations_status_retry")
op.execute(f"ALTER TABLE {schema}async_operations DROP COLUMN IF EXISTS next_retry_at")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_webhooks_bank_id")
op.execute(f"DROP TABLE IF EXISTS {schema}webhooks")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,83 +0,0 @@
"""Add CASCADE DELETE FK from async_operations and webhooks to banks.
When a bank is deleted, all its async_operations and webhooks rows are
automatically deleted by the database. This ensures that any in-flight
worker tasks detect the deletion via _check_op_alive() and abort early.
Revision ID: e5f6g7h8i9j0
Revises: d4e5f6g7h8i9
Create Date: 2026-03-11
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "e5f6g7h8i9j0"
down_revision: str | Sequence[str] | None = "d4e5f6g7h8i9"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
# Remove orphaned async_operations rows whose bank no longer exists
# (can happen because there was no FK before this migration).
op.execute(
f"""
DELETE FROM {schema}async_operations
WHERE bank_id IS NOT NULL
AND bank_id NOT IN (SELECT bank_id FROM {schema}banks)
"""
)
# Remove orphaned webhooks rows whose bank no longer exists.
op.execute(
f"""
DELETE FROM {schema}webhooks
WHERE bank_id IS NOT NULL
AND bank_id NOT IN (SELECT bank_id FROM {schema}banks)
"""
)
# Add FK with ON DELETE CASCADE so that deleting a bank automatically
# cleans up all its pending/processing operations and webhook configs.
op.execute(
f"""
ALTER TABLE {schema}async_operations
ADD CONSTRAINT fk_async_operations_bank_id
FOREIGN KEY (bank_id) REFERENCES {schema}banks(bank_id)
ON DELETE CASCADE
"""
)
op.execute(
f"""
ALTER TABLE {schema}webhooks
ADD CONSTRAINT fk_webhooks_bank_id
FOREIGN KEY (bank_id) REFERENCES {schema}banks(bank_id)
ON DELETE CASCADE
"""
)
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}async_operations DROP CONSTRAINT IF EXISTS fk_async_operations_bank_id")
op.execute(f"ALTER TABLE {schema}webhooks DROP CONSTRAINT IF EXISTS fk_webhooks_bank_id")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,67 +0,0 @@
"""chunk_fk_cascade_delete
Revision ID: f6g7h8i9j0k1
Revises: e5f6g7h8i9j0
Create Date: 2026-03-16 00:00:00.000000
"""
from collections.abc import Sequence
from alembic import op
from hindsight_api.alembic._dialect import run_for_dialect
# revision identifiers, used by Alembic.
revision: str = "f6g7h8i9j0k1"
down_revision: str | Sequence[str] | None = "e5f6g7h8i9j0"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _pg_upgrade() -> None:
"""Change memory_units.chunk_id FK from SET NULL to CASCADE.
When a document is deleted the CASCADE reaches chunks first; with SET NULL
the memory_units rows survived with chunk_id = NULL, leaving ghost records.
Switching to CASCADE ensures they are removed together with their chunk.
"""
from alembic import context
schema = context.config.get_main_option("target_schema")
schema_prefix = f'"{schema}".' if schema else ""
# Use raw SQL with IF EXISTS so this is safe on schemas where the FK was
# already dropped or never existed under this name.
op.execute(f"ALTER TABLE {schema_prefix}memory_units DROP CONSTRAINT IF EXISTS memory_units_chunk_fkey")
# Use a DO block so the ADD is also idempotent: if the FK already exists (e.g.
# the schema was provisioned after the base migration already added it) the
# duplicate_object exception is swallowed rather than failing the migration.
op.execute(
f"""
DO $$ BEGIN
ALTER TABLE {schema_prefix}memory_units
ADD CONSTRAINT memory_units_chunk_fkey
FOREIGN KEY (chunk_id)
REFERENCES {schema_prefix}chunks (chunk_id)
ON DELETE CASCADE;
EXCEPTION
WHEN duplicate_object THEN NULL;
END $$;
"""
)
def _pg_downgrade() -> None:
"""Revert to SET NULL behaviour."""
op.drop_constraint("memory_units_chunk_fkey", "memory_units", type_="foreignkey")
op.create_foreign_key(
"memory_units_chunk_fkey", "memory_units", "chunks", ["chunk_id"], ["chunk_id"], ondelete="SET NULL"
)
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,43 +0,0 @@
"""Add http_config JSONB column to webhooks table.
Stores HTTP delivery configuration (method, timeout, headers, params) as a
single JSONB column rather than separate columns.
Revision ID: f7g8h9i0j1k2
Revises: e4f5a6b7c8d9
Create Date: 2026-03-04
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "f7g8h9i0j1k2"
down_revision: str | Sequence[str] | None = "e4f5a6b7c8d9"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}webhooks ADD COLUMN IF NOT EXISTS http_config JSONB NOT NULL DEFAULT '{{}}'")
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}webhooks DROP COLUMN IF EXISTS http_config")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,93 +0,0 @@
"""remove_opinion_fact_type
Revision ID: g2h3i4j5k6l7
Revises: f1a2b3c4d5e6
Create Date: 2026-04-02
Remove the deprecated 'opinion' fact type: drop opinion-specific indexes,
update CHECK constraints, delete any remaining opinion rows, and drop the
confidence_score column (was only used for opinions, always NULL otherwise).
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
# revision identifiers, used by Alembic.
revision: str = "g2h3i4j5k6l7"
down_revision: str | Sequence[str] | None = "f1a2b3c4d5e6"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
# 1. Delete any remaining opinion rows
op.execute(f"DELETE FROM {schema}memory_units WHERE fact_type = 'opinion'")
# 2. Drop opinion-specific indexes
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_opinion_confidence")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_opinion_date")
# 3. Drop confidence_score constraints and column (only used for opinions, always NULL otherwise)
op.execute(f"ALTER TABLE {schema}memory_units DROP CONSTRAINT IF EXISTS confidence_score_fact_type_check")
op.execute(f"ALTER TABLE {schema}memory_units DROP CONSTRAINT IF EXISTS memory_units_confidence_score_check")
op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS confidence_score")
# 4. Replace fact_type CHECK constraint
op.execute(f"ALTER TABLE {schema}memory_units DROP CONSTRAINT IF EXISTS memory_units_fact_type_check")
op.execute(
f"ALTER TABLE {schema}memory_units ADD CONSTRAINT memory_units_fact_type_check "
f"CHECK (fact_type IN ('world', 'experience', 'observation'))"
)
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
# Restore confidence_score column
op.execute(f"ALTER TABLE {schema}memory_units ADD COLUMN IF NOT EXISTS confidence_score float")
op.execute(
f"ALTER TABLE {schema}memory_units ADD CONSTRAINT memory_units_confidence_score_check "
f"CHECK (confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0))"
)
op.execute(
f"ALTER TABLE {schema}memory_units ADD CONSTRAINT confidence_score_fact_type_check "
f"CHECK ((fact_type = 'opinion' AND confidence_score IS NOT NULL) OR "
f"(fact_type = 'observation') OR "
f"(fact_type NOT IN ('opinion', 'observation') AND confidence_score IS NULL))"
)
# Restore original fact_type CHECK constraint (with opinion)
op.execute(f"ALTER TABLE {schema}memory_units DROP CONSTRAINT IF EXISTS memory_units_fact_type_check")
op.execute(
f"ALTER TABLE {schema}memory_units ADD CONSTRAINT memory_units_fact_type_check "
f"CHECK (fact_type IN ('world', 'experience', 'opinion', 'observation'))"
)
# Recreate opinion indexes
op.execute(
f"CREATE INDEX idx_memory_units_opinion_confidence ON {schema}memory_units "
f"(bank_id, confidence_score DESC) WHERE fact_type = 'opinion'"
)
op.execute(
f"CREATE INDEX idx_memory_units_opinion_date ON {schema}memory_units "
f"(bank_id, event_date DESC) WHERE fact_type = 'opinion'"
)
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,81 +0,0 @@
"""backsweep_orphan_memory_units
Two-pass cleanup of memory_units rows that were never removed by earlier bugs:
Pass 1 — any fact_type, bank gone:
memory_units whose bank_id no longer exists in banks. These accumulate when
a bank is deleted without a proper cascade (no FK from memory_units to banks
exists in the schema).
Pass 2 — observations only, all sources gone:
observation rows whose bank still exists but every source_memory_id points
to a deleted memory unit. These were left behind before PR #580 fixed the
chunk FK cascade and before delete_document() called
_delete_stale_observations_for_memories.
Revision ID: g7h8i9j0k1l2
Revises: f6g7h8i9j0k1
Create Date: 2026-03-16
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "g7h8i9j0k1l2"
down_revision: str | Sequence[str] | None = "f6g7h8i9j0k1"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
mu = f"{schema}memory_units"
banks = f"{schema}banks"
# Pass 1: delete all memory_units (any fact_type) whose bank no longer exists.
# There is no FK from memory_units to banks, so these never cascade away.
op.execute(
f"""
DELETE FROM {mu}
WHERE NOT EXISTS (
SELECT 1 FROM {banks} b WHERE b.bank_id = {mu}.bank_id
)
"""
)
# Pass 2: delete orphaned observations whose bank still exists but every
# source_memory_id refers to a now-deleted memory unit (or the array is
# empty). Observations with at least one surviving source are left alone.
op.execute(
f"""
DELETE FROM {mu} orphan
WHERE orphan.fact_type = 'observation'
AND NOT EXISTS (
SELECT 1
FROM {mu} src
WHERE src.id = ANY(orphan.source_memory_ids)
AND src.bank_id = orphan.bank_id
)
"""
)
def _pg_downgrade() -> None:
# Deleted rows cannot be restored.
pass
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,52 +0,0 @@
"""Merge 3 migration heads and add unit_entities composite index
Revision ID: h3i4j5k6l7m8
Revises: a4b5c6d7e8f9, g2h3i4j5k6l7
Create Date: 2026-04-07
Merges three unmerged migration heads into one, and adds a composite index
(entity_id, unit_id) on unit_entities for index-only scans in the LATERAL
entity expansion query.
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "h3i4j5k6l7m8"
down_revision: str | Sequence[str] | None = ("a4b5c6d7e8f9", "g2h3i4j5k6l7")
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
# Composite index enables index-only scans for entity_id -> unit_id lookups
op.execute(
f"CREATE INDEX IF NOT EXISTS idx_unit_entities_entity_unit ON {schema}unit_entities (entity_id, unit_id)"
)
# Drop the now-redundant single-column index
op.execute(f"DROP INDEX IF EXISTS {schema}idx_unit_entities_entity")
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"DROP INDEX IF EXISTS {schema}idx_unit_entities_entity_unit")
# Restore the single-column index
op.execute(f"CREATE INDEX IF NOT EXISTS idx_unit_entities_entity ON {schema}unit_entities (entity_id)")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,49 +0,0 @@
"""Add 'cancelled' to async_operations status check constraint
Revision ID: i4j5k6l7m8n9
Revises: d5y6z7a8b9c0
Create Date: 2026-04-23
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "i4j5k6l7m8n9"
down_revision: str | Sequence[str] | None = "d5y6z7a8b9c0"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}async_operations DROP CONSTRAINT IF EXISTS async_operations_status_check")
op.execute(
f"ALTER TABLE {schema}async_operations ADD CONSTRAINT async_operations_status_check "
f"CHECK (status IN ('pending', 'processing', 'completed', 'failed', 'cancelled'))"
)
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}async_operations DROP CONSTRAINT IF EXISTS async_operations_status_check")
op.execute(
f"ALTER TABLE {schema}async_operations ADD CONSTRAINT async_operations_status_check "
f"CHECK (status IN ('pending', 'processing', 'completed', 'failed'))"
)
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,38 +0,0 @@
"""No-op: observation_sources table is Oracle-only
Originally created the observation_sources junction table for all backends,
but PG uses native array ops on the source_memory_ids column (faster at scale).
Oracle creates this table in the o1a2b3c4d5e6 baseline migration instead.
Kept as a no-op to preserve the Alembic revision chain.
Revision ID: k6l7m8n9o0p1
Revises: i4j5k6l7m8n9
Create Date: 2026-04-24
"""
from collections.abc import Sequence
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "k6l7m8n9o0p1"
down_revision: str | Sequence[str] | None = "i4j5k6l7m8n9"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _pg_upgrade() -> None:
# PG uses source_memory_ids[] array on memory_units — no junction table.
pass
def _pg_downgrade() -> None:
pass
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,31 +0,0 @@
"""Merge divergent heads from deferrable FK and cooccurrence backfill
Revision ID: m3rg3h3ad5f6
Revises: 9f8e7d6c5b4a, b5d4e3f2a1c9
Create Date: 2026-05-04
"""
from collections.abc import Sequence
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "m3rg3h3ad5f6"
down_revision: tuple[str, ...] = ("9f8e7d6c5b4a", "b5d4e3f2a1c9")
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _pg_upgrade() -> None:
pass
def _pg_downgrade() -> None:
pass
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -1,445 +0,0 @@
"""oracle_baseline
Brings a fresh Oracle 23ai database up to the current schema in a single step.
PostgreSQL is a no-op here — the prior 59 revisions already build the PG schema
incrementally; this revision just closes the loop so both dialects share a
single head from this point on.
After this migration ships, *every* new revision must fill both the ``_pg_*``
and ``_oracle_*`` slots (or explicitly leave one ``None``); a CI check enforces
that.
Tables mirror the PostgreSQL schema but use Oracle-native types:
- UUID -> RAW(16) DEFAULT SYS_GUID()
- TEXT / large VARCHAR -> CLOB
- JSONB -> CLOB with IS JSON CHECK
- BOOLEAN -> NUMBER(1)
- FLOAT -> BINARY_DOUBLE
- VARCHAR[] -> CLOB (JSON array stored as string)
- BYTEA -> BLOB
- vector(384) -> VECTOR(384, FLOAT32) (Oracle 23ai native)
Revision ID: o1a2b3c4d5e6
Revises: k6l7m8n9o0p1
Create Date: 2026-04-29
"""
from collections.abc import Sequence
from alembic import op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "o1a2b3c4d5e6"
down_revision: str | Sequence[str] | None = "k6l7m8n9o0p1"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
# ---------------------------------------------------------------------------
# Tables — created in dependency order
# ---------------------------------------------------------------------------
_TABLES: tuple[str, ...] = (
"""
CREATE TABLE IF NOT EXISTS banks (
bank_id VARCHAR2(256) NOT NULL,
internal_id RAW(16) DEFAULT SYS_GUID() NOT NULL,
name VARCHAR2(512),
disposition CLOB DEFAULT '{"skepticism":3,"literalism":3,"empathy":3}' NOT NULL
CONSTRAINT banks_disposition_json CHECK (disposition IS JSON),
mission CLOB,
personality CLOB DEFAULT '{}' NOT NULL
CONSTRAINT banks_personality_json CHECK (personality IS JSON),
config CLOB DEFAULT '{}' NOT NULL
CONSTRAINT banks_config_json CHECK (config IS JSON),
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
CONSTRAINT pk_banks PRIMARY KEY (bank_id),
CONSTRAINT banks_internal_id_unique UNIQUE (internal_id)
)
""",
"""
CREATE TABLE IF NOT EXISTS documents (
id VARCHAR2(512) NOT NULL,
bank_id VARCHAR2(256) NOT NULL,
original_text CLOB,
content_hash VARCHAR2(128),
metadata CLOB DEFAULT '{}' NOT NULL
CONSTRAINT docs_metadata_json CHECK (metadata IS JSON),
retain_params CLOB CONSTRAINT docs_retain_params_json CHECK (retain_params IS JSON OR retain_params IS NULL),
file_storage_key VARCHAR2(512),
file_original_name VARCHAR2(512),
file_content_type VARCHAR2(256),
tags CLOB DEFAULT '[]' NOT NULL,
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
CONSTRAINT pk_documents PRIMARY KEY (id, bank_id),
CONSTRAINT fk_documents_bank FOREIGN KEY (bank_id) REFERENCES banks(bank_id) ON DELETE CASCADE
)
""",
"""
CREATE TABLE IF NOT EXISTS chunks (
chunk_id VARCHAR2(512) NOT NULL,
document_id VARCHAR2(512) NOT NULL,
bank_id VARCHAR2(256) NOT NULL,
chunk_index NUMBER(10) NOT NULL,
chunk_text CLOB NOT NULL,
content_hash VARCHAR2(128),
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
CONSTRAINT pk_chunks PRIMARY KEY (chunk_id),
CONSTRAINT fk_chunks_document FOREIGN KEY (document_id, bank_id)
REFERENCES documents(id, bank_id) ON DELETE CASCADE
)
""",
# memory_units uses automatic list partitioning on bank_id at create time —
# no post-create ALTER required (we used to do that for legacy installs).
"""
CREATE TABLE IF NOT EXISTS memory_units (
id RAW(16) DEFAULT SYS_GUID() NOT NULL,
bank_id VARCHAR2(256) NOT NULL,
document_id VARCHAR2(512),
chunk_id VARCHAR2(512),
text CLOB NOT NULL,
embedding VECTOR(384, FLOAT32),
context CLOB,
event_date TIMESTAMP WITH TIME ZONE NOT NULL,
occurred_start TIMESTAMP WITH TIME ZONE,
occurred_end TIMESTAMP WITH TIME ZONE,
mentioned_at TIMESTAMP WITH TIME ZONE,
fact_type VARCHAR2(64) DEFAULT 'world' NOT NULL,
confidence_score BINARY_DOUBLE,
access_count NUMBER(10) DEFAULT 0 NOT NULL,
consolidated_at TIMESTAMP WITH TIME ZONE,
observation_scopes CLOB CONSTRAINT mu_obs_scopes_json CHECK (observation_scopes IS JSON OR observation_scopes IS NULL),
tags CLOB DEFAULT '[]' NOT NULL,
metadata CLOB DEFAULT '{}' NOT NULL
CONSTRAINT mu_metadata_json CHECK (metadata IS JSON),
proof_count NUMBER(10) DEFAULT 1,
source_memory_ids CLOB,
history CLOB DEFAULT '[]'
CONSTRAINT mu_history_json CHECK (history IS JSON OR history IS NULL),
text_signals CLOB,
consolidation_failed_at TIMESTAMP WITH TIME ZONE,
search_vector CLOB,
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
CONSTRAINT pk_memory_units PRIMARY KEY (id),
CONSTRAINT fk_mu_document FOREIGN KEY (document_id, bank_id)
REFERENCES documents(id, bank_id) ON DELETE CASCADE,
CONSTRAINT fk_mu_chunk FOREIGN KEY (chunk_id)
REFERENCES chunks(chunk_id) ON DELETE SET NULL,
CONSTRAINT chk_mu_fact_type CHECK (fact_type IN ('world', 'experience', 'observation')),
CONSTRAINT chk_mu_confidence CHECK (
confidence_score IS NULL
OR (confidence_score >= 0.0 AND confidence_score <= 1.0)
)
)
PARTITION BY LIST (bank_id) AUTOMATIC
(PARTITION p_default VALUES ('__default__'))
""",
"""
CREATE TABLE IF NOT EXISTS entities (
id RAW(16) DEFAULT SYS_GUID() NOT NULL,
bank_id VARCHAR2(256) NOT NULL,
canonical_name VARCHAR2(512) NOT NULL,
metadata CLOB DEFAULT '{}' NOT NULL
CONSTRAINT ent_metadata_json CHECK (metadata IS JSON),
first_seen TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
last_seen TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
mention_count NUMBER(10) DEFAULT 1 NOT NULL,
CONSTRAINT pk_entities PRIMARY KEY (id)
)
""",
"""
CREATE TABLE IF NOT EXISTS unit_entities (
unit_id RAW(16) NOT NULL,
entity_id RAW(16) NOT NULL,
CONSTRAINT pk_unit_entities PRIMARY KEY (unit_id, entity_id),
CONSTRAINT fk_ue_unit FOREIGN KEY (unit_id) REFERENCES memory_units(id) ON DELETE CASCADE,
CONSTRAINT fk_ue_entity FOREIGN KEY (entity_id) REFERENCES entities(id) ON DELETE CASCADE
)
""",
"""
CREATE TABLE IF NOT EXISTS entity_cooccurrences (
entity_id_1 RAW(16) NOT NULL,
entity_id_2 RAW(16) NOT NULL,
cooccurrence_count NUMBER(10) DEFAULT 1 NOT NULL,
last_cooccurred TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
CONSTRAINT pk_entity_cooccurrences PRIMARY KEY (entity_id_1, entity_id_2),
CONSTRAINT fk_ec_entity1 FOREIGN KEY (entity_id_1) REFERENCES entities(id) ON DELETE CASCADE,
CONSTRAINT fk_ec_entity2 FOREIGN KEY (entity_id_2) REFERENCES entities(id) ON DELETE CASCADE
)
""",
"""
CREATE TABLE IF NOT EXISTS memory_links (
from_unit_id RAW(16) NOT NULL,
to_unit_id RAW(16) NOT NULL,
link_type VARCHAR2(64) NOT NULL,
entity_id RAW(16),
bank_id VARCHAR2(256),
weight BINARY_DOUBLE DEFAULT 1.0 NOT NULL,
source_memory_ids CLOB,
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
CONSTRAINT fk_ml_from FOREIGN KEY (from_unit_id) REFERENCES memory_units(id) ON DELETE CASCADE,
CONSTRAINT fk_ml_to FOREIGN KEY (to_unit_id) REFERENCES memory_units(id) ON DELETE CASCADE,
CONSTRAINT fk_ml_entity FOREIGN KEY (entity_id) REFERENCES entities(id) ON DELETE CASCADE,
CONSTRAINT chk_ml_link_type CHECK (
link_type IN ('temporal', 'semantic', 'entity', 'causes', 'caused_by', 'enables', 'prevents')
),
CONSTRAINT chk_ml_weight CHECK (weight >= 0.0 AND weight <= 1.0)
)
""",
"""
CREATE TABLE IF NOT EXISTS mental_models (
id VARCHAR2(256) NOT NULL,
bank_id VARCHAR2(256) NOT NULL,
subtype VARCHAR2(32) NOT NULL,
name VARCHAR2(256) NOT NULL,
description CLOB NOT NULL,
source_query CLOB,
content CLOB,
embedding VECTOR(384, FLOAT32),
entity_id RAW(16),
observations CLOB DEFAULT '{"observations":[]}' NOT NULL
CONSTRAINT mm_obs_json CHECK (observations IS JSON),
links CLOB,
tags CLOB DEFAULT '[]' NOT NULL,
max_tokens NUMBER(10) DEFAULT 2048 NOT NULL,
"trigger" CLOB DEFAULT '{"refresh_after_consolidation":false}' NOT NULL
CONSTRAINT mm_trigger_json CHECK ("trigger" IS JSON),
structured_content CLOB CONSTRAINT mm_sc_json CHECK (structured_content IS JSON OR structured_content IS NULL),
last_refreshed_source_query CLOB,
reflect_response CLOB CONSTRAINT mm_reflect_resp_json CHECK (reflect_response IS JSON OR reflect_response IS NULL),
history CLOB DEFAULT '[]' NOT NULL
CONSTRAINT mm_history_json CHECK (history IS JSON),
last_refreshed_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
last_updated TIMESTAMP WITH TIME ZONE,
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
CONSTRAINT pk_mental_models PRIMARY KEY (id, bank_id),
CONSTRAINT fk_mm_bank FOREIGN KEY (bank_id) REFERENCES banks(bank_id) ON DELETE CASCADE,
CONSTRAINT fk_mm_entity FOREIGN KEY (entity_id) REFERENCES entities(id) ON DELETE SET NULL,
CONSTRAINT chk_mm_subtype CHECK (subtype IN ('directive', 'pinned'))
)
""",
"""
CREATE TABLE IF NOT EXISTS directives (
id RAW(16) DEFAULT SYS_GUID() NOT NULL,
bank_id VARCHAR2(256) NOT NULL,
name VARCHAR2(256) NOT NULL,
content CLOB NOT NULL,
priority NUMBER(10) DEFAULT 0 NOT NULL,
is_active NUMBER(1) DEFAULT 1 NOT NULL,
tags CLOB DEFAULT '[]' NOT NULL,
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
CONSTRAINT pk_directives PRIMARY KEY (id),
CONSTRAINT fk_dir_bank FOREIGN KEY (bank_id) REFERENCES banks(bank_id) ON DELETE CASCADE
)
""",
"""
CREATE TABLE IF NOT EXISTS async_operations (
operation_id RAW(16) DEFAULT SYS_GUID() NOT NULL,
bank_id VARCHAR2(256) NOT NULL,
operation_type VARCHAR2(128) NOT NULL,
status VARCHAR2(32) DEFAULT 'pending' NOT NULL,
worker_id VARCHAR2(256),
claimed_at TIMESTAMP WITH TIME ZONE,
retry_count NUMBER(10) DEFAULT 0 NOT NULL,
next_retry_at TIMESTAMP WITH TIME ZONE,
task_payload CLOB CONSTRAINT ao_payload_json CHECK (task_payload IS JSON OR task_payload IS NULL),
result_metadata CLOB DEFAULT '{}' NOT NULL
CONSTRAINT ao_result_json CHECK (result_metadata IS JSON),
error_message CLOB,
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
completed_at TIMESTAMP WITH TIME ZONE,
CONSTRAINT pk_async_operations PRIMARY KEY (operation_id),
CONSTRAINT fk_ao_bank FOREIGN KEY (bank_id) REFERENCES banks(bank_id) ON DELETE CASCADE,
CONSTRAINT chk_ao_status CHECK (status IN ('pending', 'processing', 'completed', 'failed', 'cancelled'))
)
""",
"""
CREATE TABLE IF NOT EXISTS webhooks (
id RAW(16) DEFAULT SYS_GUID() NOT NULL,
bank_id VARCHAR2(256) NOT NULL,
url VARCHAR2(2048) NOT NULL,
secret VARCHAR2(512),
event_types CLOB DEFAULT '[]' NOT NULL,
http_config CLOB DEFAULT '{}' NOT NULL
CONSTRAINT wh_http_config_json CHECK (http_config IS JSON),
enabled NUMBER(1) DEFAULT 1 NOT NULL,
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
CONSTRAINT pk_webhooks PRIMARY KEY (id),
CONSTRAINT fk_wh_bank FOREIGN KEY (bank_id) REFERENCES banks(bank_id) ON DELETE CASCADE
)
""",
"""
CREATE TABLE IF NOT EXISTS file_storage (
storage_key VARCHAR2(512) NOT NULL,
data BLOB NOT NULL,
CONSTRAINT pk_file_storage PRIMARY KEY (storage_key)
)
""",
"""
CREATE TABLE IF NOT EXISTS audit_log (
id RAW(16) DEFAULT SYS_GUID() NOT NULL,
action VARCHAR2(128) NOT NULL,
transport VARCHAR2(64) NOT NULL,
bank_id VARCHAR2(256),
started_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
ended_at TIMESTAMP WITH TIME ZONE,
request CLOB CONSTRAINT al_request_json CHECK (request IS JSON OR request IS NULL),
response CLOB CONSTRAINT al_response_json CHECK (response IS JSON OR response IS NULL),
metadata CLOB DEFAULT '{}' NOT NULL
CONSTRAINT al_metadata_json CHECK (metadata IS JSON),
CONSTRAINT pk_audit_log PRIMARY KEY (id)
)
""",
"""
CREATE TABLE IF NOT EXISTS observation_sources (
observation_id RAW(16) NOT NULL,
source_id RAW(16) NOT NULL,
CONSTRAINT pk_observation_sources PRIMARY KEY (observation_id, source_id),
CONSTRAINT fk_obs_src_observation FOREIGN KEY (observation_id)
REFERENCES memory_units(id) ON DELETE CASCADE
)
""",
)
# ---------------------------------------------------------------------------
# B-tree indexes
# ---------------------------------------------------------------------------
_INDEXES: tuple[str, ...] = (
# documents
"CREATE INDEX idx_docs_bank_id ON documents(bank_id)",
"CREATE INDEX idx_docs_content_hash ON documents(content_hash)",
# chunks
"CREATE INDEX idx_chunks_document_id ON chunks(document_id)",
"CREATE INDEX idx_chunks_bank_id ON chunks(bank_id)",
# memory_units
"CREATE INDEX idx_mu_bank_id ON memory_units(bank_id)",
"CREATE INDEX idx_mu_document_id ON memory_units(document_id)",
"CREATE INDEX idx_mu_chunk_id ON memory_units(chunk_id)",
"CREATE INDEX idx_mu_event_date ON memory_units(event_date DESC)",
"CREATE INDEX idx_mu_bank_date ON memory_units(bank_id, event_date DESC)",
"CREATE INDEX idx_mu_access_count ON memory_units(access_count DESC)",
"CREATE INDEX idx_mu_fact_type ON memory_units(fact_type)",
"CREATE INDEX idx_mu_bank_fact_type ON memory_units(bank_id, fact_type)",
"CREATE INDEX idx_mu_bank_type_date ON memory_units(bank_id, fact_type, event_date DESC)",
# entities
"CREATE INDEX idx_ent_bank_id ON entities(bank_id)",
"CREATE INDEX idx_ent_canonical_name ON entities(canonical_name)",
"CREATE INDEX idx_ent_bank_name ON entities(bank_id, canonical_name)",
"CREATE UNIQUE INDEX idx_ent_bank_lower_name ON entities(bank_id, LOWER(canonical_name))",
# unit_entities
"CREATE INDEX idx_ue_unit ON unit_entities(unit_id)",
"CREATE INDEX idx_ue_entity ON unit_entities(entity_id)",
# entity_cooccurrences
"CREATE INDEX idx_ec_entity1 ON entity_cooccurrences(entity_id_1)",
"CREATE INDEX idx_ec_entity2 ON entity_cooccurrences(entity_id_2)",
"CREATE INDEX idx_ec_count ON entity_cooccurrences(cooccurrence_count DESC)",
# memory_links — function-based unique index uses NVL with the nil UUID raw
# to handle nullable entity_id (matches PG idx_memory_links_unique).
"CREATE UNIQUE INDEX idx_memory_links_unique ON memory_links("
"from_unit_id, to_unit_id, link_type, "
"NVL(entity_id, HEXTORAW('00000000000000000000000000000000')))",
"CREATE INDEX idx_ml_from_unit ON memory_links(from_unit_id)",
"CREATE INDEX idx_ml_to_unit ON memory_links(to_unit_id)",
"CREATE INDEX idx_ml_entity ON memory_links(entity_id)",
"CREATE INDEX idx_ml_link_type ON memory_links(link_type)",
"CREATE INDEX idx_ml_bank_id ON memory_links(bank_id)",
# directives
"CREATE INDEX idx_dir_bank_id ON directives(bank_id)",
"CREATE INDEX idx_dir_bank_active ON directives(bank_id, is_active)",
# mental_models
"CREATE INDEX idx_mm_bank_id ON mental_models(bank_id)",
"CREATE INDEX idx_mm_subtype ON mental_models(bank_id, subtype)",
"CREATE INDEX idx_mm_entity_id ON mental_models(entity_id)",
# async_operations
"CREATE INDEX idx_ao_bank_id ON async_operations(bank_id)",
"CREATE INDEX idx_ao_status ON async_operations(status)",
"CREATE INDEX idx_ao_bank_status ON async_operations(bank_id, status)",
"CREATE INDEX idx_ao_status_retry ON async_operations(status, next_retry_at)",
# webhooks
"CREATE INDEX idx_wh_bank_id ON webhooks(bank_id)",
# audit_log
"CREATE INDEX idx_al_action_started ON audit_log(action, started_at DESC)",
"CREATE INDEX idx_al_bank_started ON audit_log(bank_id, started_at DESC)",
"CREATE INDEX idx_al_started ON audit_log(started_at DESC)",
# observation_sources
"CREATE INDEX idx_obs_sources_source_id ON observation_sources(source_id, observation_id)",
)
_VECTOR_INDEX = (
"CREATE VECTOR INDEX idx_mu_embedding_hnsw ON memory_units(embedding) "
"ORGANIZATION NEIGHBOR PARTITIONS "
"DISTANCE COSINE "
"WITH TARGET ACCURACY 95"
)
# Oracle Text (CTXSYS.CONTEXT) — ``SYNC (ON COMMIT)`` makes it auto-update
# without a maintenance job. Doubled single quotes for the embedded literal.
_TEXT_INDEX = (
"BEGIN "
"EXECUTE IMMEDIATE '"
"CREATE INDEX idx_mu_content_text ON memory_units(text) "
"INDEXTYPE IS CTXSYS.CONTEXT "
"PARAMETERS (''SYNC (ON COMMIT)'')"
"'; "
"EXCEPTION WHEN OTHERS THEN "
"IF SQLCODE = -955 THEN NULL; ELSE RAISE; END IF; "
"END;"
)
def _execute_ignoring_955(sql: str) -> None:
"""Run a CREATE statement and swallow ORA-00955 (object already exists).
Wraps the statement in PL/SQL so the exception handler runs server-side —
no round-trip cost for the common case.
"""
block = (
"BEGIN "
"EXECUTE IMMEDIATE :stmt; "
"EXCEPTION WHEN OTHERS THEN "
"IF SQLCODE = -955 THEN NULL; ELSE RAISE; END IF; "
"END;"
)
op.get_bind().exec_driver_sql(block, {"stmt": sql.strip()})
def _oracle_upgrade() -> None:
bind = op.get_bind()
# Tolerate concurrent DDL instead of failing immediately (ORA-00054).
bind.exec_driver_sql("ALTER SESSION SET DDL_LOCK_TIMEOUT = 30")
for ddl in _TABLES:
_execute_ignoring_955(ddl)
for idx in _INDEXES:
_execute_ignoring_955(idx)
# Hindsight on Oracle requires 23ai with VECTOR support (ASSM tablespace)
# and the CTXSYS package for full-text. Both index creations must succeed
# — the migration fails hard if either feature is unavailable, by design.
# We only swallow ORA-00955 (object already exists) so reruns are safe.
_execute_ignoring_955(_VECTOR_INDEX)
bind.exec_driver_sql(_TEXT_INDEX)
def _oracle_downgrade() -> None:
# Baseline downgrades aren't supported — dropping every table here would
# destroy customer data. Use point-in-time recovery instead.
raise NotImplementedError("Cannot downgrade past the Oracle baseline.")
def upgrade() -> None:
run_for_dialect(oracle=_oracle_upgrade)
def downgrade() -> None:
run_for_dialect(oracle=_oracle_downgrade)
@@ -1,64 +0,0 @@
"""Database URL normalization.
Hindsight accepts SQLAlchemy-style URLs like ``postgresql+asyncpg://...?ssl=require``
for its async engine, but the same string cannot be handed directly to synchronous
SQLAlchemy (psycopg2) or to :func:`asyncpg.create_pool`, which both expect a
libpq-compatible URL (``postgresql://...?sslmode=require``).
:func:`to_libpq_url` performs that translation. It is idempotent and safe to
apply to URLs that are already libpq-compatible, to the ``pg0`` embedded-PG
marker, or to any non-PostgreSQL string (returned unchanged).
"""
from __future__ import annotations
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
_ASYNCPG_SCHEMES = ("postgresql+asyncpg", "postgres+asyncpg")
_POSTGRES_SCHEMES = ("postgresql", "postgres") + _ASYNCPG_SCHEMES
def is_oracle_url(url: str) -> bool:
"""True if ``url`` is an Oracle SQLAlchemy URL (``oracle`` or ``oracle+oracledb``)."""
if not url or "://" not in url:
return False
return urlsplit(url).scheme.startswith("oracle")
def to_libpq_url(url: str) -> str:
"""Normalize a PostgreSQL URL for libpq-style consumers.
Accepts a SQLAlchemy URL (``postgresql+asyncpg://...``) or a plain libpq
URL and returns a form suitable for:
- :func:`sqlalchemy.create_engine` (sync / psycopg2)
- :func:`asyncpg.create_pool`
Transformations:
- ``postgresql+asyncpg`` / ``postgres+asyncpg`` / ``postgres`` → ``postgresql``
- Query param ``ssl=<mode>`` → ``sslmode=<mode>`` (SQLAlchemy's asyncpg
dialect uses ``ssl=``; libpq uses ``sslmode=``)
Any non-PostgreSQL input (e.g. the ``pg0`` embedded-PG marker, a sqlite
URL, an empty string) is returned unchanged. Already-normalized URLs are
returned unchanged.
"""
if not url or "://" not in url:
return url
parts = urlsplit(url)
if parts.scheme not in _POSTGRES_SCHEMES:
return url
new_scheme = "postgresql"
new_query_pairs = [
("sslmode", v) if k == "ssl" else (k, v) for k, v in parse_qsl(parts.query, keep_blank_values=True)
]
new_query = urlencode(new_query_pairs)
if new_scheme == parts.scheme and new_query == parts.query:
return url
return urlunsplit((new_scheme, parts.netloc, parts.path, new_query, parts.fragment))
@@ -1,207 +0,0 @@
"""Audit logging for feature usage tracking.
Provides fire-and-forget audit logging of all mutating and core operations
(retain, recall, reflect, bank CRUD, etc.) across HTTP, MCP, and system transports.
"""
from __future__ import annotations
import asyncio
import json
import logging
import uuid
from collections.abc import Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
from ..engine.db_utils import acquire_with_retry
logger = logging.getLogger(__name__)
@dataclass
class AuditEntry:
"""A single audit log entry."""
action: str
transport: str # "http", "mcp", "system"
bank_id: str | None = None
started_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
ended_at: datetime | None = None
request: dict[str, Any] | None = None
response: dict[str, Any] | None = None
metadata: dict[str, Any] = field(default_factory=dict)
def _json_default(obj: Any) -> str:
"""JSON serializer for objects not serializable by default."""
if isinstance(obj, datetime):
return obj.isoformat()
if isinstance(obj, uuid.UUID):
return str(obj)
if isinstance(obj, bytes):
return "<bytes>"
if isinstance(obj, set):
return list(obj)
return str(obj)
def _safe_json(data: Any) -> str | None:
"""Serialize data to JSON string, returning None on failure."""
if data is None:
return None
try:
return json.dumps(data, default=_json_default)
except Exception:
logger.debug("Failed to serialize audit data", exc_info=True)
return None
_SWEEP_INTERVAL_SECONDS = 3600 # Run retention sweep every hour
class AuditLogger:
"""Fire-and-forget audit log writer with optional retention sweep."""
def __init__(
self,
pool_getter: Callable[[], Any],
schema_getter: Callable[[], str],
enabled: bool,
allowed_actions: list[str],
retention_days: int = -1,
) -> None:
self._pool_getter = pool_getter
self._schema_getter = schema_getter
self._enabled = enabled
self._allowed_actions: frozenset[str] | None = frozenset(allowed_actions) if allowed_actions else None
self._retention_days = retention_days
self._sweep_task: asyncio.Task | None = None
def is_enabled(self, action: str) -> bool:
"""Check if audit logging is enabled for this action."""
if not self._enabled:
return False
if self._allowed_actions is not None:
return action in self._allowed_actions
return True
def log_fire_and_forget(self, entry: AuditEntry) -> None:
"""Schedule an audit write as a background task."""
if not self.is_enabled(entry.action):
return
try:
asyncio.create_task(self._safe_log(entry))
except RuntimeError:
# No running event loop (e.g. during shutdown)
logger.debug("Cannot schedule audit log write: no running event loop")
async def _safe_log(self, entry: AuditEntry) -> None:
"""Write audit entry to DB. Errors are logged, never raised."""
pool = self._pool_getter()
if pool is None:
logger.debug("Audit log skipped: pool not available")
return
try:
schema = self._schema_getter()
table = f"{schema}.audit_log"
async with acquire_with_retry(pool, max_retries=1) as conn:
await conn.execute(
f"""
INSERT INTO {table}
(id, action, transport, bank_id, started_at, ended_at, request, response, metadata)
VALUES
($1, $2, $3, $4, $5, $6, $7::jsonb, $8::jsonb, $9::jsonb)
""",
uuid.uuid4(),
entry.action,
entry.transport,
entry.bank_id,
entry.started_at,
entry.ended_at,
_safe_json(entry.request),
_safe_json(entry.response),
_safe_json(entry.metadata) or "{}",
)
except Exception as e:
logger.warning(f"Audit log write failed for action={entry.action}: {e}")
def start_retention_sweep(self) -> None:
"""Start the periodic retention sweep if retention is configured."""
if self._retention_days <= 0 or not self._enabled:
return
try:
self._sweep_task = asyncio.create_task(self._sweep_loop())
except RuntimeError:
logger.debug("Cannot start retention sweep: no running event loop")
async def stop_retention_sweep(self) -> None:
"""Stop the periodic retention sweep."""
if self._sweep_task and not self._sweep_task.done():
self._sweep_task.cancel()
try:
await self._sweep_task
except asyncio.CancelledError:
pass
self._sweep_task = None
async def _sweep_loop(self) -> None:
"""Periodically delete audit log entries older than retention_days."""
while True:
await self._run_sweep()
await asyncio.sleep(_SWEEP_INTERVAL_SECONDS)
async def _run_sweep(self) -> None:
"""Delete expired audit log entries. Concurrent-safe via row-level deletes."""
pool = self._pool_getter()
if pool is None:
return
try:
schema = self._schema_getter()
table = f"{schema}.audit_log"
async with acquire_with_retry(pool, max_retries=1) as conn:
result = await conn.execute(
f"DELETE FROM {table} WHERE started_at < NOW() - INTERVAL '{self._retention_days} days'"
)
if result and result != "DELETE 0":
logger.info(f"Audit log retention sweep: {result}")
except Exception as e:
logger.warning(f"Audit log retention sweep failed: {e}")
@asynccontextmanager
async def audit_context(
audit_logger: AuditLogger | None,
action: str,
transport: str,
bank_id: str | None = None,
request: dict[str, Any] | None = None,
metadata: dict[str, Any] | None = None,
):
"""Async context manager that times the operation and writes audit on exit.
Usage:
async with audit_context(logger, "retain", "http", bank_id, request_dict) as entry:
result = await do_work()
entry.response = result_dict
"""
if audit_logger is None or not audit_logger.is_enabled(action):
entry = AuditEntry(action=action, transport=transport, bank_id=bank_id)
yield entry
return
entry = AuditEntry(
action=action,
transport=transport,
bank_id=bank_id,
started_at=datetime.now(timezone.utc),
request=request,
metadata=metadata or {},
)
try:
yield entry
finally:
entry.ended_at = datetime.now(timezone.utc)
audit_logger.log_fire_and_forget(entry)
@@ -1,104 +0,0 @@
"""Prompts for the consolidation engine."""
# Default mission when no bank-specific mission is set
_DEFAULT_MISSION = "Track every detail: names, numbers, dates, places, and relationships. Prefer specifics over abstractions, never generalise."
# Processing rules — always present regardless of mission
_PROCESSING_RULES = """Processing rules (always apply):
1. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), etc. Never merge different facets into one observation.
2. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
3. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
4. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
5. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
6. SAME FACET → UPDATE, NOT CREATE: a new count supersedes the old count — UPDATE the existing count observation, don't create a second one. If there's an existing observation for the same specific facet, always UPDATE it rather than creating a duplicate.
7. PRESERVE HISTORY: observations that record significant events (sold, died, moved, changed) are important history — never DELETE them. Only delete an observation when it is restated identically or truly meaningless. Be very conservative with deletes.
8. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country""Sweden"), UPDATE to embed the resolved value.
9. NEVER merge observations about different people or unrelated topics."""
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
_BATCH_DATA_SECTION = """
NEW FACTS:
{facts_text}
EXISTING OBSERVATIONS (JSON array, pooled from recalls across all facts above):
{observations_text}
Each observation includes:
- id: unique identifier for updating
- text: the observation content
- proof_count: number of supporting memories
- occurred_start/occurred_end: temporal range of source facts
- source_memories: array of supporting facts with their text and dates
Compare the facts against existing observations:
- Same facet as an existing observation → UPDATE it (observation_id + source_fact_ids)
- New facet with durable knowledge → CREATE a new observation (source_fact_ids)
- Cross-reference facts within the batch: a later fact may resolve a vague reference in an earlier one
- Purely ephemeral facts → omit them unless the MISSION above explicitly targets such data (e.g. timestamped events, session state, screen content)"""
# Output format — JSON braces escaped as {{ }} so .format() leaves them literal
_BATCH_OUTPUT_FORMAT = """
Output a JSON object with three arrays.
## EXAMPLE
Input facts:
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Alice said she's exhausted from the project deadlines | Involving: Alice (occurred_start=2024-01-20, mentioned_at=2024-01-20)
Good observation text — clean prose, no metadata, each fact tracked distinctly:
"Alice works long hours, often past midnight."
"Alice feels exhausted from project deadlines."
Bad observation text — NEVER do this (verbatim copy of fact text with metadata):
"Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)"
Observation text rules:
- Write clean prose — NEVER copy raw fact lines or their metadata (temporal fields, "Involving:", "When:" labels, UUIDs).
- Parenthesized metadata like (occurred_start=...) and pipe-separated labels like "| Involving: ..." are fact formatting — strip them entirely from observation text.
- How many observations to create and how much to aggregate is driven by the MISSION above.
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890"]}}, {{"text": "Alice feels exhausted from project deadlines.", "source_fact_ids": ["b2c3d4e5-f6a7-8901-bcde-f12345678901"]}}],
"updates": [{{"text": "Alice works at Acme Corp as a senior engineer", "observation_id": "c3d4e5f6-a7b8-9012-cdef-123456789012", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"]}}],
"deletes": [{{"observation_id": "e5f6a7b8-c9d0-1234-efab-345678901234"}}]}}
Rules:
- "source_fact_ids": copy the EXACT UUID strings shown in brackets [uuid] from NEW FACTS — never use integers or positions.
- "observation_id": copy the EXACT "id" UUID string from EXISTING OBSERVATIONS.
- One create/update may reference multiple facts when they jointly support the observation.
- "deletes": only when an observation is directly superseded or contradicted by new facts.
- Do NOT include "tags" — handled automatically.
- Return {{"creates": [], "updates": [], "deletes": []}} if nothing durable is found."""
def build_batch_consolidation_prompt(
observations_mission: str | None = None,
observation_capacity_note: str | None = None,
) -> str:
"""
Build the consolidation prompt for batch mode (multiple facts per LLM call).
The mission defines *what* to track (customisable per bank).
Processing rules and output format are always present regardless of mission.
"""
mission = observations_mission or _DEFAULT_MISSION
capacity_section = ""
if observation_capacity_note:
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n{observation_capacity_note}"
return (
"You are a memory consolidation system. Synthesize facts into observations "
"and merge with existing observations when appropriate.\n\n"
f"## MISSION\n{mission}{capacity_section}\n\n"
f"{_PROCESSING_RULES}" + _BATCH_DATA_SECTION + _BATCH_OUTPUT_FORMAT
)
@@ -1,82 +0,0 @@
"""Database backend abstraction layer.
Provides a uniform interface over different database drivers (asyncpg, oracledb, etc.)
so that business logic is decoupled from any specific database platform.
Usage:
from hindsight_api.engine.db import create_database_backend, DatabaseBackend
backend = create_database_backend("postgresql")
await backend.initialize(dsn="postgresql://...")
async with backend.acquire() as conn:
rows = await conn.fetch("SELECT ...")
"""
from .base import DatabaseBackend, DatabaseConnection
from .ops import DataAccessOps
from .result import ResultRow
__all__ = [
"DataAccessOps",
"DatabaseBackend",
"DatabaseConnection",
"ResultRow",
"create_data_access_ops",
"create_database_backend",
]
def _get_backend_class(backend_type: str) -> type[DatabaseBackend]:
"""Resolve backend class by name using lazy imports."""
if backend_type == "postgresql":
from .postgresql import PostgreSQLBackend
return PostgreSQLBackend
if backend_type == "oracle":
from .oracle import OracleBackend
return OracleBackend
raise ValueError(f"Unknown database backend: {backend_type!r}. Supported: 'postgresql', 'oracle'.")
def _get_ops_class(backend_type: str) -> type[DataAccessOps]:
"""Resolve ops class by name using lazy imports."""
if backend_type == "postgresql":
from .ops_postgresql import PostgreSQLOps
return PostgreSQLOps
if backend_type == "oracle":
from .ops_oracle import OracleOps
return OracleOps
raise ValueError(f"Unknown data access ops: {backend_type!r}. Supported: 'postgresql', 'oracle'.")
def create_database_backend(backend_type: str) -> DatabaseBackend:
"""Factory: create a DatabaseBackend by name.
Args:
backend_type: One of "postgresql" or "oracle".
Returns:
An uninitialized DatabaseBackend instance.
Raises:
ValueError: If backend_type is not recognized.
"""
return _get_backend_class(backend_type)()
def create_data_access_ops(backend_type: str) -> DataAccessOps:
"""Factory: create a DataAccessOps by backend name.
Args:
backend_type: One of "postgresql" or "oracle".
Returns:
A DataAccessOps instance.
Raises:
ValueError: If backend_type is not recognized.
"""
return _get_ops_class(backend_type)()
@@ -1,342 +0,0 @@
"""Abstract base classes for database backend abstraction.
Defines the interfaces that all database backends (PostgreSQL, Oracle, etc.)
must implement. Business logic depends only on these interfaces.
"""
import json
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
# TYPE_CHECKING-only import to avoid circular import at runtime.
# DataAccessOps lives in ops.py which imports nothing from base.py,
# so the cycle is: base -> ops (type-only) and ops -> (nothing from base).
from typing import TYPE_CHECKING, Any
from .result import ResultRow
if TYPE_CHECKING:
from .ops import DataAccessOps
class DatabaseConnection(ABC):
"""Wraps a single connection from the pool.
Provides a uniform interface over asyncpg.Connection, oracledb cursor, etc.
Methods mirror asyncpg's connection API for minimal migration friction.
"""
@property
def backend_type(self) -> str:
"""Return ``"postgresql"`` or ``"oracle"``."""
return "postgresql"
def parse_json(self, value: Any) -> Any:
"""Parse a JSON column value into a Python object.
PG (asyncpg) returns JSON columns as strings that need json.loads().
Oracle returns them as pre-parsed dicts/lists (via OracleConnection
row conversion). This method normalizes both to Python objects.
"""
if value is None:
return None
if isinstance(value, str):
try:
return json.loads(value)
except (json.JSONDecodeError, TypeError):
return value
# Already a dict/list (Oracle pre-parses JSON columns)
return value
async def bulk_insert_from_arrays(
self,
table: str,
columns: list[str],
arrays: list[list],
*,
column_types: list[str] | None = None,
returning: str | None = None,
) -> list[ResultRow] | str:
"""Insert multiple rows from parallel arrays.
Default implementation uses ``INSERT ... SELECT * FROM unnest(...)``
(PostgreSQL). Oracle overrides this with ``executemany``.
Args:
table: Fully-qualified table name.
columns: Column names matching the arrays.
arrays: Parallel lists of values, one per column.
column_types: PG type suffixes for unnest casting (e.g. ``["text[]", "uuid[]"]``).
Ignored by backends that don't use unnest.
returning: Optional column expression for a RETURNING clause.
Returns:
If *returning* is set, a list of ResultRow; otherwise a status string.
"""
# Default: PostgreSQL unnest path
col_list = ", ".join(columns)
n_cols = len(columns)
types = column_types or ["text[]"] * n_cols
unnest_args = ", ".join(f"${i + 1}::{types[i]}" for i in range(n_cols))
query = f"INSERT INTO {table} ({col_list}) SELECT * FROM unnest({unnest_args})"
if returning:
query += f" RETURNING {returning}"
return await self.fetch(query, *arrays)
result = await self.execute(query, *arrays)
return result
@abstractmethod
@asynccontextmanager
async def transaction(self) -> AsyncIterator["DatabaseConnection"]:
"""Start a transaction (or savepoint if already in a transaction).
Yields:
Self — the same connection, now inside a transaction scope.
On clean exit the transaction is committed; on exception it is rolled back.
"""
... # pragma: no cover
yield # type: ignore[misc]
@abstractmethod
async def execute(self, query: str, *args: Any, timeout: float | None = None) -> str:
"""Execute a query and return a status string (e.g. 'INSERT 0 1').
Args:
query: SQL query with dialect-appropriate placeholders.
*args: Positional bind parameters.
timeout: Optional statement timeout in seconds.
Returns:
Command status string.
"""
...
@abstractmethod
async def executemany(self, query: str, args: list[tuple[Any, ...]], *, timeout: float | None = None) -> None:
"""Execute a query for each set of arguments.
Args:
query: SQL query with dialect-appropriate placeholders.
args: List of argument tuples, one per execution.
timeout: Optional statement timeout in seconds.
"""
...
@abstractmethod
async def fetch(self, query: str, *args: Any, timeout: float | None = None) -> list[ResultRow]:
"""Execute a query and return all rows.
Args:
query: SQL query with dialect-appropriate placeholders.
*args: Positional bind parameters.
timeout: Optional statement timeout in seconds.
Returns:
List of ResultRow objects.
"""
...
@abstractmethod
async def fetchrow(self, query: str, *args: Any, timeout: float | None = None) -> ResultRow | None:
"""Execute a query and return a single row (or None).
Args:
query: SQL query with dialect-appropriate placeholders.
*args: Positional bind parameters.
timeout: Optional statement timeout in seconds.
Returns:
A single ResultRow, or None if no rows match.
"""
...
@abstractmethod
async def fetchval(self, query: str, *args: Any, column: int = 0, timeout: float | None = None) -> Any:
"""Execute a query and return a single value from the first row.
Args:
query: SQL query with dialect-appropriate placeholders.
*args: Positional bind parameters.
column: Column index to return (default 0).
timeout: Optional statement timeout in seconds.
Returns:
The value from the specified column of the first row, or None.
"""
...
async def copy_records_to_table(
self,
table_name: str,
*,
records: list[tuple[Any, ...]],
columns: list[str],
timeout: float | None = None,
) -> None:
"""Bulk-load records into a table.
Default implementation uses executemany INSERT. Backends with native
bulk-load support (e.g. asyncpg COPY) should override for performance.
"""
cols = ", ".join(columns)
placeholders = ", ".join(f"${i + 1}" for i in range(len(columns)))
query = f"INSERT INTO {table_name} ({cols}) VALUES ({placeholders})"
await self.executemany(query, list(records))
class DatabaseBackend(ABC):
"""Database pool lifecycle and connection acquisition.
Manages the connection pool and provides context managers for
acquiring connections and running transactions.
The ``ops`` property provides backend-specific data access operations
(the Strategy pattern — like Django's ``connection.ops``). All business
logic should use ``backend.ops`` instead of creating DataAccessOps
instances directly.
"""
_ops_instance: "DataAccessOps | None" = None
# -- Backend capabilities --------------------------------------------
# Subclasses override these to advertise what the platform supports.
# Callers use these instead of checking ``config.database_backend``.
@property
def backend_type(self) -> str:
"""Return ``"postgresql"`` or ``"oracle"``."""
return "postgresql"
@property
def ops(self) -> "DataAccessOps":
"""Backend-specific data access operations (cached).
Follows the Django pattern: ``connection.ops`` provides the
operations handler for the current backend. Created lazily on
first access and cached for the lifetime of the backend.
"""
if self._ops_instance is None:
from . import create_data_access_ops
self._ops_instance = create_data_access_ops(self.backend_type)
return self._ops_instance
@property
def supports_partial_indexes(self) -> bool:
"""Can CREATE INDEX … WHERE <predicate>."""
return True
@property
def supports_bm25(self) -> bool:
"""Has BM25 / tsvector full-text search."""
return True
@property
def supports_unnest(self) -> bool:
"""Supports ``unnest()`` for expanding arrays into rows."""
return True
@property
def supports_pg_trgm(self) -> bool:
"""Platform *might* have pg_trgm (must still be checked at runtime)."""
return True
@property
def supports_worker_poller(self) -> bool:
"""Whether this backend supports the async WorkerPoller.
WorkerPoller is backend-agnostic (uses DatabaseBackend.acquire()).
All current backends (PostgreSQL, Oracle) support it.
"""
return True
def normalize_schema(self, schema: str | None) -> str | None:
"""Normalize a schema name for this backend.
Returns the schema as-is by default. Oracle overrides this to
convert ``"public"`` (a PG-specific default) to ``None`` (use the
connecting user's default schema).
"""
return schema
def run_migrations(self, dsn: str, *, schema: str | None = None) -> None:
"""Run database migrations for this backend.
PG uses Alembic migrations. Oracle uses its own idempotent DDL runner.
Subclasses must override this method.
"""
raise NotImplementedError(f"{type(self).__name__} must implement run_migrations()")
def create_task_backend(self, *, pool_getter: Any = None, schema_getter: Any = None) -> Any:
"""Create the task backend for this database.
All backends use BrokerTaskBackend for async worker/poller execution.
"""
from ..task_backend import BrokerTaskBackend
return BrokerTaskBackend(pool_getter=pool_getter, schema_getter=schema_getter)
@abstractmethod
async def initialize(
self,
dsn: str,
*,
min_size: int = 5,
max_size: int = 20,
command_timeout: float = 300,
acquire_timeout: float = 30,
statement_cache_size: int = 0,
init_callback: Any | None = None,
) -> None:
"""Create the connection pool.
Args:
dsn: Database connection string.
min_size: Minimum number of connections in the pool.
max_size: Maximum number of connections in the pool.
command_timeout: Default command timeout in seconds.
acquire_timeout: Timeout for acquiring a connection from the pool.
statement_cache_size: Size of the prepared-statement cache (0 to disable).
init_callback: Optional async callback invoked on each new connection.
"""
...
@abstractmethod
async def shutdown(self) -> None:
"""Close the connection pool and release all resources."""
...
@abstractmethod
@asynccontextmanager
async def acquire(self) -> AsyncIterator[DatabaseConnection]:
"""Acquire a connection from the pool.
Yields:
A DatabaseConnection wrapper.
"""
... # pragma: no cover
yield # type: ignore[misc]
@abstractmethod
@asynccontextmanager
async def transaction(self) -> AsyncIterator[DatabaseConnection]:
"""Acquire a connection and start a transaction.
The transaction is committed on clean exit, rolled back on exception.
Yields:
A DatabaseConnection wrapper inside a transaction.
"""
... # pragma: no cover
yield # type: ignore[misc]
@abstractmethod
def get_pool(self) -> Any:
"""Return the underlying raw pool object.
Escape hatch for gradual migration — callers that still need direct
pool access (e.g. asyncpg-specific features) can use this during
the transition period.
"""
...
@@ -1,437 +0,0 @@
"""Abstract base class for backend-specific data access operations.
SQLDialect handles SQL *fragment* generation (param placeholders, JSON ops, vector
distance, etc.) — stateless, no I/O.
DataAccessOps handles multi-statement *execution* patterns that differ between
backends (unnest batch insert vs executemany, LATERAL fan-out vs per-row query,
DISTINCT ON vs GROUP BY workarounds, etc.). Methods receive a DatabaseConnection
and execute complete operations.
This eliminates scattered ``if backend_type == "postgresql"`` conditionals from
business logic. Adding a new backend (e.g. Neon, Databricks) means implementing
this ABC — consumer code never checks the backend directly.
Follows the Strategy pattern (Fowler's "Replace Conditional with Polymorphism")
and mirrors Django's ``DatabaseOperations`` architecture.
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any
from uuid import UUID
from .base import DatabaseConnection
from .result import ResultRow
@dataclass
class TagListingParts:
"""Backend-specific SQL fragments for the tag listing query."""
tag_source: str
non_empty_check: str
tag_col: str
bank_prefix: str
class DataAccessOps(ABC):
"""Backend-specific multi-statement data access operations.
Each method encapsulates a complete DB operation that differs
in execution strategy between backends.
"""
@property
def uses_observation_sources_table(self) -> bool:
"""Whether this backend uses the observation_sources junction table.
PG uses native array ops (source_memory_ids column) for reads and
skips junction table writes. Oracle uses the junction table for both.
"""
return True # Default: use junction table (Oracle)
# -- Bulk insert operations ------------------------------------------
@abstractmethod
async def bulk_upsert_chunks(
self,
conn: DatabaseConnection,
table: str,
chunk_ids: list[str],
document_ids: list[str],
bank_ids: list[str],
chunk_texts: list[str],
chunk_indices: list[int],
content_hashes: list[str],
) -> None:
"""Bulk upsert chunks with ON CONFLICT handling.
PG uses INSERT ... SELECT FROM unnest() with ON CONFLICT DO UPDATE.
Non-PG uses bulk_insert_from_arrays (executemany).
"""
...
@abstractmethod
async def insert_facts_batch(
self,
conn: DatabaseConnection,
bank_id: str,
fact_texts: list[str],
embeddings: list[str],
event_dates: list,
occurred_starts: list,
occurred_ends: list,
mentioned_ats: list,
contexts: list[str],
fact_types: list[str],
metadata_jsons: list[str],
chunk_ids: list,
document_ids: list,
tags_list: list[str],
observation_scopes_list: list,
text_signals_list: list,
text_search_extension: str = "native",
) -> list[str]:
"""Batch-insert facts, returning IDs.
PG uses INSERT ... SELECT FROM unnest() with RETURNING.
Non-PG inserts row-by-row with individual RETURNING.
"""
...
@abstractmethod
async def bulk_insert_links(
self,
conn: DatabaseConnection,
table: str,
sorted_links: list[tuple],
bank_id: str,
nil_entity_uuid: str,
exists_clause: str,
chunk_size: int = 5000,
) -> None:
"""Bulk insert memory_links with conflict handling.
PG uses INSERT ... SELECT FROM unnest() with chunking.
Non-PG uses executemany.
"""
...
@abstractmethod
async def bulk_insert_entities(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
entity_names: list[str],
entity_dates: list,
) -> dict[str, str]:
"""Bulk insert entities with ON CONFLICT DO NOTHING, returning id-by-lowercase-name.
PG uses INSERT ... SELECT FROM unnest() with RETURNING.
Non-PG inserts row-by-row then SELECTs.
"""
...
@abstractmethod
async def fetch_missing_entity_ids(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
missing_names: list[str],
) -> list[ResultRow]:
"""Fetch entity IDs for names that conflicted during insert.
PG uses unnest + JOIN.
Non-PG queries each name individually.
"""
...
@abstractmethod
async def bulk_insert_unit_entities(
self,
conn: DatabaseConnection,
table: str,
unit_ids: list,
entity_ids: list,
) -> None:
"""Bulk insert unit_entities links with ON CONFLICT DO NOTHING.
PG uses INSERT ... SELECT FROM unnest().
Non-PG uses executemany.
"""
...
# -- LATERAL / fan-out queries ---------------------------------------
@abstractmethod
async def fetch_entity_unit_fanout(
self,
conn: DatabaseConnection,
ue_table: str,
entity_id_list: list[UUID],
limit_per_entity: int,
) -> list[ResultRow]:
"""Fetch unit_ids for a list of entities with per-entity row cap.
PG uses unnest + CROSS JOIN LATERAL with LIMIT.
Non-PG queries each entity individually.
"""
...
@abstractmethod
async def fetch_unit_dates(
self,
conn: DatabaseConnection,
mu_table: str,
unit_ids: list[str],
) -> list[ResultRow]:
"""Fetch event_date/fact_type for a list of unit IDs.
PG uses ANY($1) array binding.
Non-PG queries each unit individually.
"""
...
@abstractmethod
async def fetch_temporal_neighbors(
self,
conn: DatabaseConnection,
mu_table: str,
bank_id: str,
lateral_unit_ids: list,
lateral_event_dates: list,
lateral_fact_types: list,
half_limit: int,
batch_size: int = 500,
) -> list[ResultRow]:
"""Fetch temporal neighbors using bidirectional index scan.
PG uses unnest + CROSS JOIN LATERAL for batched bidirectional scan.
Non-PG queries each unit individually with backward/forward scans.
"""
...
# -- CTE builders for graph retrieval --------------------------------
@abstractmethod
def build_entity_expansion_cte(
self,
mu_table: str,
ue_table: str,
per_entity_limit: int,
) -> str:
"""Build entity expansion CTE for link expansion retrieval.
PG uses DISTINCT ON with CROSS JOIN LATERAL and GROUP BY.
Non-PG splits into entity_scores subquery then JOINs for full columns
(can't GROUP BY CLOB).
"""
...
@abstractmethod
def build_semantic_causal_cte(
self,
ml_table: str,
mu_table: str,
) -> str:
"""Build semantic + causal expansion CTEs.
PG uses DISTINCT ON for deduplication.
Non-PG computes MAX(weight) in subquery then JOINs for full columns.
"""
...
@abstractmethod
async def expand_observations(
self,
conn: DatabaseConnection,
mu_table: str,
ue_table: str,
ml_table: str,
seed_ids: list,
budget: int,
per_entity_limit: int,
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
"""Observation-specific graph expansion.
PG uses native array ops (source_memory_ids column) for performance.
Oracle uses the observation_sources junction table.
"""
...
# -- Tag listing -----------------------------------------------------
@abstractmethod
def build_tag_listing_parts(self, mu_table: str) -> TagListingParts:
"""Build SQL fragments for the tag listing query.
PG uses unnest(tags) to expand the VARCHAR[] column.
Non-PG uses CROSS APPLY JSON_TABLE on the CLOB column.
"""
...
# -- Bank index management -------------------------------------------
@abstractmethod
async def create_bank_vector_indexes(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
internal_id: str,
index_clause: str,
fact_types: dict[str, str],
) -> None:
"""Create per-bank partial vector indexes.
PG creates per-(bank, fact_type) partial indexes.
Non-PG is a no-op (uses global index).
"""
...
@abstractmethod
async def drop_bank_vector_indexes(
self,
conn: DatabaseConnection,
schema: str,
internal_id: str,
fact_types: dict[str, str],
) -> None:
"""Drop per-bank partial vector indexes.
PG drops per-(bank, fact_type) indexes.
Non-PG is a no-op.
"""
...
# -- Entity resolution strategy routing ------------------------------
@abstractmethod
def get_entity_resolution_strategy(self) -> str:
"""Return the fuzzy entity matching strategy name.
PG uses "trigram" (pg_trgm).
Non-PG uses "oracle_fuzzy" (UTL_MATCH) or falls back to "full".
"""
...
# -- Webhook operations ------------------------------------------------
@abstractmethod
async def create_webhook(
self,
conn: DatabaseConnection,
table: str,
webhook_id: Any,
bank_id: str,
url: str,
secret: str | None,
event_types: list[str],
enabled: bool,
http_config_json: str,
) -> ResultRow | None:
"""Insert a webhook row and return the created row."""
...
@abstractmethod
async def list_webhooks_for_bank(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
) -> list[ResultRow]:
"""List all webhooks for a bank, ordered by created_at."""
...
@abstractmethod
async def get_webhooks_for_dispatch(
self,
conn: DatabaseConnection,
webhook_table: str,
bank_id: str,
) -> list[ResultRow]:
"""Get enabled webhooks matching a bank (bank-specific + global NULL rows)."""
...
@abstractmethod
async def update_webhook(
self,
conn: DatabaseConnection,
table: str,
webhook_id: Any,
bank_id: str,
set_clauses: list[str],
params: list[Any],
) -> ResultRow | None:
"""Update a webhook and return the updated row, or None if not found."""
...
@abstractmethod
async def delete_webhook(
self,
conn: DatabaseConnection,
table: str,
webhook_id: Any,
bank_id: str,
) -> bool:
"""Delete a webhook. Returns True if a row was deleted."""
...
@abstractmethod
async def list_webhook_deliveries(
self,
conn: DatabaseConnection,
ops_table: str,
webhook_id: str,
bank_id: str,
limit: int,
cursor: str | None,
) -> list[ResultRow]:
"""List webhook delivery operations for a specific webhook, newest first."""
...
@abstractmethod
async def insert_webhook_delivery_task(
self,
conn: DatabaseConnection,
ops_table: str,
operation_id: Any,
bank_id: str,
payload_json: str,
timestamp: Any,
) -> None:
"""Insert a webhook delivery task into async_operations."""
...
# -- Task claiming operations ------------------------------------------
@abstractmethod
async def claim_tasks(
self,
conn: DatabaseConnection,
table: str,
worker_id: str,
reserved_limits: dict[str, int],
shared_limit: int,
) -> list[ResultRow]:
"""Claim pending tasks from the async_operations table.
PG implementation can use NOT EXISTS + FOR UPDATE SKIP LOCKED in one query.
Oracle implementation uses two-step claims (query busy banks first, then
claim excluding them) to avoid ORA-02014.
Returns claimed rows with operation_id, operation_type, task_payload, retry_count.
The caller is responsible for building ClaimedTask objects.
"""
...
# -- Shared helpers (concrete) -----------------------------------------
def _get_mu_table(self) -> str:
"""Get the fully-qualified memory_units table name."""
from ..schema import fq_table
return fq_table("memory_units")
@@ -1,938 +0,0 @@
"""Oracle 23ai implementation of DataAccessOps.
Uses executemany, per-row queries, JSON_TABLE, and ROW_NUMBER() workarounds
for Oracle-specific syntax requirements (no unnest, no DISTINCT ON, CLOB
columns can't appear in GROUP BY).
"""
import json
import uuid as uuid_mod
from datetime import UTC, datetime
from typing import Any
from uuid import UUID
from .base import DatabaseConnection
from .ops import DataAccessOps, TagListingParts
from .result import DictResultRow as ResultRow
class OracleOps(DataAccessOps):
"""Oracle-specific data access operations."""
async def bulk_upsert_chunks(
self,
conn: DatabaseConnection,
table: str,
chunk_ids: list[str],
document_ids: list[str],
bank_ids: list[str],
chunk_texts: list[str],
chunk_indices: list[int],
content_hashes: list[str],
) -> None:
# Oracle's thin-client executemany with array binds is already well-optimized —
# it batches network round-trips into a single call, so INSERT ALL or other
# patterns would not provide a meaningful improvement.
await conn.bulk_insert_from_arrays(
table,
["chunk_id", "document_id", "bank_id", "chunk_text", "chunk_index", "content_hash"],
[
chunk_ids,
document_ids,
bank_ids,
chunk_texts,
chunk_indices,
content_hashes,
],
column_types=["text[]", "text[]", "text[]", "text[]", "integer[]", "text[]"],
)
async def insert_facts_batch(
self,
conn: DatabaseConnection,
bank_id: str,
fact_texts: list[str],
embeddings: list[str],
event_dates: list,
occurred_starts: list,
occurred_ends: list,
mentioned_ats: list,
contexts: list[str],
fact_types: list[str],
metadata_jsons: list[str],
chunk_ids: list,
document_ids: list,
tags_list: list[str],
observation_scopes_list: list,
text_signals_list: list,
text_search_extension: str = "native",
) -> list[str]:
table = self._get_mu_table()
# Generate UUIDs client-side so we can use executemany (single network
# round-trip) instead of N individual INSERT+RETURNING calls.
unit_ids = [str(uuid_mod.uuid4()) for _ in range(len(fact_texts))]
rows_data = []
for i in range(len(fact_texts)):
tags_value = json.loads(tags_list[i]) if tags_list[i] else []
rows_data.append(
(
unit_ids[i],
bank_id,
fact_texts[i],
embeddings[i],
event_dates[i],
occurred_starts[i],
occurred_ends[i],
mentioned_ats[i],
contexts[i],
fact_types[i],
metadata_jsons[i],
chunk_ids[i],
document_ids[i],
tags_value,
observation_scopes_list[i],
text_signals_list[i],
)
)
await conn.executemany(
f"""
INSERT INTO {table} (id, bank_id, text, embedding, event_date, occurred_start,
occurred_end, mentioned_at, context, fact_type, metadata, chunk_id, document_id,
tags, observation_scopes, text_signals)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
""",
rows_data,
)
return unit_ids
async def bulk_insert_links(
self,
conn: DatabaseConnection,
table: str,
sorted_links: list[tuple],
bank_id: str,
nil_entity_uuid: str,
exists_clause: str,
chunk_size: int = 5000,
) -> None:
# The backend rewrites ON CONFLICT DO NOTHING for duplicate suppression.
# WHERE EXISTS checks are intentionally skipped: executemany does not support
# correlated subqueries in this form, and callers guarantee unit validity.
from_ids = [lnk[0] for lnk in sorted_links]
to_ids = [lnk[1] for lnk in sorted_links]
types = [lnk[2] for lnk in sorted_links]
weights = [lnk[3] for lnk in sorted_links]
entity_ids = [lnk[4] for lnk in sorted_links]
await conn.executemany(
f"""
INSERT INTO {table}
(from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id)
VALUES ($1, $2, $3, $4, $5, $6)
ON CONFLICT (from_unit_id, to_unit_id, link_type,
COALESCE(entity_id, '{nil_entity_uuid}'::uuid))
DO NOTHING
""",
[(from_ids[i], to_ids[i], types[i], weights[i], entity_ids[i], bank_id) for i in range(len(sorted_links))],
)
async def bulk_insert_entities(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
entity_names: list[str],
entity_dates: list,
) -> dict[str, str]:
# Row-by-row insert with duplicate suppression.
# Can't use RETURNING with ON CONFLICT DO NOTHING reliably,
# so INSERT (ignoring dups) then SELECT all IDs at the end.
id_by_name: dict[str, str] = {}
for name, event_date in zip(entity_names, entity_dates):
ts = event_date if event_date else datetime.now(UTC)
await conn.execute(
f"""
INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count)
VALUES ($1, $2, $3, $3, 0)
ON CONFLICT (bank_id, LOWER(canonical_name)) DO NOTHING
""",
bank_id,
name,
ts,
)
# Now SELECT all the entities we just inserted (or that already existed)
for name in entity_names:
row = await conn.fetchrow(
f"""
SELECT id, LOWER(canonical_name) AS name_lower
FROM {table}
WHERE bank_id = $1 AND LOWER(canonical_name) = LOWER($2)
""",
bank_id,
name,
)
if row:
id_by_name[row["name_lower"]] = row["id"]
return id_by_name
async def fetch_missing_entity_ids(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
missing_names: list[str],
) -> list[ResultRow]:
# Query each missing entity individually
results: list[ResultRow] = []
for orig_name in missing_names:
row = await conn.fetchrow(
f"""
SELECT id, LOWER(canonical_name) AS name_lower
FROM {table}
WHERE bank_id = $1 AND LOWER(canonical_name) = LOWER($2)
""",
bank_id,
orig_name,
)
if row:
# Wrap in a dict-like to include input_name for downstream compat
results.append(row)
return results
async def bulk_insert_unit_entities(
self,
conn: DatabaseConnection,
table: str,
unit_ids: list,
entity_ids: list,
) -> None:
await conn.executemany(
f"""
INSERT INTO {table} (unit_id, entity_id)
VALUES ($1, $2)
ON CONFLICT DO NOTHING
""",
list(zip(unit_ids, entity_ids)),
)
async def fetch_entity_unit_fanout(
self,
conn: DatabaseConnection,
ue_table: str,
entity_id_list: list[UUID],
limit_per_entity: int,
) -> list[ResultRow]:
# Query each entity individually
rows: list[ResultRow] = []
for eid in entity_id_list:
entity_rows = await conn.fetch(
f"""
SELECT $1 AS entity_id, ue.unit_id
FROM {ue_table} ue
WHERE ue.entity_id = $1
ORDER BY ue.unit_id DESC
LIMIT $2
""",
eid,
limit_per_entity,
)
rows.extend(entity_rows)
return rows
async def fetch_unit_dates(
self,
conn: DatabaseConnection,
mu_table: str,
unit_ids: list[str],
) -> list[ResultRow]:
# No ANY() array binding; query each unit individually
rows: list[ResultRow] = []
for uid in unit_ids:
row = await conn.fetchrow(
f"""
SELECT id, event_date, fact_type
FROM {mu_table}
WHERE id = $1
""",
uid,
)
if row:
rows.append(row)
return rows
async def fetch_temporal_neighbors(
self,
conn: DatabaseConnection,
mu_table: str,
bank_id: str,
lateral_unit_ids: list,
lateral_event_dates: list,
lateral_fact_types: list,
half_limit: int,
batch_size: int = 500,
) -> list[ResultRow]:
# Per-unit queries (no unnest/LATERAL in Oracle).
# Fetch up to half_limit in each direction, then combine and keep the
# half_limit closest overall via ROW_NUMBER — matching the PG behavior.
rows: list[ResultRow] = []
for uid, edate, ftype in zip(lateral_unit_ids, lateral_event_dates, lateral_fact_types):
uid_str = str(uid) if not isinstance(uid, str) else uid
unit_rows = await conn.fetch(
f"""
SELECT from_id, id, event_date, time_diff_hours FROM (
SELECT combined.*, ROW_NUMBER() OVER (ORDER BY combined.time_diff_hours) AS rn
FROM (
SELECT * FROM (
SELECT $1 AS from_id, mu.id, mu.event_date,
ABS(EXTRACT(DAY FROM (mu.event_date - $2)) * 24
+ EXTRACT(HOUR FROM (mu.event_date - $2))) AS time_diff_hours
FROM {mu_table} mu
WHERE mu.bank_id = $4
AND mu.fact_type = $3
AND mu.event_date <= $2
AND mu.id != $6
ORDER BY mu.event_date DESC
FETCH FIRST $5 ROWS ONLY
) bwd
UNION ALL
SELECT * FROM (
SELECT $1 AS from_id, mu.id, mu.event_date,
ABS(EXTRACT(DAY FROM (mu.event_date - $2)) * 24
+ EXTRACT(HOUR FROM (mu.event_date - $2))) AS time_diff_hours
FROM {mu_table} mu
WHERE mu.bank_id = $4
AND mu.fact_type = $3
AND mu.event_date > $2
AND mu.id != $6
ORDER BY mu.event_date ASC
FETCH FIRST $5 ROWS ONLY
) fwd
) combined
) ranked
WHERE rn <= $5
""",
uid_str,
edate,
ftype,
bank_id,
half_limit,
uid,
)
rows.extend(unit_rows)
return rows
def build_entity_expansion_cte(
self,
mu_table: str,
ue_table: str,
per_entity_limit: int,
) -> str:
# Oracle: can't GROUP BY CLOB columns (text, context).
# Restructure: count entities per unit_id in a subquery, then join to get full columns.
return f"""
seed_entities AS (
SELECT DISTINCT ue.entity_id
FROM {ue_table} ue
WHERE ue.unit_id = ANY($1::uuid[])
),
entity_scores AS (
SELECT t.unit_id, COUNT(DISTINCT se.entity_id) AS score
FROM seed_entities se
CROSS JOIN LATERAL (
SELECT ue_target.unit_id
FROM {ue_table} ue_target
WHERE ue_target.entity_id = se.entity_id
AND ue_target.unit_id != ALL($1::uuid[])
ORDER BY ue_target.unit_id DESC
FETCH FIRST {per_entity_limit} ROWS ONLY
) t
GROUP BY t.unit_id
),
entity_expanded AS (
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
es.score, 'entity' AS source
FROM entity_scores es
JOIN {mu_table} mu ON mu.id = es.unit_id
WHERE mu.fact_type = $2
ORDER BY es.score DESC
FETCH FIRST $3 ROWS ONLY
)"""
def build_semantic_causal_cte(
self,
ml_table: str,
mu_table: str,
) -> str:
# Non-PG: can't GROUP BY CLOB columns, no DISTINCT ON.
# Restructure semantic: compute max weight per id, then join for full columns.
return f"""
sem_scores AS (
SELECT id, MAX(weight) AS score
FROM (
SELECT mu.id, ml.weight
FROM {ml_table} ml
JOIN {mu_table} mu ON mu.id = ml.to_unit_id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic'
AND mu.fact_type = $2
AND mu.id != ALL($1::uuid[])
UNION ALL
SELECT mu.id, ml.weight
FROM {ml_table} ml
JOIN {mu_table} mu ON mu.id = ml.from_unit_id
WHERE ml.to_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic'
AND mu.fact_type = $2
AND mu.id != ALL($1::uuid[])
) sem_raw
GROUP BY id
),
semantic_expanded AS (
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
ss.score, 'semantic' AS source
FROM sem_scores ss
JOIN {mu_table} mu ON mu.id = ss.id
ORDER BY ss.score DESC
FETCH FIRST $3 ROWS ONLY
),
causal_ranked AS (
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
ml.weight AS score,
'causal' AS source,
ROW_NUMBER() OVER (PARTITION BY mu.id ORDER BY ml.weight DESC) AS rn_
FROM {ml_table} ml
JOIN {mu_table} mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
AND mu.fact_type = $2
),
causal_expanded AS (
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags, proof_count, score, source
FROM causal_ranked WHERE rn_ = 1
ORDER BY score DESC
FETCH FIRST $3 ROWS ONLY
)"""
async def expand_observations(
self,
conn: DatabaseConnection,
mu_table: str,
ue_table: str,
ml_table: str,
seed_ids: list,
budget: int,
per_entity_limit: int,
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
import logging
logger = logging.getLogger(__name__)
# Entity expansion via observation_sources junction table.
# Previously used JSON_TABLE to explode source_memory_ids CLOB. The junction
# table approach uses standard SQL joins, identical to the PG backend.
from ..schema import fq_table
obs_sources_table = fq_table("observation_sources")
entity_rows = await conn.fetch(
f"""
WITH seed_sources AS (
SELECT DISTINCT os.source_id
FROM {obs_sources_table} os
WHERE os.observation_id = ANY($1::uuid[])
),
source_entities AS (
SELECT DISTINCT ue_seed.entity_id
FROM seed_sources ss
JOIN {ue_table} ue_seed ON ue_seed.unit_id = ss.source_id
),
connected_sources AS (
SELECT DISTINCT t.unit_id AS source_id
FROM source_entities se
CROSS JOIN LATERAL (
SELECT ue_target.unit_id
FROM {ue_table} ue_target
WHERE ue_target.entity_id = se.entity_id
ORDER BY ue_target.unit_id DESC
FETCH FIRST {per_entity_limit} ROWS ONLY
) t
WHERE NOT EXISTS (
SELECT 1 FROM seed_sources ss WHERE ss.source_id = t.unit_id
)
)
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
(SELECT COUNT(*)
FROM {obs_sources_table} os2
WHERE os2.observation_id = mu.id
AND os2.source_id IN (SELECT source_id FROM connected_sources)
) AS score
FROM {mu_table} mu
WHERE mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
AND EXISTS (
SELECT 1 FROM {obs_sources_table} os3
WHERE os3.observation_id = mu.id
AND os3.source_id IN (SELECT source_id FROM connected_sources)
)
ORDER BY score DESC
FETCH FIRST $2 ROWS ONLY
""",
seed_ids,
budget,
)
logger.debug(f"[LinkExpansion] observation graph (Oracle): found {len(entity_rows)} connected observations")
# Semantic + causal for observations (Oracle path)
# Avoids GROUP BY CLOB and DISTINCT ON — mirrors _expand_world_facts Oracle strategy.
sem_causal_rows = await conn.fetch(
f"""
WITH sem_scores AS (
SELECT id, MAX(weight) AS score
FROM (
SELECT mu.id, ml.weight
FROM {ml_table} ml JOIN {mu_table} mu ON mu.id = ml.to_unit_id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
UNION ALL
SELECT mu.id, ml.weight
FROM {ml_table} ml JOIN {mu_table} mu ON mu.id = ml.from_unit_id
WHERE ml.to_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
) sem_raw
GROUP BY id
),
semantic_expanded AS (
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
ss.score, 'semantic' AS source
FROM sem_scores ss
JOIN {mu_table} mu ON mu.id = ss.id
ORDER BY ss.score DESC
FETCH FIRST $2 ROWS ONLY
),
causal_ranked AS (
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, mu.proof_count, ml.weight AS score,
'causal' AS source,
ROW_NUMBER() OVER (PARTITION BY mu.id ORDER BY ml.weight DESC) AS rn_
FROM {ml_table} ml
JOIN {mu_table} mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
AND mu.fact_type = 'observation'
),
causal_expanded AS (
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags, proof_count, score, source
FROM causal_ranked WHERE rn_ = 1
ORDER BY score DESC
FETCH FIRST $2 ROWS ONLY
)
SELECT * FROM semantic_expanded
UNION ALL
SELECT * FROM causal_expanded
""",
seed_ids,
budget,
)
semantic_rows = [r for r in sem_causal_rows if r["source"] == "semantic"]
causal_rows = [r for r in sem_causal_rows if r["source"] == "causal"]
return list(entity_rows), semantic_rows, causal_rows
def build_tag_listing_parts(self, mu_table: str) -> TagListingParts:
return TagListingParts(
tag_source=(
f"{mu_table} mu CROSS APPLY JSON_TABLE(mu.tags, '$[*]' COLUMNS (tag VARCHAR2(256) PATH '$')) jt"
),
non_empty_check="AND mu.tags IS NOT NULL AND DBMS_LOB.GETLENGTH(mu.tags) > 2",
tag_col="jt.tag",
bank_prefix="mu.",
)
async def create_bank_vector_indexes(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
internal_id: str,
index_clause: str,
fact_types: dict[str, str],
) -> None:
# Oracle 23ai supports HNSW vector indexes but does NOT support partial
# indexes (WHERE clause on CREATE INDEX for vector indexes). Uses a single
# global HNSW index with ORGANIZATION NEIGHBOR PARTITIONS created during
# migrations. memory_units is partitioned by LIST (bank_id) AUTOMATIC,
# so Oracle creates partitions per bank on INSERT and the optimizer can
# prune partitions on bank_id-scoped queries.
return
async def drop_bank_vector_indexes(
self,
conn: DatabaseConnection,
schema: str,
internal_id: str,
fact_types: dict[str, str],
) -> None:
# Oracle uses a single global vector index (no per-bank indexes to drop).
return
def get_entity_resolution_strategy(self) -> str:
return "oracle_fuzzy"
# -- Webhook operations ------------------------------------------------
async def create_webhook(
self,
conn,
table,
webhook_id,
bank_id,
url,
secret,
event_types,
enabled,
http_config_json,
):
return await conn.fetchrow(
f"""
INSERT INTO {table}
(id, bank_id, url, secret, event_types, enabled, http_config, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, NOW(), NOW())
RETURNING id, bank_id, url, secret, event_types, enabled,
http_config::text, created_at::text, updated_at::text
""",
webhook_id,
bank_id,
url,
secret,
event_types,
enabled,
http_config_json,
)
async def list_webhooks_for_bank(self, conn, table, bank_id):
return await conn.fetch(
f"""
SELECT id, bank_id, url, secret, event_types, enabled,
http_config::text, created_at::text, updated_at::text
FROM {table}
WHERE bank_id = $1
ORDER BY created_at
""",
bank_id,
)
async def get_webhooks_for_dispatch(self, conn, webhook_table, bank_id):
return await conn.fetch(
f"""
SELECT id, bank_id, url, secret, event_types, enabled, http_config::text
FROM {webhook_table}
WHERE (bank_id = $1 OR bank_id IS NULL) AND enabled = true
""",
bank_id,
)
async def update_webhook(self, conn, table, webhook_id, bank_id, set_clauses, params):
set_clauses_with_ts = set_clauses + ["updated_at = NOW()"]
return await conn.fetchrow(
f"""
UPDATE {table}
SET {", ".join(set_clauses_with_ts)}
WHERE id = $1 AND bank_id = $2
RETURNING id, bank_id, url, secret, event_types, enabled,
http_config::text, created_at::text, updated_at::text
""",
*params,
)
async def delete_webhook(self, conn, table, webhook_id, bank_id):
result = await conn.execute(
f"DELETE FROM {table} WHERE id = $1 AND bank_id = $2",
webhook_id,
bank_id,
)
return int(result.split()[-1]) > 0 if result else False
async def list_webhook_deliveries(self, conn, ops_table, webhook_id, bank_id, limit, cursor):
fetch_limit = limit + 1
if cursor:
return await conn.fetch(
f"""
SELECT operation_id, status, retry_count, next_retry_at::text,
error_message, task_payload, result_metadata::text, created_at::text, updated_at::text
FROM {ops_table}
WHERE operation_type = 'webhook_delivery'
AND bank_id = $1
AND task_payload->>'webhook_id' = $2
AND created_at < $3::timestamptz
ORDER BY created_at DESC
LIMIT $4
""",
bank_id,
webhook_id,
cursor,
fetch_limit,
)
return await conn.fetch(
f"""
SELECT operation_id, status, retry_count, next_retry_at::text,
error_message, task_payload, result_metadata::text, created_at::text, updated_at::text
FROM {ops_table}
WHERE operation_type = 'webhook_delivery'
AND bank_id = $1
AND task_payload->>'webhook_id' = $2
ORDER BY created_at DESC
LIMIT $3
""",
bank_id,
webhook_id,
fetch_limit,
)
async def insert_webhook_delivery_task(self, conn, ops_table, operation_id, bank_id, payload_json, timestamp):
await conn.execute(
f"""
INSERT INTO {ops_table}
(operation_id, bank_id, operation_type, status, task_payload, result_metadata, created_at, updated_at)
VALUES ($1, $2, 'webhook_delivery', 'pending', $3::jsonb, '{{}}'::jsonb, $4, $4)
""",
operation_id,
bank_id,
payload_json,
timestamp,
)
# -- Task claiming operations ------------------------------------------
async def claim_tasks(self, conn, table, worker_id, reserved_limits, shared_limit):
"""Oracle two-step claiming to avoid ORA-02014 with NOT EXISTS + FOR UPDATE."""
all_rows = []
claimed_ids = []
# --- Phase 1: claim from reserved pools ---
for op_type, limit in reserved_limits.items():
if limit <= 0:
continue
if op_type == "consolidation":
# Two-step: find busy banks first, then claim excluding them
busy_banks = await conn.fetch(
f"""
SELECT DISTINCT bank_id FROM {table}
WHERE operation_type = 'consolidation' AND status = 'processing'
""",
)
busy_bank_ids = [r["bank_id"] for r in busy_banks]
if busy_bank_ids:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
AND bank_id != ALL($1::text[])
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
""",
busy_bank_ids,
limit,
)
else:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
""",
limit,
)
else:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = $1
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
""",
op_type,
limit,
)
for row in rows:
claimed_ids.append(row["operation_id"])
all_rows.append(row)
# --- Phase 2: claim from shared pool ---
remaining_shared = shared_limit
if remaining_shared > 0:
# 2a. Non-consolidation tasks
if claimed_ids:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type != 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
AND operation_id != ALL($1::uuid[])
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
""",
claimed_ids,
remaining_shared,
)
else:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type != 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
""",
remaining_shared,
)
for row in rows:
claimed_ids.append(row["operation_id"])
all_rows.append(row)
remaining_shared -= len(rows)
# 2b. Consolidation tasks (with bank-serialization)
if remaining_shared > 0:
busy_banks_2 = await conn.fetch(
f"""
SELECT DISTINCT bank_id FROM {table}
WHERE operation_type = 'consolidation' AND status = 'processing'
""",
)
busy_bank_ids_2 = [r["bank_id"] for r in busy_banks_2]
if claimed_ids:
if busy_bank_ids_2:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
AND operation_id != ALL($1::uuid[])
AND bank_id != ALL($2::text[])
ORDER BY created_at
LIMIT $3
FOR UPDATE SKIP LOCKED
""",
claimed_ids,
busy_bank_ids_2,
remaining_shared,
)
else:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
AND operation_id != ALL($1::uuid[])
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
""",
claimed_ids,
remaining_shared,
)
else:
if busy_bank_ids_2:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
AND bank_id != ALL($1::text[])
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
""",
busy_bank_ids_2,
remaining_shared,
)
else:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
""",
remaining_shared,
)
for row in rows:
claimed_ids.append(row["operation_id"])
all_rows.append(row)
if not all_rows:
return []
# Mark all claimed rows as processing
operation_ids = [row["operation_id"] for row in all_rows]
await conn.execute(
f"""
UPDATE {table}
SET status = 'processing', worker_id = $1, claimed_at = now(), updated_at = now()
WHERE operation_id = ANY($2)
""",
worker_id,
operation_ids,
)
return all_rows
@@ -1,942 +0,0 @@
"""PostgreSQL implementation of DataAccessOps.
Uses unnest(), LATERAL, DISTINCT ON, and native array operations for
efficient batch operations.
"""
import json
from datetime import UTC, datetime
from typing import Any
from uuid import UUID
from .base import DatabaseConnection
from .ops import DataAccessOps, TagListingParts
from .result import ResultRow
class PostgreSQLOps(DataAccessOps):
"""PostgreSQL-specific data access operations using unnest and LATERAL."""
@property
def uses_observation_sources_table(self) -> bool:
return False # PG uses native array ops on source_memory_ids
async def bulk_upsert_chunks(
self,
conn: DatabaseConnection,
table: str,
chunk_ids: list[str],
document_ids: list[str],
bank_ids: list[str],
chunk_texts: list[str],
chunk_indices: list[int],
content_hashes: list[str],
) -> None:
await conn.execute(
f"""
INSERT INTO {table} (chunk_id, document_id, bank_id, chunk_text, chunk_index, content_hash)
SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[], $5::integer[], $6::text[])
ON CONFLICT (chunk_id) DO UPDATE SET
chunk_text = EXCLUDED.chunk_text,
chunk_index = EXCLUDED.chunk_index,
content_hash = EXCLUDED.content_hash
""",
chunk_ids,
document_ids,
bank_ids,
chunk_texts,
chunk_indices,
content_hashes,
)
async def insert_facts_batch(
self,
conn: DatabaseConnection,
bank_id: str,
fact_texts: list[str],
embeddings: list[str],
event_dates: list,
occurred_starts: list,
occurred_ends: list,
mentioned_ats: list,
contexts: list[str],
fact_types: list[str],
metadata_jsons: list[str],
chunk_ids: list,
document_ids: list,
tags_list: list[str],
observation_scopes_list: list,
text_signals_list: list,
text_search_extension: str = "native",
) -> list[str]:
from ...config import get_config
config = get_config()
table = self._get_mu_table()
if config.text_search_extension == "vchord":
query = f"""
WITH input_data AS (
SELECT * FROM unnest(
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id, tags_json,
observation_scopes_json, text_signals)
)
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id, tags,
observation_scopes, text_signals, search_vector)
SELECT
$1,
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id,
COALESCE(
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
'{{}}'::varchar[]
),
observation_scopes_json,
text_signals,
tokenize(
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, ''),
'llmlingua2'
)::bm25_catalog.bm25vector
FROM input_data
RETURNING id
"""
else:
query = f"""
WITH input_data AS (
SELECT * FROM unnest(
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id, tags_json,
observation_scopes_json, text_signals)
)
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id, tags,
observation_scopes, text_signals)
SELECT
$1,
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id,
COALESCE(
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
'{{}}'::varchar[]
),
observation_scopes_json,
text_signals
FROM input_data
RETURNING id
"""
results = await conn.fetch(
query,
bank_id,
fact_texts,
embeddings,
event_dates,
occurred_starts,
occurred_ends,
mentioned_ats,
contexts,
fact_types,
metadata_jsons,
chunk_ids,
document_ids,
tags_list,
observation_scopes_list,
text_signals_list,
)
return [str(row["id"]) for row in results]
async def bulk_insert_links(
self,
conn: DatabaseConnection,
table: str,
sorted_links: list[tuple],
bank_id: str,
nil_entity_uuid: str,
exists_clause: str,
chunk_size: int = 5000,
) -> None:
from_ids = [lnk[0] for lnk in sorted_links]
to_ids = [lnk[1] for lnk in sorted_links]
types = [lnk[2] for lnk in sorted_links]
weights = [lnk[3] for lnk in sorted_links]
entity_ids = [lnk[4] for lnk in sorted_links]
for chunk_start in range(0, len(sorted_links), chunk_size):
chunk_end = min(chunk_start + chunk_size, len(sorted_links))
await conn.execute(
f"""
INSERT INTO {table}
(from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id)
SELECT f, t, tp, w, e, $6
FROM unnest($1::uuid[], $2::uuid[], $3::text[], $4::float8[], $5::uuid[])
AS t(f, t, tp, w, e)
{exists_clause}
ON CONFLICT (from_unit_id, to_unit_id, link_type,
COALESCE(entity_id, '{nil_entity_uuid}'::uuid))
DO NOTHING
""",
from_ids[chunk_start:chunk_end],
to_ids[chunk_start:chunk_end],
types[chunk_start:chunk_end],
weights[chunk_start:chunk_end],
entity_ids[chunk_start:chunk_end],
bank_id,
timeout=300,
)
async def bulk_insert_entities(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
entity_names: list[str],
entity_dates: list,
) -> dict[str, str]:
inserted_rows = await conn.fetch(
f"""
INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count)
SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), 0
FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date)
ON CONFLICT (bank_id, LOWER(canonical_name))
DO NOTHING
RETURNING id, LOWER(canonical_name) AS name_lower
""",
bank_id,
entity_names,
entity_dates,
)
return {row["name_lower"]: row["id"] for row in inserted_rows}
async def fetch_missing_entity_ids(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
missing_names: list[str],
) -> list[ResultRow]:
return await conn.fetch(
f"""
SELECT e.id, LOWER(e.canonical_name) AS name_lower, inputs.input_name
FROM {table} e
JOIN (
SELECT LOWER(n) AS input_name_lower, n AS input_name
FROM unnest($2::text[]) AS n
) AS inputs ON LOWER(e.canonical_name) = inputs.input_name_lower
WHERE e.bank_id = $1
""",
bank_id,
missing_names,
)
async def bulk_insert_unit_entities(
self,
conn: DatabaseConnection,
table: str,
unit_ids: list,
entity_ids: list,
) -> None:
await conn.execute(
f"""
INSERT INTO {table} (unit_id, entity_id)
SELECT u, e FROM unnest($1::uuid[], $2::uuid[]) AS t(u, e)
ON CONFLICT DO NOTHING
""",
unit_ids,
entity_ids,
)
async def fetch_entity_unit_fanout(
self,
conn: DatabaseConnection,
ue_table: str,
entity_id_list: list[UUID],
limit_per_entity: int,
) -> list[ResultRow]:
return await conn.fetch(
f"""
SELECT e.entity_id, n.unit_id
FROM unnest($1::uuid[]) AS e(entity_id)
CROSS JOIN LATERAL (
SELECT ue.unit_id
FROM {ue_table} ue
WHERE ue.entity_id = e.entity_id
ORDER BY ue.unit_id DESC
LIMIT $2
) n
""",
entity_id_list,
limit_per_entity,
)
async def fetch_unit_dates(
self,
conn: DatabaseConnection,
mu_table: str,
unit_ids: list[str],
) -> list[ResultRow]:
return await conn.fetch(
f"""
SELECT id, event_date, fact_type
FROM {mu_table}
WHERE id::text = ANY($1)
""",
unit_ids,
)
async def fetch_temporal_neighbors(
self,
conn: DatabaseConnection,
mu_table: str,
bank_id: str,
lateral_unit_ids: list,
lateral_event_dates: list,
lateral_fact_types: list,
half_limit: int,
batch_size: int = 500,
) -> list[ResultRow]:
rows: list[ResultRow] = []
for start in range(0, len(lateral_unit_ids), batch_size):
end = min(start + batch_size, len(lateral_unit_ids))
# Exact v0.5.6 query shape: src.unit_id::text AS from_id,
# combined.*, ABS(EXTRACT(...)), ROW_NUMBER PARTITION BY src.unit_id.
batch_rows = await conn.fetch(
f"""
SELECT from_id, id, event_date, time_diff_hours FROM (
SELECT src.unit_id::text AS from_id, combined.*,
ROW_NUMBER() OVER (
PARTITION BY src.unit_id
ORDER BY combined.time_diff_hours
) AS rn
FROM unnest($1::uuid[], $2::timestamptz[], $3::text[])
AS src(unit_id, event_date, fact_type)
CROSS JOIN LATERAL (
(SELECT mu.id, mu.event_date,
ABS(EXTRACT(EPOCH FROM mu.event_date - src.event_date)) / 3600.0 AS time_diff_hours
FROM {mu_table} mu
WHERE mu.bank_id = $4
AND mu.fact_type = src.fact_type
AND mu.event_date <= src.event_date
AND mu.id != src.unit_id
ORDER BY mu.event_date DESC
LIMIT $5)
UNION ALL
(SELECT mu.id, mu.event_date,
ABS(EXTRACT(EPOCH FROM mu.event_date - src.event_date)) / 3600.0 AS time_diff_hours
FROM {mu_table} mu
WHERE mu.bank_id = $4
AND mu.fact_type = src.fact_type
AND mu.event_date > src.event_date
AND mu.id != src.unit_id
ORDER BY mu.event_date ASC
LIMIT $5)
) combined
) ranked
WHERE rn <= $5
""",
lateral_unit_ids[start:end],
lateral_event_dates[start:end],
lateral_fact_types[start:end],
bank_id,
half_limit,
)
rows.extend(batch_rows)
return rows
def build_entity_expansion_cte(
self,
mu_table: str,
ue_table: str,
per_entity_limit: int,
) -> str:
return f"""
seed_entities AS (
SELECT DISTINCT ue.entity_id
FROM {ue_table} ue
WHERE ue.unit_id = ANY($1::uuid[])
),
entity_expanded AS (
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
COUNT(DISTINCT se.entity_id)::float AS score,
'entity'::text AS source
FROM seed_entities se
CROSS JOIN LATERAL (
SELECT ue_target.unit_id
FROM {ue_table} ue_target
WHERE ue_target.entity_id = se.entity_id
AND ue_target.unit_id != ALL($1::uuid[])
ORDER BY ue_target.unit_id DESC
LIMIT {per_entity_limit}
) t
JOIN {mu_table} mu ON mu.id = t.unit_id
WHERE mu.fact_type = $2
GROUP BY mu.id
ORDER BY score DESC
LIMIT $3
)"""
def build_semantic_causal_cte(
self,
ml_table: str,
mu_table: str,
) -> str:
# Exact v0.5.6 query shape: GROUP BY + MAX(weight) for semantic,
# DISTINCT ON for causal.
return f"""
semantic_expanded AS (
SELECT
id, text, context, event_date, occurred_start,
occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags, proof_count,
MAX(weight) AS score,
'semantic'::text AS source
FROM (
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
ml.weight
FROM {ml_table} ml
JOIN {mu_table} mu ON mu.id = ml.to_unit_id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic'
AND mu.fact_type = $2
AND mu.id != ALL($1::uuid[])
UNION ALL
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
ml.weight
FROM {ml_table} ml
JOIN {mu_table} mu ON mu.id = ml.from_unit_id
WHERE ml.to_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic'
AND mu.fact_type = $2
AND mu.id != ALL($1::uuid[])
) sem_raw
GROUP BY id, text, context, event_date, occurred_start,
occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags, proof_count
ORDER BY score DESC
LIMIT $3
),
causal_expanded AS (
SELECT DISTINCT ON (mu.id)
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
ml.weight AS score,
'causal'::text AS source
FROM {ml_table} ml
JOIN {mu_table} mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
AND mu.fact_type = $2
ORDER BY mu.id, ml.weight DESC
LIMIT $3
)"""
async def expand_observations(
self,
conn: DatabaseConnection,
mu_table: str,
ue_table: str,
ml_table: str,
seed_ids: list,
budget: int,
per_entity_limit: int,
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
# v0.5.6 array ops: unnest, &&, COUNT(DISTINCT) on source_memory_ids.
from ..schema import fq_table
entity_rows = await conn.fetch(
f"""
WITH seed_sources AS (
SELECT DISTINCT unnest(source_memory_ids) AS source_id
FROM {mu_table}
WHERE id = ANY($1::uuid[])
AND source_memory_ids IS NOT NULL
),
source_entities AS (
SELECT DISTINCT ue_seed.entity_id
FROM seed_sources ss
JOIN {ue_table} ue_seed ON ue_seed.unit_id = ss.source_id
),
connected_sources AS (
SELECT DISTINCT t.unit_id AS source_id
FROM source_entities se
CROSS JOIN LATERAL (
SELECT ue_target.unit_id
FROM {ue_table} ue_target
WHERE ue_target.entity_id = se.entity_id
ORDER BY ue_target.unit_id DESC
LIMIT {per_entity_limit}
) t
WHERE NOT EXISTS (
SELECT 1 FROM seed_sources ss WHERE ss.source_id = t.unit_id
)
),
connected_array AS (
SELECT array_agg(source_id) AS source_ids FROM connected_sources
)
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
(SELECT COUNT(DISTINCT s) FROM unnest(mu.source_memory_ids) s WHERE s = ANY(ca.source_ids))::float AS score
FROM {mu_table} mu, connected_array ca
WHERE mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
AND ca.source_ids IS NOT NULL
AND mu.source_memory_ids && ca.source_ids
ORDER BY score DESC
LIMIT $2
""",
seed_ids,
budget,
)
# Exact v0.5.6 query shape: GROUP BY + MAX(weight) for semantic,
# DISTINCT ON for causal, hardcoded to fact_type='observation'.
sem_causal_rows = await conn.fetch(
f"""
WITH semantic_expanded AS (
SELECT
id, text, context, event_date, occurred_start,
occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags, proof_count,
MAX(weight) AS score,
'semantic'::text AS source
FROM (
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, mu.proof_count, ml.weight
FROM {ml_table} ml JOIN {mu_table} mu ON mu.id = ml.to_unit_id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
UNION ALL
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, mu.proof_count, ml.weight
FROM {ml_table} ml JOIN {mu_table} mu ON mu.id = ml.from_unit_id
WHERE ml.to_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
) sem_raw
GROUP BY id, text, context, event_date, occurred_start, occurred_end,
mentioned_at, fact_type, document_id, chunk_id, tags, proof_count
ORDER BY score DESC LIMIT $2
),
causal_expanded AS (
SELECT DISTINCT ON (mu.id)
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, mu.proof_count, ml.weight AS score, 'causal'::text AS source
FROM {ml_table} ml JOIN {mu_table} mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
AND mu.fact_type = 'observation'
ORDER BY mu.id, ml.weight DESC LIMIT $2
)
SELECT * FROM semantic_expanded
UNION ALL
SELECT * FROM causal_expanded
""",
seed_ids,
budget,
)
semantic_rows = [r for r in sem_causal_rows if r["source"] == "semantic"]
causal_rows = [r for r in sem_causal_rows if r["source"] == "causal"]
return list(entity_rows), semantic_rows, causal_rows
def build_tag_listing_parts(self, mu_table: str) -> TagListingParts:
return TagListingParts(
tag_source=f"{mu_table}, unnest(tags) AS tag",
non_empty_check="AND tags IS NOT NULL AND tags != '{}'",
tag_col="tag",
bank_prefix="",
)
async def create_bank_vector_indexes(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
internal_id: str,
index_clause: str,
fact_types: dict[str, str],
) -> None:
escaped = bank_id.replace("'", "''")
for ft, suffix in fact_types.items():
uid = str(internal_id).replace("-", "")[:16]
idx = f"idx_mu_emb_{suffix}_{uid}"
await conn.execute(
f"CREATE INDEX IF NOT EXISTS {idx} "
f"ON {table} {index_clause} "
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped}'"
)
async def drop_bank_vector_indexes(
self,
conn: DatabaseConnection,
schema: str,
internal_id: str,
fact_types: dict[str, str],
) -> None:
for ft, suffix in fact_types.items():
uid = str(internal_id).replace("-", "")[:16]
idx = f"idx_mu_emb_{suffix}_{uid}"
await conn.execute(f"DROP INDEX IF EXISTS {schema}.{idx}")
def get_entity_resolution_strategy(self) -> str:
return "trigram"
# -- Webhook operations ------------------------------------------------
async def create_webhook(
self,
conn,
table,
webhook_id,
bank_id,
url,
secret,
event_types,
enabled,
http_config_json,
):
return await conn.fetchrow(
f"""
INSERT INTO {table}
(id, bank_id, url, secret, event_types, enabled, http_config, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, NOW(), NOW())
RETURNING id, bank_id, url, secret, event_types, enabled,
http_config::text, created_at::text, updated_at::text
""",
webhook_id,
bank_id,
url,
secret,
event_types,
enabled,
http_config_json,
)
async def list_webhooks_for_bank(self, conn, table, bank_id):
return await conn.fetch(
f"""
SELECT id, bank_id, url, secret, event_types, enabled,
http_config::text, created_at::text, updated_at::text
FROM {table}
WHERE bank_id = $1
ORDER BY created_at
""",
bank_id,
)
async def get_webhooks_for_dispatch(self, conn, webhook_table, bank_id):
return await conn.fetch(
f"""
SELECT id, bank_id, url, secret, event_types, enabled, http_config::text
FROM {webhook_table}
WHERE (bank_id = $1 OR bank_id IS NULL) AND enabled = true
""",
bank_id,
)
async def update_webhook(self, conn, table, webhook_id, bank_id, set_clauses, params):
set_clauses_with_ts = set_clauses + ["updated_at = NOW()"]
return await conn.fetchrow(
f"""
UPDATE {table}
SET {", ".join(set_clauses_with_ts)}
WHERE id = $1 AND bank_id = $2
RETURNING id, bank_id, url, secret, event_types, enabled,
http_config::text, created_at::text, updated_at::text
""",
*params,
)
async def delete_webhook(self, conn, table, webhook_id, bank_id):
result = await conn.execute(
f"DELETE FROM {table} WHERE id = $1 AND bank_id = $2",
webhook_id,
bank_id,
)
return int(result.split()[-1]) > 0 if result else False
async def list_webhook_deliveries(self, conn, ops_table, webhook_id, bank_id, limit, cursor):
fetch_limit = limit + 1
if cursor:
return await conn.fetch(
f"""
SELECT operation_id, status, retry_count, next_retry_at::text,
error_message, task_payload, result_metadata::text, created_at::text, updated_at::text
FROM {ops_table}
WHERE operation_type = 'webhook_delivery'
AND bank_id = $1
AND task_payload->>'webhook_id' = $2
AND created_at < $3::timestamptz
ORDER BY created_at DESC
LIMIT $4
""",
bank_id,
webhook_id,
cursor,
fetch_limit,
)
return await conn.fetch(
f"""
SELECT operation_id, status, retry_count, next_retry_at::text,
error_message, task_payload, result_metadata::text, created_at::text, updated_at::text
FROM {ops_table}
WHERE operation_type = 'webhook_delivery'
AND bank_id = $1
AND task_payload->>'webhook_id' = $2
ORDER BY created_at DESC
LIMIT $3
""",
bank_id,
webhook_id,
fetch_limit,
)
async def insert_webhook_delivery_task(self, conn, ops_table, operation_id, bank_id, payload_json, timestamp):
await conn.execute(
f"""
INSERT INTO {ops_table}
(operation_id, bank_id, operation_type, status, task_payload, result_metadata, created_at, updated_at)
VALUES ($1, $2, 'webhook_delivery', 'pending', $3::jsonb, '{{}}'::jsonb, $4, $4)
""",
operation_id,
bank_id,
payload_json,
timestamp,
)
# -- Task claiming operations ------------------------------------------
async def claim_tasks(self, conn, table, worker_id, reserved_limits, shared_limit):
all_rows = []
claimed_ids = []
# --- Phase 1: claim from reserved pools ---
for op_type, limit in reserved_limits.items():
if limit <= 0:
continue
if op_type == "consolidation":
busy_banks = await conn.fetch(
f"""
SELECT DISTINCT bank_id FROM {table}
WHERE operation_type = 'consolidation' AND status = 'processing'
""",
)
busy_bank_ids = [r["bank_id"] for r in busy_banks]
if busy_bank_ids:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
AND bank_id != ALL($1::text[])
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
""",
busy_bank_ids,
limit,
)
else:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
""",
limit,
)
else:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = $1
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
""",
op_type,
limit,
)
for row in rows:
claimed_ids.append(row["operation_id"])
all_rows.append(row)
# --- Phase 2: claim from shared pool ---
remaining_shared = shared_limit
if remaining_shared > 0:
# 2a. Non-consolidation tasks
if claimed_ids:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type != 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
AND operation_id != ALL($1::uuid[])
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
""",
claimed_ids,
remaining_shared,
)
else:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type != 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
""",
remaining_shared,
)
for row in rows:
claimed_ids.append(row["operation_id"])
all_rows.append(row)
remaining_shared -= len(rows)
# 2b. Consolidation tasks (with bank-serialization)
if remaining_shared > 0:
busy_banks_2 = await conn.fetch(
f"""
SELECT DISTINCT bank_id FROM {table}
WHERE operation_type = 'consolidation' AND status = 'processing'
""",
)
busy_bank_ids_2 = [r["bank_id"] for r in busy_banks_2]
if claimed_ids:
if busy_bank_ids_2:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
AND operation_id != ALL($1::uuid[])
AND bank_id != ALL($2::text[])
ORDER BY created_at
LIMIT $3
FOR UPDATE SKIP LOCKED
""",
claimed_ids,
busy_bank_ids_2,
remaining_shared,
)
else:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
AND operation_id != ALL($1::uuid[])
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
""",
claimed_ids,
remaining_shared,
)
else:
if busy_bank_ids_2:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
AND bank_id != ALL($1::text[])
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
""",
busy_bank_ids_2,
remaining_shared,
)
else:
rows = await conn.fetch(
f"""
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
""",
remaining_shared,
)
for row in rows:
claimed_ids.append(row["operation_id"])
all_rows.append(row)
if not all_rows:
return []
# Mark all claimed rows as processing
operation_ids = [row["operation_id"] for row in all_rows]
await conn.execute(
f"""
UPDATE {table}
SET status = 'processing', worker_id = $1, claimed_at = now(), updated_at = now()
WHERE operation_id = ANY($2)
""",
worker_id,
operation_ids,
)
return all_rows
File diff suppressed because it is too large Load Diff
@@ -1,130 +0,0 @@
"""PostgreSQL backend implementation using asyncpg.
Wraps asyncpg's pool and connection objects behind the DatabaseBackend
and DatabaseConnection interfaces. Returns raw asyncpg.Record objects
from fetch/fetchrow — they satisfy the ResultRow protocol natively in C,
avoiding Python-level wrapping overhead (~570K __getitem__ calls per
20-query benchmark → measurable regression when wrapped).
"""
import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Any
import asyncpg # noqa: F401
from .base import DatabaseBackend, DatabaseConnection
logger = logging.getLogger(__name__)
class PostgresConnection(DatabaseConnection):
"""DatabaseConnection wrapper around an asyncpg.Connection."""
__slots__ = ("_conn",)
def __init__(self, conn: asyncpg.Connection) -> None:
self._conn = conn
@asynccontextmanager
async def transaction(self) -> AsyncIterator["PostgresConnection"]:
async with self._conn.transaction():
yield self
async def execute(self, query: str, *args: Any, timeout: float | None = None) -> str:
return await self._conn.execute(query, *args, timeout=timeout)
async def executemany(self, query: str, args: list[tuple[Any, ...]], *, timeout: float | None = None) -> None:
await self._conn.executemany(query, args, timeout=timeout)
async def fetch(self, query: str, *args: Any, timeout: float | None = None) -> list:
# Return raw asyncpg.Record objects — they satisfy the ResultRow
# protocol natively (key access, .keys(), .get(), etc.) with zero
# Python wrapping overhead.
return await self._conn.fetch(query, *args, timeout=timeout)
async def fetchrow(self, query: str, *args: Any, timeout: float | None = None):
# Return raw asyncpg.Record — no wrapping needed.
return await self._conn.fetchrow(query, *args, timeout=timeout)
async def fetchval(self, query: str, *args: Any, column: int = 0, timeout: float | None = None) -> Any:
return await self._conn.fetchval(query, *args, column=column, timeout=timeout)
async def copy_records_to_table(
self,
table_name: str,
*,
records: list[tuple[Any, ...]],
columns: list[str],
timeout: float | None = None,
) -> None:
"""Use asyncpg's native COPY for fast bulk loading."""
await self._conn.copy_records_to_table(table_name, records=records, columns=columns, timeout=timeout)
class PostgreSQLBackend(DatabaseBackend):
"""DatabaseBackend implementation wrapping an asyncpg connection pool."""
def run_migrations(self, dsn: str, *, schema: str | None = None) -> None:
"""Run Alembic migrations for PostgreSQL."""
from ...config import get_config
from ...migrations import run_migrations
config = get_config()
run_migrations(dsn, schema=schema, migration_database_url=config.migration_database_url)
def __init__(self) -> None:
self._pool: asyncpg.Pool | None = None
async def initialize(
self,
dsn: str,
*,
min_size: int = 5,
max_size: int = 20,
command_timeout: float = 300,
acquire_timeout: float = 30,
statement_cache_size: int = 0,
init_callback: Any | None = None,
) -> None:
self._pool = await asyncpg.create_pool(
dsn,
min_size=min_size,
max_size=max_size,
command_timeout=command_timeout,
statement_cache_size=statement_cache_size,
timeout=acquire_timeout,
init=init_callback,
)
logger.info(
f"PostgreSQL pool created (min={min_size}, max={max_size}, "
f"cmd_timeout={command_timeout}s, acquire_timeout={acquire_timeout}s)"
)
async def shutdown(self) -> None:
if self._pool is not None:
await self._pool.close()
self._pool = None
logger.info("PostgreSQL pool closed")
@asynccontextmanager
async def acquire(self) -> AsyncIterator[PostgresConnection]:
pool = self._ensure_pool()
async with pool.acquire() as conn:
yield PostgresConnection(conn)
@asynccontextmanager
async def transaction(self) -> AsyncIterator[PostgresConnection]:
pool = self._ensure_pool()
async with pool.acquire() as conn:
async with conn.transaction():
yield PostgresConnection(conn)
def get_pool(self) -> asyncpg.Pool:
return self._ensure_pool()
def _ensure_pool(self) -> asyncpg.Pool:
if self._pool is None:
raise RuntimeError("PostgreSQLBackend is not initialized. Call initialize() first.")
return self._pool
@@ -1,104 +0,0 @@
"""Uniform row interface over heterogeneous database drivers.
ResultRow is a Protocol that describes the dict-like access pattern all
database rows must support. asyncpg.Record already satisfies this protocol
natively (key-based access, .keys(), .values(), etc.) so the PostgreSQL
backend returns raw Records — zero wrapping overhead.
Only backends whose native row type does NOT satisfy the protocol (e.g.,
Oracle named-tuple rows) need the concrete DictResultRow wrapper.
"""
from typing import Any, Protocol, runtime_checkable
@runtime_checkable
class ResultRow(Protocol):
"""Dict-like row interface returned by all database operations."""
def __getitem__(self, key: str | int) -> Any: ...
def get(self, key: str, default: Any = None) -> Any: ...
def keys(self) -> Any: ...
def values(self) -> Any: ...
def items(self) -> Any: ...
def __contains__(self, key: str) -> bool: ...
def __len__(self) -> int: ...
class DictResultRow:
"""Concrete wrapper for backends whose native rows lack dict-like access.
Used by Oracle (named-tuple rows, plain dicts) and tests.
asyncpg.Record satisfies ResultRow natively — do NOT wrap it.
"""
__slots__ = ("_data",)
def __init__(self, data: Any) -> None:
object.__setattr__(self, "_data", data)
def __getitem__(self, key: str | int) -> Any:
data = object.__getattribute__(self, "_data")
return data[key]
def __getattr__(self, key: str) -> Any:
data = object.__getattribute__(self, "_data")
if isinstance(data, dict):
try:
return data[key]
except KeyError:
raise AttributeError(key) from None
try:
return data[key]
except (KeyError, TypeError):
raise AttributeError(key) from None
def get(self, key: str, default: Any = None) -> Any:
try:
return self[key]
except (KeyError, IndexError):
return default
def keys(self) -> list[str]:
data = object.__getattribute__(self, "_data")
if isinstance(data, dict):
return list(data.keys())
if hasattr(data, "keys"):
return list(data.keys())
return []
def values(self) -> list[Any]:
data = object.__getattribute__(self, "_data")
if isinstance(data, dict):
return list(data.values())
if hasattr(data, "values"):
return list(data.values())
return []
def items(self) -> list[tuple[str, Any]]:
data = object.__getattribute__(self, "_data")
if isinstance(data, dict):
return list(data.items())
if hasattr(data, "items"):
return list(data.items())
return list(zip(self.keys(), self.values()))
def __repr__(self) -> str:
data = object.__getattribute__(self, "_data")
return f"DictResultRow({data!r})"
def __contains__(self, key: str) -> bool:
data = object.__getattribute__(self, "_data")
if isinstance(data, dict):
return key in data
if hasattr(data, "keys"):
return key in data.keys()
return False
def __len__(self) -> int:
data = object.__getattribute__(self, "_data")
return len(data)
def __bool__(self) -> bool:
data = object.__getattribute__(self, "_data")
return bool(data)
@@ -1,164 +0,0 @@
"""
Database utility functions for connection management with retry logic.
"""
import asyncio
import logging
import time
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Any
logger = logging.getLogger(__name__)
# Default retry configuration for database operations
DEFAULT_MAX_RETRIES = 3
DEFAULT_BASE_DELAY = 0.5 # seconds
DEFAULT_MAX_DELAY = 5.0 # seconds
# Retryable exception types (checked by class name to avoid hard imports)
_RETRYABLE_EXCEPTION_NAMES = frozenset(
{
"InterfaceError",
"ConnectionDoesNotExistError",
"TooManyConnectionsError",
"DeadlockDetectedError",
}
)
def _is_oracle_deadlock(exc: BaseException) -> bool:
"""Check if an exception is an Oracle ORA-00060 deadlock."""
try:
import oracledb # type: ignore[import-not-found]
except ImportError:
return False
if isinstance(exc, oracledb.DatabaseError) and exc.args:
err = exc.args[0]
return getattr(err, "code", None) == 60 # ORA-00060
return False
def _is_retryable(exc: BaseException) -> bool:
"""Check if an exception is retryable (transient connection issue)."""
if isinstance(exc, (OSError, ConnectionError, asyncio.TimeoutError)):
return True
if type(exc).__name__ in _RETRYABLE_EXCEPTION_NAMES:
return True
return _is_oracle_deadlock(exc)
async def retry_with_backoff(
func,
max_retries: int = DEFAULT_MAX_RETRIES,
base_delay: float = DEFAULT_BASE_DELAY,
max_delay: float = DEFAULT_MAX_DELAY,
):
"""
Execute an async function with exponential backoff retry.
Args:
func: Async function to execute
max_retries: Maximum number of retry attempts
base_delay: Initial delay between retries (seconds)
max_delay: Maximum delay between retries (seconds)
Returns:
Result of the function
Raises:
The last exception if all retries fail
"""
last_exception = None
for attempt in range(max_retries + 1):
try:
return await func()
except Exception as e:
if not _is_retryable(e):
raise
last_exception = e
if attempt < max_retries:
delay = min(base_delay * (2**attempt), max_delay)
if type(e).__name__ == "DeadlockDetectedError" or _is_oracle_deadlock(e):
logger.warning(
"Deadlock detected during parallel document processing — "
"this is expected and will resolve automatically "
f"(attempt {attempt + 1}/{max_retries + 1}, retrying in {delay:.1f}s)"
)
else:
logger.warning(
f"Database operation failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
f"Retrying in {delay:.1f}s..."
)
await asyncio.sleep(delay)
else:
logger.error(f"Database operation failed after {max_retries + 1} attempts: {e}")
raise last_exception
@asynccontextmanager
async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MAX_RETRIES) -> AsyncIterator[Any]:
"""
Async context manager to acquire a database connection with retry logic.
Accepts either a DatabaseBackend or a raw asyncpg.Pool for backward compatibility.
Usage:
async with acquire_with_retry(backend) as conn:
await conn.execute(...)
Args:
backend_or_pool: A DatabaseBackend instance or asyncpg.Pool
max_retries: Maximum number of retry attempts
Yields:
A DatabaseConnection (if backend) or asyncpg.Connection (if pool)
"""
from .db.base import DatabaseBackend
if isinstance(backend_or_pool, DatabaseBackend) or getattr(backend_or_pool, "_wraps_backend", False):
# Use the backend's acquire context manager with retry
start = time.time()
last_exception = None
for attempt in range(max_retries + 1):
try:
async with backend_or_pool.acquire() as conn:
acquire_time = time.time() - start
if acquire_time > 0.05:
logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s")
yield conn
return
except Exception as e:
if not _is_retryable(e):
raise
last_exception = e
if attempt < max_retries:
delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
logger.warning(
f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
f"Retrying in {delay:.1f}s..."
)
await asyncio.sleep(delay)
else:
logger.error(f"Database acquire failed after {max_retries + 1} attempts: {e}")
raise last_exception
else:
# Legacy path: raw asyncpg.Pool
pool = backend_or_pool
start = time.time()
async def acquire():
return await pool.acquire()
conn = await retry_with_backoff(acquire, max_retries=max_retries)
acquire_time = time.time() - start
if acquire_time > 0.05:
pool_size = pool.get_size()
pool_free = pool.get_idle_size()
logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s | size={pool_size}, idle={pool_free}")
try:
yield conn
finally:
await pool.release(conn)
@@ -1,144 +0,0 @@
"""
MLX implementation of jina-reranker-v3 for Apple Silicon.
This file is adapted from the official model repository:
https://huggingface.co/jinaai/jina-reranker-v3-mlx/blob/main/rerank.py
License: CC BY-NC 4.0 (contact Jina AI for commercial usage)
Changes from upstream:
- Removed the __main__ example block
- Type annotations added to public methods
- top_n parameter added to rerank() (upstream only exposed it implicitly)
"""
import numpy as np
class _MLPProjector:
def __init__(self):
import mlx.nn as nn
self.linear1 = nn.Linear(1024, 512, bias=False)
self.linear2 = nn.Linear(512, 512, bias=False)
def __call__(self, x):
import mlx.nn as nn
x = self.linear1(x)
x = nn.relu(x)
x = self.linear2(x)
return x
def _load_projector(projector_path: str) -> _MLPProjector:
import mlx.core as mx
from safetensors import safe_open
projector = _MLPProjector()
with safe_open(projector_path, framework="numpy") as f:
projector.linear1.weight = mx.array(f.get_tensor("linear1.weight"))
projector.linear2.weight = mx.array(f.get_tensor("linear2.weight"))
return projector
def _sanitize(text: str, special_tokens: dict[str, str]) -> str:
for token in special_tokens.values():
text = text.replace(token, "")
return text
def _format_prompt(query: str, docs: list[str], special_tokens: dict[str, str]) -> str:
query = _sanitize(query, special_tokens)
docs = [_sanitize(d, special_tokens) for d in docs]
doc_token = special_tokens["doc_embed_token"]
query_token = special_tokens["query_embed_token"]
prefix = (
"<|im_start|>system\n"
"You are a search relevance expert who can determine a ranking of the passages based on how relevant they are to the query. "
"If the query is a question, how relevant a passage is depends on how well it answers the question. "
"If not, try to analyze the intent of the query and assess how well each passage satisfies the intent. "
"If an instruction is provided, you should follow the instruction when determining the ranking."
"<|im_end|>\n<|im_start|>user\n"
)
suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
body = (
f"I will provide you with {len(docs)} passages, each indicated by a numerical identifier. "
f"Rank the passages based on their relevance to query: {query}\n"
)
body += "\n".join(f'<passage id="{i}">\n{doc}{doc_token}\n</passage>' for i, doc in enumerate(docs))
body += f"\n<query>\n{query}{query_token}\n</query>"
return prefix + body + suffix
class MLXReranker:
"""
MLX-accelerated jina-reranker-v3 for Apple Silicon.
Loads the model from a local directory (use huggingface_hub.snapshot_download
to fetch jinaai/jina-reranker-v3-mlx if you don't have it already).
"""
_SPECIAL_TOKENS = {
"query_embed_token": "<|rerank_token|>",
"doc_embed_token": "<|embed_token|>",
}
_DOC_TOKEN_ID = 151670
_QUERY_TOKEN_ID = 151671
def __init__(self, model_path: str, projector_path: str):
from mlx_lm import load
self.model, self.tokenizer = load(model_path)
self.model.eval()
self.projector = _load_projector(projector_path)
def rerank(self, query: str, documents: list[str], top_n: int | None = None) -> list[dict]:
"""
Rank documents by relevance to a query.
Returns a list of dicts with keys: document, relevance_score, index.
Sorted by descending relevance_score.
"""
import mlx.core as mx
prompt = _format_prompt(query, documents, self._SPECIAL_TOKENS)
input_ids = self.tokenizer.encode(prompt)
hidden_states = self.model.model([input_ids])[0] # [seq_len, hidden_size]
input_ids_np = np.array(input_ids)
query_positions = np.where(input_ids_np == self._QUERY_TOKEN_ID)[0]
doc_positions = np.where(input_ids_np == self._DOC_TOKEN_ID)[0]
if len(query_positions) == 0:
raise ValueError("Query embed token not found in prompt")
if len(doc_positions) == 0:
raise ValueError("Document embed tokens not found in prompt")
query_hidden = mx.expand_dims(hidden_states[int(query_positions[0])], axis=0)
doc_hidden = mx.stack([hidden_states[int(p)] for p in doc_positions])
query_emb = self.projector(query_hidden) # [1, 512]
doc_emb = self.projector(doc_hidden) # [num_docs, 512]
query_exp = mx.broadcast_to(mx.expand_dims(query_emb, 0), (1, len(documents), 512))
doc_exp = mx.expand_dims(doc_emb, 0)
scores = mx.sum(doc_exp * query_exp, axis=-1) / (
mx.sqrt(mx.sum(doc_exp * doc_exp, axis=-1)) * mx.sqrt(mx.sum(query_exp * query_exp, axis=-1))
) # [1, num_docs]
scores_np = np.array(scores[0])
order = np.argsort(scores_np)[::-1]
n = min(top_n, len(documents)) if top_n is not None else len(documents)
return [
{
"document": documents[order[i]],
"relevance_score": float(scores_np[order[i]]),
"index": int(order[i]),
}
for i in range(n)
]
@@ -1,130 +0,0 @@
"""File parser implementations."""
import logging
from dataclasses import dataclass
from .base import FileParser, UnsupportedFileTypeError
from .iris import IrisParser
from .llama_parse import LlamaParseParser
from .markitdown import MarkitdownParser
__all__ = [
"FileParser",
"UnsupportedFileTypeError",
"IrisParser",
"LlamaParseParser",
"MarkitdownParser",
"FileParserRegistry",
"ConvertResult",
]
@dataclass
class ConvertResult:
"""Result of a successful file conversion."""
content: str
parser_name: str
logger = logging.getLogger(__name__)
class FileParserRegistry:
"""Registry for file parsers with auto-detection."""
def __init__(self):
"""Initialize empty parser registry."""
self._parsers: dict[str, FileParser] = {}
def register(self, parser: FileParser):
"""
Register a parser.
Args:
parser: FileParser instance
"""
self._parsers[parser.name()] = parser
def get_parser(
self,
name: str | None,
filename: str,
content_type: str | None = None,
) -> FileParser:
"""
Get parser by name or auto-detect.
Args:
name: Parser name (e.g., "markitdown") or None for auto-detect
filename: File name for auto-detection
content_type: MIME type (optional)
Returns:
FileParser instance
Raises:
ValueError: If no suitable parser found
"""
if name:
# Explicit parser requested — return it directly, let the parser
# raise UnsupportedFileTypeError from convert() if needed
if name not in self._parsers:
raise ValueError(f"Parser '{name}' not found. Available: {list(self._parsers.keys())}")
return self._parsers[name]
# Auto-detect parser
for parser in self._parsers.values():
if parser.supports(filename, content_type):
return parser
raise ValueError(f"No parser found for {filename}. Available parsers: {list(self._parsers.keys())}")
async def convert_with_fallback(
self,
parsers: list[str],
file_data: bytes,
filename: str,
content_type: str | None = None,
) -> ConvertResult:
"""
Try each parser in order, falling back on failure or empty content.
Moves to the next parser if the current one raises UnsupportedFileTypeError
or returns empty content. Any other exception (RuntimeError, network error,
etc.) also triggers a fallback so the chain is exhausted before failing.
Args:
parsers: Ordered list of parser names to try
file_data: Raw file bytes
filename: Original filename
content_type: MIME type (optional)
Returns:
ConvertResult with the parsed content and the name of the parser that succeeded
Raises:
ValueError: If a parser name is not registered
RuntimeError: If all parsers fail or return empty content
"""
last_error: Exception | None = None
for name in parsers:
parser = self.get_parser(name, filename, content_type)
try:
content = await parser.convert(file_data, filename)
if content and content.strip():
return ConvertResult(content=content, parser_name=name)
logger.warning(f"Parser '{name}' returned empty content for '{filename}', trying next")
last_error = RuntimeError(f"Parser '{name}' returned no content for '{filename}'")
except UnsupportedFileTypeError as e:
logger.warning(f"Parser '{name}' does not support '{filename}', trying next: {e}")
last_error = e
except Exception as e:
logger.warning(f"Parser '{name}' failed for '{filename}', trying next: {e}")
last_error = e
raise last_error or RuntimeError(f"No parsers available for '{filename}'")
def list_parsers(self) -> list[str]:
"""Get list of registered parser names."""
return list(self._parsers.keys())
@@ -1,125 +0,0 @@
"""LlamaParse parser implementation using the LlamaIndex Cloud parsing API."""
import asyncio
import logging
import mimetypes
import time
import httpx
from .base import FileParser, UnsupportedFileTypeError
logger = logging.getLogger(__name__)
_LLAMA_PARSE_BASE_URL = "https://api.cloud.llamaindex.ai/api/parsing"
_DEFAULT_POLL_INTERVAL = 2.0 # seconds
_DEFAULT_TIMEOUT = 300.0 # seconds
# HTTP status codes that indicate the file type is not supported.
# Other 4xx codes (401, 403, 429, etc.) are operational errors, not file-type issues.
_UNSUPPORTED_FILE_STATUS_CODES = {400, 415, 422}
class LlamaParseParser(FileParser):
"""
LlamaParse file parser using LlamaIndex's hosted parsing service.
Uploads files to the LlamaParse API, polls until the parse job completes,
and returns the resulting markdown. The API determines which file types
are supported — UnsupportedFileTypeError is raised if the file is rejected.
"""
def __init__(
self,
api_key: str,
poll_interval: float = _DEFAULT_POLL_INTERVAL,
timeout: float = _DEFAULT_TIMEOUT,
):
"""
Initialize llama_parse parser.
Args:
api_key: LlamaCloud API key (typically starts with "llx-")
poll_interval: Seconds between status poll requests (default: 2)
timeout: Maximum seconds to wait for parsing (default: 300)
"""
self._api_key = api_key
self._poll_interval = poll_interval
self._timeout = timeout
self._auth_headers = {"Authorization": f"Bearer {api_key}"}
self._client = httpx.AsyncClient(timeout=httpx.Timeout(30.0, read=120.0))
async def convert(self, file_data: bytes, filename: str) -> str:
"""
Parse file to markdown using the LlamaParse API.
Raises:
UnsupportedFileTypeError: If the LlamaParse API rejects the file type
RuntimeError: If parsing fails for another reason
"""
content_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
# Step 1: Upload file and start parse job
upload_resp = await self._client.post(
f"{_LLAMA_PARSE_BASE_URL}/upload",
headers=self._auth_headers,
# Ensure file_data is plain bytes (storage backends may return obstore.Bytes)
files={"file": (filename, bytes(file_data), content_type)},
)
_raise_for_status(upload_resp, filename, "upload")
job_id: str = upload_resp.json()["id"]
# Step 2: Poll job status until SUCCESS or ERROR
deadline = time.monotonic() + self._timeout
while True:
status_resp = await self._client.get(
f"{_LLAMA_PARSE_BASE_URL}/job/{job_id}",
headers=self._auth_headers,
)
_raise_for_status(status_resp, filename, "poll job status")
status_data = status_resp.json()
status = status_data.get("status")
if status == "SUCCESS":
break
if status in ("ERROR", "CANCELLED"):
error = status_data.get("error_code") or status_data.get("error") or "unknown error"
raise RuntimeError(f"LlamaParse job failed for '{filename}': {error}")
if time.monotonic() >= deadline:
raise RuntimeError(f"LlamaParse job timed out after {self._timeout}s for '{filename}'")
await asyncio.sleep(self._poll_interval)
# Step 3: Fetch the markdown result
result_resp = await self._client.get(
f"{_LLAMA_PARSE_BASE_URL}/job/{job_id}/result/markdown",
headers=self._auth_headers,
)
_raise_for_status(result_resp, filename, "fetch markdown result")
markdown = result_resp.json().get("markdown")
if not markdown:
raise RuntimeError(f"No content extracted from '{filename}'")
return markdown
def name(self) -> str:
"""Get parser name."""
return "llama_parse"
def _raise_for_status(response: httpx.Response, filename: str, step: str) -> None:
"""
Raise an appropriate error for HTTP errors.
Raises UnsupportedFileTypeError for 400/415/422 (file rejected by the API).
Raises RuntimeError for all other errors (auth, rate-limit, server errors).
"""
if not response.is_error:
return
body = response.text or "<empty>"
msg = (
f"LlamaParse API error during {step} for '{filename}': {response.status_code} {response.reason_phrase}{body}"
)
if response.status_code in _UNSUPPORTED_FILE_STATUS_CODES:
raise UnsupportedFileTypeError(msg)
raise RuntimeError(msg)
@@ -1,385 +0,0 @@
"""
LiteLLM LLM provider for universal model support.
This provider enables using 100+ LLM providers via the LiteLLM SDK, including:
- AWS Bedrock (bedrock/anthropic.claude-3-5-sonnet-...)
- Azure OpenAI (azure/gpt-4o)
- Together AI (together_ai/meta-llama/...)
- Any other LiteLLM-supported provider
Uses litellm.acompletion() for async chat completions.
Authentication for cloud providers (e.g., AWS Bedrock via boto3 credential chain)
is handled automatically by LiteLLM.
"""
import asyncio
import json
import logging
import time
from typing import Any
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.metrics import get_metrics_collector
from hindsight_api.worker.stage import set_stage
logger = logging.getLogger(__name__)
class LiteLLMLLM(LLMInterface):
"""
LLM provider using the LiteLLM SDK for universal model support.
Supports any model accessible via litellm.acompletion(), including AWS Bedrock,
Azure OpenAI, Together AI, Fireworks AI, and more.
Model names follow LiteLLM conventions with provider prefixes:
- bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0
- azure/gpt-4o
- together_ai/meta-llama/Llama-3-70b-chat-hf
- fireworks_ai/accounts/fireworks/models/llama-v3p1-70b-instruct
"""
def __init__(
self,
provider: str,
api_key: str,
base_url: str,
model: str,
reasoning_effort: str = "low",
timeout: float = 300.0,
**kwargs: Any,
):
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
self.timeout = timeout
self._litellm: Any = None
try:
import litellm
self._litellm = litellm
# Suppress LiteLLM's verbose logging
litellm.suppress_debug_info = True # type: ignore[assignment]
# Drop unsupported params instead of raising errors (e.g. tool_choice on some Bedrock models)
litellm.drop_params = True # type: ignore[assignment]
logging.getLogger("LiteLLM").setLevel(logging.WARNING)
logger.info(f"LiteLLM SDK initialized for model: {self.model}")
except ImportError as e:
raise RuntimeError("LiteLLM SDK not installed. Run: uv add litellm or pip install litellm") from e
async def verify_connection(self) -> None:
try:
test_messages = [{"role": "user", "content": "test"}]
await self.call(
messages=test_messages,
max_completion_tokens=50,
temperature=0.0,
scope="verification",
max_retries=0,
)
logger.info("LiteLLM connection verified successfully")
except OutputTooLongError:
# Truncation is fine for verification — it means the connection works
logger.info("LiteLLM connection verified successfully (response truncated)")
except Exception as e:
logger.error(f"LiteLLM connection verification failed: {e}")
raise RuntimeError(f"Failed to verify LiteLLM connection: {e}") from e
def _build_common_kwargs(
self,
messages: list[dict[str, Any]],
max_completion_tokens: int | None = None,
temperature: float | None = None,
) -> dict[str, Any]:
"""Build common kwargs for litellm calls."""
kwargs: dict[str, Any] = {
"model": self.model,
"messages": messages,
"timeout": self.timeout,
}
if self.api_key:
kwargs["api_key"] = self.api_key
if self.base_url:
kwargs["api_base"] = self.base_url
if max_completion_tokens is not None:
kwargs["max_completion_tokens"] = max_completion_tokens
if temperature is not None:
kwargs["temperature"] = temperature
return kwargs
async def call(
self,
messages: list[dict[str, str]],
response_format: Any | None = None,
max_completion_tokens: int | None = None,
temperature: float | None = None,
scope: str = "memory",
max_retries: int = 10,
initial_backoff: float = 1.0,
max_backoff: float = 60.0,
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
) -> Any:
start_time = time.time()
call_kwargs = self._build_common_kwargs(messages, max_completion_tokens, temperature)
# Add JSON schema response format if provided
if response_format is not None and hasattr(response_format, "model_json_schema"):
schema = response_format.model_json_schema()
call_kwargs["response_format"] = {
"type": "json_schema",
"json_schema": {
"name": response_format.__name__ if hasattr(response_format, "__name__") else "response",
"schema": schema,
"strict": strict_schema,
},
}
last_exception = None
for attempt in range(max_retries + 1):
if attempt > 0:
set_stage(f"llm.litellm.{scope}.attempt={attempt + 1}/{max_retries + 1}")
try:
response = await self._litellm.acompletion(**call_kwargs)
content = response.choices[0].message.content or ""
finish_reason = response.choices[0].finish_reason
# Check for length-limited output
if finish_reason == "length":
raise OutputTooLongError("LiteLLM response was truncated due to token limit")
if response_format is not None:
# Strip markdown code fences if present
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:
json_data = json.loads(content)
if skip_validation:
result = json_data
else:
result = response_format.model_validate(json_data)
else:
result = content
# Extract usage
input_tokens = getattr(response.usage, "prompt_tokens", 0) or 0
output_tokens = getattr(response.usage, "completion_tokens", 0) or 0
total_tokens = input_tokens + output_tokens
# Record metrics
duration = time.time() - start_time
metrics = get_metrics_collector()
metrics.record_llm_call(
provider=self.provider,
model=self.model,
scope=scope,
duration=duration,
input_tokens=input_tokens,
output_tokens=output_tokens,
success=True,
)
# Record trace span
from hindsight_api.tracing import _serialize_for_span, get_span_recorder
span_recorder = get_span_recorder()
span_recorder.record_llm_call(
provider=self.provider,
model=self.model,
scope=scope,
messages=messages,
response_content=_serialize_for_span(result),
input_tokens=input_tokens,
output_tokens=output_tokens,
duration=duration,
finish_reason=finish_reason,
error=None,
)
if duration > 10.0:
logger.info(
f"slow llm call: scope={scope}, model={self.provider}/{self.model}, "
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
f"time={duration:.3f}s"
)
if return_usage:
token_usage = TokenUsage(
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=total_tokens,
)
return result, token_usage
return result
except OutputTooLongError:
raise
except json.JSONDecodeError as e:
last_exception = e
if attempt < max_retries:
logger.warning("LiteLLM returned invalid JSON, retrying...")
backoff = min(initial_backoff * (2**attempt), max_backoff)
await asyncio.sleep(backoff)
continue
else:
logger.error(f"LiteLLM returned invalid JSON after {max_retries + 1} attempts")
raise
except Exception as e:
error_str = str(e).lower()
# Fast fail on auth errors
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
logger.error(f"LiteLLM auth error, not retrying: {e}")
raise
last_exception = e
if attempt < max_retries:
# Retry on rate limits, connection errors, server errors
is_retryable = any(
keyword in error_str
for keyword in ("rate", "limit", "timeout", "connection", "500", "502", "503", "529")
)
if is_retryable:
backoff = min(initial_backoff * (2**attempt), max_backoff)
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
await asyncio.sleep(backoff + jitter)
continue
logger.error(f"LiteLLM API error after {attempt + 1} attempts: {e}")
raise
if last_exception:
raise last_exception
raise RuntimeError("LiteLLM call failed after all retries")
async def call_with_tools(
self,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]],
max_completion_tokens: int | None = None,
temperature: float | None = None,
scope: str = "tools",
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
start_time = time.time()
call_kwargs = self._build_common_kwargs(messages, max_completion_tokens, temperature)
call_kwargs["tools"] = tools
call_kwargs["tool_choice"] = tool_choice
last_exception = None
for attempt in range(max_retries + 1):
if attempt > 0:
set_stage(f"llm.litellm.tools.attempt={attempt + 1}/{max_retries + 1}")
try:
response = await self._litellm.acompletion(**call_kwargs)
message = response.choices[0].message
content = message.content
finish_reason = response.choices[0].finish_reason
# Extract tool calls
tool_calls: list[LLMToolCall] = []
if message.tool_calls:
for tc in message.tool_calls:
arguments = tc.function.arguments
if isinstance(arguments, str):
arguments = json.loads(arguments)
tool_calls.append(
LLMToolCall(
id=tc.id,
name=tc.function.name,
arguments=arguments,
)
)
# Extract usage
input_tokens = getattr(response.usage, "prompt_tokens", 0) or 0
output_tokens = getattr(response.usage, "completion_tokens", 0) or 0
# Record metrics
duration = time.time() - start_time
metrics = get_metrics_collector()
metrics.record_llm_call(
provider=self.provider,
model=self.model,
scope=scope,
duration=duration,
input_tokens=input_tokens,
output_tokens=output_tokens,
success=True,
)
# Record trace span
from hindsight_api.tracing import get_span_recorder
span_recorder = get_span_recorder()
tool_calls_dict = (
[{"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in tool_calls]
if tool_calls
else None
)
span_recorder.record_llm_call(
provider=self.provider,
model=self.model,
scope=scope,
messages=messages,
response_content=content,
input_tokens=input_tokens,
output_tokens=output_tokens,
duration=duration,
finish_reason=finish_reason,
error=None,
tool_calls=tool_calls_dict,
)
return LLMToolCallResult(
content=content,
tool_calls=tool_calls,
finish_reason=finish_reason or ("tool_calls" if tool_calls else "stop"),
input_tokens=input_tokens,
output_tokens=output_tokens,
)
except Exception as e:
error_str = str(e).lower()
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
raise
last_exception = e
if attempt < max_retries:
is_retryable = any(
keyword in error_str
for keyword in ("rate", "limit", "timeout", "connection", "500", "502", "503", "529")
)
if is_retryable:
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
continue
logger.error(f"LiteLLM tool call error after {attempt + 1} attempts: {e}")
raise
if last_exception:
raise last_exception
raise RuntimeError("LiteLLM tool call failed after all retries")
async def cleanup(self) -> None:
"""Clean up resources."""
pass
@@ -1,428 +0,0 @@
"""
Built-in llama.cpp LLM provider for fully offline operation.
Manages a llama-cpp-python server as a subprocess, downloads GGUF models
from HuggingFace on first use, and delegates inference to the OpenAI-compatible API.
Usage:
HINDSIGHT_API_LLM_PROVIDER=llamacpp
HINDSIGHT_API_LLAMACPP_MODEL_PATH=~/.hindsight/models/gemma-4-E2B-it-Q4_K_M.gguf
HINDSIGHT_API_LLAMACPP_GPU_LAYERS=-1 # -1 = all layers on GPU
HINDSIGHT_API_LLAMACPP_CONTEXT_SIZE=8192
"""
import asyncio
import logging
import os
import signal
import socket
import subprocess
import sys
import time
from pathlib import Path
from typing import Any
from hindsight_api.engine.llm_interface import LLMInterface
from hindsight_api.engine.response_models import LLMToolCallResult
logger = logging.getLogger(__name__)
# Default GGUF model for offline mode
DEFAULT_LLAMACPP_HF_REPO = "bartowski/google_gemma-4-E2B-it-GGUF"
DEFAULT_LLAMACPP_HF_FILENAME = "google_gemma-4-E2B-it-Q4_K_M.gguf"
DEFAULT_LLAMACPP_MODEL_ALIAS = "gemma-4-e2b-it"
MODELS_DIR = Path.home() / ".hindsight" / "models"
# Singleton server instance — shared across all LlamaCppLLM instances
# (retain, reflect, consolidation each create their own LLMProvider,
# but they should all share one llama.cpp server process)
_shared_server: "LlamaCppServer | None" = None
_shared_server_lock = asyncio.Lock()
def _find_free_port() -> int:
"""Find a free TCP port on localhost."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def _download_default_model() -> Path:
"""Download the default GGUF model from HuggingFace if not already cached.
Returns:
Path to the downloaded GGUF file.
"""
try:
from huggingface_hub import hf_hub_download
except ImportError:
raise ImportError(
"huggingface-hub is required for automatic model download. "
"Install with: pip install 'hindsight-api-slim[local-llm]'"
)
MODELS_DIR.mkdir(parents=True, exist_ok=True)
target = MODELS_DIR / DEFAULT_LLAMACPP_HF_FILENAME
if target.exists():
logger.info(f"Using cached model: {target}")
return target
logger.info(
f"Downloading {DEFAULT_LLAMACPP_HF_FILENAME} from {DEFAULT_LLAMACPP_HF_REPO} (~3.5 GB, first run only)..."
)
downloaded = hf_hub_download(
repo_id=DEFAULT_LLAMACPP_HF_REPO,
filename=DEFAULT_LLAMACPP_HF_FILENAME,
local_dir=str(MODELS_DIR),
)
logger.info(f"Model downloaded: {downloaded}")
return Path(downloaded)
def _resolve_model_path(model_path: str | None) -> Path:
"""Resolve the model path, downloading the default if needed.
Args:
model_path: Explicit path to a GGUF file, or None to use the default.
Returns:
Resolved Path to the GGUF file.
"""
if model_path:
p = Path(model_path).expanduser()
if not p.exists():
raise FileNotFoundError(
f"GGUF model not found: {p}\n"
f"Set HINDSIGHT_API_LLAMACPP_MODEL_PATH to a valid .gguf file, "
f"or remove the setting to auto-download the default model."
)
return p
return _download_default_model()
class LlamaCppServer:
"""Manages a llama-cpp-python OpenAI-compatible server as a subprocess."""
def __init__(
self,
model_path: Path,
port: int,
gpu_layers: int = -1,
context_size: int = 8192,
chat_format: str | None = None,
extra_args: str | None = None,
):
self.model_path = model_path
self.port = port
self.gpu_layers = gpu_layers
self.context_size = context_size
self.chat_format = chat_format
self.extra_args = extra_args
self._process: subprocess.Popen | None = None
@property
def base_url(self) -> str:
return f"http://127.0.0.1:{self.port}/v1"
async def start(self) -> None:
"""Start the llama.cpp server subprocess."""
cmd = [
sys.executable,
"-m",
"llama_cpp.server",
"--model",
str(self.model_path),
"--host",
"127.0.0.1",
"--port",
str(self.port),
"--n_gpu_layers",
str(self.gpu_layers),
"--n_ctx",
str(self.context_size),
"--flash_attn",
"true",
"--n_batch",
"2048",
# Prompt cache: reuse KV cache for repeated system prompts
"--cache",
"true",
]
# Only pass chat_format if explicitly set (most GGUF models have it embedded)
if self.chat_format:
cmd.extend(["--chat_format", self.chat_format])
# User-provided extra args (e.g. "--type_k 1 --type_v 1 --n_threads 8")
if self.extra_args:
cmd.extend(self.extra_args.split())
logger.info(f"Starting llama.cpp server: {' '.join(cmd)}")
# Write stderr to a log file to avoid pipe buffer deadlock
# (llama.cpp outputs a lot of model metadata on stderr during loading)
self._log_path = MODELS_DIR / "llamacpp_server.log"
self._log_file = open(self._log_path, "w")
self._process = subprocess.Popen(
cmd,
stdout=subprocess.DEVNULL,
stderr=self._log_file,
# Ensure the subprocess is killed when the parent exits
preexec_fn=os.setsid if hasattr(os, "setsid") else None,
)
# Wait for the server to be ready
await self._wait_for_ready()
async def _wait_for_ready(self, timeout: float = 120.0) -> None:
"""Wait for the llama.cpp server to accept connections."""
import httpx
start = time.monotonic()
url = f"http://127.0.0.1:{self.port}/v1/models"
last_log = start
while time.monotonic() - start < timeout:
# Check if process died
if self._process and self._process.poll() is not None:
stderr = ""
try:
stderr = self._log_path.read_text()[-2000:]
except Exception:
pass
raise RuntimeError(f"llama.cpp server exited with code {self._process.returncode}.\nstderr: {stderr}")
try:
async with httpx.AsyncClient() as client:
resp = await client.get(url, timeout=5.0)
if resp.status_code == 200:
logger.info(f"llama.cpp server ready on port {self.port}")
return
except (httpx.ConnectError, httpx.TimeoutException, httpx.ConnectTimeout):
pass
# Log progress every 15s
now = time.monotonic()
if now - last_log > 15:
elapsed = int(now - start)
logger.info(f"Waiting for llama.cpp server to load model... ({elapsed}s)")
last_log = now
await asyncio.sleep(1.0)
# Timeout — read the log to help debug
stderr = ""
try:
stderr = self._log_path.read_text()[-2000:]
except Exception:
pass
raise TimeoutError(
f"llama.cpp server did not become ready within {timeout}s.\n"
f"Check model compatibility and available memory.\n"
f"Server log: {stderr}"
)
async def stop(self) -> None:
"""Stop the llama.cpp server subprocess."""
if self._process is None:
return
logger.info("Stopping llama.cpp server...")
try:
# Send SIGTERM to the process group
if hasattr(os, "killpg"):
os.killpg(os.getpgid(self._process.pid), signal.SIGTERM)
else:
self._process.terminate()
# Wait up to 10s for graceful shutdown
try:
self._process.wait(timeout=10)
except subprocess.TimeoutExpired:
if hasattr(os, "killpg"):
os.killpg(os.getpgid(self._process.pid), signal.SIGKILL)
else:
self._process.kill()
self._process.wait(timeout=5)
except (ProcessLookupError, OSError):
pass # Process already exited
finally:
self._process = None
if hasattr(self, "_log_file") and self._log_file:
self._log_file.close()
self._log_file = None
logger.info("llama.cpp server stopped")
class LlamaCppLLM(LLMInterface):
"""
Built-in llama.cpp provider.
Manages a llama-cpp-python server subprocess and delegates to OpenAICompatibleLLM
for actual inference calls. Handles model downloading and server lifecycle.
"""
def __init__(
self,
provider: str,
api_key: str,
base_url: str,
model: str,
reasoning_effort: str = "low",
model_path: str | None = None,
gpu_layers: int = -1,
context_size: int = 8192,
chat_format: str | None = None,
no_grammar: bool = False,
extra_args: str | None = None,
**kwargs: Any,
):
super().__init__(
provider=provider,
api_key=api_key or "llamacpp",
base_url=base_url or "",
model=model or DEFAULT_LLAMACPP_MODEL_ALIAS,
reasoning_effort=reasoning_effort,
)
self._model_path_str = model_path
self._gpu_layers = gpu_layers
self._context_size = context_size
self._chat_format = chat_format
self._no_grammar = no_grammar
self._extra_args = extra_args
self._server: LlamaCppServer | None = None
self._delegate: Any = None # OpenAICompatibleLLM, created after server starts
self._initialized = False
async def _ensure_initialized(self) -> None:
"""Lazy initialization: download model + start shared server on first use."""
if self._initialized:
return
global _shared_server
from .openai_compatible_llm import OpenAICompatibleLLM
async with _shared_server_lock:
if _shared_server is None:
# Resolve and potentially download the model
model_path = _resolve_model_path(self._model_path_str)
logger.info(f"Using GGUF model: {model_path}")
# Start the shared llama.cpp server
port = _find_free_port()
_shared_server = LlamaCppServer(
model_path=model_path,
port=port,
gpu_layers=self._gpu_layers,
context_size=self._context_size,
chat_format=self._chat_format,
extra_args=self._extra_args,
)
await _shared_server.start()
self._server = _shared_server
# Create the delegate that talks to the shared server's OpenAI-compatible API
if self._no_grammar:
logger.info("Grammar enforcement disabled (HINDSIGHT_API_LLAMACPP_NO_GRAMMAR=true)")
self._delegate = OpenAICompatibleLLM(
provider="llamacpp",
api_key="llamacpp",
base_url=self._server.base_url,
model=self.model,
reasoning_effort=self.reasoning_effort,
)
self._initialized = True
async def verify_connection(self) -> None:
"""Verify the llama.cpp server is running and can generate text."""
await self._ensure_initialized()
# Make a simple test call to verify the model can actually generate
await self._delegate.call(
messages=[{"role": "user", "content": "Say 'ok'"}],
max_completion_tokens=10,
max_retries=2,
initial_backoff=0.5,
max_backoff=2.0,
scope="verification",
)
logger.info("llama.cpp LLM verification passed")
async def call(
self,
messages: list[dict[str, str]],
response_format: Any | None = None,
max_completion_tokens: int | None = None,
temperature: float | None = None,
scope: str = "memory",
max_retries: int = 10,
initial_backoff: float = 1.0,
max_backoff: float = 60.0,
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
) -> Any:
"""Delegate call to the OpenAI-compatible API."""
await self._ensure_initialized()
return await self._delegate.call(
messages=messages,
response_format=response_format,
max_completion_tokens=max_completion_tokens,
temperature=temperature,
scope=scope,
max_retries=max_retries,
initial_backoff=initial_backoff,
max_backoff=max_backoff,
skip_validation=skip_validation,
strict_schema=strict_schema,
return_usage=return_usage,
)
async def call_with_tools(
self,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]],
max_completion_tokens: int | None = None,
temperature: float | None = None,
scope: str = "tools",
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""Delegate tool calls to the OpenAI-compatible API."""
await self._ensure_initialized()
return await self._delegate.call_with_tools(
messages=messages,
tools=tools,
max_completion_tokens=max_completion_tokens,
temperature=temperature,
scope=scope,
max_retries=max_retries,
initial_backoff=initial_backoff,
max_backoff=max_backoff,
tool_choice=tool_choice,
)
async def cleanup(self) -> None:
"""Stop the shared llama.cpp server."""
global _shared_server
if self._delegate:
await self._delegate.cleanup()
self._delegate = None
# Stop the shared server (only the first cleanup call actually stops it)
async with _shared_server_lock:
if _shared_server is not None:
await _shared_server.stop()
_shared_server = None
self._server = None
self._initialized = False
@@ -1,78 +0,0 @@
"""
No-op LLM provider for chunk-only storage mode.
When the LLM provider is set to "none", the system operates without any LLM dependency.
Retain uses chunks mode (no fact extraction), and reflect/consolidation are disabled.
This provider acts as a safety net — if any code path unexpectedly tries to call the LLM,
it raises a clear error instead of a confusing connection failure.
"""
import logging
from typing import Any
from ..llm_interface import LLMInterface
from ..response_models import LLMToolCallResult
logger = logging.getLogger(__name__)
class LLMNotAvailableError(Exception):
"""Raised when an operation requires an LLM but the provider is set to 'none'."""
pass
class NoneLLM(LLMInterface):
"""
No-op LLM provider that rejects all LLM calls.
Used when HINDSIGHT_API_LLM_PROVIDER=none to run Hindsight as a chunk store
with semantic search but without LLM-based features (fact extraction, reflect,
consolidation).
"""
async def verify_connection(self) -> None:
"""No-op — no LLM connection to verify."""
logger.debug("NoneLLM: no LLM connection to verify (provider=none)")
async def call(
self,
messages: list[dict[str, str]],
response_format: Any | None = None,
max_completion_tokens: int | None = None,
temperature: float | None = None,
scope: str = "memory",
max_retries: int = 10,
initial_backoff: float = 1.0,
max_backoff: float = 60.0,
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
) -> Any:
"""Raise LLMNotAvailableError — no LLM is configured."""
raise LLMNotAvailableError(
"LLM provider is set to 'none'. This operation requires an LLM. "
"Set HINDSIGHT_API_LLM_PROVIDER to a real provider (e.g., openai, anthropic, gemini)."
)
async def call_with_tools(
self,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]],
max_completion_tokens: int | None = None,
temperature: float | None = None,
scope: str = "tools",
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""Raise LLMNotAvailableError — no LLM is configured."""
raise LLMNotAvailableError(
"LLM provider is set to 'none'. This operation requires an LLM. "
"Set HINDSIGHT_API_LLM_PROVIDER to a real provider (e.g., openai, anthropic, gemini)."
)
async def cleanup(self) -> None:
"""No-op — nothing to clean up."""
pass
@@ -1,307 +0,0 @@
"""Delta operations for structured mental models.
The LLM's job during a delta refresh is to emit a list of these operations,
each targeting an existing section (by id) or referencing a position relative
to one. ``apply_operations`` validates and applies each op in turn against a
copy of the document; invalid ops (unknown ``section_id``, out-of-range
``block_index``, malformed payloads) are dropped with a debug-friendly reason.
Sections and blocks not mentioned by any op are physically copied through
unchanged — there is no LLM-mediated re-emission of unchanged text, so prose
drift is structurally impossible.
Why operations and not "output the new structured doc":
- "Output the new doc" still asks the LLM to *generate* every section's
blocks, including ones it didn't intend to modify, which gives it the same
opportunity to drift.
- Operations make the no-change case mechanical: zero ops → identical doc.
- Operations are auditable: each refresh produces a log of exactly what
changed, useful for debugging the LLM's behaviour and explaining diffs.
Failure modes are by design conservative: an operation list that fails to
parse against the Pydantic schema, or an LLM that returns invalid ops, results
in zero changes — the document stays as-is. The structure can only get better
or stay the same per refresh, never get worse.
"""
from __future__ import annotations
import logging
from typing import Annotated, Any, Literal, Union
from pydantic import BaseModel, ConfigDict, Field
from .structured_doc import (
Block,
Section,
StructuredDocument,
make_unique_id,
slugify_heading,
)
logger = logging.getLogger(__name__)
# Op payloads ---------------------------------------------------------------
class _OpBase(BaseModel):
model_config = ConfigDict(extra="forbid")
class AppendBlockOp(_OpBase):
"""Add a new block at the end of an existing section."""
op: Literal["append_block"] = "append_block"
section_id: str
block: Block
class InsertBlockOp(_OpBase):
"""Insert a new block at ``index`` in an existing section.
``index`` may equal ``len(section.blocks)`` (append) but not be greater.
"""
op: Literal["insert_block"] = "insert_block"
section_id: str
index: int = Field(ge=0)
block: Block
class ReplaceBlockOp(_OpBase):
"""Replace the block at ``index`` of an existing section."""
op: Literal["replace_block"] = "replace_block"
section_id: str
index: int = Field(ge=0)
block: Block
class RemoveBlockOp(_OpBase):
"""Remove the block at ``index`` of an existing section."""
op: Literal["remove_block"] = "remove_block"
section_id: str
index: int = Field(ge=0)
class AddSectionOp(_OpBase):
"""Add a brand-new section.
``after_section_id`` is optional; when omitted the new section is appended
at the end. ``new_id`` is optional; when omitted we slugify the heading
and disambiguate against existing IDs.
"""
op: Literal["add_section"] = "add_section"
heading: str
level: int = Field(default=2, ge=1, le=6)
blocks: list[Block] = Field(default_factory=list)
after_section_id: str | None = None
new_id: str | None = None
class RemoveSectionOp(_OpBase):
"""Remove an entire section by id."""
op: Literal["remove_section"] = "remove_section"
section_id: str
class ReplaceSectionBlocksOp(_OpBase):
"""Replace all blocks of a section in one go.
Used when most of a section's contents are stale and rebuilding it as a
unit is clearer than emitting many block-level ops. The section's heading
and id are preserved.
"""
op: Literal["replace_section_blocks"] = "replace_section_blocks"
section_id: str
blocks: list[Block] = Field(default_factory=list)
class RenameSectionOp(_OpBase):
"""Rename a section's heading. The id is unchanged so future ops still resolve."""
op: Literal["rename_section"] = "rename_section"
section_id: str
new_heading: str
Operation = Annotated[
Union[
AppendBlockOp,
InsertBlockOp,
ReplaceBlockOp,
RemoveBlockOp,
AddSectionOp,
RemoveSectionOp,
ReplaceSectionBlocksOp,
RenameSectionOp,
],
Field(discriminator="op"),
]
class DeltaOperationList(BaseModel):
"""Container for the operations produced by an LLM delta call."""
model_config = ConfigDict(extra="forbid")
operations: list[Operation] = Field(default_factory=list)
# Application ---------------------------------------------------------------
class AppliedDelta(BaseModel):
"""Outcome of applying a list of operations to a document."""
model_config = ConfigDict(extra="forbid")
document: StructuredDocument
applied: list[dict[str, Any]] = Field(default_factory=list)
skipped: list[dict[str, Any]] = Field(default_factory=list)
@property
def changed(self) -> bool:
return len(self.applied) > 0
def _op_summary(op: Operation) -> dict[str, Any]:
"""Compact dict suitable for the audit trail."""
data = op.model_dump()
return {k: v for k, v in data.items() if k != "block" and k != "blocks"} | {
"op": data["op"],
}
def apply_operations(
doc: StructuredDocument,
operations: list[Operation],
) -> AppliedDelta:
"""Apply a list of operations to a document, returning a new document.
The original document is never mutated. Invalid operations (unknown
section, out-of-range index, name collision when adding a section) are
skipped and recorded in ``skipped`` with a ``reason`` string.
"""
new_doc = doc.model_copy(deep=True)
applied: list[dict[str, Any]] = []
skipped: list[dict[str, Any]] = []
def skip(op: Operation, reason: str) -> None:
entry = _op_summary(op)
entry["reason"] = reason
skipped.append(entry)
logger.debug(f"[STRUCTURED_DELTA] skipping op {entry}")
for op in operations:
if isinstance(op, AppendBlockOp):
section = new_doc.section_by_id(op.section_id)
if section is None:
skip(op, f"unknown section_id: {op.section_id}")
continue
section.blocks.append(op.block)
applied.append(_op_summary(op))
continue
if isinstance(op, InsertBlockOp):
section = new_doc.section_by_id(op.section_id)
if section is None:
skip(op, f"unknown section_id: {op.section_id}")
continue
if op.index > len(section.blocks):
skip(
op,
f"index out of range: {op.index} > {len(section.blocks)}",
)
continue
section.blocks.insert(op.index, op.block)
applied.append(_op_summary(op))
continue
if isinstance(op, ReplaceBlockOp):
section = new_doc.section_by_id(op.section_id)
if section is None:
skip(op, f"unknown section_id: {op.section_id}")
continue
if op.index >= len(section.blocks):
skip(
op,
f"index out of range: {op.index} >= {len(section.blocks)}",
)
continue
section.blocks[op.index] = op.block
applied.append(_op_summary(op))
continue
if isinstance(op, RemoveBlockOp):
section = new_doc.section_by_id(op.section_id)
if section is None:
skip(op, f"unknown section_id: {op.section_id}")
continue
if op.index >= len(section.blocks):
skip(
op,
f"index out of range: {op.index} >= {len(section.blocks)}",
)
continue
section.blocks.pop(op.index)
applied.append(_op_summary(op))
continue
if isinstance(op, AddSectionOp):
existing_ids = {s.id for s in new_doc.sections}
base_id = op.new_id or slugify_heading(op.heading)
section_id = make_unique_id(base_id, existing_ids)
new_section = Section(
id=section_id,
heading=op.heading,
level=op.level,
blocks=list(op.blocks),
)
if op.after_section_id is None:
new_doc.sections.append(new_section)
else:
idx = new_doc.section_index(op.after_section_id)
if idx is None:
skip(op, f"unknown after_section_id: {op.after_section_id}")
continue
new_doc.sections.insert(idx + 1, new_section)
entry = _op_summary(op)
entry["assigned_id"] = section_id
applied.append(entry)
continue
if isinstance(op, RemoveSectionOp):
idx = new_doc.section_index(op.section_id)
if idx is None:
skip(op, f"unknown section_id: {op.section_id}")
continue
new_doc.sections.pop(idx)
applied.append(_op_summary(op))
continue
if isinstance(op, ReplaceSectionBlocksOp):
section = new_doc.section_by_id(op.section_id)
if section is None:
skip(op, f"unknown section_id: {op.section_id}")
continue
section.blocks = list(op.blocks)
applied.append(_op_summary(op))
continue
if isinstance(op, RenameSectionOp):
section = new_doc.section_by_id(op.section_id)
if section is None:
skip(op, f"unknown section_id: {op.section_id}")
continue
section.heading = op.new_heading
applied.append(_op_summary(op))
continue
skip(op, f"unhandled op type: {type(op).__name__}") # pragma: no cover
return AppliedDelta(document=new_doc, applied=applied, skipped=skipped)
@@ -1,301 +0,0 @@
"""Structured representation of a mental model document.
Why this exists
---------------
Storing mental models as raw markdown forces every refresh to round-trip prose
through an LLM, which then drifts on stylistic details (numbered vs bulleted
lists, casing, separator lines, paraphrasing) even when instructed to preserve
content byte-for-byte. The intrinsic mechanism of an LLM is to *generate* the
next token from a gestalt of the input — not to copy tokens verbatim — so any
"preserve unchanged content" instruction is fundamentally a soft constraint.
The fix is to give the LLM no opportunity to drift on unchanged content. We
keep an authoritative structured representation of the document; the markdown
shown to users is a deterministic render of that structure. Delta refreshes
emit *operations* against the structure (see ``delta_ops.py``); sections and
blocks not mentioned by any operation are physically untouched.
Schema (v1)
-----------
A document is an ordered list of ``Section``s. Each section has:
- ``id`` : stable slug derived from ``heading`` (used as the operation
target across refreshes; surviving renames is a separate
concern handled by an explicit ``rename`` op).
- ``heading``: the markdown heading text (without the ``#`` prefix).
- ``level`` : 1 (``#``) … 6 (``######``). Default 2.
- ``blocks``: ordered list of typed blocks — paragraph, bullet_list,
ordered_list, code.
The schema is intentionally narrow: it covers what real mental-model documents
actually contain (the kind a coding agent writes for itself or a user writes as
a "skill" doc). Tables, images, and raw HTML are out of scope until needed.
"""
from __future__ import annotations
import re
from typing import Annotated, Literal, Union
from pydantic import BaseModel, ConfigDict, Field
# Blocks ---------------------------------------------------------------------
class ParagraphBlock(BaseModel):
model_config = ConfigDict(extra="forbid")
type: Literal["paragraph"] = "paragraph"
text: str
class BulletListBlock(BaseModel):
model_config = ConfigDict(extra="forbid")
type: Literal["bullet_list"] = "bullet_list"
items: list[str] = Field(default_factory=list)
class OrderedListBlock(BaseModel):
model_config = ConfigDict(extra="forbid")
type: Literal["ordered_list"] = "ordered_list"
items: list[str] = Field(default_factory=list)
class CodeBlock(BaseModel):
model_config = ConfigDict(extra="forbid")
type: Literal["code"] = "code"
language: str = ""
text: str
Block = Annotated[
Union[ParagraphBlock, BulletListBlock, OrderedListBlock, CodeBlock],
Field(discriminator="type"),
]
# Section / Document ---------------------------------------------------------
class Section(BaseModel):
model_config = ConfigDict(extra="forbid")
id: str
heading: str
level: int = Field(default=2, ge=1, le=6)
blocks: list[Block] = Field(default_factory=list)
class StructuredDocument(BaseModel):
"""Top-level structured representation of a mental model."""
model_config = ConfigDict(extra="forbid")
version: Literal[1] = 1
sections: list[Section] = Field(default_factory=list)
def section_by_id(self, section_id: str) -> Section | None:
for s in self.sections:
if s.id == section_id:
return s
return None
def section_index(self, section_id: str) -> int | None:
for i, s in enumerate(self.sections):
if s.id == section_id:
return i
return None
# Slug helpers ---------------------------------------------------------------
_SLUG_RX = re.compile(r"[^a-z0-9]+")
def slugify_heading(heading: str) -> str:
"""Stable, deterministic slug from a heading.
"Stop Conditions" -> "stop-conditions"
"Inputs and Context" -> "inputs-and-context"
"""
slug = _SLUG_RX.sub("-", heading.strip().lower()).strip("-")
return slug or "section"
def make_unique_id(base: str, existing: set[str]) -> str:
"""Disambiguate by appending -2, -3, … if the slug is already in use."""
if base not in existing:
return base
i = 2
while f"{base}-{i}" in existing:
i += 1
return f"{base}-{i}"
# Renderer -------------------------------------------------------------------
def render_block(block: Block) -> str:
"""Render a single block to markdown. No trailing newline."""
if isinstance(block, ParagraphBlock):
return block.text.rstrip()
if isinstance(block, BulletListBlock):
return "\n".join(f"- {item.rstrip()}" for item in block.items)
if isinstance(block, OrderedListBlock):
return "\n".join(f"{i + 1}. {item.rstrip()}" for i, item in enumerate(block.items))
if isinstance(block, CodeBlock):
fence_lang = block.language or ""
return f"```{fence_lang}\n{block.text}\n```"
raise TypeError(f"Unknown block type: {type(block)!r}")
def render_section(section: Section) -> str:
"""Render a section: heading + blank line + blocks separated by blank lines."""
parts = ["#" * section.level + " " + section.heading.strip()]
for block in section.blocks:
parts.append("") # blank line before each block
parts.append(render_block(block))
return "\n".join(parts)
def render_document(doc: StructuredDocument) -> str:
"""Render the whole document. Sections separated by a single blank line.
The output is byte-stable: same structured input always produces the same
markdown, modulo the inherent ordering of sections/blocks/items.
"""
if not doc.sections:
return ""
return "\n\n".join(render_section(s) for s in doc.sections) + "\n"
# Parser ---------------------------------------------------------------------
#
# The parser is intentionally lenient: it accepts the markdown produced by
# our own renderer (round-trip-safe) and the markdown an LLM tends to produce
# for mental-model documents. It is *not* a general CommonMark parser — it
# does not need to be. When it cannot classify a block it falls back to a
# paragraph so that no content is silently dropped.
_HEADING_RX = re.compile(r"^(#{1,6})\s+(.+?)\s*$")
_BULLET_RX = re.compile(r"^\s*[-*+]\s+(.*)$")
_ORDERED_RX = re.compile(r"^\s*\d+[.)]\s+(.*)$")
_FENCE_RX = re.compile(r"^```([A-Za-z0-9_+-]*)\s*$")
def _strip_separators(lines: list[str]) -> list[str]:
"""Drop horizontal-rule lines (`---`, `***`) used as section separators.
Our renderer never emits these, but LLM output frequently includes them
between sections; treating them as blank lines avoids parsing them as
paragraphs.
"""
return ["" if re.fullmatch(r"\s*([-*_])\1{2,}\s*", line) else line for line in lines]
def _split_blocks(lines: list[str]) -> list[list[str]]:
"""Group consecutive non-blank lines into block chunks."""
chunks: list[list[str]] = []
current: list[str] = []
in_fence = False
for line in lines:
if _FENCE_RX.match(line):
current.append(line)
in_fence = not in_fence
continue
if in_fence:
current.append(line)
continue
if line.strip() == "":
if current:
chunks.append(current)
current = []
else:
current.append(line)
if current:
chunks.append(current)
return chunks
def _parse_block(chunk: list[str]) -> Block:
"""Parse a single non-empty chunk into a block."""
if chunk and _FENCE_RX.match(chunk[0]):
m = _FENCE_RX.match(chunk[0])
lang = m.group(1) if m else ""
body_lines = chunk[1:]
if body_lines and _FENCE_RX.match(body_lines[-1]):
body_lines = body_lines[:-1]
return CodeBlock(language=lang, text="\n".join(body_lines))
if all(_BULLET_RX.match(line) for line in chunk):
items = []
for line in chunk:
m = _BULLET_RX.match(line)
assert m is not None
items.append(m.group(1).strip())
return BulletListBlock(items=items)
if all(_ORDERED_RX.match(line) for line in chunk):
items = []
for line in chunk:
m = _ORDERED_RX.match(line)
assert m is not None
items.append(m.group(1).strip())
return OrderedListBlock(items=items)
return ParagraphBlock(text=" ".join(line.strip() for line in chunk).strip())
def parse_markdown(markdown: str) -> StructuredDocument:
"""Best-effort parse of a markdown document into the structured schema.
Sections are introduced by ATX headings (``#``..``######``). Anything
before the first heading is wrapped into an implicit "Overview" section
so we never silently drop user content. Section IDs are unique slugs of
their headings.
"""
raw_lines = (markdown or "").splitlines()
lines = _strip_separators(raw_lines)
sections: list[Section] = []
used_ids: set[str] = set()
pending: list[str] = []
current: Section | None = None
def flush_pending_into(section: Section) -> None:
if not pending:
return
for chunk in _split_blocks(pending):
section.blocks.append(_parse_block(chunk))
pending.clear()
for line in lines:
m = _HEADING_RX.match(line)
if m:
if current is not None:
flush_pending_into(current)
sections.append(current)
elif pending:
# Content before the first heading: wrap in implicit section.
base = "overview"
section_id = make_unique_id(base, used_ids)
used_ids.add(section_id)
implicit = Section(id=section_id, heading="Overview", level=2)
flush_pending_into(implicit)
sections.append(implicit)
level = len(m.group(1))
heading = m.group(2).strip()
section_id = make_unique_id(slugify_heading(heading), used_ids)
used_ids.add(section_id)
current = Section(id=section_id, heading=heading, level=level)
else:
pending.append(line)
if current is not None:
flush_pending_into(current)
sections.append(current)
elif pending:
base = "overview"
section_id = make_unique_id(base, used_ids)
used_ids.add(section_id)
implicit = Section(id=section_id, heading="Overview", level=2)
flush_pending_into(implicit)
sections.append(implicit)
return StructuredDocument(sections=sections)
@@ -1,138 +0,0 @@
"""
Chunk storage for retain pipeline.
Handles storage of document chunks in the database.
"""
import hashlib
import logging
from dataclasses import dataclass
from ..memory_engine import fq_table
from .types import ChunkMetadata
logger = logging.getLogger(__name__)
def compute_chunk_hash(chunk_text: str) -> str:
"""Compute SHA256 hash of chunk text for delta comparison."""
return hashlib.sha256(chunk_text.encode()).hexdigest()
@dataclass
class ExistingChunk:
"""Represents a chunk already stored in the database."""
chunk_id: str
chunk_index: int
content_hash: str | None
async def load_existing_chunks(conn, bank_id: str, document_id: str) -> list[ExistingChunk]:
"""
Load existing chunk metadata for a document.
Returns list of ExistingChunk with chunk_id, chunk_index, and content_hash.
"""
rows = await conn.fetch(
f"""
SELECT chunk_id, chunk_index, content_hash
FROM {fq_table("chunks")}
WHERE document_id = $1 AND bank_id = $2
ORDER BY chunk_index
""",
document_id,
bank_id,
)
return [
ExistingChunk(
chunk_id=row["chunk_id"],
chunk_index=row["chunk_index"],
content_hash=row["content_hash"],
)
for row in rows
]
async def delete_chunks_by_ids(conn, chunk_ids: list[str]) -> None:
"""
Delete specific chunks by their IDs.
This cascades to memory_units (via FK with CASCADE delete)
and their links.
"""
if not chunk_ids:
return
await conn.execute(
f"DELETE FROM {fq_table('chunks')} WHERE chunk_id = ANY($1::text[])",
chunk_ids,
)
async def store_chunks_batch(
conn, bank_id: str, document_id: str, chunks: list[ChunkMetadata], ops=None
) -> dict[int, str]:
"""
Store document chunks in the database.
Args:
conn: Database connection
bank_id: Bank identifier
document_id: Document identifier
chunks: List of ChunkMetadata objects
ops: DataAccessOps instance (from backend.ops)
Returns:
Dictionary mapping global chunk index to chunk_id
"""
if not chunks:
return {}
# Prepare chunk data for batch insert
chunk_ids = []
chunk_texts = []
chunk_indices = []
content_hashes = []
chunk_id_map = {}
for chunk in chunks:
chunk_id = f"{bank_id}_{document_id}_{chunk.chunk_index}"
chunk_ids.append(chunk_id)
chunk_texts.append(chunk.chunk_text)
chunk_indices.append(chunk.chunk_index)
content_hashes.append(compute_chunk_hash(chunk.chunk_text))
chunk_id_map[chunk.chunk_index] = chunk_id
# Batch upsert all chunks. ON CONFLICT makes this idempotent: re-submitting
# a retain under the same document_id may produce chunk_ids that already exist.
# Overwriting is the correct behavior per document_id grouping semantics.
await ops.bulk_upsert_chunks(
conn,
fq_table("chunks"),
chunk_ids,
[document_id] * len(chunk_texts),
[bank_id] * len(chunk_texts),
chunk_texts,
chunk_indices,
content_hashes,
)
return chunk_id_map
def map_facts_to_chunks(facts_chunk_indices: list[int], chunk_id_map: dict[int, str]) -> list[str | None]:
"""
Map fact chunk indices to chunk IDs.
Args:
facts_chunk_indices: List of chunk indices for each fact
chunk_id_map: Dictionary mapping chunk index to chunk_id
Returns:
List of chunk_ids (same length as facts_chunk_indices)
"""
chunk_ids = []
for chunk_idx in facts_chunk_indices:
chunk_id = chunk_id_map.get(chunk_idx)
chunk_ids.append(chunk_id)
return chunk_ids
@@ -1,166 +0,0 @@
"""
Entity processing for retain pipeline.
Handles entity extraction, resolution, and link creation for stored facts.
"""
import logging
from . import link_utils
from .types import EntityLink, ProcessedFact
logger = logging.getLogger(__name__)
def _prepare_facts_for_entity_processing(
facts: list[ProcessedFact],
user_entities_per_content: dict[int, list[dict]] | None = None,
) -> tuple[list[str], list, list[list[dict]]]:
"""
Extract fact texts, dates, and merged entity lists from ProcessedFact objects.
Returns:
Tuple of (fact_texts, fact_dates, entities_per_fact)
"""
user_entities_per_content = user_entities_per_content or {}
fact_texts = [fact.fact_text for fact in facts]
fact_dates = [fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at for fact in facts]
entities_per_fact = []
for fact in facts:
llm_entities = [{"text": entity.name, "type": "CONCEPT"} for entity in (fact.entities or [])]
user_entities = user_entities_per_content.get(fact.content_index, [])
seen_texts = {e["text"].lower() for e in llm_entities}
for user_entity in user_entities:
if user_entity["text"].lower() not in seen_texts:
llm_entities.append(
{
"text": user_entity["text"],
"type": user_entity.get("type", "CONCEPT"),
}
)
seen_texts.add(user_entity["text"].lower())
entities_per_fact.append(llm_entities)
return fact_texts, fact_dates, entities_per_fact
async def resolve_entities(
entity_resolver,
conn,
bank_id: str,
unit_ids: list[str],
facts: list[ProcessedFact],
log_buffer: list[str] = None,
user_entities_per_content: dict[int, list[dict]] = None,
entity_labels: list | None = None,
) -> tuple[list[str], list[tuple], dict[str, list[str]]]:
"""
Phase 1: Resolve entity names to canonical IDs (read-heavy).
Should be called on a SEPARATE connection OUTSIDE the main write transaction
to avoid holding the transaction open during expensive trigram scans.
Args:
entity_resolver: EntityResolver instance
conn: Database connection (separate from the main write transaction)
bank_id: Bank identifier
unit_ids: Placeholder unit IDs (used only for grouping)
facts: List of ProcessedFact objects
log_buffer: Optional buffer for detailed logging
user_entities_per_content: Dict mapping content_index to user-provided entities
entity_labels: Optional entity label taxonomy
Returns:
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids)
to pass to build_entity_links().
"""
if not unit_ids or not facts:
return [], [], {}
if len(unit_ids) != len(facts):
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})")
fact_texts, fact_dates, entities_per_fact = _prepare_facts_for_entity_processing(facts, user_entities_per_content)
return await link_utils.resolve_entities_only(
entity_resolver,
conn,
bank_id,
unit_ids,
fact_texts,
"", # context (not used in current implementation)
fact_dates,
entities_per_fact,
log_buffer,
entity_labels=entity_labels,
)
async def build_entity_links(
entity_resolver,
conn,
bank_id: str,
unit_ids: list[str],
resolved_entity_ids: list[str],
entity_to_unit: list[tuple],
unit_to_entity_ids: dict[str, list[str]],
log_buffer: list[str] = None,
skip_unit_entities_insert: bool = False,
ops=None,
) -> list[EntityLink]:
"""
Build entity links for UI graph visualization.
Queries unit_entities to find shared entities between new and existing units,
then generates EntityLink objects. When called from Phase 3 (post-transaction),
set skip_unit_entities_insert=True since unit_entities were already inserted
in Phase 2.
Args:
entity_resolver: EntityResolver instance
conn: Database connection
bank_id: Bank identifier
unit_ids: Actual unit IDs (must already be inserted in the DB)
resolved_entity_ids: From resolve_entities()
entity_to_unit: From resolve_entities()
unit_to_entity_ids: From resolve_entities()
log_buffer: Optional buffer for detailed logging
skip_unit_entities_insert: Skip unit_entities INSERT (already done in Phase 2)
ops: DataAccessOps instance (from backend.ops)
Returns:
List of EntityLink objects for batch insertion
"""
return await link_utils.build_entity_links_from_resolved(
entity_resolver,
conn,
bank_id,
unit_ids,
resolved_entity_ids,
entity_to_unit,
unit_to_entity_ids,
log_buffer,
skip_unit_entities_insert=skip_unit_entities_insert,
ops=ops,
)
async def insert_entity_links_batch(conn, entity_links: list[EntityLink], bank_id: str, ops=None) -> None:
"""
Insert entity links in batch.
Args:
conn: Database connection
entity_links: List of EntityLink objects
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
ops: DataAccessOps instance (from backend.ops)
"""
if not entity_links:
return
await link_utils.insert_entity_links_batch(conn, entity_links, bank_id, ops=ops)
@@ -1,442 +0,0 @@
"""
Fact storage for retain pipeline.
Handles insertion of facts into the database.
"""
import json
import logging
import uuid
from datetime import datetime
from ...config import get_config
from ..memory_engine import fq_table
from .bank_utils import DEFAULT_DISPOSITION, create_bank_vector_indexes
from .fact_extraction import _sanitize_text
from .types import ProcessedFact
logger = logging.getLogger(__name__)
async def get_document_content(
conn,
bank_id: str,
document_id: str,
) -> str | None:
"""Fetch the original_text of an existing document.
Returns None if the document does not exist.
"""
row = await conn.fetchval(
f"SELECT original_text FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2",
document_id,
bank_id,
)
return row
async def insert_facts_batch(
conn, bank_id: str, facts: list[ProcessedFact], document_id: str | None = None, ops=None
) -> list[str]:
"""
Insert facts into the database in batch.
Args:
conn: Database connection
bank_id: Bank identifier
facts: List of ProcessedFact objects to insert
document_id: Optional document ID to associate with facts
Returns:
List of unit IDs (UUIDs as strings) for the inserted facts
"""
if not facts:
return []
# Prepare data for batch insert
fact_texts = []
embeddings = []
event_dates = []
occurred_starts = []
occurred_ends = []
mentioned_ats = []
contexts = []
fact_types = []
metadata_jsons = []
chunk_ids = []
document_ids = []
tags_list = []
observation_scopes_list = []
text_signals_list = []
for fact in facts:
fact_texts.append(_sanitize_text(fact.fact_text))
# Convert embedding to string for asyncpg vector type
embeddings.append(str(fact.embedding))
# event_date: Use occurred_start if available, otherwise use mentioned_at
# This maintains backward compatibility while handling None occurred_start
event_dates.append(fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at)
occurred_starts.append(fact.occurred_start)
occurred_ends.append(fact.occurred_end)
mentioned_ats.append(fact.mentioned_at)
contexts.append(_sanitize_text(fact.context))
fact_types.append(fact.fact_type)
metadata_jsons.append(json.dumps(fact.metadata))
chunk_ids.append(fact.chunk_id)
# Use per-fact document_id if available, otherwise fallback to batch-level document_id
document_ids.append(fact.document_id if fact.document_id else document_id)
# Convert tags to JSON string for proper batch insertion (PostgreSQL unnest doesn't handle 2D arrays well)
tags_list.append(json.dumps(fact.tags if fact.tags else []))
# observation_scopes: stored as JSONB (string or 2D array), None if not provided
observation_scopes_list.append(
json.dumps(fact.observation_scopes) if fact.observation_scopes is not None else None
)
# Build text_signals: entity names + date tokens for enriched BM25 indexing
signal_parts = []
if fact.entities:
signal_parts.extend(e.name for e in fact.entities)
if fact.occurred_start:
try:
signal_parts.append(fact.occurred_start.strftime("%B %d %Y").lstrip("0").replace(" 0", " "))
except (ValueError, AttributeError):
pass
if fact.occurred_end and fact.occurred_end != fact.occurred_start:
try:
signal_parts.append(fact.occurred_end.strftime("%B %d %Y").lstrip("0").replace(" 0", " "))
except (ValueError, AttributeError):
pass
text_signals_list.append(" ".join(signal_parts) if signal_parts else None)
# Batch insert all facts — delegates to DataAccessOps which handles
# unnest (PG) vs row-by-row (Oracle) transparently.
config = get_config()
return await ops.insert_facts_batch(
conn,
bank_id,
fact_texts,
embeddings,
event_dates,
occurred_starts,
occurred_ends,
mentioned_ats,
contexts,
fact_types,
metadata_jsons,
chunk_ids,
document_ids,
tags_list,
observation_scopes_list,
text_signals_list,
text_search_extension=config.text_search_extension,
)
async def ensure_bank_exists(conn, bank_id: str, ops=None) -> None:
"""
Ensure bank exists in the database.
Creates bank with default values if it doesn't exist.
Args:
conn: Database connection
bank_id: Bank identifier
"""
# Generate internal_id here so we control the value and can use it
# immediately for HNSW index creation without a RETURNING round-trip.
internal_id = uuid.uuid4()
inserted = await conn.fetchval(
f"""
INSERT INTO {fq_table("banks")} (bank_id, disposition, mission, internal_id)
VALUES ($1, $2::jsonb, $3, $4)
ON CONFLICT (bank_id) DO NOTHING
RETURNING bank_id
""",
bank_id,
json.dumps(DEFAULT_DISPOSITION),
"",
internal_id,
)
if inserted:
# Fresh insert — create per-bank vector indexes
await create_bank_vector_indexes(conn, bank_id, str(internal_id), ops=ops)
async def delete_stale_observations_for_memories(
conn,
bank_id: str,
fact_ids: "list[str | uuid.UUID]",
ops=None,
) -> int:
"""Delete observations whose source memories are about to be removed.
Mirrors the cleanup performed by ``MemoryEngine.delete_document`` so that
every code path that removes ``memory_units`` also removes the
observations derived from them. Without this, ingesting a fresh version
of a document via the retain pipeline (which does a full-replace
``DELETE FROM documents`` cascade) used to leave orphan observations
pointing at memory IDs that no longer existed.
For each observation referencing any of ``fact_ids``:
1. Delete the observation row (its text is stale once even one source
memory disappears).
2. Reset ``consolidated_at = NULL`` on the surviving source memories so
they get re-consolidated under fresh observations on the next run.
Must be called within an active transaction, before the source memories
are deleted.
Returns the number of observations deleted.
"""
if not fact_ids:
return 0
fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in fact_ids]
if ops is not None and not ops.uses_observation_sources_table:
# PG: use native array overlap operator
affected_obs = await conn.fetch(
f"""
SELECT id, source_memory_ids
FROM {fq_table("memory_units")}
WHERE bank_id = $1
AND fact_type = 'observation'
AND source_memory_ids && $2::uuid[]
""",
bank_id,
fact_uuids,
)
else:
# Oracle / default: use observation_sources junction table
affected_obs = await conn.fetch(
f"""
SELECT mu.id, mu.source_memory_ids
FROM {fq_table("memory_units")} mu
WHERE mu.bank_id = $1
AND mu.fact_type = 'observation'
AND EXISTS (
SELECT 1 FROM {fq_table("observation_sources")} os
WHERE os.observation_id = mu.id
AND os.source_id = ANY($2::uuid[])
)
""",
bank_id,
fact_uuids,
)
if not affected_obs:
return 0
deleted_set = {str(uid) for uid in fact_uuids}
obs_ids = [obs["id"] for obs in affected_obs]
seen_remaining: set[str] = set()
remaining_source_ids: list[uuid.UUID] = []
for obs in affected_obs:
for src_id in obs["source_memory_ids"] or []:
src_str = str(src_id)
if src_str not in deleted_set and src_str not in seen_remaining:
remaining_source_ids.append(src_id)
seen_remaining.add(src_str)
await conn.execute(
f"DELETE FROM {fq_table('memory_units')} WHERE id = ANY($1::uuid[])",
obs_ids,
)
if remaining_source_ids:
await conn.execute(
f"""
UPDATE {fq_table("memory_units")}
SET consolidated_at = NULL
WHERE id = ANY($1::uuid[])
AND fact_type IN ('experience', 'world')
""",
remaining_source_ids,
)
logger.info(
f"[OBSERVATIONS] Deleted {len(obs_ids)} observations, reset {len(remaining_source_ids)} "
f"source memories for re-consolidation in bank {bank_id}"
)
return len(obs_ids)
async def handle_document_tracking(
conn,
bank_id: str,
document_id: str,
combined_content: str,
is_first_batch: bool,
retain_params: dict | None = None,
document_tags: list[str] | None = None,
ops=None,
) -> None:
"""
Handle document tracking in the database (full-replace mode).
Deletes the existing document (cascading to all units and links) on the
first batch, then inserts the new document record.
Args:
conn: Database connection
bank_id: Bank identifier
document_id: Document identifier
combined_content: Combined content text from all content items
is_first_batch: Whether this is the first batch (for chunked operations)
retain_params: Optional parameters passed during retain (context, event_date, etc.)
document_tags: Optional list of tags to associate with the document
ops: Backend-specific DataAccessOps. Required by the inner
``delete_stale_observations_for_memories`` call to choose the PG
(native array) vs Oracle (junction table) read path. Defaults to
None so older callers don't break, but the PG branch is only
taken when ops is non-None — pass ``pool.ops`` from the caller.
"""
import hashlib
# Sanitize and calculate content hash
combined_content = _sanitize_text(combined_content) or ""
content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
# Delete old document first (cascades to units and links).
# Only delete on the first batch to avoid deleting data we just inserted.
# Before the cascade, fan out to delete observations derived from the
# outgoing memory_units — otherwise the FK ON DELETE CASCADE removes the
# source memory_units but leaves observation rows pointing at IDs that
# no longer exist (consolidated_at on co-source memories also stays
# frozen). Same cleanup the explicit ``delete_document`` API performs.
preserved_created_at = None
if is_first_batch:
existing_unit_rows = await conn.fetch(
f"""
SELECT id FROM {fq_table("memory_units")}
WHERE document_id = $1 AND fact_type IN ('experience', 'world')
""",
document_id,
)
existing_unit_ids = [row["id"] for row in existing_unit_rows]
if existing_unit_ids:
invalidated = await delete_stale_observations_for_memories(conn, bank_id, existing_unit_ids, ops=ops)
if invalidated:
logger.info(
f"[RETAIN] Document {document_id} re-ingested: invalidated "
f"{invalidated} observation(s) derived from {len(existing_unit_ids)} outgoing memory_units"
)
# Explicitly delete memory_units by document_id BEFORE deleting the
# document row. The CASCADE from documents→chunks→memory_units only
# catches units that have a non-NULL chunk_id FK. Units with chunk_id=NULL
# (e.g. from partial writes or edge cases) would survive the cascade.
# This explicit delete ensures complete cleanup.
await conn.execute(
f"DELETE FROM {fq_table('memory_units')} WHERE document_id = $1 AND bank_id = $2",
document_id,
bank_id,
)
# Capture created_at before deletion so re-ingestion preserves it.
preserved_created_at = await conn.fetchval(
f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING created_at",
document_id,
bank_id,
)
# Insert document (or update if exists from concurrent operations)
await _upsert_document_row(
conn,
bank_id,
document_id,
combined_content,
content_hash,
retain_params,
document_tags,
preserved_created_at=preserved_created_at,
)
async def upsert_document_metadata(
conn,
bank_id: str,
document_id: str,
combined_content: str,
retain_params: dict | None = None,
document_tags: list[str] | None = None,
) -> None:
"""
Update document metadata without deleting existing facts/chunks.
Used by delta retain: the document row is upserted but chunks and
memory_units are managed separately at the chunk level.
"""
import hashlib
combined_content = _sanitize_text(combined_content) or ""
content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
await _upsert_document_row(conn, bank_id, document_id, combined_content, content_hash, retain_params, document_tags)
async def _upsert_document_row(
conn,
bank_id: str,
document_id: str,
combined_content: str,
content_hash: str,
retain_params: dict | None = None,
document_tags: list[str] | None = None,
preserved_created_at: datetime | None = None,
) -> None:
"""Insert or update a document row.
When ``preserved_created_at`` is provided, it is used for ``created_at`` on
INSERT so that re-ingesting a document (which deletes + inserts the row)
keeps the original creation timestamp. ``updated_at`` is always set to
``NOW()`` on both INSERT and the ON CONFLICT UPDATE branch.
"""
await conn.execute(
f"""
INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, retain_params, tags, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, $6, COALESCE($7, NOW()), NOW())
ON CONFLICT (id, bank_id) DO UPDATE
SET original_text = EXCLUDED.original_text,
content_hash = EXCLUDED.content_hash,
retain_params = EXCLUDED.retain_params,
tags = EXCLUDED.tags,
updated_at = NOW()
""",
document_id,
bank_id,
combined_content,
content_hash,
json.dumps(retain_params) if retain_params else None,
document_tags or [],
preserved_created_at,
)
async def update_memory_units_tags(
conn,
bank_id: str,
document_id: str,
tags: list[str],
) -> int:
"""
Update tags on all memory_units belonging to a document.
Used during delta retain to propagate tag changes to unchanged facts.
Returns:
Number of memory units updated.
"""
result = await conn.execute(
f"""
UPDATE {fq_table("memory_units")}
SET tags = $3, updated_at = NOW()
WHERE bank_id = $1 AND document_id = $2
""",
bank_id,
document_id,
tags or [],
)
# result is a status string like "UPDATE 5"
try:
return int(result.split()[-1])
except (ValueError, IndexError):
return 0
@@ -1,994 +0,0 @@
"""
Link creation utilities for temporal, semantic, and entity links.
"""
import logging
import time
from datetime import UTC, datetime, timedelta
from uuid import UUID
from ..memory_engine import fq_table
from .types import EntityLink
logger = logging.getLogger(__name__)
# Sentinel UUID used in the unique index to represent NULL entity_id
_NIL_ENTITY_UUID = "00000000-0000-0000-0000-000000000000"
# Maximum number of temporal links to keep per unit (from_unit_id).
# Retrieval only reads top 10-20 per unit via LATERAL join, so keeping
# more is wasted storage and write amplification.
MAX_TEMPORAL_LINKS_PER_UNIT = 20
def _cap_links_per_unit(links: list[tuple], max_per_unit: int = MAX_TEMPORAL_LINKS_PER_UNIT) -> list[tuple]:
"""Keep only the top-N links per from_unit_id, ranked by weight descending.
Args:
links: List of (from_unit_id, to_unit_id, link_type, weight, entity_id) tuples.
max_per_unit: Maximum number of links to retain per from_unit_id.
Returns:
Filtered list of link tuples.
"""
if not links:
return links
# Group by from_unit_id (index 0)
groups: dict[str, list[tuple]] = {}
for link in links:
key = str(link[0])
if key not in groups:
groups[key] = []
groups[key].append(link)
# For each group, sort by weight (index 3) descending and keep top N
result: list[tuple] = []
for group_links in groups.values():
group_links.sort(key=lambda lnk: lnk[3], reverse=True)
result.extend(group_links[:max_per_unit])
return result
async def _bulk_insert_links(
conn,
links: list[tuple],
bank_id: str = "",
chunk_size: int = 5000,
skip_exists_check: bool = False,
ops=None,
) -> None:
"""Bulk-insert links using sorted INSERT FROM unnest().
Sorting by (from_unit_id, to_unit_id) ensures all concurrent transactions
acquire index locks in the same order, eliminating circular-wait deadlocks.
Args:
conn: Database connection (must be inside a transaction).
links: List of (from_unit_id, to_unit_id, link_type, weight, entity_id) tuples.
bank_id: Bank identifier stored on memory_links for fast filtering.
chunk_size: Max rows per INSERT statement to avoid query timeouts on
very large tables (100M+ rows).
skip_exists_check: Skip WHERE EXISTS checks on memory_units. Use when
all referenced unit IDs are guaranteed to exist (e.g., within
the same transaction that inserted them).
ops: DataAccessOps instance for backend-specific bulk operations.
"""
if not links:
return
# Sort by (from_unit_id, to_unit_id) to guarantee consistent lock ordering
# across concurrent transactions — prevents deadlocks.
sorted_links = sorted(links, key=lambda lnk: (str(lnk[0]), str(lnk[1])))
exists_clause = ""
if not skip_exists_check:
exists_clause = (
f"WHERE EXISTS (SELECT 1 FROM {fq_table('memory_units')} mu WHERE mu.id = f)"
f" AND EXISTS (SELECT 1 FROM {fq_table('memory_units')} mu WHERE mu.id = t)"
)
await ops.bulk_insert_links(
conn,
fq_table("memory_links"),
sorted_links,
bank_id,
_NIL_ENTITY_UUID,
exists_clause,
chunk_size,
)
def _normalize_datetime(dt):
"""Normalize datetime to be timezone-aware (UTC) for consistent comparison."""
if dt is None:
return None
if dt.tzinfo is None:
# Naive datetime - assume UTC
return dt.replace(tzinfo=UTC)
return dt
def compute_temporal_links(
new_units: dict,
candidates: list,
time_window_hours: int = 24,
) -> list:
"""
Compute temporal links between new units and candidate neighbors.
This is a pure function that takes query results and returns link tuples,
making it easy to test without database access.
Args:
new_units: Dict mapping unit_id (str) to event_date (datetime)
candidates: List of dicts with 'id' and 'event_date' keys (candidate neighbors)
time_window_hours: Time window in hours for temporal links
Returns:
List of tuples: (from_unit_id, to_unit_id, 'temporal', weight, None)
"""
if not new_units:
return []
links = []
for unit_id, unit_event_date in new_units.items():
# Units without event_date can't form temporal links
if unit_event_date is None:
continue
# Normalize unit_event_date for consistent comparison
unit_event_date_norm = _normalize_datetime(unit_event_date)
# Calculate time window bounds with overflow protection
try:
time_lower = unit_event_date_norm - timedelta(hours=time_window_hours)
except OverflowError:
time_lower = datetime.min.replace(tzinfo=UTC)
try:
time_upper = unit_event_date_norm + timedelta(hours=time_window_hours)
except OverflowError:
time_upper = datetime.max.replace(tzinfo=UTC)
# Filter candidates within this unit's time window
matching_neighbors = [
(row["id"], row["event_date"])
for row in candidates
if time_lower <= _normalize_datetime(row["event_date"]) <= time_upper
][:10] # Limit to top 10
for recent_id, recent_event_date in matching_neighbors:
# Calculate temporal proximity weight
time_diff_hours = abs(
(unit_event_date_norm - _normalize_datetime(recent_event_date)).total_seconds() / 3600
)
weight = max(0.3, 1.0 - (time_diff_hours / time_window_hours))
links.append((unit_id, str(recent_id), "temporal", weight, None))
return _cap_links_per_unit(links)
def compute_temporal_query_bounds(
new_units: dict,
time_window_hours: int = 24,
) -> tuple:
"""
Compute the min/max date bounds for querying temporal neighbors.
Args:
new_units: Dict mapping unit_id (str) to event_date (datetime)
time_window_hours: Time window in hours
Returns:
Tuple of (min_date, max_date) with overflow protection
"""
if not new_units:
return None, None
# Normalize all dates to be timezone-aware to avoid comparison issues
# Filter out None values — units without event_date can't form temporal links
all_dates = [_normalize_datetime(d) for d in new_units.values() if d is not None]
if not all_dates:
return None, None
try:
min_date = min(all_dates) - timedelta(hours=time_window_hours)
except OverflowError:
min_date = datetime.min.replace(tzinfo=UTC)
try:
max_date = max(all_dates) + timedelta(hours=time_window_hours)
except OverflowError:
max_date = datetime.max.replace(tzinfo=UTC)
return min_date, max_date
def _log(log_buffer, message, level="info"):
"""Helper to log to buffer if available, otherwise use logger.
Args:
log_buffer: Buffer to append messages to (for main output)
message: The log message
level: 'info', 'debug', 'warning', or 'error'. Debug messages are not added to buffer.
"""
if level == "debug":
# Debug messages only go to logger, not to buffer
logger.debug(message)
return
if log_buffer is not None:
log_buffer.append(message)
else:
if level == "info":
logger.info(message)
else:
logger.log(logging.WARNING if level == "warning" else logging.ERROR, message)
def _prepare_entities_for_resolution(
unit_ids: list[str],
sentences: list[str],
fact_dates: list,
llm_entities: list[list[dict]],
log_buffer: list[str] = None,
) -> tuple[list[dict], list[list[dict]], list[tuple]]:
"""
Convert LLM entities into the flat format expected by entity resolver.
Returns:
Tuple of (all_entities_flat, all_entities, entity_to_unit) where:
- all_entities_flat: flat list of entity dicts ready for resolve_entities_batch
- all_entities: per-unit formatted entity lists
- entity_to_unit: maps flat index to (unit_id, local_index, fact_date)
"""
substep_start = time.time()
all_entities = []
for entity_list in llm_entities:
formatted_entities = []
for ent in entity_list:
if hasattr(ent, "text"):
formatted_entities.append({"text": ent.text, "type": "CONCEPT"})
elif isinstance(ent, dict):
formatted_entities.append({"text": ent.get("text", ""), "type": ent.get("type", "CONCEPT")})
all_entities.append(formatted_entities)
total_entities = sum(len(ents) for ents in all_entities)
_log(
log_buffer,
f" [6.1] Process LLM entities: {total_entities} entities from {len(sentences)} facts in {time.time() - substep_start:.3f}s",
level="debug",
)
substep_start = time.time()
all_entities_flat = []
entity_to_unit: list[tuple] = []
for unit_id, entities, fact_date in zip(unit_ids, all_entities, fact_dates):
if not entities:
continue
for local_idx, entity in enumerate(entities):
all_entities_flat.append(
{
"text": entity["text"],
"type": entity["type"],
"nearby_entities": entities,
}
)
entity_to_unit.append((unit_id, local_idx, fact_date))
_log(
log_buffer,
f" [6.2.1] Prepare entities: {len(all_entities_flat)} entities in {time.time() - substep_start:.3f}s",
level="debug",
)
# Attach per-entity dates
for idx, (_unit_id, _local_idx, fact_date) in enumerate(entity_to_unit):
all_entities_flat[idx]["event_date"] = fact_date
return all_entities_flat, all_entities, entity_to_unit
async def resolve_entities_only(
entity_resolver,
conn,
bank_id: str,
unit_ids: list[str],
sentences: list[str],
context: str,
fact_dates: list,
llm_entities: list[list[dict]],
log_buffer: list[str] = None,
entity_labels: list | None = None,
) -> tuple[list[str], list[tuple], dict[str, list[str]]]:
"""
Phase 1 of entity processing: resolve entity names to canonical IDs.
Runs the expensive read-heavy trigram search, co-occurrence fetch, and scoring
OUTSIDE the main write transaction. Also INSERTs new entities (idempotent
DO NOTHING) so that IDs are available for the subsequent write phase.
Args:
entity_resolver: EntityResolver instance
conn: Database connection (separate from the main write transaction)
bank_id: Bank identifier
unit_ids: Placeholder unit IDs (used only for grouping, not yet inserted)
sentences: Fact texts
context: Context string
fact_dates: Per-fact dates
llm_entities: Per-fact entity lists from LLM extraction
log_buffer: Optional logging buffer
entity_labels: Optional entity label taxonomy
Returns:
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids) where:
- resolved_entity_ids: list of entity IDs in same order as flattened entities
- entity_to_unit: maps flat index to (unit_id, local_index, fact_date)
- unit_to_entity_ids: maps unit_id to list of resolved entity IDs
"""
all_entities_flat, _all_entities, entity_to_unit = _prepare_entities_for_resolution(
unit_ids, sentences, fact_dates, llm_entities, log_buffer
)
if not all_entities_flat:
_log(log_buffer, " [6.2] Entity resolution (batched): 0 entities", level="debug")
return [], [], {}
step_start = time.time()
resolved_entity_ids = await entity_resolver.resolve_entities_batch(
bank_id=bank_id,
entities_data=all_entities_flat,
context=context,
unit_event_date=None,
conn=conn,
entity_labels=entity_labels,
)
_log(
log_buffer,
f" [6.2.2] Resolve entities: {len(all_entities_flat)} entities in single batch in {time.time() - step_start:.3f}s",
level="debug",
)
# Build unit_to_entity_ids mapping
unit_to_entity_ids: dict[str, list[str]] = {}
for idx, (unit_id, _local_idx, _fact_date) in enumerate(entity_to_unit):
if unit_id not in unit_to_entity_ids:
unit_to_entity_ids[unit_id] = []
unit_to_entity_ids[unit_id].append(resolved_entity_ids[idx])
_log(
log_buffer,
f" [6.2] Entity resolution (batched): {len(all_entities_flat)} entities resolved in {time.time() - step_start:.3f}s",
level="debug",
)
return resolved_entity_ids, entity_to_unit, unit_to_entity_ids
async def build_entity_links_from_resolved(
entity_resolver,
conn,
bank_id: str,
unit_ids: list[str],
resolved_entity_ids: list[str],
entity_to_unit: list[tuple],
unit_to_entity_ids: dict[str, list[str]],
log_buffer: list[str] = None,
skip_unit_entities_insert: bool = False,
ops=None,
) -> list["EntityLink"]:
"""
Build entity links between units that share entities.
Queries unit_entities to find which existing units share entities with the
new units, then generates EntityLink objects for UI graph visualization.
Args:
entity_resolver: EntityResolver instance
conn: Database connection
bank_id: Bank identifier
unit_ids: Actual unit IDs (must already be inserted in the DB)
resolved_entity_ids: Entity IDs from resolve_entities_only
entity_to_unit: Mapping from resolve_entities_only
unit_to_entity_ids: Mapping from resolve_entities_only
log_buffer: Optional logging buffer
skip_unit_entities_insert: If True, skip unit_entities INSERT (already done in Phase 2)
Returns:
List of EntityLink objects for batch insertion
"""
if not resolved_entity_ids:
return []
if not skip_unit_entities_insert:
# Insert unit-entity links (used in fallback path where Phase 2 didn't do this)
substep_start = time.time()
unit_entity_pairs = []
for idx, (unit_id, _local_idx, fact_date) in enumerate(entity_to_unit):
# Propagate the unit's fact_date so entity_cooccurrences.last_cooccurred
# reflects the event timeline, not the ingest moment.
unit_entity_pairs.append((unit_id, resolved_entity_ids[idx], fact_date))
await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
_log(
log_buffer,
f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_start:.3f}s",
level="debug",
)
# Build entity links between units that share entities
substep_start = time.time()
all_entity_ids = set()
for entity_ids_list in unit_to_entity_ids.values():
all_entity_ids.update(entity_ids_list)
_log(log_buffer, f" [6.3] Creating entity links for {len(all_entity_ids)} unique entities...", level="debug")
MAX_LINKS_PER_ENTITY = 10
entity_to_units = {}
if all_entity_ids:
query_start = time.time()
import uuid
entity_id_list = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in all_entity_ids]
limit_per_entity = MAX_LINKS_PER_ENTITY + len(unit_ids) # room for new units + existing cap
rows = await ops.fetch_entity_unit_fanout(
conn,
fq_table("unit_entities"),
entity_id_list,
limit_per_entity,
)
_log(
log_buffer,
f" [6.3.1] Query unit_entities (LATERAL): {len(rows)} rows in {time.time() - query_start:.3f}s",
level="debug",
)
group_start = time.time()
for row in rows:
entity_id = row["entity_id"]
if entity_id not in entity_to_units:
entity_to_units[entity_id] = []
entity_to_units[entity_id].append(row["unit_id"])
_log(log_buffer, f" [6.3.2] Group by entity_id: {time.time() - group_start:.3f}s", level="debug")
link_gen_start = time.time()
links: list[EntityLink] = []
new_unit_set = set(unit_ids)
def to_uuid(val) -> UUID:
return UUID(val) if isinstance(val, str) else val
for entity_id, units_with_entity in entity_to_units.items():
entity_uuid = to_uuid(entity_id)
new_units = [u for u in units_with_entity if str(u) in new_unit_set or u in new_unit_set]
existing_units = [u for u in units_with_entity if str(u) not in new_unit_set and u not in new_unit_set]
new_units_to_link = new_units[-MAX_LINKS_PER_ENTITY:] if len(new_units) > MAX_LINKS_PER_ENTITY else new_units
for i, unit_id_1 in enumerate(new_units_to_link):
for unit_id_2 in new_units_to_link[i + 1 :]:
links.append(
EntityLink(from_unit_id=to_uuid(unit_id_1), to_unit_id=to_uuid(unit_id_2), entity_id=entity_uuid)
)
links.append(
EntityLink(from_unit_id=to_uuid(unit_id_2), to_unit_id=to_uuid(unit_id_1), entity_id=entity_uuid)
)
existing_to_link = existing_units[-MAX_LINKS_PER_ENTITY:]
for new_unit in new_units:
for existing_unit in existing_to_link:
links.append(
EntityLink(from_unit_id=to_uuid(new_unit), to_unit_id=to_uuid(existing_unit), entity_id=entity_uuid)
)
links.append(
EntityLink(from_unit_id=to_uuid(existing_unit), to_unit_id=to_uuid(new_unit), entity_id=entity_uuid)
)
_log(log_buffer, f" [6.3.3] Generate {len(links)} links: {time.time() - link_gen_start:.3f}s", level="debug")
_log(
log_buffer,
f" [6.3] Entity link creation: {len(links)} links for {len(all_entity_ids)} unique entities in {time.time() - substep_start:.3f}s",
level="debug",
)
return links
async def create_temporal_links_batch_per_fact(
conn,
bank_id: str,
unit_ids: list[str],
time_window_hours: int = 24,
log_buffer: list[str] = None,
ops=None,
) -> int:
"""
Create temporal links for multiple units, each with their own event_date.
Queries the event_date for each unit from the database and creates temporal
links based on individual dates (supports per-fact dating).
Args:
conn: Database connection
bank_id: Bank identifier
unit_ids: List of unit IDs
time_window_hours: Time window in hours for temporal links
log_buffer: Optional buffer for logging
Returns:
Number of temporal links created
"""
if not unit_ids:
return 0
try:
import time as time_mod
# Get the event_date for each new unit
fetch_dates_start = time_mod.time()
rows = await ops.fetch_unit_dates(conn, fq_table("memory_units"), unit_ids)
new_units = {str(row["id"]): (row["event_date"], row["fact_type"]) for row in rows}
_log(
log_buffer,
f" [7.1] Fetch event_dates for {len(unit_ids)} units: {time_mod.time() - fetch_dates_start:.3f}s",
)
# Use LATERAL push-down to fetch only top-N temporal neighbors per new unit,
# avoiding transfer of the entire time-window result set (could be 50k+ rows).
fetch_neighbors_start = time_mod.time()
# Build arrays of new unit IDs, event dates, and fact types for the LATERAL query
new_unit_entries = [(uid, edate, ftype) for uid, (edate, ftype) in new_units.items() if edate is not None]
if new_unit_entries:
import uuid as uuid_mod
lateral_unit_ids = [
uuid_mod.UUID(uid) if isinstance(uid, str) else uid for uid in [e[0] for e in new_unit_entries]
]
lateral_event_dates = [_normalize_datetime(e[1]) for e in new_unit_entries]
lateral_fact_types = [e[2] for e in new_unit_entries]
# Bidirectional index scan: instead of scanning all units in the 24h
# window (O(N) — 164k rows at scale) and sorting by proximity, we scan
# the nearest K units in each direction using the B-tree index on
# (bank_id, fact_type, event_date). This reads only 2×K rows per probe
# regardless of bank size — 120x faster at 164k units (0.6ms vs 74ms).
TEMPORAL_LATERAL_BATCH = 500
half_limit = MAX_TEMPORAL_LINKS_PER_UNIT # fetch K in each direction, take top K combined
mu = fq_table("memory_units")
# Bidirectional index scan: instead of scanning all units in the 24h
# window (O(N) — 164k rows at scale) and sorting by proximity, we scan
# the nearest K units in each direction using the B-tree index on
# (bank_id, fact_type, event_date). This reads only 2×K rows per probe
# regardless of bank size — 120x faster at 164k units (0.6ms vs 74ms).
rows = await ops.fetch_temporal_neighbors(
conn,
mu,
bank_id,
lateral_unit_ids,
lateral_event_dates,
lateral_fact_types,
half_limit,
batch_size=TEMPORAL_LATERAL_BATCH,
)
else:
rows = []
_log(
log_buffer,
f" [7.2] Fetch {len(rows)} candidate neighbors (LATERAL): {time_mod.time() - fetch_neighbors_start:.3f}s",
)
# Build links directly from the LATERAL results (already per-unit limited)
link_gen_start = time_mod.time()
links = []
for row in rows:
time_diff_h = float(row["time_diff_hours"])
weight = max(0.3, 1.0 - (time_diff_h / time_window_hours))
links.append((row["from_id"], str(row["id"]), "temporal", weight, None))
# Also compute temporal links WITHIN the new batch (new units to each other)
if len(new_units) > 1:
# Convert new_units dict to candidate format for within-batch linking
new_unit_items = list(new_units.items())
for i, (unit_id, (event_date, fact_type)) in enumerate(new_unit_items):
if event_date is None:
continue # Skip units without event_date for temporal linking
unit_event_date_norm = _normalize_datetime(event_date)
# Compare with other new units (only those after this one to avoid duplicates)
for j in range(i + 1, len(new_unit_items)):
other_id, (other_event_date, other_fact_type) = new_unit_items[j]
if other_event_date is None:
continue # Skip units without event_date
if fact_type != other_fact_type:
continue # Only link facts of the same type
other_event_date_norm = _normalize_datetime(other_event_date)
# Check if within time window
time_diff_hours = abs((unit_event_date_norm - other_event_date_norm).total_seconds() / 3600)
if time_diff_hours <= time_window_hours:
weight = max(0.3, 1.0 - (time_diff_hours / time_window_hours))
# Create bidirectional links
links.append((unit_id, other_id, "temporal", weight, None))
links.append((other_id, unit_id, "temporal", weight, None))
# Cap temporal links per unit to avoid write amplification;
# retrieval only reads top 10-20 per unit anyway.
links = _cap_links_per_unit(links)
_log(log_buffer, f" [7.3] Generate {len(links)} temporal links: {time_mod.time() - link_gen_start:.3f}s")
if links:
insert_start = time_mod.time()
await _bulk_insert_links(conn, links, bank_id=bank_id, skip_exists_check=True, ops=ops)
_log(log_buffer, f" [7.4] Insert {len(links)} temporal links: {time_mod.time() - insert_start:.3f}s")
return len(links)
except Exception as e:
logger.error(f"Failed to create temporal links: {str(e)}")
import traceback
traceback.print_exc()
raise
async def compute_semantic_links_ann(
conn,
bank_id: str,
unit_ids: list[str],
embeddings: list[list[float]],
fact_types: list[str] | None = None,
top_k: int = 50,
threshold: float = 0.7,
log_buffer: list[str] = None,
) -> list[tuple]:
"""
Phase 1: ANN search for semantic neighbors among existing units.
Runs on a separate connection OUTSIDE the write transaction to avoid
holding locks during expensive HNSW index probes. Uses a temp table +
LATERAL join to batch all probes in a single query.
Queries are split by fact_type so PostgreSQL uses the per-bank partial
HNSW indexes (idx_mu_emb_worl_*, idx_mu_emb_expr_*). Without the
fact_type filter, the planner falls back to sequential scan (~50x slower).
Args:
conn: Database connection (separate from write transaction, autocommit)
bank_id: Bank identifier
unit_ids: Placeholder unit IDs (real IDs not yet created)
embeddings: Embedding vectors for each unit
fact_types: Per-unit fact types (same length as unit_ids). Used to
query only the matching HNSW index per seed.
top_k: Max neighbors per unit
threshold: Minimum cosine similarity
log_buffer: Optional logging buffer
Returns:
List of (from_id, to_id, "semantic", similarity, None) tuples
where from_id uses placeholder IDs.
"""
if not unit_ids or not embeddings:
return []
import time as time_mod
ann_start = time_mod.time()
links = []
logger.debug(f"[ANN] Starting: {len(unit_ids)} seeds, top_k={top_k}")
# Build per-unit fact_types (default to 'world' if not provided)
if fact_types is None:
fact_types = ["world"] * len(unit_ids)
# No exclude_uuids — large exclusion lists (8k+ UUIDs) force PostgreSQL to
# sequential-scan every HNSW probe result against the array, destroying
# performance (67s for 8k seeds). Self-links are harmless (ON CONFLICT DO
# NOTHING handles duplicates in memory_links).
#
# The entire CREATE TEMP TABLE → COPY → SELECT sequence MUST run inside a
# single transaction. Callers may connect through pgBouncer in `transaction`
# pool mode, in which case the backend is only pinned to the client for the
# duration of a transaction. Outside a transaction, pgBouncer can rebind
# the client to a different backend between statements, and the temp table
# (which is session-scoped to its creating backend) becomes invisible.
# The observed failure mode was an intermittent
# `relation "_ann_seeds" does not exist` on the second statement.
#
# Using ON COMMIT DROP + SET LOCAL also means we don't have to remember to
# manually drop the temp table or reset hnsw.ef_search — the transaction
# end handles both.
rows: list = []
async with conn.transaction():
# Transaction-local ef_search. Default 400 is tuned for recall precision
# but at 164k units each HNSW probe takes 94ms. ef_search=60 gives 2.7ms
# per probe (35x faster) with sufficient accuracy for top-50 semantic
# link creation. SET LOCAL auto-reverts at commit, so we don't pollute
# the pool for subsequent recall queries.
await conn.execute("SET LOCAL hnsw.ef_search = 60")
t_setup = time_mod.time()
await conn.execute("CREATE TEMP TABLE _ann_seeds (unit_id text, emb_text text, fact_type text) ON COMMIT DROP")
records = [
(uid, emb if isinstance(emb, str) else str(emb), ft)
for uid, emb, ft in zip(unit_ids, embeddings, fact_types)
]
await conn.copy_records_to_table("_ann_seeds", records=records, columns=["unit_id", "emb_text", "fact_type"])
logger.debug(f"[ANN] Temp table setup: {time_mod.time() - t_setup:.3f}s ({len(records)} seeds)")
# Run one ANN query per fact_type so each uses the right HNSW index.
active_types = set(fact_types)
for fact_type in active_types:
t_query = time_mod.time()
seed_count = sum(1 for ft in fact_types if ft == fact_type)
logger.debug(f"[ANN] Querying fact_type={fact_type}: {seed_count} seeds")
ft_rows = await conn.fetch(
f"""
SELECT s.unit_id AS from_id,
n.id::text AS to_id,
n.similarity
FROM _ann_seeds s
CROSS JOIN LATERAL (
SELECT mu.id,
1 - (mu.embedding <=> s.emb_text::vector) AS similarity
FROM {fq_table("memory_units")} mu
WHERE mu.bank_id = $1
AND mu.fact_type = $2
AND mu.embedding IS NOT NULL
ORDER BY mu.embedding <=> s.emb_text::vector
LIMIT $3
) n
WHERE s.fact_type = $2
""",
bank_id,
fact_type,
top_k,
)
logger.debug(f"[ANN] fact_type={fact_type}: {len(ft_rows)} rows in {time_mod.time() - t_query:.3f}s")
rows.extend(ft_rows)
# Transaction commits here. _ann_seeds is dropped (ON COMMIT DROP).
# hnsw.ef_search reverts (SET LOCAL).
for row in rows:
sim = float(min(1.0, max(0.0, row["similarity"])))
if sim >= threshold:
links.append((row["from_id"], row["to_id"], "semantic", sim, None))
_log(
log_buffer,
f" [8.1] ANN search (Phase 1): {len(unit_ids)} units → {len(links)} links in {time_mod.time() - ann_start:.3f}s",
)
return links
def compute_semantic_links_within_batch(
unit_ids: list[str],
embeddings: list[list[float]],
top_k: int = 50,
threshold: float = 0.7,
) -> list[tuple]:
"""
Compute semantic links between units within the same batch (no DB needed).
Uses numpy dot product on embeddings already in memory — instant.
Args:
unit_ids: Unit IDs (real IDs from insert_facts_batch)
embeddings: Embedding vectors
top_k: Max neighbors per unit
threshold: Minimum cosine similarity
Returns:
List of (from_id, to_id, "semantic", similarity, None) tuples
"""
if len(unit_ids) < 2:
return []
import numpy as np
links = []
new_embeddings_matrix = np.array(embeddings)
for i, unit_id in enumerate(unit_ids):
other_indices = [j for j in range(len(unit_ids)) if j != i]
if not other_indices:
continue
other_embeddings = new_embeddings_matrix[other_indices]
similarities = np.dot(other_embeddings, new_embeddings_matrix[i])
above_threshold = np.where(similarities >= threshold)[0]
if len(above_threshold) > 0:
sorted_local_indices = above_threshold[np.argsort(-similarities[above_threshold])][:top_k]
for local_idx in sorted_local_indices:
other_idx = other_indices[local_idx]
other_id = unit_ids[other_idx]
similarity = float(min(1.0, max(0.0, similarities[local_idx])))
links.append((unit_id, other_id, "semantic", similarity, None))
return links
async def create_semantic_links_batch(
conn,
bank_id: str,
unit_ids: list[str],
embeddings: list[list[float]],
top_k: int = 50,
threshold: float = 0.7,
log_buffer: list[str] = None,
pre_computed_ann_links: list[tuple] | None = None,
ops=None,
) -> int:
"""
Phase 2: Create semantic links (within-batch + pre-computed ANN results).
Within-batch similarities are computed in Python (numpy, instant).
ANN results from Phase 1 are passed in via pre_computed_ann_links and
inserted alongside the within-batch links.
Args:
conn: Database connection (inside write transaction)
bank_id: Bank identifier
unit_ids: Real unit IDs (from insert_facts_batch)
embeddings: Embedding vectors
top_k: Max neighbors per unit
threshold: Minimum cosine similarity
log_buffer: Optional logging buffer
pre_computed_ann_links: ANN results from Phase 1 (already remapped to real IDs)
Returns:
Number of semantic links created
"""
if not unit_ids or not embeddings:
return 0
try:
import time as time_mod
all_links = []
# Within-batch similarities (numpy, no DB)
batch_start = time_mod.time()
within_batch_links = compute_semantic_links_within_batch(unit_ids, embeddings, top_k, threshold)
all_links.extend(within_batch_links)
_log(
log_buffer,
f" [8.1] Within-batch semantic: {len(within_batch_links)} links in {time_mod.time() - batch_start:.3f}s",
)
# Add pre-computed ANN links from Phase 1
if pre_computed_ann_links:
all_links.extend(pre_computed_ann_links)
_log(
log_buffer,
f" [8.2] Pre-computed ANN: {len(pre_computed_ann_links)} links",
)
if all_links:
insert_start = time_mod.time()
await _bulk_insert_links(conn, all_links, bank_id=bank_id, ops=ops)
_log(
log_buffer, f" [8.3] Insert {len(all_links)} semantic links: {time_mod.time() - insert_start:.3f}s"
)
return len(all_links)
except Exception as e:
logger.error(f"Failed to create semantic links: {str(e)}")
import traceback
traceback.print_exc()
raise
async def insert_entity_links_batch(conn, links: list[EntityLink], bank_id: str, chunk_size: int = 5000, ops=None):
"""
Bulk-insert entity links via sorted INSERT FROM unnest().
Args:
conn: Database connection
links: List of EntityLink objects
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
chunk_size: Number of rows per INSERT chunk (default 5000)
"""
if not links:
return
import time as time_mod
total_start = time_mod.time()
tuples = [(link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id) for link in links]
await _bulk_insert_links(conn, tuples, bank_id=bank_id, chunk_size=chunk_size, ops=ops)
logger.debug(
f" [9.TOTAL] Entity links batch insert ({len(tuples)} rows): {time_mod.time() - total_start:.3f}s"
)
async def create_causal_links_batch(
conn,
bank_id: str,
unit_ids: list[str],
causal_relations_per_fact: list[list[dict]],
ops=None,
) -> int:
"""
Create causal links between facts based on LLM-extracted causal relationships.
Args:
conn: Database connection
unit_ids: List of unit IDs (in same order as causal_relations_per_fact)
causal_relations_per_fact: List of causal relations for each fact.
Each element is a list of dicts with:
- target_fact_index: Index into unit_ids for the target fact
- relation_type: "caused_by"
Returns:
Number of causal links created
Causal link type:
- "caused_by": This fact was caused by the target fact
"""
if not unit_ids or not causal_relations_per_fact:
return 0
try:
import time as time_mod
create_start = time_mod.time()
# Build links list
links = []
for fact_idx, causal_relations in enumerate(causal_relations_per_fact):
if not causal_relations:
continue
from_unit_id = unit_ids[fact_idx]
for relation in causal_relations:
target_idx = relation["target_fact_index"]
relation_type = relation["relation_type"]
# Validate relation_type - only "caused_by" is supported (DB constraint)
valid_types = {"caused_by"}
if relation_type not in valid_types:
logger.error(
f"Invalid relation_type '{relation_type}' (type: {type(relation_type).__name__}) "
f"from fact {fact_idx}. Must be one of: {valid_types}. "
f"Relation data: {relation}"
)
continue
# Validate target index
if target_idx < 0 or target_idx >= len(unit_ids):
logger.warning(f"Invalid target_fact_index {target_idx} in causal relation from fact {fact_idx}")
continue
to_unit_id = unit_ids[target_idx]
# Don't create self-links
if from_unit_id == to_unit_id:
continue
links.append((from_unit_id, to_unit_id, relation_type, 1.0, None))
if links:
insert_start = time_mod.time()
await _bulk_insert_links(conn, links, bank_id=bank_id, skip_exists_check=True, ops=ops)
logger.debug(f" [10.1] Insert {len(links)} causal links: {time_mod.time() - insert_start:.3f}s")
return len(links)
except Exception as e:
logger.error(f"Failed to create causal links: {str(e)}")
import traceback
traceback.print_exc()
raise
File diff suppressed because it is too large Load Diff
@@ -1,41 +0,0 @@
"""
Centralized schema-qualified table name helpers.
Single source of truth for producing ``"schema".table_name`` references
that respect both the active schema context and the database backend.
"""
from ..config import get_config
def _is_oracle() -> bool:
"""Return True when the configured database backend is Oracle."""
return get_config().database_backend == "oracle"
def fq_table(table_name: str) -> str:
"""Get fully-qualified table name using the current schema context.
On Oracle the schema is set at the session level (``ALTER SESSION SET
CURRENT_SCHEMA``), so we return the bare table name. On PostgreSQL
we prefix with the schema from :func:`memory_engine.get_current_schema`.
"""
if _is_oracle():
return table_name
from .memory_engine import get_current_schema
return f"{get_current_schema()}.{table_name}"
def fq_table_explicit(table: str, schema: str | None = None) -> str:
"""Get fully-qualified table name with an explicit schema override.
Used by modules that don't rely on the context-variable schema
(e.g. task_backend, worker poller) and instead pass the schema
explicitly.
"""
if _is_oracle():
return table
if schema:
return f'"{schema}".{table}'
return table
@@ -1,70 +0,0 @@
"""
Graph retrieval strategies for memory recall.
This module provides an abstraction for graph-based memory retrieval,
allowing different algorithms to be swapped without changing the rest
of the recall pipeline.
"""
import logging
from abc import ABC, abstractmethod
from datetime import datetime
from .tags import TagGroup, TagsMatch
from .types import GraphRetrievalTimings, RetrievalResult
logger = logging.getLogger(__name__)
class GraphRetriever(ABC):
"""
Abstract base class for graph-based memory retrieval.
Implementations traverse the memory graph (entity links, temporal links,
causal links) to find relevant facts that might not be found by
semantic or keyword search alone.
"""
@property
@abstractmethod
def name(self) -> str:
"""Return identifier for this retrieval strategy (e.g., 'link_expansion')."""
pass
@abstractmethod
async def retrieve(
self,
pool,
query_embedding_str: str,
bank_id: str,
fact_type: str,
budget: int,
query_text: str | None = None,
semantic_seeds: list[RetrievalResult] | None = None,
temporal_seeds: list[RetrievalResult] | None = None,
adjacency=None, # TypedAdjacency, optional pre-loaded graph
tags: list[str] | None = None, # Visibility scope tags for filtering
tags_match: TagsMatch = "any", # How to match tags: 'any' (OR) or 'all' (AND)
tag_groups: list[TagGroup] | None = None, # Compound boolean tag filter groups
created_after: datetime | None = None, # Only include memory_units created after this time
created_before: datetime | None = None, # Only include memory_units created before this time
) -> tuple[list[RetrievalResult], GraphRetrievalTimings | None]:
"""
Retrieve relevant facts via graph traversal.
Args:
pool: Database connection pool
query_embedding_str: Query embedding as string (for finding entry points)
bank_id: Memory bank identifier
fact_type: Fact type to filter ('world', 'experience', 'observation')
budget: Maximum number of nodes to explore/return
query_text: Original query text (optional, for some strategies)
semantic_seeds: Pre-computed semantic entry points (from semantic retrieval)
temporal_seeds: Pre-computed temporal entry points (from temporal retrieval)
adjacency: Pre-loaded typed adjacency graph (optional)
tags: Optional list of tags for visibility filtering (OR matching)
Returns:
Tuple of (List of RetrievalResult with activation scores, optional timing info)
"""
pass

Some files were not shown because too many files have changed in this diff Show More