Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0ac7a4c957 | ||
|
|
cc374efbe5 | ||
|
|
528a15db00 | ||
|
|
a8926d5749 |
@@ -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"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -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.
|
||||
+2
-44
@@ -2,7 +2,7 @@
|
||||
# Copy this file to .env and fill in your values
|
||||
|
||||
# LLM Configuration (Required)
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, zai, volcano
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
@@ -20,21 +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: z.ai configuration (Zhipu GLM series, https://z.ai)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=zai
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-zai-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=glm-4.5-flash # or glm-4.5-air for the paid tier
|
||||
|
||||
# Example: LM Studio local configuration (Qwen 2.5 32B recommended)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=lmstudio
|
||||
# HINDSIGHT_API_LLM_API_KEY=lmstudio
|
||||
@@ -54,8 +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_READ_DATABASE_URL= # Optional read-replica URL. When set, recall queries (semantic, BM25, graph, temporal) flow through a separate pool against this URL, offloading the primary. Typically points to a read-only endpoint (CNPG's <cluster>-ro service or Aurora reader endpoint).
|
||||
# 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)
|
||||
@@ -65,25 +48,12 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale # Auto-detects pg_diskann on Azure
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default), "tei", "openai", "cohere", "google", "openrouter", "litellm", or "litellm-sdk"
|
||||
# Provider: "local" (default) or "tei" (HuggingFace Text Embeddings Inference)
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
# For local provider:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
|
||||
# Optional for China network / restricted HF access:
|
||||
# HF_ENDPOINT=https://hf-mirror.com
|
||||
# For TEI provider:
|
||||
# HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://localhost:8080
|
||||
# For OpenAI-compatible embeddings:
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=sk-xxxx
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL=text-embedding-3-small
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL=https://api.openai.com/v1
|
||||
#
|
||||
# IMPORTANT: Embedding keys require provider-specific names:
|
||||
# HINDSIGHT_API_EMBEDDINGS_{PROVIDER}_{PARAMETER}
|
||||
# (for example, HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL).
|
||||
#
|
||||
# DeepSeek note: DeepSeek is supported for LLM calls, but not for embeddings.
|
||||
# If using DeepSeek as LLM provider, keep embeddings on local/openai/cohere/google/etc.
|
||||
|
||||
# Reranker Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default) or "tei" (HuggingFace Text Embeddings Inference)
|
||||
@@ -107,15 +77,3 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Custom service name and environment (optional, defaults: hindsight-api, development)
|
||||
# HINDSIGHT_API_OTEL_SERVICE_NAME=hindsight-production
|
||||
# HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT=production
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Control Plane (Optional)
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# Dataplane API URL - where the CP proxies requests to
|
||||
# HINDSIGHT_CP_DATAPLANE_API_URL=http://localhost:8888
|
||||
|
||||
# Optional: Require a shared access key to view the Control Plane UI.
|
||||
# When set, visitors see a login page and must enter the key before
|
||||
# accessing the dashboard or any /api/* routes (except /api/health).
|
||||
# HINDSIGHT_CP_ACCESS_KEY=your-shared-secret-key
|
||||
|
||||
@@ -1,6 +0,0 @@
|
||||
version: 2
|
||||
updates:
|
||||
- package-ecosystem: "github-actions"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
@@ -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
|
||||
|
||||
@@ -1,200 +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_conversations:
|
||||
description: "LoComo conversation IDs (space-separated). Blank = curated set (conv-26 conv-30 conv-43)."
|
||||
type: string
|
||||
default: ""
|
||||
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: Run perf-test
|
||||
run: |
|
||||
SUITE_ARG=""
|
||||
if [ -n "${{ inputs.suite }}" ]; then
|
||||
SUITE_ARG="--suite ${{ inputs.suite }}"
|
||||
fi
|
||||
./scripts/benchmarks/run-perf-test.sh \
|
||||
--scale ${{ inputs.scale || 'large' }} \
|
||||
$SUITE_ARG \
|
||||
--output perf-results.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
|
||||
|
||||
# Publish enriched results (perf JSON + commit metadata) to the dashboard
|
||||
# repo's gh-pages branch. The static site at
|
||||
# https://vectorize-io.github.io/hindsight-continuous-performance-monitor/
|
||||
# reads data/index.json + data/<run>.json and renders charts client-side.
|
||||
- name: Publish to dashboard
|
||||
if: github.event_name == 'schedule' || github.event_name == 'push' || github.event_name == 'workflow_dispatch'
|
||||
env:
|
||||
PERF_DASHBOARD_TOKEN: ${{ secrets.PERF_DASHBOARD_TOKEN }}
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: ./scripts/benchmarks/publish-perf-results.sh hindsight-dev/perf-results.json
|
||||
|
||||
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
|
||||
# Curated 3-conversation subset (best/middle/worst by accuracy on the
|
||||
# last successful full run): conv-26 (best), conv-30 (middle), conv-43
|
||||
# (worst). Excludes conv-44, the bank with the largest unconsolidated
|
||||
# set that has been pushing scheduled runs over the per-bank
|
||||
# _wait_for_consolidation timeout. Override via workflow_dispatch with
|
||||
# the locomo_conversations input.
|
||||
run: |
|
||||
CONVERSATIONS="${{ inputs.locomo_conversations }}"
|
||||
if [ -z "$CONVERSATIONS" ]; then
|
||||
CONVERSATIONS="conv-26 conv-30 conv-43"
|
||||
fi
|
||||
uv run python hindsight-dev/benchmarks/locomo/locomo_benchmark.py \
|
||||
--wait-consolidation \
|
||||
--conversation $CONVERSATIONS
|
||||
|
||||
- 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
|
||||
|
||||
- name: Publish LoComo to dashboard
|
||||
if: success() && (github.event_name == 'schedule' || github.event_name == 'push' || github.event_name == 'workflow_dispatch')
|
||||
env:
|
||||
PERF_DASHBOARD_TOKEN: ${{ secrets.PERF_DASHBOARD_TOKEN }}
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: ./scripts/benchmarks/publish-locomo-results.sh hindsight-dev/benchmarks/locomo/results/benchmark_results.json
|
||||
@@ -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 --provenance 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 }}
|
||||
@@ -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
-109
@@ -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
|
||||
@@ -284,7 +393,7 @@ jobs:
|
||||
asset_name: hindsight-linux-arm64
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
@@ -302,7 +411,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 }}
|
||||
@@ -314,7 +423,6 @@ jobs:
|
||||
permissions:
|
||||
contents: read
|
||||
packages: write
|
||||
id-token: write
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
@@ -344,7 +452,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
|
||||
@@ -358,13 +466,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 }}
|
||||
@@ -376,7 +484,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: |
|
||||
@@ -392,7 +500,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
|
||||
@@ -411,8 +519,7 @@ jobs:
|
||||
|
||||
# Build multi-platform and push to release tags
|
||||
- name: Build and push release images
|
||||
id: build
|
||||
uses: docker/build-push-action@v7
|
||||
uses: docker/build-push-action@v6
|
||||
with:
|
||||
context: .
|
||||
file: docker/standalone/Dockerfile
|
||||
@@ -423,31 +530,6 @@ jobs:
|
||||
tags: ${{ steps.meta.outputs.tags }}
|
||||
labels: ${{ steps.meta.outputs.labels }}
|
||||
|
||||
- name: Install cosign
|
||||
uses: sigstore/cosign-installer@v3
|
||||
|
||||
- name: Sign published images
|
||||
env:
|
||||
TAGS: ${{ steps.meta.outputs.tags }}
|
||||
DIGEST: ${{ steps.build.outputs.digest }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
refs=()
|
||||
while IFS= read -r tag; do
|
||||
[[ -z "${tag}" ]] && continue
|
||||
refs+=("${tag}@${DIGEST}")
|
||||
done <<< "${TAGS}"
|
||||
cosign sign --yes "${refs[@]}"
|
||||
|
||||
- name: Verify signature on primary tag
|
||||
env:
|
||||
IMAGE: ghcr.io/${{ github.repository_owner }}/${{ matrix.image_name }}
|
||||
DIGEST: ${{ steps.build.outputs.digest }}
|
||||
run: |
|
||||
cosign verify "${IMAGE}@${DIGEST}" \
|
||||
--certificate-identity-regexp "^https://github\.com/${{ github.repository }}/\.github/workflows/release\.yml@.*" \
|
||||
--certificate-oidc-issuer https://token.actions.githubusercontent.com
|
||||
|
||||
release-helm-chart:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
@@ -455,10 +537,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'
|
||||
|
||||
@@ -475,7 +557,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
|
||||
@@ -483,67 +565,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
|
||||
@@ -553,20 +641,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
|
||||
@@ -574,7 +665,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
|
||||
|
||||
@@ -1,71 +0,0 @@
|
||||
name: Sign published images
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
version:
|
||||
description: 'Version to sign (without leading v, e.g. 0.6.0)'
|
||||
required: true
|
||||
type: string
|
||||
default: '0.6.0'
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
packages: write
|
||||
id-token: write
|
||||
|
||||
jobs:
|
||||
sign:
|
||||
name: Sign ${{ matrix.image }}:${{ inputs.version }}${{ matrix.suffix }}
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
include:
|
||||
- { image: hindsight-api, suffix: '' }
|
||||
- { image: hindsight-api, suffix: '-slim' }
|
||||
- { image: hindsight-control-plane, suffix: '' }
|
||||
- { image: hindsight, suffix: '' }
|
||||
- { image: hindsight, suffix: '-slim' }
|
||||
|
||||
steps:
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v4
|
||||
|
||||
- name: Log in to GHCR
|
||||
uses: docker/login-action@v4
|
||||
with:
|
||||
registry: ghcr.io
|
||||
username: ${{ github.actor }}
|
||||
password: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
- name: Install cosign
|
||||
uses: sigstore/cosign-installer@v3
|
||||
|
||||
- name: Resolve image digest
|
||||
id: resolve
|
||||
env:
|
||||
IMAGE: ghcr.io/${{ github.repository_owner }}/${{ matrix.image }}
|
||||
TAG: ${{ inputs.version }}${{ matrix.suffix }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
DIGEST=$(docker buildx imagetools inspect "${IMAGE}:${TAG}" --format '{{json .Manifest.Digest}}' | tr -d '"')
|
||||
if [[ -z "${DIGEST}" || "${DIGEST}" != sha256:* ]]; then
|
||||
echo "Failed to resolve digest for ${IMAGE}:${TAG} (got: ${DIGEST})" >&2
|
||||
exit 1
|
||||
fi
|
||||
echo "Resolved ${IMAGE}:${TAG} -> ${DIGEST}"
|
||||
echo "ref=${IMAGE}@${DIGEST}" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Sign image
|
||||
env:
|
||||
REF: ${{ steps.resolve.outputs.ref }}
|
||||
run: cosign sign --yes "${REF}"
|
||||
|
||||
- name: Verify signature
|
||||
env:
|
||||
REF: ${{ steps.resolve.outputs.ref }}
|
||||
run: |
|
||||
cosign verify "${REF}" \
|
||||
--certificate-identity-regexp "^https://github\.com/${{ github.repository }}/\.github/workflows/sign-images\.yml@.*" \
|
||||
--certificate-oidc-issuer https://token.actions.githubusercontent.com
|
||||
+180
-2290
File diff suppressed because it is too large
Load Diff
+1
-2
@@ -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
|
||||
|
||||
@@ -1,7 +0,0 @@
|
||||
{
|
||||
"semi": true,
|
||||
"singleQuote": false,
|
||||
"tabWidth": 2,
|
||||
"trailingComma": "es5",
|
||||
"printWidth": 100
|
||||
}
|
||||
@@ -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,58 +133,25 @@ 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
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -7,12 +7,10 @@
|
||||
[](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml)
|
||||
[](https://join.slack.com/t/hindsight-space/shared_invite/zt-3nhbm4w29-LeSJ5Ixi6j8PdiYOCPlOgg)
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
[](https://gitcgr.com/vectorize-io/hindsight)
|
||||

|
||||

|
||||
<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>
|
||||
|
||||
---
|
||||
@@ -30,7 +28,7 @@ It eliminates the shortcomings of alternative techniques such as RAG and knowled
|
||||
|
||||
Hindsight is the most accurate agent memory system ever tested according to benchmark performance. It has achieved state-of-the-art performance on the LongMemEval benchmark, widely used to assess memory system performance across a variety of conversational AI scenarios. The current reported performance of Hindsight and other agent memory solutions as of January 2026 is shown here:
|
||||
|
||||

|
||||

|
||||
|
||||
The benchmark performance data for Hindsight has been independently reproduced by research collaborators at the Virginia Tech [Sanghani Center for Artificial Intelligence and Data Analytics](https://sanghani.cs.vt.edu/) and The Washington Post. Other scores are self-reported by software vendors.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,90 +0,0 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with AlloyDB Omni and ScaNN
|
||||
# Uses Google's free AlloyDB Omni container image: https://hub.docker.com/r/google/alloydbomni
|
||||
#
|
||||
# Usage:
|
||||
# docker compose -f docker/docker-compose/alloydb/docker-compose.yaml up -d
|
||||
#
|
||||
# Make sure to set the required environment variables before running:
|
||||
# - HINDSIGHT_DB_PASSWORD: password for the AlloyDB Omni/PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see the hindsight service below)
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_VERSION: AlloyDB Omni image tag (default: 17)
|
||||
# - HINDSIGHT_DB_USER: database user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: database name (default: hindsight_db)
|
||||
|
||||
services:
|
||||
db:
|
||||
image: google/alloydbomni:${HINDSIGHT_DB_VERSION:-17}
|
||||
container_name: hindsight-db-alloydb
|
||||
restart: always
|
||||
ports:
|
||||
- "5438:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- alloydb_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
alloydb-init:
|
||||
image: google/alloydbomni:${HINDSIGHT_DB_VERSION:-17}
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command:
|
||||
- bash
|
||||
- -c
|
||||
- |
|
||||
echo 'Waiting for AlloyDB Omni to be ready...'
|
||||
until pg_isready -h hindsight-db-alloydb -p 5432 -U ${HINDSIGHT_DB_USER:-hindsight_user}; do
|
||||
echo 'AlloyDB Omni is unavailable - sleeping'
|
||||
sleep 2
|
||||
done
|
||||
echo 'AlloyDB Omni is ready - creating ${HINDSIGHT_DB_NAME:-hindsight_db} database'
|
||||
psql -h hindsight-db-alloydb -p 5432 -U ${HINDSIGHT_DB_USER:-hindsight_user} -c 'CREATE DATABASE ${HINDSIGHT_DB_NAME:-hindsight_db};' 2>/dev/null || echo 'Database already exists'
|
||||
echo 'Creating vector and alloydb_scann extensions'
|
||||
psql -h hindsight-db-alloydb -p 5432 -U ${HINDSIGHT_DB_USER:-hindsight_user} -d ${HINDSIGHT_DB_NAME:-hindsight_db} -c 'CREATE EXTENSION IF NOT EXISTS vector;'
|
||||
psql -h hindsight-db-alloydb -p 5432 -U ${HINDSIGHT_DB_USER:-hindsight_user} -d ${HINDSIGHT_DB_NAME:-hindsight_db} -c 'CREATE EXTENSION IF NOT EXISTS alloydb_scann CASCADE;'
|
||||
echo 'Database and extensions created successfully'
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: scann
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: native
|
||||
|
||||
depends_on:
|
||||
db:
|
||||
condition: service_started
|
||||
alloydb-init:
|
||||
condition: service_completed_successfully
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
alloydb_data:
|
||||
@@ -1,34 +0,0 @@
|
||||
# Example: custom Hindsight image with non-default local models baked in.
|
||||
#
|
||||
# Use this pattern in production when you run a non-default embedder or
|
||||
# reranker. Baking models into the image removes the runtime dependency on
|
||||
# HuggingFace and lets the container registry handle caching per node, so
|
||||
# you don't need a model-cache PVC.
|
||||
#
|
||||
# Built on top of the slim image so only the deps and models you actually
|
||||
# use end up in the final image.
|
||||
FROM ghcr.io/vectorize-io/hindsight:latest-slim
|
||||
|
||||
# Install the local-ml deps required to load sentence-transformers /
|
||||
# cross-encoder models at runtime. Pinned ranges mirror hindsight-api-slim's
|
||||
# `local-ml` extra in hindsight-api-slim/pyproject.toml. Use `uv pip
|
||||
# install` against the image's venv explicitly: the slim image's venv was
|
||||
# created by `uv sync` and does not ship its own `pip`, so a bare
|
||||
# `pip install` would fall back to user site-packages and not be visible
|
||||
# to the runtime python.
|
||||
RUN uv pip install --python /app/api/.venv/bin/python --no-cache \
|
||||
'sentence-transformers>=3.3.0' \
|
||||
'transformers>=4.53.0' \
|
||||
'torch>=2.6.0'
|
||||
|
||||
# Pre-download the models you want to use. Replace these with your own.
|
||||
# The defaults bundled in the full image are BAAI/bge-small-en-v1.5 and
|
||||
# cross-encoder/ms-marco-MiniLM-L-6-v2; here we pick multilingual variants
|
||||
# as a concrete non-default example.
|
||||
ARG EMBEDDER=sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2
|
||||
ARG RERANKER=cross-encoder/mmarco-mMiniLMv2-L12-H384-v1
|
||||
ENV HF_HUB_DOWNLOAD_TIMEOUT=600
|
||||
RUN python -c "\
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder; \
|
||||
SentenceTransformer('${EMBEDDER}'); \
|
||||
CrossEncoder('${RERANKER}')"
|
||||
@@ -1,81 +0,0 @@
|
||||
# Hindsight with Custom Local Models
|
||||
|
||||
Example Docker Compose setup that builds a Hindsight image with **non-default
|
||||
local embedder and reranker models baked in at build time**.
|
||||
|
||||
This is the recommended pattern for production when you use a non-default
|
||||
local model: the container registry caches model layers per node, pod
|
||||
startup is deterministic, and you don't need a model-cache PVC (or any
|
||||
runtime dependency on HuggingFace).
|
||||
|
||||
## When to use this
|
||||
|
||||
- You override `HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL` or
|
||||
`HINDSIGHT_API_RERANKER_LOCAL_MODEL` to a non-default model.
|
||||
- You want pod startup to be deterministic and offline-capable.
|
||||
- You'd otherwise reach for a Helm `modelCache` PVC just to avoid
|
||||
re-downloading models.
|
||||
|
||||
If you're using the **default** local models, the published full image
|
||||
(`ghcr.io/vectorize-io/hindsight:latest`) already bakes them in — you don't
|
||||
need this example.
|
||||
|
||||
If you're using **external** providers (TEI, OpenAI, Cohere, ...) for
|
||||
embeddings and reranking, use the slim image directly — no models are
|
||||
needed in the image.
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
|
||||
docker compose -f docker/docker-compose/custom-models/docker-compose.yaml up --build
|
||||
```
|
||||
|
||||
- API: http://localhost:8888
|
||||
- Control Plane: http://localhost:9999
|
||||
|
||||
## Using your own models
|
||||
|
||||
Override the build args to bake different models:
|
||||
|
||||
```bash
|
||||
docker compose -f docker/docker-compose/custom-models/docker-compose.yaml build \
|
||||
--build-arg EMBEDDER=your-org/your-embedder \
|
||||
--build-arg RERANKER=your-org/your-reranker
|
||||
```
|
||||
|
||||
Then update the matching `HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL` and
|
||||
`HINDSIGHT_API_RERANKER_LOCAL_MODEL` values in `docker-compose.yaml` so the
|
||||
runtime points at the same model IDs.
|
||||
|
||||
## Verifying the models are baked in
|
||||
|
||||
`docker-compose.yaml` sets `HF_HUB_OFFLINE=1` and `TRANSFORMERS_OFFLINE=1`
|
||||
so that any attempt to download a model at runtime fails loudly instead of
|
||||
silently re-downloading. If the container starts and serves recall queries
|
||||
with these set, the models are correctly baked in.
|
||||
|
||||
You can also inspect the image directly:
|
||||
|
||||
```bash
|
||||
docker run --rm --entrypoint sh hindsight-custom-models-hindsight \
|
||||
-c 'ls ~/.cache/huggingface/hub/'
|
||||
```
|
||||
|
||||
## Why not a model-cache PVC?
|
||||
|
||||
The Helm chart exposes an optional `api.persistence.modelCache` PVC for
|
||||
caching downloaded models across pod restarts. Compared to baking models
|
||||
into the image:
|
||||
|
||||
- A PVC adds storage cost — one PVC per worker replica with
|
||||
`volumeClaimTemplates`.
|
||||
- `ReadWriteOnce` (the default) pins pods to a node.
|
||||
- The PVC needs lifecycle management on `helm uninstall` / `helm upgrade`
|
||||
— without `helm.sh/resource-policy: keep` it is deleted on uninstall;
|
||||
with it, storage keeps billing forever until manually cleaned up.
|
||||
- Pod startup still depends on HuggingFace being reachable on first run.
|
||||
|
||||
Image layers, by contrast, are pulled once per node and cached for free by
|
||||
the container runtime, with no orphaned-storage cleanup story.
|
||||
@@ -1,44 +0,0 @@
|
||||
name: hindsight-custom-models
|
||||
# Example: run a custom Hindsight image with non-default local models baked
|
||||
# in at build time, so pod startup does not depend on HuggingFace at runtime.
|
||||
#
|
||||
# Quick start:
|
||||
# export OPENAI_API_KEY=sk-xxx
|
||||
# docker compose -f docker/docker-compose/custom-models/docker-compose.yaml up --build
|
||||
#
|
||||
# Required environment variables:
|
||||
# - OPENAI_API_KEY (or configure another LLM provider via HINDSIGHT_API_LLM_*)
|
||||
|
||||
services:
|
||||
hindsight:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
# Override at build time to bake different models:
|
||||
# docker compose build --build-arg EMBEDDER=your-org/your-embedder
|
||||
args:
|
||||
EMBEDDER: sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2
|
||||
RERANKER: cross-encoder/mmarco-mMiniLMv2-L12-H384-v1
|
||||
container_name: hindsight-custom-models
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Point Hindsight at the models baked into the image above.
|
||||
HINDSIGHT_API_EMBEDDINGS_PROVIDER: local
|
||||
HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL: sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2
|
||||
HINDSIGHT_API_RERANKER_PROVIDER: local
|
||||
HINDSIGHT_API_RERANKER_LOCAL_MODEL: cross-encoder/mmarco-mMiniLMv2-L12-H384-v1
|
||||
|
||||
# Fail fast if a model is missing from the image instead of silently
|
||||
# falling back to a HuggingFace download at runtime.
|
||||
HF_HUB_OFFLINE: "1"
|
||||
TRANSFORMERS_OFFLINE: "1"
|
||||
volumes:
|
||||
- pg_data:/home/hindsight/.pg0
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -51,8 +51,6 @@ services:
|
||||
|
||||
# Control Plane config
|
||||
HINDSIGHT_CP_DATAPLANE_API_URL: http://localhost:8888
|
||||
# Optional: Require a shared access key for Control Plane UI access
|
||||
# HINDSIGHT_CP_ACCESS_KEY: your-secret-key
|
||||
volumes:
|
||||
# Persist embedded pg0 database
|
||||
- hindsight_data:/app/data
|
||||
|
||||
@@ -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 .
|
||||
@@ -172,10 +175,6 @@ USER hindsight
|
||||
# "Permission denied" error when mounting a fresh root-owned volume.
|
||||
RUN mkdir -p /home/hindsight/.pg0
|
||||
|
||||
# Make /home/hindsight traversable when running with --user UID:GID overrides
|
||||
# (default 0700 blocks traversal by non-owner UIDs needed for bind-mount ownership matching)
|
||||
RUN chmod 755 /home/hindsight
|
||||
|
||||
ENV PATH="/app/api/.venv/bin:${PATH}"
|
||||
|
||||
# Pre-download tiktoken encoding (ALWAYS - required for token counting even in air-gapped envs)
|
||||
@@ -332,10 +331,6 @@ USER hindsight
|
||||
# "Permission denied" error when mounting a fresh root-owned volume.
|
||||
RUN mkdir -p /home/hindsight/.pg0
|
||||
|
||||
# Make /home/hindsight traversable when running with --user UID:GID overrides
|
||||
# (default 0700 blocks traversal by non-owner UIDs needed for bind-mount ownership matching)
|
||||
RUN chmod 755 /home/hindsight
|
||||
|
||||
ENV PATH="/app/api/.venv/bin:${PATH}"
|
||||
|
||||
# Pre-download tiktoken encoding (ALWAYS - required for token counting even in air-gapped envs)
|
||||
|
||||
@@ -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,63 +71,6 @@ 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=()
|
||||
|
||||
@@ -190,7 +111,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 +137,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 $?
|
||||
|
||||
@@ -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,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.6.2
|
||||
appVersion: "0.6.2"
|
||||
version: 0.4.16
|
||||
appVersion: "0.4.16"
|
||||
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 }}
|
||||
|
||||
@@ -67,39 +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.
|
||||
#
|
||||
# For production, prefer baking models into a custom image instead of
|
||||
# enabling this PVC: image layers are pulled once per node and cached
|
||||
# for free, while a PVC adds storage cost, pins pods to a node
|
||||
# (ReadWriteOnce), and needs lifecycle management on uninstall/upgrade.
|
||||
# See docs: developer/installation#bundling-custom-models-in-a-custom-image
|
||||
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"
|
||||
@@ -173,35 +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 — one PVC per
|
||||
# replica. For production, prefer baking models into a custom image; see
|
||||
# api.persistence.modelCache above and docs:
|
||||
# developer/installation#bundling-custom-models-in-a-custom-image
|
||||
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: {}
|
||||
|
||||
|
||||
@@ -1,4 +0,0 @@
|
||||
node_modules
|
||||
dist
|
||||
*.tgz
|
||||
.DS_Store
|
||||
@@ -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
|
||||
@@ -1,57 +0,0 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.6.2",
|
||||
"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"
|
||||
}
|
||||
}
|
||||
@@ -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",
|
||||
]);
|
||||
});
|
||||
});
|
||||
@@ -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}`];
|
||||
}
|
||||
@@ -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";
|
||||
@@ -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),
|
||||
};
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -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}`
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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"]
|
||||
}
|
||||
@@ -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,
|
||||
});
|
||||
@@ -1,8 +0,0 @@
|
||||
import { defineConfig } from "vitest/config";
|
||||
|
||||
export default defineConfig({
|
||||
test: {
|
||||
include: ["src/**/*.test.ts"],
|
||||
environment: "node",
|
||||
},
|
||||
});
|
||||
@@ -1,33 +0,0 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=61"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.6.2"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim==0.6.2",
|
||||
"hindsight-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"
|
||||
@@ -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
|
||||
```
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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,170 +0,0 @@
|
||||
"""Shared PostgreSQL vector-extension dispatch helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.engine import Connection
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Extensions a user can set via HINDSIGHT_API_VECTOR_EXTENSION.
|
||||
CONFIGURABLE_EXTENSIONS = ("pgvector", "pgvectorscale", "vchord", "scann")
|
||||
|
||||
# Extensions detect_vector_extension() can return. pg_diskann is a runtime-only
|
||||
# resolution from a configured "pgvectorscale" backend on Azure (uses a different
|
||||
# WITH clause), never a value the user sets directly.
|
||||
RESOLVED_EXTENSIONS = (*CONFIGURABLE_EXTENSIONS, "pg_diskann")
|
||||
|
||||
# Backwards-compatible alias for older imports.
|
||||
VALID_EXTENSIONS = CONFIGURABLE_EXTENSIONS
|
||||
|
||||
SCANN_MIN_ROWS_FOR_AUTO_INDEX = 10_000
|
||||
|
||||
|
||||
_EXTENSION_NAMES = {
|
||||
"pgvector": "vector",
|
||||
"pgvectorscale": "vectorscale",
|
||||
"vchord": "vchord",
|
||||
"scann": "alloydb_scann",
|
||||
}
|
||||
|
||||
_INDEX_USING_CLAUSES = {
|
||||
"pgvector": "USING hnsw (embedding vector_cosine_ops)",
|
||||
"pgvectorscale": "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)",
|
||||
"pg_diskann": "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)",
|
||||
"vchord": "USING vchordrq (embedding vector_l2_ops)",
|
||||
"scann": "USING scann (embedding cosine) WITH (mode = 'AUTO')",
|
||||
}
|
||||
|
||||
_INDEX_TYPE_KEYWORDS = {
|
||||
"pgvector": "hnsw",
|
||||
"pgvectorscale": "diskann",
|
||||
"pg_diskann": "diskann",
|
||||
"vchord": "vchordrq",
|
||||
"scann": "scann",
|
||||
}
|
||||
|
||||
_EXTENSION_INSTALL_SQL = {
|
||||
"pgvector": ("CREATE EXTENSION IF NOT EXISTS vector",),
|
||||
"pgvectorscale": (
|
||||
"CREATE EXTENSION IF NOT EXISTS vector",
|
||||
"CREATE EXTENSION IF NOT EXISTS vectorscale CASCADE",
|
||||
),
|
||||
"vchord": ("CREATE EXTENSION IF NOT EXISTS vchord CASCADE",),
|
||||
"scann": (
|
||||
"CREATE EXTENSION IF NOT EXISTS vector",
|
||||
"CREATE EXTENSION IF NOT EXISTS alloydb_scann CASCADE",
|
||||
),
|
||||
}
|
||||
|
||||
_INSTALL_HINTS = {
|
||||
"pgvector": "CREATE EXTENSION vector;",
|
||||
"pgvectorscale": "CREATE EXTENSION vector; then CREATE EXTENSION vectorscale CASCADE; (or pg_diskann on Azure)",
|
||||
"vchord": "CREATE EXTENSION vchord CASCADE;",
|
||||
"scann": "CREATE EXTENSION vector; then CREATE EXTENSION alloydb_scann CASCADE;",
|
||||
}
|
||||
|
||||
|
||||
def validate_extension(name: str) -> str:
|
||||
"""Return a normalized configurable vector extension name or raise.
|
||||
|
||||
Used at the user-facing config boundary; pg_diskann is rejected here because
|
||||
it is a detection-time alias, never a value the user sets directly.
|
||||
"""
|
||||
ext = name.lower()
|
||||
if ext not in CONFIGURABLE_EXTENSIONS:
|
||||
valid = ", ".join(CONFIGURABLE_EXTENSIONS)
|
||||
raise ValueError(f"Invalid vector_extension: {name}. Must be one of: {valid}")
|
||||
return ext
|
||||
|
||||
|
||||
def _normalize_resolved(name: str) -> str:
|
||||
"""Normalize either a user-configurable or detect-time extension name."""
|
||||
ext = name.lower()
|
||||
if ext not in RESOLVED_EXTENSIONS:
|
||||
valid = ", ".join(RESOLVED_EXTENSIONS)
|
||||
raise ValueError(f"Unknown vector extension: {name}. Must be one of: {valid}")
|
||||
return ext
|
||||
|
||||
|
||||
def pg_extension_name(ext: str) -> str:
|
||||
"""Return the PostgreSQL extension name for a configured vector backend."""
|
||||
return _EXTENSION_NAMES[validate_extension(ext)]
|
||||
|
||||
|
||||
def index_using_clause(ext: str) -> str:
|
||||
"""Return the CREATE INDEX USING clause for the vector backend."""
|
||||
return _INDEX_USING_CLAUSES[_normalize_resolved(ext)]
|
||||
|
||||
|
||||
def index_type_keyword(ext: str) -> str:
|
||||
"""Return the keyword that identifies this index type in pg_indexes.indexdef."""
|
||||
return _INDEX_TYPE_KEYWORDS[_normalize_resolved(ext)]
|
||||
|
||||
|
||||
def minimum_rows_for_index(ext: str) -> int:
|
||||
"""Return the minimum populated embedding rows before creating this index type."""
|
||||
return SCANN_MIN_ROWS_FOR_AUTO_INDEX if _normalize_resolved(ext) == "scann" else 0
|
||||
|
||||
|
||||
def should_defer_index_creation(ext: str, row_count: int) -> bool:
|
||||
"""Return True when index creation should wait for more embeddings."""
|
||||
minimum_rows = minimum_rows_for_index(ext)
|
||||
return minimum_rows > 0 and row_count < minimum_rows
|
||||
|
||||
|
||||
def uses_per_bank_vector_indexes(ext: str) -> bool:
|
||||
"""Return whether the backend should create per-bank partial vector indexes."""
|
||||
return _normalize_resolved(ext) != "scann"
|
||||
|
||||
|
||||
def bootstrap_extension(conn: Connection, ext: str) -> None:
|
||||
"""Install the configured vector extension and any prerequisites if possible."""
|
||||
normalized = validate_extension(ext)
|
||||
for statement in _EXTENSION_INSTALL_SQL[normalized]:
|
||||
conn.execute(text(statement))
|
||||
|
||||
|
||||
def detect_vector_extension(conn: Connection, vector_extension: str = "pgvector") -> str:
|
||||
"""Validate the configured vector extension exists and return the index backend."""
|
||||
configured_ext = validate_extension(vector_extension)
|
||||
|
||||
if configured_ext == "pgvectorscale":
|
||||
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
|
||||
if not pgvector_check:
|
||||
raise RuntimeError(
|
||||
"DiskANN (pgvectorscale/pg_diskann) requires pgvector to be installed. "
|
||||
f"Install it with: {_INSTALL_HINTS['pgvectorscale']}"
|
||||
)
|
||||
|
||||
vectorscale_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vectorscale'")).scalar()
|
||||
pg_diskann_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_diskann'")).scalar()
|
||||
|
||||
if vectorscale_check:
|
||||
logger.debug("Using vector extension: pgvectorscale (DiskANN)")
|
||||
return "pgvectorscale"
|
||||
if pg_diskann_check:
|
||||
logger.debug("Using vector extension: pg_diskann (Azure DiskANN)")
|
||||
return "pg_diskann"
|
||||
|
||||
raise RuntimeError(
|
||||
"Configured vector extension 'pgvectorscale' not found. Install either:\n"
|
||||
" - pgvectorscale (open source): CREATE EXTENSION vectorscale CASCADE;\n"
|
||||
" - pg_diskann (Azure): CREATE EXTENSION pg_diskann CASCADE;"
|
||||
)
|
||||
|
||||
extension_name = pg_extension_name(configured_ext)
|
||||
extension_check = conn.execute(
|
||||
text("SELECT 1 FROM pg_extension WHERE extname = :extension_name"),
|
||||
{"extension_name": extension_name},
|
||||
).scalar()
|
||||
if not extension_check:
|
||||
raise RuntimeError(
|
||||
f"Configured vector extension '{configured_ext}' not found. "
|
||||
f"Install it with: {_INSTALL_HINTS[configured_ext]}"
|
||||
)
|
||||
|
||||
logger.debug("Using configured vector extension: %s", configured_ext)
|
||||
return configured_ext
|
||||
@@ -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()
|
||||
-55
@@ -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)
|
||||
-117
@@ -1,117 +0,0 @@
|
||||
"""Repair mental_models.subtype on databases stuck at m3rg3h3ad5f6
|
||||
|
||||
Three production deployments reported `column "subtype" of relation
|
||||
"mental_models" does not exist` on `create_mental_model` even after their
|
||||
container reported `Database migrations completed successfully` and
|
||||
`alembic_version` advanced to `m3rg3h3ad5f6` (see issue #1553, #1553#1
|
||||
confirmations from @4Lienau and @khanhduyvt0101).
|
||||
|
||||
Both `h3c4d5e6f7g8_mental_models_v4` and `d5y6z7a8b9c0_backfill_mental_models_subtype`
|
||||
were meant to ensure `subtype` exists, but on databases that came through the
|
||||
`reflections -> mental_models` rename chain *and* whose alembic_version
|
||||
advanced past `d5y6z7a8b9c0` along an alternate path during the divergent-heads
|
||||
reorganization, neither column-add actually fired. The result is a head-tagged
|
||||
database with a v3-shaped `mental_models` table missing six columns:
|
||||
``subtype``, ``description``, ``entity_id``, ``observations``, ``links``,
|
||||
``last_updated``.
|
||||
|
||||
This migration sits at the current head (`m3rg3h3ad5f6`) so every affected
|
||||
deployment will pick it up on next container start. It mirrors the column-add
|
||||
block from `d5y6z7a8b9c0_backfill_mental_models_subtype` using
|
||||
``ADD COLUMN IF NOT EXISTS`` so it is a no-op on databases where the columns
|
||||
are already present.
|
||||
|
||||
Revision ID: 86f7a033d372
|
||||
Revises: m3rg3h3ad5f6
|
||||
Create Date: 2026-05-14
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "86f7a033d372"
|
||||
down_revision: str | Sequence[str] | None = "m3rg3h3ad5f6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
"""Idempotently ensure mental_models has the v4 column set.
|
||||
|
||||
Safe to re-apply on databases that already received the columns via
|
||||
`h3c4d5e6f7g8_mental_models_v4` or `d5y6z7a8b9c0_backfill_mental_models_subtype` —
|
||||
every column-add uses ``IF NOT EXISTS`` and the constraint is recreated
|
||||
from scratch with the canonical v4 allowlist.
|
||||
"""
|
||||
schema = _pg_schema_prefix()
|
||||
bare_schema = schema.strip(".").strip('"') if schema else ""
|
||||
schema_clause = f"AND table_schema = '{bare_schema}'" if bare_schema else ""
|
||||
|
||||
# Wrapped in a DO block so the existence check skips databases that
|
||||
# predate the reflections -> mental_models rename chain (no table to
|
||||
# repair). On those, every ALTER below would error.
|
||||
op.execute(
|
||||
f"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF EXISTS (
|
||||
SELECT 1 FROM information_schema.tables
|
||||
WHERE table_name = 'mental_models'
|
||||
{schema_clause}
|
||||
) THEN
|
||||
-- Add the six v4 columns idempotently.
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD COLUMN IF NOT EXISTS subtype VARCHAR(32) NOT NULL DEFAULT 'structural';
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD COLUMN IF NOT EXISTS description TEXT NOT NULL DEFAULT '';
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD COLUMN IF NOT EXISTS entity_id UUID;
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD COLUMN IF NOT EXISTS observations JSONB DEFAULT '{{"observations": []}}'::jsonb;
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD COLUMN IF NOT EXISTS links VARCHAR[];
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD COLUMN IF NOT EXISTS last_updated TIMESTAMP WITH TIME ZONE;
|
||||
|
||||
-- Recreate the CHECK constraint with the canonical v4 allowlist.
|
||||
-- Existing rows with subtype = 'directive' (possible on databases
|
||||
-- that ran the o0j1k2l3m4n5 directive-only path) are rewritten to
|
||||
-- 'structural' first so the constraint add succeeds.
|
||||
UPDATE {schema}mental_models SET subtype = 'structural' WHERE subtype = 'directive';
|
||||
|
||||
ALTER TABLE {schema}mental_models DROP CONSTRAINT IF EXISTS ck_mental_models_subtype;
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD CONSTRAINT ck_mental_models_subtype
|
||||
CHECK (subtype IN ('structural', 'emergent', 'pinned', 'learned'));
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_mental_models_subtype
|
||||
ON {schema}mental_models(bank_id, subtype);
|
||||
END IF;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
"""No-op: dropping these columns would corrupt v4 application code."""
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# PG-only: Oracle's baseline (o1a2b3c4d5e6) creates mental_models with its
|
||||
# own subtype shape (chk_mm_subtype IN ('directive', 'pinned')) and a
|
||||
# different table topology, so this PG-shaped repair does not apply.
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-50
@@ -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)
|
||||
-128
@@ -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)
|
||||
-48
@@ -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)
|
||||
-62
@@ -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)
|
||||
-161
@@ -1,161 +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) or scann. ScaNN uses global vector indexes because empty or tiny
|
||||
per-bank indexes cannot be built safely on AlloyDB.
|
||||
"""
|
||||
|
||||
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 _validate_extension(name: str) -> str:
|
||||
ext = name.lower()
|
||||
if ext not in {"pgvector", "pgvectorscale", "vchord", "scann"}:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {ext}. Must be 'pgvector', 'vchord', 'pgvectorscale', or 'scann'"
|
||||
)
|
||||
return ext
|
||||
|
||||
|
||||
def _index_type_keyword(ext: str) -> str:
|
||||
if ext == "pgvectorscale":
|
||||
return "diskann"
|
||||
if ext == "vchord":
|
||||
return "vchordrq"
|
||||
if ext == "scann":
|
||||
return "scann"
|
||||
return "hnsw"
|
||||
|
||||
|
||||
def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pgvectorscale":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
ext = _validate_extension(os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector"))
|
||||
if ext in {"pgvector", "scann"}:
|
||||
return
|
||||
target = _index_type_keyword(ext)
|
||||
|
||||
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(ext)
|
||||
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)
|
||||
ext = _validate_extension(os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector"))
|
||||
if ext in {"pgvector", "scann"}:
|
||||
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)
|
||||
-42
@@ -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)
|
||||
-54
@@ -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)
|
||||
-91
@@ -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)
|
||||
-76
@@ -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)
|
||||
-58
@@ -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)
|
||||
-63
@@ -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)
|
||||
-168
@@ -1,168 +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. For non-ScaNN backends, drops the global vector index (competes with
|
||||
per-bank partial indexes)
|
||||
3. For non-ScaNN backends, 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.
|
||||
- AlloyDB ScaNN uses global vector indexes with filtered vector search instead
|
||||
because empty or tiny per-bank indexes cannot be built safely.
|
||||
"""
|
||||
|
||||
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 _configured_vector_extension() -> str:
|
||||
ext = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
|
||||
if ext not in {"pgvector", "pgvectorscale", "vchord", "scann"}:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {ext}. Must be 'pgvector', 'vchord', 'pgvectorscale', or 'scann'"
|
||||
)
|
||||
return ext
|
||||
|
||||
|
||||
def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pgvectorscale":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
|
||||
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 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)")
|
||||
|
||||
ext = _configured_vector_extension()
|
||||
if ext == "scann":
|
||||
# ScaNN should keep/use a global vector index. Per-bank partial indexes
|
||||
# are created while banks are empty and can fail AlloyDB's ScaNN build
|
||||
# requirements, so this migration leaves vector index reconciliation to
|
||||
# runtime ensure_vector_extension once enough rows exist.
|
||||
return
|
||||
|
||||
# 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(ext)
|
||||
|
||||
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 _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)
|
||||
-67
@@ -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)
|
||||
-49
@@ -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)
|
||||
-83
@@ -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)
|
||||
-67
@@ -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)
|
||||
-93
@@ -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)
|
||||
-81
@@ -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)
|
||||
-52
@@ -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)
|
||||
-49
@@ -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)
|
||||
-38
@@ -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)
|
||||
-31
@@ -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,185 +0,0 @@
|
||||
"""
|
||||
Daemon mode support for Hindsight API.
|
||||
|
||||
Provides idle timeout for running as a background daemon.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import IO
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default daemon configuration
|
||||
DEFAULT_DAEMON_PORT = 8888
|
||||
DEFAULT_IDLE_TIMEOUT = 0 # 0 = no auto-exit (hindsight-embed passes its own timeout)
|
||||
|
||||
# Allow override via environment variable for profile-specific logs
|
||||
DAEMON_LOG_PATH = Path(os.getenv("HINDSIGHT_API_DAEMON_LOG", str(Path.home() / ".hindsight" / "daemon.log")))
|
||||
|
||||
# Internal env var: set by daemonize() in the re-exec'd child so the child
|
||||
# skips re-exec and just redirects stdio. Also set by hindsight-embed's
|
||||
# DaemonEmbedManager so the daemon launched via Popen skips re-exec entirely
|
||||
# (hindsight-embed's Popen already provides a clean, detached process).
|
||||
ENV_DAEMON_CHILD = "_HINDSIGHT_DAEMON_CHILD"
|
||||
|
||||
|
||||
class IdleTimeoutMiddleware:
|
||||
"""ASGI middleware that tracks activity and exits after idle timeout."""
|
||||
|
||||
def __init__(self, app, idle_timeout: int = DEFAULT_IDLE_TIMEOUT):
|
||||
self.app = app
|
||||
self.idle_timeout = idle_timeout
|
||||
self.last_activity = time.time()
|
||||
self._checker_task = None
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
# Update activity timestamp on each request
|
||||
self.last_activity = time.time()
|
||||
await self.app(scope, receive, send)
|
||||
|
||||
def start_idle_checker(self):
|
||||
"""Start the background task that checks for idle timeout."""
|
||||
self._checker_task = asyncio.create_task(self._check_idle())
|
||||
|
||||
async def _check_idle(self):
|
||||
"""Background task that exits the process after idle timeout."""
|
||||
# If idle_timeout is 0, don't auto-exit
|
||||
if self.idle_timeout <= 0:
|
||||
return
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(30) # Check every 30 seconds
|
||||
idle_time = time.time() - self.last_activity
|
||||
if idle_time > self.idle_timeout:
|
||||
logger.info(f"Idle timeout reached ({self.idle_timeout}s), shutting down daemon")
|
||||
# Give a moment for any in-flight requests
|
||||
await asyncio.sleep(1)
|
||||
# Send SIGTERM to ourselves to trigger graceful shutdown
|
||||
import signal
|
||||
|
||||
os.kill(os.getpid(), signal.SIGTERM)
|
||||
|
||||
|
||||
def _detach_popen_kwargs(log_handle: "IO[bytes]") -> dict:
|
||||
"""Cross-platform kwargs to spawn a subprocess detached from the caller.
|
||||
|
||||
On POSIX, ``start_new_session=True`` calls ``setsid(2)`` so the child
|
||||
survives the parent's terminal. On Windows we use
|
||||
``DETACHED_PROCESS | CREATE_NEW_PROCESS_GROUP``.
|
||||
|
||||
``log_handle`` receives the child's stdout/stderr so output never leaks
|
||||
into the parent's terminal.
|
||||
"""
|
||||
if platform.system() == "Windows":
|
||||
detached_process = getattr(subprocess, "DETACHED_PROCESS", 0)
|
||||
create_new_process_group = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
|
||||
return {
|
||||
"creationflags": detached_process | create_new_process_group,
|
||||
"stdin": subprocess.DEVNULL,
|
||||
"stdout": log_handle,
|
||||
"stderr": subprocess.STDOUT,
|
||||
"close_fds": True,
|
||||
}
|
||||
return {
|
||||
"start_new_session": True,
|
||||
"stdin": subprocess.DEVNULL,
|
||||
"stdout": log_handle,
|
||||
"stderr": log_handle,
|
||||
}
|
||||
|
||||
|
||||
def _redirect_stdio_to_log() -> None:
|
||||
"""Redirect stdin/stdout/stderr to the daemon log file.
|
||||
|
||||
Called in the daemon child process after re-exec.
|
||||
"""
|
||||
DAEMON_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
|
||||
with open(os.devnull, "r") as devnull:
|
||||
os.dup2(devnull.fileno(), sys.stdin.fileno())
|
||||
|
||||
log_fd = open(DAEMON_LOG_PATH, "a")
|
||||
os.dup2(log_fd.fileno(), sys.stdout.fileno())
|
||||
os.dup2(log_fd.fileno(), sys.stderr.fileno())
|
||||
|
||||
|
||||
def daemonize():
|
||||
"""Detach the current process into a background daemon.
|
||||
|
||||
Uses ``subprocess.Popen`` (which maps to ``posix_spawn`` on macOS) to
|
||||
re-exec the current command in a detached session. This replaces the
|
||||
traditional double-fork pattern because ``os.fork()`` without ``exec()``
|
||||
corrupts Apple framework state (XPC, Metal/MPS, ObjC runtime) on macOS,
|
||||
causing SIGBUS crashes when PyTorch uses the MPS backend.
|
||||
|
||||
The function has two code paths controlled by the ``_HINDSIGHT_DAEMON_CHILD``
|
||||
environment variable:
|
||||
|
||||
* **Parent** (env var not set): re-exec the same command via Popen with
|
||||
``start_new_session=True``, stripping ``--daemon`` from argv and setting
|
||||
``_HINDSIGHT_DAEMON_CHILD=1``. Then ``sys.exit(0)``.
|
||||
* **Child** (env var set): redirect stdio to the daemon log file and return.
|
||||
No fork, no re-exec.
|
||||
|
||||
On Windows there is no fork model: the spawning parent is expected to
|
||||
detach us via ``CREATE_NEW_PROCESS_GROUP | DETACHED_PROCESS`` and to
|
||||
redirect stdout/stderr to ``HINDSIGHT_API_DAEMON_LOG`` before exec.
|
||||
We still ensure the log directory exists.
|
||||
"""
|
||||
if sys.platform == "win32":
|
||||
DAEMON_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
return
|
||||
|
||||
# If we are already the daemon child (re-exec'd by a previous daemonize()
|
||||
# call, or launched by hindsight-embed with the env var set), just redirect
|
||||
# stdio and return — no re-exec needed.
|
||||
if os.environ.get(ENV_DAEMON_CHILD) == "1":
|
||||
_redirect_stdio_to_log()
|
||||
return
|
||||
|
||||
# --- Parent path: re-exec ourselves as a detached background process ---
|
||||
|
||||
# Build child command: same Python, same module entry point, all args
|
||||
# except --daemon (replaced by the env var).
|
||||
child_args = [a for a in sys.argv[1:] if a != "--daemon"]
|
||||
cmd = [sys.executable, "-m", "hindsight_api.main"] + child_args
|
||||
|
||||
env = os.environ.copy()
|
||||
env[ENV_DAEMON_CHILD] = "1"
|
||||
env["HINDSIGHT_API_DAEMON_LOG"] = str(DAEMON_LOG_PATH)
|
||||
|
||||
DAEMON_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(DAEMON_LOG_PATH, "ab") as log_handle:
|
||||
subprocess.Popen(cmd, env=env, **_detach_popen_kwargs(log_handle))
|
||||
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
def check_daemon_running(port: int = DEFAULT_DAEMON_PORT) -> bool:
|
||||
"""Check if a daemon is running and responsive on the given port."""
|
||||
import socket
|
||||
|
||||
try:
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
sock.settimeout(1)
|
||||
result = sock.connect_ex(("127.0.0.1", port))
|
||||
sock.close()
|
||||
return result == 0
|
||||
except Exception:
|
||||
return False
|
||||
@@ -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,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
|
||||
@@ -1,181 +0,0 @@
|
||||
"""Registry of optional PostgreSQL server-side routines.
|
||||
|
||||
Some hot paths (currently only the worker poller's per-cycle scan) can be
|
||||
sped up by a server-side PL/pgSQL routine that the API never installs
|
||||
itself — operators install it out-of-band (e.g. a Helm hook in
|
||||
hindsight-cloud) when they want the optimisation. When the routine is not
|
||||
installed, callers must fall back to a pure-Python implementation.
|
||||
|
||||
This module centralises that pattern so we don't sprinkle ad-hoc
|
||||
``try / except`` blocks (which silently log a server-side error on every
|
||||
call) around the codebase. Each registered entry carries:
|
||||
|
||||
* a ``schema`` and ``name`` (used to probe ``pg_proc``)
|
||||
* a ``contract`` describing the expected signature and return shape
|
||||
|
||||
Bodies are deliberately not stored here. Hindsight never installs these
|
||||
routines, so a body checked into this repo would (a) drift from whatever
|
||||
operators actually deploy and (b) imply ownership we don't have. The
|
||||
contract is the entire API surface: any operator-supplied implementation
|
||||
that satisfies it is interchangeable.
|
||||
|
||||
Probe behaviour:
|
||||
|
||||
* On first ``is_installed()`` call per backend instance we issue a single
|
||||
``SELECT EXISTS(...) FROM pg_proc`` and cache the boolean result in
|
||||
memory for the life of the process.
|
||||
* No TTL: if an operator installs a routine on a running cluster, workers
|
||||
pick it up only after restart. This is intentional — these routines are
|
||||
expected to be installed once at deploy time, and a probe-per-poll would
|
||||
defeat the optimisation.
|
||||
* Non-PostgreSQL backends short-circuit to ``False`` without touching the
|
||||
database, so callers can use the same code path for Oracle.
|
||||
|
||||
PostgreSQL terminology note: ``CREATE FUNCTION ... RETURNS SETOF`` defines
|
||||
a *function* (invoked via ``SELECT``); ``CREATE PROCEDURE`` defines a
|
||||
*procedure* (invoked via ``CALL``). The SQL-standard umbrella term
|
||||
covering both is *routine*, and the system catalog (``pg_proc``) stores
|
||||
both. The module name uses "routine" so a future ``CREATE PROCEDURE``
|
||||
entry slots in without a rename.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .base import DatabaseBackend, DatabaseConnection
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OptionalRoutine:
|
||||
"""One optional server-side routine Hindsight may call when present.
|
||||
|
||||
Hindsight never installs these — operators do, out-of-band. The fields
|
||||
here describe the contract the API expects; any implementation that
|
||||
matches is interchangeable.
|
||||
|
||||
Attributes:
|
||||
name: Unqualified routine name as it appears in ``pg_proc.proname``.
|
||||
schema: Schema the routine lives in (matched against
|
||||
``pg_namespace.nspname``).
|
||||
contract: Free-form description of the expected signature,
|
||||
arguments, return type, and any semantic constraints. Read
|
||||
this before installing a custom implementation.
|
||||
"""
|
||||
|
||||
name: str
|
||||
schema: str
|
||||
contract: str
|
||||
|
||||
|
||||
# Registry of known optional routines.
|
||||
#
|
||||
# Add new entries here when a hot path grows a server-side optimisation.
|
||||
# Document the *contract* — not the implementation — so operator-supplied
|
||||
# variants stay interchangeable and we don't pretend to own SQL we never
|
||||
# install.
|
||||
SCHEMAS_WITH_PENDING_WORK = OptionalRoutine(
|
||||
name="schemas_with_pending_work",
|
||||
schema="public",
|
||||
contract="""
|
||||
Signature: public.schemas_with_pending_work() RETURNS SETOF text
|
||||
|
||||
Called by the worker poller on every cycle to find schemas with
|
||||
claimable async_operations rows. The poller then runs FOR UPDATE
|
||||
SKIP LOCKED only against the returned schemas, so an empty set means
|
||||
"nothing to do, skip the expensive claim query".
|
||||
|
||||
A schema is "claimable" iff at least one row matches:
|
||||
status = 'pending' AND task_payload IS NOT NULL
|
||||
in that schema's ``async_operations`` table.
|
||||
|
||||
Required semantics:
|
||||
* No arguments.
|
||||
* Returns a set of schema names (``text``); each must match a real
|
||||
``pg_namespace.nspname``. The poller passes them straight into
|
||||
the claim query — anything that isn't a valid schema will fail.
|
||||
* Operators choose the search scope (e.g. ``tenant_%`` only, or
|
||||
include ``public``). A schema omitted from the scan will *never*
|
||||
be serviced by the poller, so the implementation must cover
|
||||
every schema that holds an ``async_operations`` table in that
|
||||
deployment.
|
||||
* Should be cheap and idempotent — called every poll cycle (~30s).
|
||||
|
||||
Fallback when the routine is absent: per-schema ``EXISTS`` queries
|
||||
from Python (~4ms per schema). The server-side path is a single-
|
||||
round-trip optimisation worth ~200ms in deployments with thousands
|
||||
of tenant schemas; everything else works correctly without it.
|
||||
""",
|
||||
)
|
||||
|
||||
_REGISTRY: dict[str, OptionalRoutine] = {
|
||||
SCHEMAS_WITH_PENDING_WORK.name: SCHEMAS_WITH_PENDING_WORK,
|
||||
}
|
||||
|
||||
|
||||
class OptionalRoutines:
|
||||
"""Per-backend cache of which optional routines are installed.
|
||||
|
||||
One instance per long-lived consumer (e.g. one per ``WorkerPoller``).
|
||||
Probes ``pg_proc`` lazily on first lookup and caches the result in
|
||||
memory until the process restarts.
|
||||
"""
|
||||
|
||||
def __init__(self, backend: DatabaseBackend) -> None:
|
||||
self._backend = backend
|
||||
self._cache: dict[str, bool] = {}
|
||||
|
||||
async def is_installed(self, conn: DatabaseConnection, routine_name: str) -> bool:
|
||||
"""Return True iff *routine_name* exists in ``pg_proc``.
|
||||
|
||||
On non-PostgreSQL backends always returns False without issuing a
|
||||
query. Result is memoised for the life of this instance.
|
||||
"""
|
||||
if self._backend.backend_type != "postgresql":
|
||||
return False
|
||||
|
||||
cached = self._cache.get(routine_name)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
routine = _REGISTRY.get(routine_name)
|
||||
if routine is None:
|
||||
raise KeyError(f"Unknown optional routine: {routine_name!r}")
|
||||
|
||||
exists = await conn.fetchval(
|
||||
"SELECT EXISTS(SELECT 1 FROM pg_proc p "
|
||||
"JOIN pg_namespace n ON p.pronamespace = n.oid "
|
||||
"WHERE n.nspname = $1 AND p.proname = $2)",
|
||||
routine.schema,
|
||||
routine.name,
|
||||
)
|
||||
installed = bool(exists)
|
||||
self._cache[routine_name] = installed
|
||||
if installed:
|
||||
logger.info(
|
||||
"Optional PG routine %s.%s detected — using server-side path",
|
||||
routine.schema,
|
||||
routine.name,
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
"Optional PG routine %s.%s not installed — using fallback path",
|
||||
routine.schema,
|
||||
routine.name,
|
||||
)
|
||||
return installed
|
||||
|
||||
def invalidate(self, routine_name: str | None = None) -> None:
|
||||
"""Drop cached probe results (test helper).
|
||||
|
||||
Without an argument, clears the entire cache.
|
||||
"""
|
||||
if routine_name is None:
|
||||
self._cache.clear()
|
||||
else:
|
||||
self._cache.pop(routine_name, None)
|
||||
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,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)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,433 +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"] = self._cap_max_completion_tokens(max_completion_tokens)
|
||||
if temperature is not None:
|
||||
kwargs["temperature"] = temperature
|
||||
|
||||
return kwargs
|
||||
|
||||
# ── per-model output-tokens cap (shared with Router subclass) ────────────
|
||||
# Hindsight's defaults (e.g. retain_max_completion_tokens=64000) target
|
||||
# high-capacity models. When a configured deployment supports fewer
|
||||
# completion tokens (e.g. gpt-4.1-nano caps at 32768), the call would
|
||||
# otherwise be rejected. Cap pre-emptively using LiteLLM's per-model
|
||||
# registry so things work out of the box across the supported model set.
|
||||
|
||||
def _cap_max_completion_tokens(self, value: int) -> int:
|
||||
cap = self._get_model_output_cap()
|
||||
if cap and value > cap:
|
||||
logger.debug("capping max_completion_tokens %d -> %d for model %s", value, cap, self.model)
|
||||
return cap
|
||||
return value
|
||||
|
||||
def _get_model_output_cap(self) -> int | None:
|
||||
"""Return the configured model's max output tokens, per LiteLLM's registry."""
|
||||
try:
|
||||
cap = self._litellm.get_max_tokens(self.model)
|
||||
return int(cap) if cap else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
# ── hooks for Router-style subclasses ────────────────────────────────────
|
||||
# The retry+parse loop in call() / call_with_tools() is shared by every
|
||||
# LiteLLM-backed provider. Subclasses override the small surface below to
|
||||
# swap the completion fn (direct vs Router) and rename the deployment that
|
||||
# actually answered the request.
|
||||
|
||||
@property
|
||||
def _stage_label(self) -> str:
|
||||
"""Stage breadcrumb label — overridden by subclasses (e.g. ``litellmrouter``)."""
|
||||
return "litellm"
|
||||
|
||||
async def _acompletion(self, **kwargs: Any) -> Any:
|
||||
"""Issue a chat completion. Subclasses override to route via ``litellm.Router``."""
|
||||
return await self._litellm.acompletion(**kwargs)
|
||||
|
||||
def _resolve_completion_model(self, response: Any) -> str:
|
||||
"""
|
||||
Return the model name to record in metrics/tracing.
|
||||
|
||||
For Router-backed providers this can differ from ``self.model`` — the Router
|
||||
may pick a different deployment than the primary. Default: ``self.model``.
|
||||
"""
|
||||
return self.model
|
||||
|
||||
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.{self._stage_label}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._acompletion(**call_kwargs)
|
||||
|
||||
content = response.choices[0].message.content or ""
|
||||
finish_reason = response.choices[0].finish_reason
|
||||
model_name = self._resolve_completion_model(response)
|
||||
|
||||
# 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=model_name,
|
||||
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=model_name,
|
||||
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}/{model_name}, "
|
||||
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.{self._stage_label}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._acompletion(**call_kwargs)
|
||||
|
||||
message = response.choices[0].message
|
||||
content = message.content
|
||||
finish_reason = response.choices[0].finish_reason
|
||||
model_name = self._resolve_completion_model(response)
|
||||
|
||||
# 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=model_name,
|
||||
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=model_name,
|
||||
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,167 +0,0 @@
|
||||
"""
|
||||
LiteLLM Router LLM provider — pure pass-through to ``litellm.Router``.
|
||||
|
||||
The full configuration object is forwarded verbatim. We do not translate model
|
||||
names, infer fallbacks, validate shape, or introspect Router internals:
|
||||
whatever the user puts in ``HINDSIGHT_API_LLM_LITELLMROUTER_CONFIG`` becomes
|
||||
``Router(**config)``. If the shape is wrong, LiteLLM Router raises.
|
||||
|
||||
The only Hindsight-imposed convention is that one entry in ``model_list``
|
||||
must have ``model_name: "default"`` — that's the entrypoint we issue
|
||||
completions against. Everything else (ordering, fallbacks, load-balancing,
|
||||
weighted picks, rate limits, retries, cooldowns) is whatever the user
|
||||
configures via LiteLLM's own keys.
|
||||
|
||||
See https://docs.litellm.ai/docs/routing for the supported keys (``model_list``,
|
||||
``fallbacks``, ``context_window_fallbacks``, ``num_retries``, ``cooldown_time``,
|
||||
``routing_strategy``, ``allowed_fails``, …).
|
||||
|
||||
The retry/parse/metrics loop is shared with ``LiteLLMLLM`` via inheritance:
|
||||
this class only overrides the completion fn, the call kwargs, and the model
|
||||
name reported in metrics.
|
||||
|
||||
Example ``HINDSIGHT_API_LLM_LITELLMROUTER_CONFIG``::
|
||||
|
||||
{
|
||||
"model_list": [
|
||||
{"model_name": "default", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-..."}},
|
||||
{"model_name": "fallback", "litellm_params": {"model": "anthropic/claude-sonnet-4", "api_key": "sk-ant-..."}}
|
||||
],
|
||||
"fallbacks": [{"default": ["fallback"]}],
|
||||
"num_retries": 0,
|
||||
"cooldown_time": 60
|
||||
}
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from hindsight_api.engine.providers.litellm_llm import LiteLLMLLM
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Hindsight always issues completions against this ``model_name``. Users must
|
||||
# include at least one entry with ``model_name: "default"`` in their config's
|
||||
# ``model_list``; that entry is the entrypoint, and any other entries become
|
||||
# fallback / load-balance / weighted-pool members per the user's own
|
||||
# ``fallbacks`` / ``routing_strategy`` settings.
|
||||
_ENTRYPOINT_MODEL_NAME = "default"
|
||||
|
||||
|
||||
class LiteLLMRouterLLM(LiteLLMLLM):
|
||||
"""
|
||||
LLM provider backed by ``litellm.Router``.
|
||||
|
||||
The full Router config is supplied by the caller. We pass it verbatim to
|
||||
``Router(**config)`` and route requests against the first ``model_list``
|
||||
entry's ``model_name``. Inherits the retry/parse/metrics loop from
|
||||
``LiteLLMLLM``; only the completion fn and the call kwargs differ.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
config: dict[str, Any],
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float = 300.0,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
timeout=timeout,
|
||||
**kwargs,
|
||||
)
|
||||
self.config = config
|
||||
|
||||
from litellm import Router
|
||||
|
||||
logging.getLogger("LiteLLM Router").setLevel(logging.WARNING)
|
||||
# Pure pass-through: whatever the user gave goes straight to LiteLLM Router.
|
||||
# If the shape is invalid, Router raises its own error — we don't pre-validate
|
||||
# or introspect Router internals.
|
||||
self._router = Router(**config)
|
||||
|
||||
# Pre-compute the most conservative output-tokens cap across every configured
|
||||
# deployment so a single max_completion_tokens value works no matter which
|
||||
# deployment Router picks. Uses LiteLLM's own per-model registry; unknown
|
||||
# models contribute no cap. See LiteLLMLLM._cap_max_completion_tokens.
|
||||
self._router_output_cap = self._compute_router_output_cap(config)
|
||||
|
||||
logger.info("LiteLLM Router initialized; entrypoint model_name=%r", _ENTRYPOINT_MODEL_NAME)
|
||||
|
||||
def _compute_router_output_cap(self, config: dict[str, Any]) -> int | None:
|
||||
caps: list[int] = []
|
||||
for deployment in (config.get("model_list") or []) if isinstance(config, dict) else []:
|
||||
if not isinstance(deployment, dict):
|
||||
continue
|
||||
params = deployment.get("litellm_params") or {}
|
||||
model_str = params.get("model") if isinstance(params, dict) else None
|
||||
if not model_str:
|
||||
continue
|
||||
try:
|
||||
cap = self._litellm.get_max_tokens(model_str)
|
||||
except Exception:
|
||||
cap = None
|
||||
if cap:
|
||||
caps.append(int(cap))
|
||||
return min(caps) if caps else None
|
||||
|
||||
# ── overrides for the shared retry/parse loop ───────────────────────────
|
||||
|
||||
@property
|
||||
def _stage_label(self) -> str:
|
||||
return "litellmrouter"
|
||||
|
||||
async def _acompletion(self, **kwargs: Any) -> Any:
|
||||
return await self._router.acompletion(**kwargs)
|
||||
|
||||
def _resolve_completion_model(self, response: Any) -> str:
|
||||
hidden = getattr(response, "_hidden_params", None) or {}
|
||||
return hidden.get("model") or _ENTRYPOINT_MODEL_NAME
|
||||
|
||||
def _get_model_output_cap(self) -> int | None:
|
||||
return self._router_output_cap
|
||||
|
||||
def _build_common_kwargs(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
) -> dict[str, Any]:
|
||||
# Always issue against the entrypoint group; Router handles deployment selection,
|
||||
# cross-group fallbacks, retries, cooldowns — whatever the user configured.
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": _ENTRYPOINT_MODEL_NAME,
|
||||
"messages": messages,
|
||||
}
|
||||
if max_completion_tokens is not None:
|
||||
kwargs["max_completion_tokens"] = self._cap_max_completion_tokens(max_completion_tokens)
|
||||
if temperature is not None:
|
||||
kwargs["temperature"] = temperature
|
||||
return kwargs
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
from hindsight_api.engine.llm_interface import OutputTooLongError
|
||||
|
||||
try:
|
||||
await self.call(
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
max_completion_tokens=50,
|
||||
temperature=0.0,
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
logger.info("LiteLLM Router connection verified successfully")
|
||||
except OutputTooLongError:
|
||||
logger.info("LiteLLM Router connection verified successfully (response truncated)")
|
||||
except Exception as e:
|
||||
logger.error(f"LiteLLM Router connection verification failed: {e}")
|
||||
raise RuntimeError(f"Failed to verify LiteLLM Router connection: {e}") from e
|
||||
@@ -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,17 +0,0 @@
|
||||
"""Token counting helpers for reflect prompts and agent control flow."""
|
||||
|
||||
from functools import lru_cache
|
||||
|
||||
import tiktoken
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _get_cl100k_base_encoding() -> tiktoken.Encoding:
|
||||
# tiktoken downloads this encoding on first lookup when it is not cached.
|
||||
# Keep the lookup lazy so importing hindsight_api does not depend on network access.
|
||||
return tiktoken.get_encoding("cl100k_base")
|
||||
|
||||
|
||||
def count_cl100k_tokens(text: str) -> int:
|
||||
"""Return the number of cl100k_base tokens in text."""
|
||||
return len(_get_cl100k_base_encoding().encode(text))
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user