Compare commits
78
Commits
embed-again
..
py314
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0f190d4fd3 | ||
|
|
0d8d805832 | ||
|
|
1cd836229b | ||
|
|
90ad003c46 | ||
|
|
278718dd84 | ||
|
|
093ecff48d | ||
|
|
85b9074f43 | ||
|
|
7e339e1677 | ||
|
|
dd621a69d0 | ||
|
|
7097716204 | ||
|
|
d3302c95b9 | ||
|
|
665877bb01 | ||
|
|
a43d208e93 | ||
|
|
34d9188e13 | ||
|
|
9a776e9f58 | ||
|
|
d02affd8f2 | ||
|
|
6b346925e2 | ||
|
|
63e2964a4c | ||
|
|
d5403a4b29 | ||
|
|
a24941f83b | ||
|
|
21b25fe8fe | ||
|
|
794a7435a9 | ||
|
|
038a9c2313 | ||
|
|
749478d9f9 | ||
|
|
96f0e54efa | ||
|
|
382550690a | ||
|
|
6c7f057e9d | ||
|
|
539190b69e | ||
|
|
1499ce5549 | ||
|
|
8564135b2a | ||
|
|
44d912533c | ||
|
|
35127d5f8b | ||
|
|
86c733c10e | ||
|
|
cb7ebe80bb | ||
|
|
615509011e | ||
|
|
af6bd1b5e1 | ||
|
|
579b10b53d | ||
|
|
4b57b82301 | ||
|
|
9c3fda74e2 | ||
|
|
f0cb1925ec | ||
|
|
039944cae2 | ||
|
|
ef9d3a15cb | ||
|
|
d788a55e28 | ||
|
|
c8ae82d62f | ||
|
|
27498f99d0 | ||
|
|
1530c09120 | ||
|
|
1163b1f6a6 | ||
|
|
fe88bdf704 | ||
|
|
cbb8fc6723 | ||
|
|
c33b9b8bb2 | ||
|
|
b364bc3402 | ||
|
|
35f0984b72 | ||
|
|
5dc45194c9 | ||
|
|
ff47814422 | ||
|
|
1ba70f81c8 | ||
|
|
fe15b5ec87 | ||
|
|
10e21f7302 | ||
|
|
7d3ac5ddb9 | ||
|
|
f4f86e3842 | ||
|
|
728ce13cea | ||
|
|
ecc590cb79 | ||
|
|
381c96c093 | ||
|
|
ab5e31f203 | ||
|
|
0da77ce2c9 | ||
|
|
d57e8639c5 | ||
|
|
03bf13e9e3 | ||
|
|
ff20bf9dc7 | ||
|
|
751f99a82f | ||
|
|
49ae55af03 | ||
|
|
c2ac7d0440 | ||
|
|
657fe023b2 | ||
|
|
9c95a1ac1d | ||
|
|
15540075b2 | ||
|
|
3f211f0729 | ||
|
|
8781c9fbfe | ||
|
|
12e9a3d305 | ||
|
|
c16ccc2c22 | ||
|
|
a7c094d436 |
+8
-1
@@ -2,7 +2,7 @@
|
||||
# Copy this file to .env and fill in your values
|
||||
|
||||
# LLM Configuration (Required)
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio
|
||||
# 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=o3-mini
|
||||
@@ -13,6 +13,13 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-anthropic-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=claude-sonnet-4-20250514
|
||||
|
||||
# Example: Google Vertex AI configuration
|
||||
# HINDSIGHT_API_LLM_PROVIDER=vertexai
|
||||
# HINDSIGHT_API_LLM_MODEL=google/gemini-2.0-flash-001
|
||||
# HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=your-gcp-project-id
|
||||
# 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: LM Studio local configuration (Qwen 2.5 32B recommended)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=lmstudio
|
||||
# HINDSIGHT_API_LLM_API_KEY=lmstudio
|
||||
|
||||
@@ -139,6 +139,104 @@ jobs:
|
||||
path: hindsight-clients/typescript/*.tgz
|
||||
retention-days: 1
|
||||
|
||||
release-openclaw-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/openclaw
|
||||
run: npm ci
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm run build
|
||||
|
||||
- name: Publish to npm
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
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/openclaw
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
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-control-plane:
|
||||
runs-on: ubuntu-latest
|
||||
environment: npm
|
||||
@@ -366,7 +464,7 @@ jobs:
|
||||
|
||||
create-github-release:
|
||||
runs-on: ubuntu-latest
|
||||
needs: [release-python-packages, release-typescript-client, 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-control-plane, release-rust-cli, release-docker-images, release-helm-chart]
|
||||
permissions:
|
||||
contents: write
|
||||
|
||||
@@ -389,6 +487,18 @@ jobs:
|
||||
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 Control Plane
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
@@ -430,6 +540,10 @@ jobs:
|
||||
cp artifacts/python-packages/hindsight-embed/dist/* release-assets/ || true
|
||||
# TypeScript client
|
||||
cp artifacts/typescript-client/*.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
|
||||
# Control Plane
|
||||
cp artifacts/control-plane/*.tgz release-assets/ || true
|
||||
# Rust CLI binaries
|
||||
|
||||
+163
-36
@@ -9,42 +9,11 @@ concurrency:
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
build-python-packages:
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- name: hindsight-all
|
||||
path: hindsight
|
||||
- name: hindsight-api
|
||||
path: hindsight-api
|
||||
- name: hindsight-client
|
||||
path: hindsight-clients/python
|
||||
- name: hindsight-embed
|
||||
path: hindsight-embed
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
enable-cache: true
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build ${{ matrix.name }}
|
||||
working-directory: ./${{ matrix.path }}
|
||||
run: uv build
|
||||
|
||||
build-api-python-versions:
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ['3.11', '3.12', '3.13']
|
||||
python-version: ['3.11', '3.12', '3.13', '3.14']
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
@@ -82,6 +51,52 @@ jobs:
|
||||
- name: Build TypeScript client
|
||||
run: npm run build --workspace=hindsight-clients/typescript
|
||||
|
||||
build-openclaw-integration:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '22'
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm ci
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm test
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm run build
|
||||
|
||||
build-ai-sdk-integration:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '22'
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/ai-sdk
|
||||
run: npm ci
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/ai-sdk
|
||||
run: npm test
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/ai-sdk
|
||||
run: npm run build
|
||||
|
||||
build-control-plane:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
@@ -726,9 +741,9 @@ jobs:
|
||||
test-embed:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_EMBED_LLM_PROVIDER: groq
|
||||
HINDSIGHT_EMBED_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
HINDSIGHT_EMBED_LLM_MODEL: openai/gpt-oss-20b
|
||||
HINDSIGHT_API_LLM_PROVIDER: groq
|
||||
HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
|
||||
# Prefer CPU-only PyTorch in CI
|
||||
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
|
||||
|
||||
@@ -759,10 +774,62 @@ jobs:
|
||||
${{ runner.os }}-huggingface-embed-
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Run unit and integration tests
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv run pytest tests/ -v
|
||||
|
||||
- name: Run smoke test
|
||||
working-directory: ./hindsight-embed
|
||||
run: ./test.sh
|
||||
|
||||
test-hindsight-all:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: groq
|
||||
HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
|
||||
# For test_server_integration.py compatibility
|
||||
HINDSIGHT_LLM_PROVIDER: groq
|
||||
HINDSIGHT_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
HINDSIGHT_LLM_MODEL: openai/gpt-oss-20b
|
||||
# Prefer CPU-only PyTorch in CI
|
||||
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build hindsight-all
|
||||
working-directory: ./hindsight
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight
|
||||
run: uv sync --frozen --extra test --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-all-${{ hashFiles('hindsight/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-all-
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Run unit tests
|
||||
working-directory: ./hindsight
|
||||
run: uv run pytest tests/ -v
|
||||
|
||||
test-doc-examples:
|
||||
runs-on: ubuntu-latest
|
||||
needs: test-rust-cli
|
||||
@@ -935,6 +1002,18 @@ jobs:
|
||||
working-directory: ./hindsight-dev
|
||||
run: uv run pytest upgrade_tests/ -v --tb=short
|
||||
|
||||
- name: Show upgrade test logs
|
||||
if: always()
|
||||
run: |
|
||||
echo "=== Upgrade Test Server Logs ==="
|
||||
for log in /tmp/upgrade-test-*.log; do
|
||||
if [ -f "$log" ]; then
|
||||
echo ""
|
||||
echo "--- $log ---"
|
||||
tail -500 "$log"
|
||||
fi
|
||||
done
|
||||
|
||||
verify-generated-files:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
@@ -1005,4 +1084,52 @@ jobs:
|
||||
git diff --stat
|
||||
exit 1
|
||||
fi
|
||||
echo "✓ All generated files are up to date"
|
||||
echo "✓ All generated files are up to date"
|
||||
|
||||
check-openapi-compatibility:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0 # Fetch full git history to access base branch
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
enable-cache: true
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Install hindsight-dev dependencies
|
||||
run: |
|
||||
cd hindsight-dev && uv sync --frozen --index-strategy unsafe-best-match
|
||||
|
||||
- name: Check OpenAPI compatibility with base branch
|
||||
run: |
|
||||
# Get the base branch (usually main)
|
||||
BASE_BRANCH="${{ github.base_ref }}"
|
||||
|
||||
if [ -z "$BASE_BRANCH" ]; then
|
||||
echo "⚠️ Warning: No base branch found (not a PR?). Skipping compatibility check."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "Checking OpenAPI compatibility against base branch: $BASE_BRANCH"
|
||||
|
||||
# Extract the old OpenAPI spec from base branch
|
||||
git show "origin/$BASE_BRANCH:hindsight-docs/static/openapi.json" > /tmp/old-openapi.json
|
||||
|
||||
if [ ! -s /tmp/old-openapi.json ]; then
|
||||
echo "⚠️ Warning: Could not find OpenAPI spec in base branch. Skipping compatibility check."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# Check compatibility using our tool
|
||||
cd hindsight-dev
|
||||
uv run check-openapi-compatibility /tmp/old-openapi.json ../hindsight-docs/static/openapi.json
|
||||
@@ -93,6 +93,34 @@ uv run ty check hindsight_api # Type check
|
||||
3. Run tests to ensure nothing breaks
|
||||
4. Submit a PR with a clear description of changes
|
||||
|
||||
## Release Process
|
||||
|
||||
The project uses `scripts/release.sh` for creating releases. This script automates the entire release workflow:
|
||||
|
||||
1. Bumps version in all components (API, clients, CLI, control plane, Helm)
|
||||
2. **Regenerates OpenAPI spec and client SDKs** (Python, TypeScript, Rust)
|
||||
3. Updates documentation versioning
|
||||
4. Creates a commit and git tag
|
||||
5. Pushes to GitHub (triggers CI/CD to publish packages)
|
||||
|
||||
### Usage
|
||||
|
||||
```bash
|
||||
./scripts/release.sh <version>
|
||||
```
|
||||
|
||||
**Example:**
|
||||
```bash
|
||||
./scripts/release.sh 0.5.0
|
||||
```
|
||||
|
||||
### Important for Developers
|
||||
|
||||
- During development, version bumps in `__init__.py` do NOT require client regeneration
|
||||
- Clients are only regenerated during releases
|
||||
- Do not manually run `./scripts/generate-clients.sh` unless testing generation changes
|
||||
- Client version comments will reflect the API version from the latest release
|
||||
|
||||
## Reporting Issues
|
||||
|
||||
Open an issue on GitHub with:
|
||||
|
||||
@@ -169,16 +169,34 @@ ENV PATH="/app/api/.venv/bin:${PATH}"
|
||||
|
||||
# Pre-download ML models to avoid runtime download (conditional)
|
||||
# Only runs if both PRELOAD_ML_MODELS=true AND INCLUDE_LOCAL_MODELS=true
|
||||
# Includes retry logic with exponential backoff for transient network failures
|
||||
ARG PRELOAD_ML_MODELS
|
||||
ARG INCLUDE_LOCAL_MODELS
|
||||
ENV HF_HUB_DOWNLOAD_TIMEOUT=600
|
||||
RUN if [ "$PRELOAD_ML_MODELS" = "true" ] && [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
|
||||
/app/api/.venv/bin/python -c "\
|
||||
MAX_RETRIES=3; \
|
||||
RETRY_DELAY=10; \
|
||||
for i in $(seq 1 $MAX_RETRIES); do \
|
||||
echo "Attempt $i/$MAX_RETRIES: Downloading ML models..."; \
|
||||
/app/api/.venv/bin/python -c "\
|
||||
import os; os.environ['HF_HUB_DOWNLOAD_TIMEOUT'] = '600'; \
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder; \
|
||||
print('Downloading embedding model...'); \
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5'); \
|
||||
print('Downloading cross-encoder model...'); \
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); \
|
||||
print('Models cached successfully')"; \
|
||||
print('Downloading tiktoken encoding...'); import tiktoken; tiktoken.get_encoding('cl100k_base'); \
|
||||
print('Models cached successfully')" && break; \
|
||||
if [ $i -lt $MAX_RETRIES ]; then \
|
||||
echo "Attempt $i failed, retrying in ${RETRY_DELAY}s..."; \
|
||||
sleep $RETRY_DELAY; \
|
||||
RETRY_DELAY=$((RETRY_DELAY * 2)); \
|
||||
fi; \
|
||||
done; \
|
||||
if [ $i -eq $MAX_RETRIES ] && ! /app/api/.venv/bin/python -c "from sentence_transformers import SentenceTransformer; SentenceTransformer('BAAI/bge-small-en-v1.5')" 2>/dev/null; then \
|
||||
echo "ERROR: Failed to download models after $MAX_RETRIES attempts"; \
|
||||
exit 1; \
|
||||
fi; \
|
||||
elif [ "$INCLUDE_LOCAL_MODELS" != "true" ]; then echo "Skipping ML model preload (local-models not included)"; \
|
||||
else echo "Skipping ML model preload"; fi
|
||||
|
||||
@@ -190,6 +208,10 @@ ENV HINDSIGHT_API_LOG_LEVEL=info
|
||||
ENV HINDSIGHT_ENABLE_API=true
|
||||
ENV HINDSIGHT_ENABLE_CP=false
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
# Suppress verbose transformers/HuggingFace model loading warnings
|
||||
ENV TRANSFORMERS_VERBOSITY=error
|
||||
ENV HF_HUB_VERBOSITY=error
|
||||
ENV TOKENIZERS_PARALLELISM=false
|
||||
|
||||
CMD ["/app/start-all.sh"]
|
||||
|
||||
@@ -277,16 +299,34 @@ ENV PATH="/app/api/.venv/bin:${PATH}"
|
||||
|
||||
# Pre-download ML models to avoid runtime download (conditional)
|
||||
# Only runs if both PRELOAD_ML_MODELS=true AND INCLUDE_LOCAL_MODELS=true
|
||||
# Includes retry logic with exponential backoff for transient network failures
|
||||
ARG PRELOAD_ML_MODELS
|
||||
ARG INCLUDE_LOCAL_MODELS
|
||||
ENV HF_HUB_DOWNLOAD_TIMEOUT=600
|
||||
RUN if [ "$PRELOAD_ML_MODELS" = "true" ] && [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
|
||||
/app/api/.venv/bin/python -c "\
|
||||
MAX_RETRIES=3; \
|
||||
RETRY_DELAY=10; \
|
||||
for i in $(seq 1 $MAX_RETRIES); do \
|
||||
echo "Attempt $i/$MAX_RETRIES: Downloading ML models..."; \
|
||||
/app/api/.venv/bin/python -c "\
|
||||
import os; os.environ['HF_HUB_DOWNLOAD_TIMEOUT'] = '600'; \
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder; \
|
||||
print('Downloading embedding model...'); \
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5'); \
|
||||
print('Downloading cross-encoder model...'); \
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); \
|
||||
print('Models cached successfully')"; \
|
||||
print('Downloading tiktoken encoding...'); import tiktoken; tiktoken.get_encoding('cl100k_base'); \
|
||||
print('Models cached successfully')" && break; \
|
||||
if [ $i -lt $MAX_RETRIES ]; then \
|
||||
echo "Attempt $i failed, retrying in ${RETRY_DELAY}s..."; \
|
||||
sleep $RETRY_DELAY; \
|
||||
RETRY_DELAY=$((RETRY_DELAY * 2)); \
|
||||
fi; \
|
||||
done; \
|
||||
if [ $i -eq $MAX_RETRIES ] && ! /app/api/.venv/bin/python -c "from sentence_transformers import SentenceTransformer; SentenceTransformer('BAAI/bge-small-en-v1.5')" 2>/dev/null; then \
|
||||
echo "ERROR: Failed to download models after $MAX_RETRIES attempts"; \
|
||||
exit 1; \
|
||||
fi; \
|
||||
elif [ "$INCLUDE_LOCAL_MODELS" != "true" ]; then echo "Skipping ML model preload (local-models not included)"; \
|
||||
else echo "Skipping ML model preload"; fi
|
||||
|
||||
@@ -300,6 +340,10 @@ ENV HINDSIGHT_CP_DATAPLANE_API_URL=http://localhost:8888
|
||||
ENV HINDSIGHT_ENABLE_API=true
|
||||
ENV HINDSIGHT_ENABLE_CP=true
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
# Suppress verbose transformers/HuggingFace model loading warnings
|
||||
ENV TRANSFORMERS_VERBOSITY=error
|
||||
ENV HF_HUB_VERBOSITY=error
|
||||
ENV TOKENIZERS_PARALLELISM=false
|
||||
|
||||
CMD ["/app/start-all.sh"]
|
||||
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.4.1
|
||||
appVersion: "0.4.1"
|
||||
version: 0.4.9
|
||||
appVersion: "0.4.9"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -46,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.4.1"
|
||||
__version__ = "0.4.9"
|
||||
|
||||
@@ -11,6 +11,7 @@ from collections.abc import Sequence
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from pgvector.sqlalchemy import Vector
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -23,8 +24,21 @@ depends_on: str | Sequence[str] | None = None
|
||||
def upgrade() -> None:
|
||||
"""Upgrade schema - create all tables from scratch."""
|
||||
|
||||
# Enable required extensions
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS vector")
|
||||
# Note: pgvector extension is installed globally BEFORE migrations run
|
||||
# See migrations.py:run_migrations() - this ensures the extension is available
|
||||
# to all schemas, not just the one being migrated
|
||||
|
||||
# We keep this here as a fallback for backwards compatibility
|
||||
# This may fail if user lacks permissions, which is fine if extension already exists
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS vector")
|
||||
except Exception:
|
||||
# Extension might already exist or user lacks permissions - verify it exists
|
||||
conn = op.get_bind()
|
||||
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).fetchone()
|
||||
if not result:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
|
||||
# Create banks table
|
||||
op.create_table(
|
||||
|
||||
@@ -5,6 +5,7 @@ This module provides the create_app function to create and configure
|
||||
the FastAPI application with all API endpoints.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
@@ -35,7 +36,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.engine.db_utils import acquire_with_retry
|
||||
from hindsight_api.engine.memory_engine import Budget, fq_table
|
||||
from hindsight_api.engine.memory_engine import Budget, _get_tiktoken_encoding, fq_table
|
||||
from hindsight_api.engine.reflect.observations import Observation
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES, TokenUsage
|
||||
from hindsight_api.engine.search.tags import TagsMatch
|
||||
@@ -45,6 +46,8 @@ from hindsight_api.models import RequestContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MAX_QUERY_TOKENS = 500 # Maximum tokens allowed in recall query
|
||||
|
||||
|
||||
class EntityIncludeOptions(BaseModel):
|
||||
"""Options for including entity observations in recall results."""
|
||||
@@ -92,8 +95,7 @@ class RecallRequest(BaseModel):
|
||||
query: str
|
||||
types: list[str] | None = Field(
|
||||
default=None,
|
||||
description="List of fact types to recall: 'world', 'experience', 'observation'. Defaults to world and experience if not specified. "
|
||||
"Note: 'opinion' is accepted but ignored (opinions are excluded from recall).",
|
||||
description="List of fact types to recall: 'world', 'experience', 'observation'. Defaults to world and experience if not specified.",
|
||||
)
|
||||
budget: Budget = Budget.MID
|
||||
max_tokens: int = 4096
|
||||
@@ -504,13 +506,6 @@ class ReflectRequest(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class OpinionItem(BaseModel):
|
||||
"""Model for an opinion with confidence score."""
|
||||
|
||||
text: str
|
||||
confidence: float
|
||||
|
||||
|
||||
class ReflectFact(BaseModel):
|
||||
"""A fact used in think response."""
|
||||
|
||||
@@ -529,7 +524,7 @@ class ReflectFact(BaseModel):
|
||||
|
||||
id: str | None = None
|
||||
text: str
|
||||
type: str | None = None # fact type: world, experience, opinion
|
||||
type: str | None = None # fact type: world, experience, observation
|
||||
context: str | None = None
|
||||
occurred_start: str | None = None
|
||||
occurred_end: str | None = None
|
||||
@@ -871,6 +866,7 @@ class ListDocumentsResponse(BaseModel):
|
||||
"updated_at": "2024-01-15T10:30:00Z",
|
||||
"text_length": 5420,
|
||||
"memory_unit_count": 15,
|
||||
"tags": ["user_a", "session_123"],
|
||||
}
|
||||
],
|
||||
"total": 50,
|
||||
@@ -1142,6 +1138,7 @@ class CreateMentalModelRequest(BaseModel):
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"id": "team-communication",
|
||||
"name": "Team Communication Preferences",
|
||||
"source_query": "How does the team prefer to communicate?",
|
||||
"tags": ["team"],
|
||||
@@ -1151,6 +1148,9 @@ class CreateMentalModelRequest(BaseModel):
|
||||
}
|
||||
)
|
||||
|
||||
id: str | None = Field(
|
||||
None, description="Optional custom ID for the mental model (alphanumeric lowercase with hyphens)"
|
||||
)
|
||||
name: str = Field(description="Human-readable name for the mental model")
|
||||
source_query: str = Field(description="The query to run to generate content")
|
||||
tags: list[str] = Field(default_factory=list, description="Tags for scoped visibility")
|
||||
@@ -1161,7 +1161,8 @@ class CreateMentalModelRequest(BaseModel):
|
||||
class CreateMentalModelResponse(BaseModel):
|
||||
"""Response model for mental model creation."""
|
||||
|
||||
operation_id: str = Field(description="Operation ID to track progress")
|
||||
mental_model_id: str | None = Field(None, description="ID of the created mental model")
|
||||
operation_id: str = Field(description="Operation ID to track refresh progress")
|
||||
|
||||
|
||||
class UpdateMentalModelRequest(BaseModel):
|
||||
@@ -1406,15 +1407,21 @@ def create_app(
|
||||
|
||||
# Start worker poller if enabled (standalone mode)
|
||||
if config.worker_enabled and memory._pool is not None:
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA
|
||||
|
||||
worker_id = config.worker_id or socket.gethostname()
|
||||
# Convert default schema to None for SQL compatibility (no schema prefix)
|
||||
schema = None if config.database_schema == DEFAULT_DATABASE_SCHEMA else config.database_schema
|
||||
poller = WorkerPoller(
|
||||
pool=memory._pool,
|
||||
worker_id=worker_id,
|
||||
executor=memory.execute_task,
|
||||
poll_interval_ms=config.worker_poll_interval_ms,
|
||||
batch_size=config.worker_batch_size,
|
||||
max_retries=config.worker_max_retries,
|
||||
tenant_extension=getattr(memory, "_tenant_extension", None),
|
||||
schema=schema,
|
||||
tenant_extension=memory._tenant_extension,
|
||||
max_slots=config.worker_max_slots,
|
||||
consolidation_max_slots=config.worker_consolidation_max_slots,
|
||||
)
|
||||
poller_task = asyncio.create_task(poller.run())
|
||||
logging.info(f"Worker poller started (worker_id={worker_id})")
|
||||
@@ -1707,9 +1714,7 @@ def _register_routes(app: FastAPI):
|
||||
description="Recall memory using semantic similarity and spreading activation.\n\n"
|
||||
"The type parameter is optional and must be one of:\n"
|
||||
"- `world`: General knowledge about people, places, events, and things that happen\n"
|
||||
"- `experience`: Memories about experience, conversations, actions taken, and tasks performed\n"
|
||||
"- `opinion`: The bank's formed beliefs, perspectives, and viewpoints\n\n"
|
||||
"Set `include_entities=true` to get entity observations alongside recall results.",
|
||||
"- `experience`: Memories about experience, conversations, actions taken, and tasks performed",
|
||||
operation_id="recall_memories",
|
||||
tags=["Memory"],
|
||||
)
|
||||
@@ -1722,11 +1727,18 @@ def _register_routes(app: FastAPI):
|
||||
handler_start = time.time()
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
# Validate query length to prevent expensive operations on oversized queries
|
||||
encoding = _get_tiktoken_encoding()
|
||||
query_tokens = len(encoding.encode(request.query))
|
||||
if query_tokens > MAX_QUERY_TOKENS:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Query too long: {query_tokens} tokens exceeds maximum of {MAX_QUERY_TOKENS}. Please shorten your query.",
|
||||
)
|
||||
|
||||
try:
|
||||
# Default to world and experience if not specified (exclude observation and opinion)
|
||||
# Filter out 'opinion' even if requested - opinions are excluded from recall
|
||||
# Default to world and experience if not specified (exclude observation)
|
||||
fact_types = request.types if request.types else list(VALID_RECALL_FACT_TYPES)
|
||||
fact_types = [ft for ft in fact_types if ft != "opinion"]
|
||||
|
||||
# Parse query_timestamp if provided
|
||||
question_date = None
|
||||
@@ -1838,6 +1850,15 @@ def _register_routes(app: FastAPI):
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except (asyncio.TimeoutError, TimeoutError):
|
||||
handler_duration = time.time() - handler_start
|
||||
logger.error(
|
||||
f"[RECALL TIMEOUT] bank={bank_id} handler_duration={handler_duration:.3f}s - database query timed out"
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=504,
|
||||
detail="Request timed out while searching memories. Try a shorter or more specific query.",
|
||||
)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
@@ -1858,8 +1879,7 @@ def _register_routes(app: FastAPI):
|
||||
"2. Retrieves world facts relevant to the query\n"
|
||||
"3. Retrieves existing opinions (bank's perspectives)\n"
|
||||
"4. Uses LLM to formulate a contextual answer\n"
|
||||
"5. Extracts and stores any new opinions formed\n"
|
||||
"6. Returns plain text answer, the facts used, and new opinions",
|
||||
"5. Returns plain text answer and the facts used",
|
||||
operation_id="reflect",
|
||||
tags=["Memory"],
|
||||
)
|
||||
@@ -1896,17 +1916,17 @@ def _register_routes(app: FastAPI):
|
||||
directives = []
|
||||
for fact_type, facts in core_result.based_on.items():
|
||||
if fact_type == "directives":
|
||||
# Directives have different structure (id, name, content)
|
||||
# Directives are dicts with id, name, content (not MemoryFact objects)
|
||||
for directive in facts:
|
||||
directives.append(
|
||||
ReflectDirective(
|
||||
id=directive.id,
|
||||
name=directive.name,
|
||||
content=directive.content,
|
||||
id=directive["id"],
|
||||
name=directive["name"],
|
||||
content=directive["content"],
|
||||
)
|
||||
)
|
||||
elif fact_type == "mental_models":
|
||||
# Mental models are MemoryFact with type "mental_models"
|
||||
elif fact_type == "mental-models":
|
||||
# Mental models are MemoryFact with type "mental-models" (note: hyphen, not underscore)
|
||||
for fact in facts:
|
||||
mental_models.append(
|
||||
ReflectMentalModel(
|
||||
@@ -2297,6 +2317,23 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""Get a mental model by ID."""
|
||||
try:
|
||||
# Pre-operation validation hook
|
||||
validator = app.state.memory._operation_validator
|
||||
if validator:
|
||||
from hindsight_api.extensions.operation_validator import MentalModelGetContext
|
||||
|
||||
ctx = MentalModelGetContext(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
validation = await validator.validate_mental_model_get(ctx)
|
||||
if not validation.allowed:
|
||||
raise OperationValidationError(
|
||||
validation.reason or "Operation not allowed",
|
||||
status_code=validation.status_code,
|
||||
)
|
||||
|
||||
mental_model = await app.state.memory.get_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
@@ -2304,9 +2341,31 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
if mental_model is None:
|
||||
raise HTTPException(status_code=404, detail=f"Mental model '{mental_model_id}' not found")
|
||||
|
||||
# Post-operation hook
|
||||
if validator:
|
||||
from hindsight_api.extensions.operation_validator import MentalModelGetResult
|
||||
|
||||
content = mental_model.get("content", "")
|
||||
output_tokens = len(content) // 4 if content else 0
|
||||
|
||||
result_ctx = MentalModelGetResult(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=request_context,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
try:
|
||||
await validator.on_mental_model_get_complete(result_ctx)
|
||||
except Exception as hook_err:
|
||||
logger.warning(f"Post-mental-model-get hook error (non-fatal): {hook_err}")
|
||||
|
||||
return MentalModelResponse(**mental_model)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
@@ -2331,12 +2390,30 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""Create a mental model (async - returns operation_id)."""
|
||||
try:
|
||||
# Pre-operation validation hook
|
||||
validator = app.state.memory._operation_validator
|
||||
if validator:
|
||||
from hindsight_api.extensions.operation_validator import MentalModelRefreshContext
|
||||
|
||||
ctx = MentalModelRefreshContext(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=None, # Not yet created
|
||||
request_context=request_context,
|
||||
)
|
||||
validation = await validator.validate_mental_model_refresh(ctx)
|
||||
if not validation.allowed:
|
||||
raise OperationValidationError(
|
||||
validation.reason or "Operation not allowed",
|
||||
status_code=validation.status_code,
|
||||
)
|
||||
|
||||
# 1. Create the mental model with placeholder content
|
||||
mental_model = await app.state.memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name=body.name,
|
||||
source_query=body.source_query,
|
||||
content="Generating content...",
|
||||
mental_model_id=body.id if body.id else None,
|
||||
tags=body.tags if body.tags else None,
|
||||
max_tokens=body.max_tokens,
|
||||
trigger=body.trigger.model_dump() if body.trigger else None,
|
||||
@@ -2348,11 +2425,13 @@ def _register_routes(app: FastAPI):
|
||||
mental_model_id=mental_model["id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
return CreateMentalModelResponse(operation_id=result["operation_id"])
|
||||
return CreateMentalModelResponse(mental_model_id=mental_model["id"], operation_id=result["operation_id"])
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
@@ -2375,6 +2454,23 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""Refresh a mental model by re-running its source query (async)."""
|
||||
try:
|
||||
# Pre-operation validation hook
|
||||
validator = app.state.memory._operation_validator
|
||||
if validator:
|
||||
from hindsight_api.extensions.operation_validator import MentalModelRefreshContext
|
||||
|
||||
ctx = MentalModelRefreshContext(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
validation = await validator.validate_mental_model_refresh(ctx)
|
||||
if not validation.allowed:
|
||||
raise OperationValidationError(
|
||||
validation.reason or "Operation not allowed",
|
||||
status_code=validation.status_code,
|
||||
)
|
||||
|
||||
result = await app.state.memory.submit_async_refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
@@ -2385,6 +2481,8 @@ def _register_routes(app: FastAPI):
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
|
||||
@@ -29,15 +29,26 @@ logger = logging.getLogger(__name__)
|
||||
# Default bank_id from environment variable
|
||||
DEFAULT_BANK_ID = os.environ.get("HINDSIGHT_MCP_BANK_ID", "default")
|
||||
|
||||
# MCP authentication token (optional - if set, Bearer token auth is required)
|
||||
MCP_AUTH_TOKEN = os.environ.get("HINDSIGHT_API_MCP_AUTH_TOKEN")
|
||||
|
||||
# Context variable to hold the current bank_id
|
||||
_current_bank_id: ContextVar[str | None] = ContextVar("current_bank_id", default=None)
|
||||
|
||||
# Context variable to hold the current API key (for tenant auth propagation)
|
||||
_current_api_key: ContextVar[str | None] = ContextVar("current_api_key", default=None)
|
||||
|
||||
|
||||
def get_current_bank_id() -> str | None:
|
||||
"""Get the current bank_id from context."""
|
||||
return _current_bank_id.get()
|
||||
|
||||
|
||||
def get_current_api_key() -> str | None:
|
||||
"""Get the current API key from context."""
|
||||
return _current_api_key.get()
|
||||
|
||||
|
||||
def create_mcp_server(memory: MemoryEngine) -> FastMCP:
|
||||
"""
|
||||
Create and configure the Hindsight MCP server.
|
||||
@@ -54,6 +65,7 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
|
||||
# Configure and register tools using shared module
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=get_current_bank_id,
|
||||
api_key_resolver=get_current_api_key, # Propagate API key for tenant auth
|
||||
include_bank_id_param=True, # HTTP MCP supports multi-bank via parameter
|
||||
tools=None, # All tools
|
||||
retain_fire_and_forget=False, # HTTP MCP supports sync/async modes
|
||||
@@ -65,7 +77,11 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
|
||||
|
||||
|
||||
class MCPMiddleware:
|
||||
"""ASGI middleware that extracts bank_id from header or path and sets context.
|
||||
"""ASGI middleware that handles authentication and extracts bank_id from header or path.
|
||||
|
||||
Authentication:
|
||||
If HINDSIGHT_API_MCP_AUTH_TOKEN is set, all requests must include a valid
|
||||
Authorization header with Bearer token or direct token matching the configured value.
|
||||
|
||||
Bank ID can be provided via:
|
||||
1. X-Bank-Id header (recommended for Claude Code)
|
||||
@@ -74,7 +90,7 @@ class MCPMiddleware:
|
||||
|
||||
For Claude Code, configure with:
|
||||
claude mcp add --transport http hindsight http://localhost:8888/mcp \\
|
||||
--header "X-Bank-Id: my-bank"
|
||||
--header "X-Bank-Id: my-bank" --header "Authorization: Bearer <token>"
|
||||
"""
|
||||
|
||||
def __init__(self, app, memory: MemoryEngine):
|
||||
@@ -98,6 +114,22 @@ class MCPMiddleware:
|
||||
await self.mcp_app(scope, receive, send)
|
||||
return
|
||||
|
||||
# Extract auth token from header (for tenant auth propagation)
|
||||
auth_header = self._get_header(scope, "Authorization")
|
||||
auth_token: str | None = None
|
||||
if auth_header:
|
||||
# Support both "Bearer <token>" and direct token
|
||||
auth_token = auth_header[7:].strip() if auth_header.startswith("Bearer ") else auth_header.strip()
|
||||
|
||||
# Authenticate if MCP_AUTH_TOKEN is configured
|
||||
if MCP_AUTH_TOKEN:
|
||||
if not auth_token:
|
||||
await self._send_error(send, 401, "Authorization header required")
|
||||
return
|
||||
if auth_token != MCP_AUTH_TOKEN:
|
||||
await self._send_error(send, 401, "Invalid authentication token")
|
||||
return
|
||||
|
||||
path = scope.get("path", "")
|
||||
|
||||
# Strip any mount prefix (e.g., /mcp) that FastAPI might not have stripped
|
||||
@@ -132,8 +164,10 @@ class MCPMiddleware:
|
||||
bank_id = DEFAULT_BANK_ID
|
||||
logger.debug(f"Using default bank_id: {bank_id}")
|
||||
|
||||
# Set bank_id context
|
||||
token = _current_bank_id.set(bank_id)
|
||||
# Set bank_id and api_key context
|
||||
bank_id_token = _current_bank_id.set(bank_id)
|
||||
# Store the auth token for tenant extension to validate
|
||||
api_key_token = _current_api_key.set(auth_token) if auth_token else None
|
||||
try:
|
||||
new_scope = scope.copy()
|
||||
new_scope["path"] = new_path
|
||||
@@ -152,7 +186,9 @@ class MCPMiddleware:
|
||||
|
||||
await self.mcp_app(new_scope, receive, send_wrapper)
|
||||
finally:
|
||||
_current_bank_id.reset(token)
|
||||
_current_bank_id.reset(bank_id_token)
|
||||
if api_key_token is not None:
|
||||
_current_api_key.reset(api_key_token)
|
||||
|
||||
async def _send_error(self, send, status: int, message: str):
|
||||
"""Send an error response."""
|
||||
@@ -176,6 +212,10 @@ def create_mcp_app(memory: MemoryEngine):
|
||||
"""
|
||||
Create an ASGI app that handles MCP requests.
|
||||
|
||||
Authentication:
|
||||
Set HINDSIGHT_API_MCP_AUTH_TOKEN to require Bearer token authentication.
|
||||
If not set, MCP endpoint is open (for local development).
|
||||
|
||||
Bank ID can be provided via:
|
||||
1. X-Bank-Id header: claude mcp add --transport http hindsight http://localhost:8888/mcp --header "X-Bank-Id: my-bank"
|
||||
2. URL path: /mcp/{bank_id}/
|
||||
|
||||
@@ -83,9 +83,12 @@ def print_startup_info(
|
||||
embeddings_provider: str,
|
||||
reranker_provider: str,
|
||||
mcp_enabled: bool = False,
|
||||
version: str | None = None,
|
||||
):
|
||||
"""Print styled startup information."""
|
||||
print(color_start("Starting Hindsight API..."))
|
||||
if version:
|
||||
print(f" {dim('Version:')} {color(f'v{version}', 0.1)}")
|
||||
print(f" {dim('URL:')} {color(f'http://{host}:{port}', 0.2)}")
|
||||
print(f" {dim('Database:')} {color(database_url, 0.4)}")
|
||||
print(f" {dim('LLM:')} {color(f'{llm_provider} / {llm_model}', 0.6)}")
|
||||
|
||||
@@ -26,6 +26,9 @@ ENV_LLM_API_KEY = "HINDSIGHT_API_LLM_API_KEY"
|
||||
ENV_LLM_MODEL = "HINDSIGHT_API_LLM_MODEL"
|
||||
ENV_LLM_BASE_URL = "HINDSIGHT_API_LLM_BASE_URL"
|
||||
ENV_LLM_MAX_CONCURRENT = "HINDSIGHT_API_LLM_MAX_CONCURRENT"
|
||||
ENV_LLM_MAX_RETRIES = "HINDSIGHT_API_LLM_MAX_RETRIES"
|
||||
ENV_LLM_INITIAL_BACKOFF = "HINDSIGHT_API_LLM_INITIAL_BACKOFF"
|
||||
ENV_LLM_MAX_BACKOFF = "HINDSIGHT_API_LLM_MAX_BACKOFF"
|
||||
ENV_LLM_TIMEOUT = "HINDSIGHT_API_LLM_TIMEOUT"
|
||||
ENV_LLM_GROQ_SERVICE_TIER = "HINDSIGHT_API_LLM_GROQ_SERVICE_TIER"
|
||||
|
||||
@@ -34,16 +37,31 @@ ENV_RETAIN_LLM_PROVIDER = "HINDSIGHT_API_RETAIN_LLM_PROVIDER"
|
||||
ENV_RETAIN_LLM_API_KEY = "HINDSIGHT_API_RETAIN_LLM_API_KEY"
|
||||
ENV_RETAIN_LLM_MODEL = "HINDSIGHT_API_RETAIN_LLM_MODEL"
|
||||
ENV_RETAIN_LLM_BASE_URL = "HINDSIGHT_API_RETAIN_LLM_BASE_URL"
|
||||
ENV_RETAIN_LLM_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_LLM_MAX_CONCURRENT"
|
||||
ENV_RETAIN_LLM_MAX_RETRIES = "HINDSIGHT_API_RETAIN_LLM_MAX_RETRIES"
|
||||
ENV_RETAIN_LLM_INITIAL_BACKOFF = "HINDSIGHT_API_RETAIN_LLM_INITIAL_BACKOFF"
|
||||
ENV_RETAIN_LLM_MAX_BACKOFF = "HINDSIGHT_API_RETAIN_LLM_MAX_BACKOFF"
|
||||
ENV_RETAIN_LLM_TIMEOUT = "HINDSIGHT_API_RETAIN_LLM_TIMEOUT"
|
||||
|
||||
ENV_REFLECT_LLM_PROVIDER = "HINDSIGHT_API_REFLECT_LLM_PROVIDER"
|
||||
ENV_REFLECT_LLM_API_KEY = "HINDSIGHT_API_REFLECT_LLM_API_KEY"
|
||||
ENV_REFLECT_LLM_MODEL = "HINDSIGHT_API_REFLECT_LLM_MODEL"
|
||||
ENV_REFLECT_LLM_BASE_URL = "HINDSIGHT_API_REFLECT_LLM_BASE_URL"
|
||||
ENV_REFLECT_LLM_MAX_CONCURRENT = "HINDSIGHT_API_REFLECT_LLM_MAX_CONCURRENT"
|
||||
ENV_REFLECT_LLM_MAX_RETRIES = "HINDSIGHT_API_REFLECT_LLM_MAX_RETRIES"
|
||||
ENV_REFLECT_LLM_INITIAL_BACKOFF = "HINDSIGHT_API_REFLECT_LLM_INITIAL_BACKOFF"
|
||||
ENV_REFLECT_LLM_MAX_BACKOFF = "HINDSIGHT_API_REFLECT_LLM_MAX_BACKOFF"
|
||||
ENV_REFLECT_LLM_TIMEOUT = "HINDSIGHT_API_REFLECT_LLM_TIMEOUT"
|
||||
|
||||
ENV_CONSOLIDATION_LLM_PROVIDER = "HINDSIGHT_API_CONSOLIDATION_LLM_PROVIDER"
|
||||
ENV_CONSOLIDATION_LLM_API_KEY = "HINDSIGHT_API_CONSOLIDATION_LLM_API_KEY"
|
||||
ENV_CONSOLIDATION_LLM_MODEL = "HINDSIGHT_API_CONSOLIDATION_LLM_MODEL"
|
||||
ENV_CONSOLIDATION_LLM_BASE_URL = "HINDSIGHT_API_CONSOLIDATION_LLM_BASE_URL"
|
||||
ENV_CONSOLIDATION_LLM_MAX_CONCURRENT = "HINDSIGHT_API_CONSOLIDATION_LLM_MAX_CONCURRENT"
|
||||
ENV_CONSOLIDATION_LLM_MAX_RETRIES = "HINDSIGHT_API_CONSOLIDATION_LLM_MAX_RETRIES"
|
||||
ENV_CONSOLIDATION_LLM_INITIAL_BACKOFF = "HINDSIGHT_API_CONSOLIDATION_LLM_INITIAL_BACKOFF"
|
||||
ENV_CONSOLIDATION_LLM_MAX_BACKOFF = "HINDSIGHT_API_CONSOLIDATION_LLM_MAX_BACKOFF"
|
||||
ENV_CONSOLIDATION_LLM_TIMEOUT = "HINDSIGHT_API_CONSOLIDATION_LLM_TIMEOUT"
|
||||
|
||||
ENV_EMBEDDINGS_PROVIDER = "HINDSIGHT_API_EMBEDDINGS_PROVIDER"
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL"
|
||||
@@ -90,13 +108,17 @@ ENV_MCP_LOCAL_BANK_ID = "HINDSIGHT_API_MCP_LOCAL_BANK_ID"
|
||||
ENV_MCP_INSTRUCTIONS = "HINDSIGHT_API_MCP_INSTRUCTIONS"
|
||||
ENV_MENTAL_MODEL_REFRESH_CONCURRENCY = "HINDSIGHT_API_MENTAL_MODEL_REFRESH_CONCURRENCY"
|
||||
|
||||
# Vertex AI configuration
|
||||
ENV_LLM_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID"
|
||||
ENV_LLM_VERTEXAI_REGION = "HINDSIGHT_API_LLM_VERTEXAI_REGION"
|
||||
ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY"
|
||||
|
||||
# Retain settings
|
||||
ENV_RETAIN_MAX_COMPLETION_TOKENS = "HINDSIGHT_API_RETAIN_MAX_COMPLETION_TOKENS"
|
||||
ENV_RETAIN_CHUNK_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_SIZE"
|
||||
ENV_RETAIN_EXTRACT_CAUSAL_LINKS = "HINDSIGHT_API_RETAIN_EXTRACT_CAUSAL_LINKS"
|
||||
ENV_RETAIN_EXTRACTION_MODE = "HINDSIGHT_API_RETAIN_EXTRACTION_MODE"
|
||||
ENV_RETAIN_CUSTOM_INSTRUCTIONS = "HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS"
|
||||
ENV_RETAIN_OBSERVATIONS_ASYNC = "HINDSIGHT_API_RETAIN_OBSERVATIONS_ASYNC"
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
ENV_ENABLE_OBSERVATIONS = "HINDSIGHT_API_ENABLE_OBSERVATIONS"
|
||||
@@ -121,8 +143,9 @@ ENV_WORKER_ENABLED = "HINDSIGHT_API_WORKER_ENABLED"
|
||||
ENV_WORKER_ID = "HINDSIGHT_API_WORKER_ID"
|
||||
ENV_WORKER_POLL_INTERVAL_MS = "HINDSIGHT_API_WORKER_POLL_INTERVAL_MS"
|
||||
ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES"
|
||||
ENV_WORKER_BATCH_SIZE = "HINDSIGHT_API_WORKER_BATCH_SIZE"
|
||||
ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT"
|
||||
ENV_WORKER_MAX_SLOTS = "HINDSIGHT_API_WORKER_MAX_SLOTS"
|
||||
ENV_WORKER_CONSOLIDATION_MAX_SLOTS = "HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLOTS"
|
||||
|
||||
# Reflect agent settings
|
||||
ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS"
|
||||
@@ -131,10 +154,32 @@ ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS"
|
||||
DEFAULT_DATABASE_URL = "pg0"
|
||||
DEFAULT_DATABASE_SCHEMA = "public"
|
||||
DEFAULT_LLM_PROVIDER = "openai"
|
||||
DEFAULT_LLM_MODEL = "gpt-5-mini"
|
||||
|
||||
# Provider-specific default models
|
||||
PROVIDER_DEFAULT_MODELS = {
|
||||
"openai": "o3-mini",
|
||||
"anthropic": "claude-haiku-4-5-20251001",
|
||||
"gemini": "gemini-2.5-flash",
|
||||
"groq": "openai/gpt-oss-120b",
|
||||
"ollama": "gemma3:12b",
|
||||
"lmstudio": "local-model",
|
||||
"vertexai": "gemini-2.0-flash-001",
|
||||
"openai-codex": "gpt-5.2-codex",
|
||||
"claude-code": "claude-sonnet-4-5-20250929",
|
||||
"mock": "mock-model",
|
||||
}
|
||||
DEFAULT_LLM_MODEL = "o3-mini" # Fallback if provider not in table
|
||||
DEFAULT_LLM_MAX_CONCURRENT = 32
|
||||
DEFAULT_LLM_MAX_RETRIES = 10 # Max retry attempts for LLM API calls
|
||||
DEFAULT_LLM_INITIAL_BACKOFF = 1.0 # Initial backoff in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_MAX_BACKOFF = 60.0 # Max backoff cap in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_TIMEOUT = 120.0 # seconds
|
||||
|
||||
# Vertex AI defaults
|
||||
DEFAULT_LLM_VERTEXAI_PROJECT_ID = None # Required for Vertex AI
|
||||
DEFAULT_LLM_VERTEXAI_REGION = "us-central1"
|
||||
DEFAULT_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY = None # Optional, uses ADC if not set
|
||||
|
||||
DEFAULT_EMBEDDINGS_PROVIDER = "local"
|
||||
DEFAULT_EMBEDDINGS_LOCAL_MODEL = "BAAI/bge-small-en-v1.5"
|
||||
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU = False # Force CPU mode for local embeddings (avoids MPS/XPC issues on macOS)
|
||||
@@ -179,7 +224,6 @@ DEFAULT_RETAIN_EXTRACT_CAUSAL_LINKS = True # Extract causal links between facts
|
||||
DEFAULT_RETAIN_EXTRACTION_MODE = "concise" # Extraction mode: "concise", "verbose", or "custom"
|
||||
RETAIN_EXTRACTION_MODES = ("concise", "verbose", "custom") # Allowed extraction modes
|
||||
DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS = None # Custom extraction guidelines (only used when mode="custom")
|
||||
DEFAULT_RETAIN_OBSERVATIONS_ASYNC = False # Run observation generation async (after retain completes)
|
||||
|
||||
# Observations defaults (consolidated knowledge from facts)
|
||||
DEFAULT_ENABLE_OBSERVATIONS = True # Observations enabled by default
|
||||
@@ -200,8 +244,9 @@ DEFAULT_WORKER_ENABLED = True # API runs worker by default (standalone mode)
|
||||
DEFAULT_WORKER_ID = None # Will use hostname if not specified
|
||||
DEFAULT_WORKER_POLL_INTERVAL_MS = 500 # Poll database every 500ms
|
||||
DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed
|
||||
DEFAULT_WORKER_BATCH_SIZE = 10 # Tasks to claim per poll cycle
|
||||
DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health
|
||||
DEFAULT_WORKER_MAX_SLOTS = 10 # Total concurrent tasks per worker
|
||||
DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS = 2 # Max concurrent consolidation tasks per worker
|
||||
|
||||
# Reflect agent settings
|
||||
DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response
|
||||
@@ -272,6 +317,11 @@ def _validate_extraction_mode(mode: str) -> str:
|
||||
return mode_lower
|
||||
|
||||
|
||||
def _get_default_model_for_provider(provider: str) -> str:
|
||||
"""Get the default model for a given provider."""
|
||||
return PROVIDER_DEFAULT_MODELS.get(provider.lower(), DEFAULT_LLM_MODEL)
|
||||
|
||||
|
||||
@dataclass
|
||||
class HindsightConfig:
|
||||
"""Configuration container for Hindsight API."""
|
||||
@@ -286,23 +336,46 @@ class HindsightConfig:
|
||||
llm_model: str
|
||||
llm_base_url: str | None
|
||||
llm_max_concurrent: int
|
||||
llm_max_retries: int
|
||||
llm_initial_backoff: float
|
||||
llm_max_backoff: float
|
||||
llm_timeout: float
|
||||
|
||||
# Vertex AI configuration
|
||||
llm_vertexai_project_id: str | None
|
||||
llm_vertexai_region: str
|
||||
llm_vertexai_service_account_key: str | None
|
||||
|
||||
# Per-operation LLM configuration (None = use default LLM config)
|
||||
retain_llm_provider: str | None
|
||||
retain_llm_api_key: str | None
|
||||
retain_llm_model: str | None
|
||||
retain_llm_base_url: str | None
|
||||
retain_llm_max_concurrent: int | None
|
||||
retain_llm_max_retries: int | None
|
||||
retain_llm_initial_backoff: float | None
|
||||
retain_llm_max_backoff: float | None
|
||||
retain_llm_timeout: float | None
|
||||
|
||||
reflect_llm_provider: str | None
|
||||
reflect_llm_api_key: str | None
|
||||
reflect_llm_model: str | None
|
||||
reflect_llm_base_url: str | None
|
||||
reflect_llm_max_concurrent: int | None
|
||||
reflect_llm_max_retries: int | None
|
||||
reflect_llm_initial_backoff: float | None
|
||||
reflect_llm_max_backoff: float | None
|
||||
reflect_llm_timeout: float | None
|
||||
|
||||
consolidation_llm_provider: str | None
|
||||
consolidation_llm_api_key: str | None
|
||||
consolidation_llm_model: str | None
|
||||
consolidation_llm_base_url: str | None
|
||||
consolidation_llm_max_concurrent: int | None
|
||||
consolidation_llm_max_retries: int | None
|
||||
consolidation_llm_initial_backoff: float | None
|
||||
consolidation_llm_max_backoff: float | None
|
||||
consolidation_llm_timeout: float | None
|
||||
|
||||
# Embeddings
|
||||
embeddings_provider: str
|
||||
@@ -343,7 +416,6 @@ class HindsightConfig:
|
||||
retain_extract_causal_links: bool
|
||||
retain_extraction_mode: str
|
||||
retain_custom_instructions: str | None
|
||||
retain_observations_async: bool
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations: bool
|
||||
@@ -368,8 +440,9 @@ class HindsightConfig:
|
||||
worker_id: str | None
|
||||
worker_poll_interval_ms: int
|
||||
worker_max_retries: int
|
||||
worker_batch_size: int
|
||||
worker_http_port: int
|
||||
worker_max_slots: int
|
||||
worker_consolidation_max_slots: int
|
||||
|
||||
# Reflect agent settings
|
||||
reflect_max_iterations: int
|
||||
@@ -377,30 +450,100 @@ class HindsightConfig:
|
||||
@classmethod
|
||||
def from_env(cls) -> "HindsightConfig":
|
||||
"""Create configuration from environment variables."""
|
||||
# Get provider first to determine default model
|
||||
llm_provider = os.getenv(ENV_LLM_PROVIDER, DEFAULT_LLM_PROVIDER)
|
||||
llm_model = os.getenv(ENV_LLM_MODEL) or _get_default_model_for_provider(llm_provider)
|
||||
|
||||
return cls(
|
||||
# Database
|
||||
database_url=os.getenv(ENV_DATABASE_URL, DEFAULT_DATABASE_URL),
|
||||
database_schema=os.getenv(ENV_DATABASE_SCHEMA, DEFAULT_DATABASE_SCHEMA),
|
||||
# LLM
|
||||
llm_provider=os.getenv(ENV_LLM_PROVIDER, DEFAULT_LLM_PROVIDER),
|
||||
llm_provider=llm_provider,
|
||||
llm_api_key=os.getenv(ENV_LLM_API_KEY),
|
||||
llm_model=os.getenv(ENV_LLM_MODEL, DEFAULT_LLM_MODEL),
|
||||
llm_model=llm_model,
|
||||
llm_base_url=os.getenv(ENV_LLM_BASE_URL) or None,
|
||||
llm_max_concurrent=int(os.getenv(ENV_LLM_MAX_CONCURRENT, str(DEFAULT_LLM_MAX_CONCURRENT))),
|
||||
llm_max_retries=int(os.getenv(ENV_LLM_MAX_RETRIES, str(DEFAULT_LLM_MAX_RETRIES))),
|
||||
llm_initial_backoff=float(os.getenv(ENV_LLM_INITIAL_BACKOFF, str(DEFAULT_LLM_INITIAL_BACKOFF))),
|
||||
llm_max_backoff=float(os.getenv(ENV_LLM_MAX_BACKOFF, str(DEFAULT_LLM_MAX_BACKOFF))),
|
||||
llm_timeout=float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT))),
|
||||
# Vertex AI
|
||||
llm_vertexai_project_id=os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID) or DEFAULT_LLM_VERTEXAI_PROJECT_ID,
|
||||
llm_vertexai_region=os.getenv(ENV_LLM_VERTEXAI_REGION, DEFAULT_LLM_VERTEXAI_REGION),
|
||||
llm_vertexai_service_account_key=os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY)
|
||||
or DEFAULT_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY,
|
||||
# Per-operation LLM config (None = use default)
|
||||
retain_llm_provider=os.getenv(ENV_RETAIN_LLM_PROVIDER) or None,
|
||||
retain_llm_api_key=os.getenv(ENV_RETAIN_LLM_API_KEY) or None,
|
||||
retain_llm_model=os.getenv(ENV_RETAIN_LLM_MODEL) or None,
|
||||
retain_llm_model=os.getenv(ENV_RETAIN_LLM_MODEL)
|
||||
or (
|
||||
_get_default_model_for_provider(os.getenv(ENV_RETAIN_LLM_PROVIDER))
|
||||
if os.getenv(ENV_RETAIN_LLM_PROVIDER)
|
||||
else None
|
||||
),
|
||||
retain_llm_base_url=os.getenv(ENV_RETAIN_LLM_BASE_URL) or None,
|
||||
retain_llm_max_concurrent=int(os.getenv(ENV_RETAIN_LLM_MAX_CONCURRENT))
|
||||
if os.getenv(ENV_RETAIN_LLM_MAX_CONCURRENT)
|
||||
else None,
|
||||
retain_llm_max_retries=int(os.getenv(ENV_RETAIN_LLM_MAX_RETRIES))
|
||||
if os.getenv(ENV_RETAIN_LLM_MAX_RETRIES)
|
||||
else None,
|
||||
retain_llm_initial_backoff=float(os.getenv(ENV_RETAIN_LLM_INITIAL_BACKOFF))
|
||||
if os.getenv(ENV_RETAIN_LLM_INITIAL_BACKOFF)
|
||||
else None,
|
||||
retain_llm_max_backoff=float(os.getenv(ENV_RETAIN_LLM_MAX_BACKOFF))
|
||||
if os.getenv(ENV_RETAIN_LLM_MAX_BACKOFF)
|
||||
else None,
|
||||
retain_llm_timeout=float(os.getenv(ENV_RETAIN_LLM_TIMEOUT)) if os.getenv(ENV_RETAIN_LLM_TIMEOUT) else None,
|
||||
reflect_llm_provider=os.getenv(ENV_REFLECT_LLM_PROVIDER) or None,
|
||||
reflect_llm_api_key=os.getenv(ENV_REFLECT_LLM_API_KEY) or None,
|
||||
reflect_llm_model=os.getenv(ENV_REFLECT_LLM_MODEL) or None,
|
||||
reflect_llm_model=os.getenv(ENV_REFLECT_LLM_MODEL)
|
||||
or (
|
||||
_get_default_model_for_provider(os.getenv(ENV_REFLECT_LLM_PROVIDER))
|
||||
if os.getenv(ENV_REFLECT_LLM_PROVIDER)
|
||||
else None
|
||||
),
|
||||
reflect_llm_base_url=os.getenv(ENV_REFLECT_LLM_BASE_URL) or None,
|
||||
reflect_llm_max_concurrent=int(os.getenv(ENV_REFLECT_LLM_MAX_CONCURRENT))
|
||||
if os.getenv(ENV_REFLECT_LLM_MAX_CONCURRENT)
|
||||
else None,
|
||||
reflect_llm_max_retries=int(os.getenv(ENV_REFLECT_LLM_MAX_RETRIES))
|
||||
if os.getenv(ENV_REFLECT_LLM_MAX_RETRIES)
|
||||
else None,
|
||||
reflect_llm_initial_backoff=float(os.getenv(ENV_REFLECT_LLM_INITIAL_BACKOFF))
|
||||
if os.getenv(ENV_REFLECT_LLM_INITIAL_BACKOFF)
|
||||
else None,
|
||||
reflect_llm_max_backoff=float(os.getenv(ENV_REFLECT_LLM_MAX_BACKOFF))
|
||||
if os.getenv(ENV_REFLECT_LLM_MAX_BACKOFF)
|
||||
else None,
|
||||
reflect_llm_timeout=float(os.getenv(ENV_REFLECT_LLM_TIMEOUT))
|
||||
if os.getenv(ENV_REFLECT_LLM_TIMEOUT)
|
||||
else None,
|
||||
consolidation_llm_provider=os.getenv(ENV_CONSOLIDATION_LLM_PROVIDER) or None,
|
||||
consolidation_llm_api_key=os.getenv(ENV_CONSOLIDATION_LLM_API_KEY) or None,
|
||||
consolidation_llm_model=os.getenv(ENV_CONSOLIDATION_LLM_MODEL) or None,
|
||||
consolidation_llm_model=os.getenv(ENV_CONSOLIDATION_LLM_MODEL)
|
||||
or (
|
||||
_get_default_model_for_provider(os.getenv(ENV_CONSOLIDATION_LLM_PROVIDER))
|
||||
if os.getenv(ENV_CONSOLIDATION_LLM_PROVIDER)
|
||||
else None
|
||||
),
|
||||
consolidation_llm_base_url=os.getenv(ENV_CONSOLIDATION_LLM_BASE_URL) or None,
|
||||
consolidation_llm_max_concurrent=int(os.getenv(ENV_CONSOLIDATION_LLM_MAX_CONCURRENT))
|
||||
if os.getenv(ENV_CONSOLIDATION_LLM_MAX_CONCURRENT)
|
||||
else None,
|
||||
consolidation_llm_max_retries=int(os.getenv(ENV_CONSOLIDATION_LLM_MAX_RETRIES))
|
||||
if os.getenv(ENV_CONSOLIDATION_LLM_MAX_RETRIES)
|
||||
else None,
|
||||
consolidation_llm_initial_backoff=float(os.getenv(ENV_CONSOLIDATION_LLM_INITIAL_BACKOFF))
|
||||
if os.getenv(ENV_CONSOLIDATION_LLM_INITIAL_BACKOFF)
|
||||
else None,
|
||||
consolidation_llm_max_backoff=float(os.getenv(ENV_CONSOLIDATION_LLM_MAX_BACKOFF))
|
||||
if os.getenv(ENV_CONSOLIDATION_LLM_MAX_BACKOFF)
|
||||
else None,
|
||||
consolidation_llm_timeout=float(os.getenv(ENV_CONSOLIDATION_LLM_TIMEOUT))
|
||||
if os.getenv(ENV_CONSOLIDATION_LLM_TIMEOUT)
|
||||
else None,
|
||||
# Embeddings
|
||||
embeddings_provider=os.getenv(ENV_EMBEDDINGS_PROVIDER, DEFAULT_EMBEDDINGS_PROVIDER),
|
||||
embeddings_local_model=os.getenv(ENV_EMBEDDINGS_LOCAL_MODEL, DEFAULT_EMBEDDINGS_LOCAL_MODEL),
|
||||
@@ -460,10 +603,6 @@ class HindsightConfig:
|
||||
os.getenv(ENV_RETAIN_EXTRACTION_MODE, DEFAULT_RETAIN_EXTRACTION_MODE)
|
||||
),
|
||||
retain_custom_instructions=os.getenv(ENV_RETAIN_CUSTOM_INSTRUCTIONS) or DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS,
|
||||
retain_observations_async=os.getenv(
|
||||
ENV_RETAIN_OBSERVATIONS_ASYNC, str(DEFAULT_RETAIN_OBSERVATIONS_ASYNC)
|
||||
).lower()
|
||||
== "true",
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations=os.getenv(ENV_ENABLE_OBSERVATIONS, str(DEFAULT_ENABLE_OBSERVATIONS)).lower() == "true",
|
||||
consolidation_batch_size=int(
|
||||
@@ -484,8 +623,11 @@ class HindsightConfig:
|
||||
worker_id=os.getenv(ENV_WORKER_ID) or DEFAULT_WORKER_ID,
|
||||
worker_poll_interval_ms=int(os.getenv(ENV_WORKER_POLL_INTERVAL_MS, str(DEFAULT_WORKER_POLL_INTERVAL_MS))),
|
||||
worker_max_retries=int(os.getenv(ENV_WORKER_MAX_RETRIES, str(DEFAULT_WORKER_MAX_RETRIES))),
|
||||
worker_batch_size=int(os.getenv(ENV_WORKER_BATCH_SIZE, str(DEFAULT_WORKER_BATCH_SIZE))),
|
||||
worker_http_port=int(os.getenv(ENV_WORKER_HTTP_PORT, str(DEFAULT_WORKER_HTTP_PORT))),
|
||||
worker_max_slots=int(os.getenv(ENV_WORKER_MAX_SLOTS, str(DEFAULT_WORKER_MAX_SLOTS))),
|
||||
worker_consolidation_max_slots=int(
|
||||
os.getenv(ENV_WORKER_CONSOLIDATION_MAX_SLOTS, str(DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS))
|
||||
),
|
||||
# Reflect agent settings
|
||||
reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))),
|
||||
)
|
||||
|
||||
@@ -1,11 +1,10 @@
|
||||
"""
|
||||
Daemon mode support for Hindsight API.
|
||||
|
||||
Provides idle timeout and lockfile management for running as a background daemon.
|
||||
Provides idle timeout for running as a background daemon.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import fcntl
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
@@ -15,10 +14,11 @@ from pathlib import Path
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default daemon configuration
|
||||
DEFAULT_DAEMON_PORT = 8889
|
||||
DEFAULT_DAEMON_PORT = 8888
|
||||
DEFAULT_IDLE_TIMEOUT = 0 # 0 = no auto-exit (hindsight-embed passes its own timeout)
|
||||
LOCKFILE_PATH = Path.home() / ".hindsight" / "daemon.lock"
|
||||
DAEMON_LOG_PATH = Path.home() / ".hindsight" / "daemon.log"
|
||||
|
||||
# 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")))
|
||||
|
||||
|
||||
class IdleTimeoutMiddleware:
|
||||
@@ -58,97 +58,27 @@ class IdleTimeoutMiddleware:
|
||||
os.kill(os.getpid(), signal.SIGTERM)
|
||||
|
||||
|
||||
class DaemonLock:
|
||||
"""
|
||||
File-based lock to prevent multiple daemon instances.
|
||||
|
||||
Uses fcntl.flock for atomic locking on Unix systems.
|
||||
"""
|
||||
|
||||
def __init__(self, lockfile: Path = LOCKFILE_PATH):
|
||||
self.lockfile = lockfile
|
||||
self._fd = None
|
||||
|
||||
def acquire(self) -> bool:
|
||||
"""
|
||||
Try to acquire the daemon lock.
|
||||
|
||||
Returns True if lock acquired, False if another daemon is running.
|
||||
"""
|
||||
self.lockfile.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
self._fd = open(self.lockfile, "w")
|
||||
fcntl.flock(self._fd.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
# Write PID for debugging
|
||||
self._fd.write(str(os.getpid()))
|
||||
self._fd.flush()
|
||||
return True
|
||||
except (IOError, OSError):
|
||||
# Lock is held by another process
|
||||
if self._fd:
|
||||
self._fd.close()
|
||||
self._fd = None
|
||||
return False
|
||||
|
||||
def release(self):
|
||||
"""Release the daemon lock."""
|
||||
if self._fd:
|
||||
try:
|
||||
fcntl.flock(self._fd.fileno(), fcntl.LOCK_UN)
|
||||
self._fd.close()
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
self._fd = None
|
||||
# Remove lockfile
|
||||
try:
|
||||
self.lockfile.unlink()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def is_locked(self) -> bool:
|
||||
"""Check if the lock is held by another process."""
|
||||
if not self.lockfile.exists():
|
||||
return False
|
||||
|
||||
try:
|
||||
fd = open(self.lockfile, "r")
|
||||
fcntl.flock(fd.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
# We got the lock, so no one else has it
|
||||
fcntl.flock(fd.fileno(), fcntl.LOCK_UN)
|
||||
fd.close()
|
||||
return False
|
||||
except (IOError, OSError):
|
||||
return True
|
||||
|
||||
def get_pid(self) -> int | None:
|
||||
"""Get the PID of the daemon holding the lock."""
|
||||
if not self.lockfile.exists():
|
||||
return None
|
||||
try:
|
||||
with open(self.lockfile, "r") as f:
|
||||
return int(f.read().strip())
|
||||
except (ValueError, IOError):
|
||||
return None
|
||||
|
||||
|
||||
def daemonize():
|
||||
"""
|
||||
Fork the current process into a background daemon.
|
||||
|
||||
Uses double-fork technique to properly detach from terminal.
|
||||
"""
|
||||
# First fork
|
||||
pid = os.fork()
|
||||
if pid > 0:
|
||||
# Parent exits
|
||||
sys.exit(0)
|
||||
# First fork - detach from parent
|
||||
try:
|
||||
pid = os.fork()
|
||||
if pid > 0:
|
||||
sys.exit(0)
|
||||
except OSError as e:
|
||||
sys.stderr.write(f"fork #1 failed: {e}\n")
|
||||
sys.exit(1)
|
||||
|
||||
# Create new session
|
||||
# Decouple from parent environment
|
||||
os.chdir("/")
|
||||
os.setsid()
|
||||
os.umask(0)
|
||||
|
||||
# Second fork to prevent zombie processes
|
||||
# Second fork - prevent zombie
|
||||
pid = os.fork()
|
||||
if pid > 0:
|
||||
sys.exit(0)
|
||||
@@ -181,27 +111,3 @@ def check_daemon_running(port: int = DEFAULT_DAEMON_PORT) -> bool:
|
||||
return result == 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def stop_daemon(port: int = DEFAULT_DAEMON_PORT) -> bool:
|
||||
"""Stop a running daemon by sending SIGTERM to the process."""
|
||||
lock = DaemonLock()
|
||||
pid = lock.get_pid()
|
||||
|
||||
if pid is None:
|
||||
return False
|
||||
|
||||
try:
|
||||
import signal
|
||||
|
||||
os.kill(pid, signal.SIGTERM)
|
||||
# Wait for process to exit
|
||||
for _ in range(50): # Wait up to 5 seconds
|
||||
time.sleep(0.1)
|
||||
try:
|
||||
os.kill(pid, 0) # Check if process exists
|
||||
except OSError:
|
||||
return True # Process exited
|
||||
return False
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
@@ -143,11 +143,18 @@ async def run_consolidation_job(
|
||||
"skipped": 0,
|
||||
}
|
||||
|
||||
# Track all unique tags from consolidated memories for mental model refresh filtering
|
||||
consolidated_tags: set[str] = set()
|
||||
|
||||
batch_num = 0
|
||||
last_progress_timings = {} # Track timings at last progress log
|
||||
while True:
|
||||
batch_num += 1
|
||||
batch_start = time.time()
|
||||
|
||||
# Snapshot timings at batch start for per-batch calculation
|
||||
batch_start_timings = perf.timings.copy()
|
||||
|
||||
# Fetch next batch of unconsolidated memories
|
||||
async with pool.acquire() as conn:
|
||||
t0 = time.time()
|
||||
@@ -172,6 +179,11 @@ async def run_consolidation_job(
|
||||
for memory in memories:
|
||||
mem_start = time.time()
|
||||
|
||||
# Track tags from this memory for mental model refresh filtering
|
||||
memory_tags = memory.get("tags") or []
|
||||
if memory_tags:
|
||||
consolidated_tags.update(memory_tags)
|
||||
|
||||
# Process the memory (uses its own connection internally)
|
||||
async with pool.acquire() as conn:
|
||||
result = await _process_memory(
|
||||
@@ -217,19 +229,44 @@ async def run_consolidation_job(
|
||||
elif action == "skipped":
|
||||
stats["skipped"] += 1
|
||||
|
||||
# Log progress periodically
|
||||
# Log progress periodically with timing breakdown
|
||||
if stats["memories_processed"] % 10 == 0:
|
||||
# Calculate timing deltas since last progress log
|
||||
timing_parts = []
|
||||
for key in ["recall", "llm", "embedding", "db_write"]:
|
||||
if key in perf.timings:
|
||||
delta = perf.timings[key] - last_progress_timings.get(key, 0)
|
||||
timing_parts.append(f"{key}={delta:.2f}s")
|
||||
|
||||
timing_str = f" | {', '.join(timing_parts)}" if timing_parts else ""
|
||||
logger.info(
|
||||
f"[CONSOLIDATION] bank={bank_id} progress: "
|
||||
f"{stats['memories_processed']}/{total_count} memories processed"
|
||||
f"{stats['memories_processed']}/{total_count} memories processed{timing_str}"
|
||||
)
|
||||
|
||||
# Update last progress snapshot
|
||||
last_progress_timings = perf.timings.copy()
|
||||
|
||||
batch_time = time.time() - batch_start
|
||||
perf.log(
|
||||
f"[2] Batch {batch_num}: {len(memories)} memories in {batch_time:.3f}s "
|
||||
f"(avg {batch_time / len(memories):.3f}s/memory)"
|
||||
)
|
||||
|
||||
# Log timing breakdown after each batch (delta from batch start)
|
||||
timing_parts = []
|
||||
for key in ["recall", "llm", "embedding", "db_write"]:
|
||||
if key in perf.timings:
|
||||
delta = perf.timings[key] - batch_start_timings.get(key, 0)
|
||||
timing_parts.append(f"{key}={delta:.3f}s")
|
||||
|
||||
if timing_parts:
|
||||
avg_per_memory = batch_time / len(memories) if memories else 0
|
||||
logger.info(
|
||||
f"[CONSOLIDATION] bank={bank_id} batch {batch_num}/{len(memories)} memories: "
|
||||
f"{', '.join(timing_parts)} | avg={avg_per_memory:.3f}s/memory"
|
||||
)
|
||||
|
||||
# Build summary
|
||||
perf.log(
|
||||
f"[3] Results: {stats['memories_processed']} memories -> "
|
||||
@@ -255,10 +292,12 @@ async def run_consolidation_job(
|
||||
perf.log(f"[4] Timing breakdown: {', '.join(timing_parts)}")
|
||||
|
||||
# Trigger mental model refreshes for models with refresh_after_consolidation=true
|
||||
# SECURITY: Only refresh mental models with matching tags (or all if no tags were consolidated)
|
||||
mental_models_refreshed = await _trigger_mental_model_refreshes(
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
consolidated_tags=list(consolidated_tags) if consolidated_tags else None,
|
||||
perf=perf,
|
||||
)
|
||||
stats["mental_models_refreshed"] = mental_models_refreshed
|
||||
@@ -272,15 +311,20 @@ async def _trigger_mental_model_refreshes(
|
||||
memory_engine: "MemoryEngine",
|
||||
bank_id: str,
|
||||
request_context: "RequestContext",
|
||||
consolidated_tags: list[str] | None = None,
|
||||
perf: ConsolidationPerfLog | None = None,
|
||||
) -> int:
|
||||
"""
|
||||
Trigger refreshes for mental models with refresh_after_consolidation=true.
|
||||
|
||||
SECURITY: Only triggers refresh for mental models whose tags overlap with the
|
||||
consolidated memory tags, preventing unnecessary refreshes across security boundaries.
|
||||
|
||||
Args:
|
||||
memory_engine: MemoryEngine instance
|
||||
bank_id: Bank identifier
|
||||
request_context: Request context for authentication
|
||||
consolidated_tags: Tags from memories that were consolidated (None = refresh all)
|
||||
perf: Performance logging
|
||||
|
||||
Returns:
|
||||
@@ -289,22 +333,52 @@ async def _trigger_mental_model_refreshes(
|
||||
pool = memory_engine._pool
|
||||
|
||||
# Find mental models with refresh_after_consolidation=true
|
||||
# SECURITY: Control which mental models get refreshed based on tags
|
||||
async with pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, name
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1
|
||||
AND (trigger->>'refresh_after_consolidation')::boolean = true
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
if consolidated_tags:
|
||||
# Tagged memories were consolidated - refresh:
|
||||
# 1. Mental models with overlapping tags (security boundary)
|
||||
# 2. Untagged mental models (they're "global" and available to all contexts)
|
||||
# DO NOT refresh mental models with different tags
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, name, tags
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1
|
||||
AND (trigger->>'refresh_after_consolidation')::boolean = true
|
||||
AND (
|
||||
(tags IS NOT NULL AND tags != '{{}}' AND tags && $2::varchar[])
|
||||
OR (tags IS NULL OR tags = '{{}}')
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
consolidated_tags,
|
||||
)
|
||||
else:
|
||||
# Untagged memories were consolidated - only refresh untagged mental models
|
||||
# SECURITY: Tagged mental models are NOT refreshed when untagged memories are consolidated
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, name, tags
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1
|
||||
AND (trigger->>'refresh_after_consolidation')::boolean = true
|
||||
AND (tags IS NULL OR tags = '{{}}')
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not rows:
|
||||
return 0
|
||||
|
||||
if perf:
|
||||
perf.log(f"[5] Triggering refresh for {len(rows)} mental models with refresh_after_consolidation=true")
|
||||
if consolidated_tags:
|
||||
perf.log(
|
||||
f"[5] Triggering refresh for {len(rows)} mental models with refresh_after_consolidation=true "
|
||||
f"(filtered by tags: {consolidated_tags})"
|
||||
)
|
||||
else:
|
||||
perf.log(f"[5] Triggering refresh for {len(rows)} mental models with refresh_after_consolidation=true")
|
||||
|
||||
# Submit refresh tasks for each mental model
|
||||
refreshed_count = 0
|
||||
@@ -356,7 +430,8 @@ async def _process_memory(
|
||||
memory_id = memory["id"]
|
||||
fact_tags = memory.get("tags") or []
|
||||
|
||||
# Find related observations using the full recall system (NO tag filtering)
|
||||
# Find related observations using the full recall system
|
||||
# SECURITY: Pass tags to ensure observations don't leak across security boundaries
|
||||
t0 = time.time()
|
||||
related_observations = await _find_related_observations(
|
||||
conn=conn,
|
||||
@@ -364,6 +439,7 @@ async def _process_memory(
|
||||
bank_id=bank_id,
|
||||
query=fact_text,
|
||||
request_context=request_context,
|
||||
tags=fact_tags, # Pass source memory's tags for security
|
||||
)
|
||||
if perf:
|
||||
perf.record_timing("recall", time.time() - t0)
|
||||
@@ -637,17 +713,20 @@ async def _find_related_observations(
|
||||
bank_id: str,
|
||||
query: str,
|
||||
request_context: "RequestContext",
|
||||
tags: list[str] | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Find observations related to the given query using optimized recall.
|
||||
|
||||
IMPORTANT: We do NOT filter by tags here. Consolidation needs to see ALL
|
||||
potentially related observations regardless of scope, so the LLM can
|
||||
decide on tag routing (same scope update vs cross-scope create).
|
||||
SECURITY: Filters by tags using all_strict matching to prevent cross-tenant/cross-user
|
||||
information leakage. Observations are only consolidated within the same tag scope.
|
||||
|
||||
Uses max_tokens to naturally limit observations (no artificial count limit).
|
||||
Includes source memories with dates for LLM context.
|
||||
|
||||
Args:
|
||||
tags: Optional tags to filter observations (uses all_strict matching for security)
|
||||
|
||||
Returns:
|
||||
List of related observations with their tags, source memories, and dates
|
||||
"""
|
||||
@@ -656,14 +735,19 @@ async def _find_related_observations(
|
||||
from ...config import get_config
|
||||
|
||||
config = get_config()
|
||||
|
||||
# SECURITY: Use all_strict matching if tags provided to prevent cross-scope consolidation
|
||||
tags_match = "all_strict" if tags else "any"
|
||||
|
||||
recall_result = await memory_engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
max_tokens=config.consolidation_max_tokens, # Token budget for observations (configurable)
|
||||
fact_type=["observation"], # Only retrieve observations
|
||||
request_context=request_context,
|
||||
tags=tags, # Filter by source memory's tags
|
||||
tags_match=tags_match, # Use strict matching for security
|
||||
_quiet=True, # Suppress logging
|
||||
# NO tags parameter - intentionally get ALL observations
|
||||
)
|
||||
|
||||
# If no observations returned, return empty list
|
||||
@@ -836,7 +920,14 @@ Focus on DURABLE knowledge that serves this mission, not ephemeral state.
|
||||
)
|
||||
# Parse JSON response - should be an array
|
||||
if isinstance(result, str):
|
||||
result = json.loads(result)
|
||||
# Strip markdown code fences (some models wrap JSON in ```json ... ```)
|
||||
clean = result.strip()
|
||||
if clean.startswith("```"):
|
||||
clean = clean.split("\n", 1)[1] if "\n" in clean else clean[3:]
|
||||
if clean.endswith("```"):
|
||||
clean = clean[:-3]
|
||||
clean = clean.strip()
|
||||
result = json.loads(clean)
|
||||
# Ensure result is a list
|
||||
if isinstance(result, list):
|
||||
return result
|
||||
|
||||
@@ -32,13 +32,16 @@ BAD examples:
|
||||
|
||||
## MERGE RULES (when comparing to existing observations):
|
||||
1. REDUNDANT: Same information worded differently → update existing
|
||||
2. CONTRADICTION: Opposite information about same topic → update with history (e.g., "used to X, now Y")
|
||||
3. UPDATE: New state replacing old state → update with history
|
||||
2. CONTRADICTION: Opposite information about same topic → update with temporal markers showing change
|
||||
Example: "Alex used to love pizza but now hates it" OR "Alex's pizza preference changed from love to hate"
|
||||
3. UPDATE: New state replacing old state → update showing the transition with "used to", "now", "changed from X to Y"
|
||||
|
||||
## CRITICAL RULES:
|
||||
- NEVER merge facts about DIFFERENT people
|
||||
- NEVER merge unrelated topics (food preferences vs work vs hobbies)
|
||||
- When merging contradictions, capture the CHANGE (before → after)
|
||||
- When merging contradictions, the "text" field MUST capture BOTH states with temporal markers:
|
||||
* Use "used to X, now Y" OR "changed from X to Y" OR "X but now Y"
|
||||
* DO NOT just state the new fact - you MUST show the change
|
||||
- Keep observations focused on ONE specific topic per person
|
||||
- The "text" field MUST contain durable knowledge, not ephemeral state
|
||||
- Do NOT include "tags" in output - tags are handled automatically"""
|
||||
|
||||
@@ -9,6 +9,7 @@ Configuration via environment variables - see hindsight_api.config for all env v
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
@@ -162,11 +163,28 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}")
|
||||
|
||||
self._model = CrossEncoder(
|
||||
self.model_name,
|
||||
device=device,
|
||||
model_kwargs={"low_cpu_mem_usage": False},
|
||||
)
|
||||
# Suppress verbose transformers warnings during model loading
|
||||
# This suppresses the "UNEXPECTED" warnings from CrossEncoder which are harmless
|
||||
# but look alarming to users (e.g., "embeddings.position_ids | UNEXPECTED")
|
||||
with warnings.catch_warnings():
|
||||
warnings.filterwarnings("ignore", category=UserWarning)
|
||||
warnings.filterwarnings("ignore", message=".*was not found in model state dict.*")
|
||||
warnings.filterwarnings("ignore", message=".*UNEXPECTED.*")
|
||||
|
||||
# Also suppress transformers library logging temporarily
|
||||
transformers_logger = logging.getLogger("transformers")
|
||||
original_level = transformers_logger.level
|
||||
transformers_logger.setLevel(logging.ERROR)
|
||||
|
||||
try:
|
||||
self._model = CrossEncoder(
|
||||
self.model_name,
|
||||
device=device,
|
||||
model_kwargs={"low_cpu_mem_usage": False},
|
||||
)
|
||||
finally:
|
||||
# Restore original logging level
|
||||
transformers_logger.setLevel(original_level)
|
||||
|
||||
# Initialize shared executor (limited workers naturally limits concurrency)
|
||||
if LocalSTCrossEncoder._executor is None:
|
||||
@@ -614,7 +632,7 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
return
|
||||
|
||||
try:
|
||||
from flashrank import Ranker # type: ignore[import-untyped]
|
||||
from flashrank import Ranker
|
||||
except ImportError:
|
||||
raise ImportError("flashrank is required for FlashRankCrossEncoder. Install it with: pip install flashrank")
|
||||
|
||||
@@ -641,7 +659,7 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
|
||||
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""Synchronous predict - processes each query group."""
|
||||
from flashrank import RerankRequest # type: ignore[import-untyped]
|
||||
from flashrank import RerankRequest
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
@@ -11,6 +11,7 @@ Configuration via environment variables - see hindsight_api.config for all env v
|
||||
|
||||
import logging
|
||||
import os
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import httpx
|
||||
@@ -157,11 +158,28 @@ class LocalSTEmbeddings(Embeddings):
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}")
|
||||
|
||||
self._model = SentenceTransformer(
|
||||
self.model_name,
|
||||
device=device,
|
||||
model_kwargs={"low_cpu_mem_usage": False},
|
||||
)
|
||||
# Suppress verbose transformers warnings during model loading
|
||||
# This suppresses the "UNEXPECTED" warnings from BertModel which are harmless
|
||||
# but look alarming to users (e.g., "embeddings.position_ids | UNEXPECTED")
|
||||
with warnings.catch_warnings():
|
||||
warnings.filterwarnings("ignore", category=UserWarning)
|
||||
warnings.filterwarnings("ignore", message=".*was not found in model state dict.*")
|
||||
warnings.filterwarnings("ignore", message=".*UNEXPECTED.*")
|
||||
|
||||
# Also suppress transformers library logging temporarily
|
||||
transformers_logger = logging.getLogger("transformers")
|
||||
original_level = transformers_logger.level
|
||||
transformers_logger.setLevel(logging.ERROR)
|
||||
|
||||
try:
|
||||
self._model = SentenceTransformer(
|
||||
self.model_name,
|
||||
device=device,
|
||||
model_kwargs={"low_cpu_mem_usage": False},
|
||||
)
|
||||
finally:
|
||||
# Restore original logging level
|
||||
transformers_logger.setLevel(original_level)
|
||||
|
||||
self._dimension = self._model.get_sentence_embedding_dimension()
|
||||
logger.info(f"Embeddings: local provider initialized (dim: {self._dimension})")
|
||||
@@ -545,7 +563,7 @@ class CohereEmbeddings(Embeddings):
|
||||
model=self.model,
|
||||
input_type=self.input_type,
|
||||
)
|
||||
if response.embeddings:
|
||||
if response.embeddings and isinstance(response.embeddings, list):
|
||||
self._dimension = len(response.embeddings[0])
|
||||
|
||||
logger.info(f"Embeddings: Cohere provider initialized (model: {self.model}, dim: {self._dimension})")
|
||||
|
||||
@@ -442,49 +442,6 @@ class MemoryEngineInterface(ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_entity_observations(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_id: str,
|
||||
*,
|
||||
limit: int = 10,
|
||||
request_context: "RequestContext",
|
||||
) -> list[Any]:
|
||||
"""
|
||||
Get observations for an entity.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
entity_id: The entity ID.
|
||||
limit: Maximum observations.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
List of EntityObservation objects.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def regenerate_entity_observations(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_id: str,
|
||||
entity_name: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> None:
|
||||
"""
|
||||
Regenerate observations for an entity.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
entity_id: The entity ID.
|
||||
entity_name: The entity's canonical name.
|
||||
request_context: Request context for authentication.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Statistics & Operations
|
||||
# =========================================================================
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
"""
|
||||
Abstract interface for LLM providers.
|
||||
|
||||
This module defines the interface that all LLM providers must implement,
|
||||
enabling support for multiple LLM backends (OpenAI, Anthropic, Gemini, Codex, etc.)
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
from .response_models import LLMToolCallResult, TokenUsage
|
||||
|
||||
|
||||
class LLMInterface(ABC):
|
||||
"""
|
||||
Abstract interface for LLM providers.
|
||||
|
||||
All LLM provider implementations must inherit from this class and implement
|
||||
the required methods.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
Initialize LLM provider.
|
||||
|
||||
Args:
|
||||
provider: Provider name (e.g., "openai", "codex", "anthropic", "gemini").
|
||||
api_key: API key or authentication token.
|
||||
base_url: Base URL for the API.
|
||||
model: Model name.
|
||||
reasoning_effort: Reasoning effort level for supported providers.
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
"""
|
||||
self.provider = provider.lower()
|
||||
self.api_key = api_key
|
||||
self.base_url = base_url
|
||||
self.model = model
|
||||
self.reasoning_effort = reasoning_effort
|
||||
|
||||
@abstractmethod
|
||||
async def verify_connection(self) -> None:
|
||||
"""
|
||||
Verify that the LLM provider is configured correctly by making a simple test call.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the connection test fails.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
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:
|
||||
"""
|
||||
Make an LLM API call with retry logic.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts with 'role' and 'content'.
|
||||
response_format: Optional Pydantic model for structured output.
|
||||
max_completion_tokens: Maximum tokens in response.
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Use strict JSON schema enforcement (OpenAI only).
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
|
||||
Returns:
|
||||
If return_usage=False: Parsed response if response_format is provided, otherwise text content.
|
||||
If return_usage=True: Tuple of (result, TokenUsage) with token counts.
|
||||
|
||||
Raises:
|
||||
OutputTooLongError: If output exceeds token limits.
|
||||
Exception: Re-raises API errors after retries exhausted.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
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:
|
||||
"""
|
||||
Make an LLM API call with tool/function calling support.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts. Can include tool results with role='tool'.
|
||||
tools: List of tool definitions in OpenAI format.
|
||||
max_completion_tokens: Maximum tokens in response.
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or specific function.
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (close connections, etc.)."""
|
||||
pass
|
||||
|
||||
|
||||
class OutputTooLongError(Exception):
|
||||
"""
|
||||
Bridge exception raised when LLM output exceeds token limits.
|
||||
|
||||
This wraps provider-specific errors (e.g., OpenAI's LengthFinishReasonError)
|
||||
to allow callers to handle output length issues without depending on
|
||||
provider-specific implementations.
|
||||
"""
|
||||
|
||||
pass
|
||||
File diff suppressed because it is too large
Load Diff
@@ -303,8 +303,10 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
db_url = db_url or config.database_url
|
||||
memory_llm_provider = memory_llm_provider or config.llm_provider
|
||||
memory_llm_api_key = memory_llm_api_key or config.llm_api_key
|
||||
# Ollama and mock don't require an API key
|
||||
if not memory_llm_api_key and memory_llm_provider not in ("ollama", "mock"):
|
||||
# Ollama, openai-codex, claude-code, and mock don't require an API key
|
||||
# openai-codex uses OAuth tokens from ~/.codex/auth.json
|
||||
# claude-code uses OAuth tokens from macOS Keychain
|
||||
if not memory_llm_api_key and memory_llm_provider not in ("ollama", "openai-codex", "claude-code", "mock"):
|
||||
raise ValueError("LLM API key is required. Set HINDSIGHT_API_LLM_API_KEY environment variable.")
|
||||
memory_llm_model = memory_llm_model or config.llm_model
|
||||
memory_llm_base_url = memory_llm_base_url or config.get_llm_base_url() or None
|
||||
@@ -457,7 +459,11 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
# Store operation validator extension (optional)
|
||||
self._operation_validator = operation_validator
|
||||
|
||||
# Store tenant extension (optional)
|
||||
# Store tenant extension (always set, use default if none provided)
|
||||
if tenant_extension is None:
|
||||
from ..extensions.builtin.tenant import DefaultTenantExtension
|
||||
|
||||
tenant_extension = DefaultTenantExtension(config={})
|
||||
self._tenant_extension = tenant_extension
|
||||
|
||||
async def _validate_operation(self, validation_coro) -> None:
|
||||
@@ -495,23 +501,18 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
Raises:
|
||||
AuthenticationError: If authentication fails or request_context is missing when required.
|
||||
"""
|
||||
if self._tenant_extension is None:
|
||||
_current_schema.set("public")
|
||||
return "public"
|
||||
|
||||
from hindsight_api.extensions import AuthenticationError
|
||||
|
||||
if request_context is None:
|
||||
raise AuthenticationError("RequestContext is required when tenant extension is configured")
|
||||
raise AuthenticationError("RequestContext is required")
|
||||
|
||||
# For internal/background operations (e.g., worker tasks), skip extension authentication
|
||||
# if the schema has already been set by execute_task via the _schema field.
|
||||
# For internal/background operations (e.g., worker tasks), skip extension authentication.
|
||||
# The task was already authenticated at submission time, and execute_task sets _current_schema
|
||||
# from the task's _schema field.
|
||||
if request_context.internal:
|
||||
current = _current_schema.get()
|
||||
if current and current != "public":
|
||||
return current
|
||||
return _current_schema.get()
|
||||
|
||||
# Let AuthenticationError propagate - HTTP layer will convert to 401
|
||||
# Authenticate through tenant extension (always set, may be default no-auth extension)
|
||||
tenant_context = await self._tenant_extension.authenticate(request_context)
|
||||
|
||||
_current_schema.set(tenant_context.schema_name)
|
||||
@@ -537,10 +538,15 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
f"[BATCH_RETAIN_TASK] Starting background batch retain for bank_id={bank_id}, {len(contents)} items"
|
||||
)
|
||||
|
||||
# Use internal request context for background tasks (skips tenant auth when schema is pre-set)
|
||||
# Restore tenant_id/api_key_id from task payload so downstream operations
|
||||
# (e.g., consolidation and mental model refreshes) can attribute usage.
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
internal_context = RequestContext(internal=True)
|
||||
internal_context = RequestContext(
|
||||
internal=True,
|
||||
tenant_id=task_dict.get("_tenant_id"),
|
||||
api_key_id=task_dict.get("_api_key_id"),
|
||||
)
|
||||
await self.retain_batch_async(bank_id=bank_id, contents=contents, request_context=internal_context)
|
||||
|
||||
logger.info(f"[BATCH_RETAIN_TASK] Completed background batch retain for bank_id={bank_id}")
|
||||
@@ -566,7 +572,13 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
from .consolidation import run_consolidation_job
|
||||
|
||||
internal_context = RequestContext(internal=True)
|
||||
# Restore tenant_id/api_key_id from task payload so downstream operations
|
||||
# (e.g., mental model refreshes) can attribute usage to the correct org.
|
||||
internal_context = RequestContext(
|
||||
internal=True,
|
||||
tenant_id=task_dict.get("_tenant_id"),
|
||||
api_key_id=task_dict.get("_api_key_id"),
|
||||
)
|
||||
result = await run_consolidation_job(
|
||||
memory_engine=self,
|
||||
bank_id=bank_id,
|
||||
@@ -598,7 +610,13 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
internal_context = RequestContext(internal=True)
|
||||
# Restore tenant_id/api_key_id from task payload so extensions can
|
||||
# attribute the mental_model_refresh operation to the correct org.
|
||||
internal_context = RequestContext(
|
||||
internal=True,
|
||||
tenant_id=task_dict.get("_tenant_id"),
|
||||
api_key_id=task_dict.get("_api_key_id"),
|
||||
)
|
||||
|
||||
# Get the current mental model to get source_query
|
||||
mental_model = await self.get_mental_model(bank_id, mental_model_id, request_context=internal_context)
|
||||
@@ -607,11 +625,19 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
source_query = mental_model["source_query"]
|
||||
|
||||
# SECURITY: If the mental model has tags, pass them to reflect with "all_strict" matching
|
||||
# to ensure it can only access other mental models/memories with the SAME tags.
|
||||
# This prevents cross-tenant/cross-user information leakage by excluding untagged content.
|
||||
tags = mental_model.get("tags")
|
||||
tags_match = "all_strict" if tags else "any"
|
||||
|
||||
# Run reflect to generate new content, excluding the mental model being refreshed
|
||||
reflect_result = await self.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query=source_query,
|
||||
request_context=internal_context,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
exclude_mental_model_ids=[mental_model_id],
|
||||
)
|
||||
|
||||
@@ -642,6 +668,42 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
request_context=internal_context,
|
||||
)
|
||||
|
||||
# Call post-operation hook if validator is configured
|
||||
if self._operation_validator:
|
||||
from hindsight_api.extensions.operation_validator import MentalModelRefreshResult
|
||||
|
||||
# Count facts and mental models from based_on
|
||||
facts_used = 0
|
||||
mental_models_used = 0
|
||||
if reflect_result.based_on:
|
||||
for fact_type, facts in reflect_result.based_on.items():
|
||||
if facts:
|
||||
if fact_type == "mental_models":
|
||||
mental_models_used += len(facts)
|
||||
else:
|
||||
facts_used += len(facts)
|
||||
|
||||
# Estimate tokens
|
||||
query_tokens = len(source_query) // 4 if source_query else 0
|
||||
output_tokens = len(generated_content) // 4 if generated_content else 0
|
||||
context_tokens = 0 # refresh doesn't use additional context
|
||||
|
||||
result_ctx = MentalModelRefreshResult(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=internal_context,
|
||||
query_tokens=query_tokens,
|
||||
output_tokens=output_tokens,
|
||||
context_tokens=context_tokens,
|
||||
facts_used=facts_used,
|
||||
mental_models_used=mental_models_used,
|
||||
success=True,
|
||||
)
|
||||
try:
|
||||
await self._operation_validator.on_mental_model_refresh_complete(result_ctx)
|
||||
except Exception as hook_err:
|
||||
logger.warning(f"Post-mental-model-refresh hook error (non-fatal): {hook_err}")
|
||||
|
||||
logger.info(f"[REFRESH_MENTAL_MODEL_TASK] Completed for bank_id={bank_id}, mental_model_id={mental_model_id}")
|
||||
|
||||
async def execute_task(self, task_dict: dict[str, Any]):
|
||||
@@ -789,7 +851,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
kwargs = {"name": self._pg0_instance_name}
|
||||
if self._pg0_port is not None:
|
||||
kwargs["port"] = self._pg0_port
|
||||
pg0 = EmbeddedPostgres(**kwargs) # type: ignore[invalid-argument-type] - dict kwargs
|
||||
pg0 = EmbeddedPostgres(**kwargs)
|
||||
# Check if pg0 is already running before we start it
|
||||
was_already_running = await pg0.is_running()
|
||||
self.db_url = await pg0.ensure_running()
|
||||
@@ -885,13 +947,34 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
if not self.db_url:
|
||||
raise ValueError("Database URL is required for migrations")
|
||||
logger.info("Running database migrations...")
|
||||
# Use configured database schema for migrations (defaults to "public")
|
||||
run_migrations(self.db_url, schema=get_config().database_schema)
|
||||
|
||||
# Ensure embedding column dimension matches the model's dimension
|
||||
# This is done after migrations and after embeddings.initialize()
|
||||
ensure_embedding_dimension(self.db_url, self.embeddings.dimension, schema=get_config().database_schema)
|
||||
# Migrate all schemas from the tenant extension
|
||||
# The tenant extension is the single source of truth for which schemas exist
|
||||
logger.info("Running database migrations...")
|
||||
try:
|
||||
tenants = await self._tenant_extension.list_tenants()
|
||||
if tenants:
|
||||
logger.info(f"Running migrations on {len(tenants)} schema(s)...")
|
||||
for tenant in tenants:
|
||||
schema = tenant.schema
|
||||
if schema:
|
||||
try:
|
||||
run_migrations(self.db_url, schema=schema)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to migrate schema {schema}: {e}")
|
||||
logger.info("Schema migrations completed")
|
||||
|
||||
# Ensure embedding column dimension matches the model's dimension
|
||||
# This is done after migrations and after embeddings.initialize()
|
||||
for tenant in tenants:
|
||||
schema = tenant.schema
|
||||
if schema:
|
||||
try:
|
||||
ensure_embedding_dimension(self.db_url, self.embeddings.dimension, schema=schema)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to ensure embedding dimension for schema {schema}: {e}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to run schema migrations: {e}")
|
||||
|
||||
logger.info(f"Connecting to PostgreSQL at {self.db_url}")
|
||||
|
||||
@@ -1175,15 +1258,15 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
context: Context about when/why this memory was formed
|
||||
event_date: When the event occurred (defaults to now)
|
||||
document_id: Optional document ID for tracking (always upserts if document already exists)
|
||||
fact_type_override: Override fact type ('world', 'experience', 'opinion')
|
||||
confidence_score: Confidence score for opinions (0.0 to 1.0)
|
||||
fact_type_override: Override fact type ('world', 'experience')
|
||||
confidence_score: Confidence score (0.0 to 1.0)
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
List of created unit IDs
|
||||
"""
|
||||
# Build content dict
|
||||
content_dict: RetainContentDict = {"content": content, "context": context} # type: ignore[typeddict-item] - building incrementally
|
||||
content_dict: RetainContentDict = {"content": content, "context": context}
|
||||
if event_date:
|
||||
content_dict["event_date"] = event_date
|
||||
if document_id:
|
||||
@@ -1231,8 +1314,8 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
- "document_id" (optional): Document ID for this specific content item
|
||||
document_id: **DEPRECATED** - Use "document_id" key in each content dict instead.
|
||||
Applies the same document_id to ALL content items that don't specify their own.
|
||||
fact_type_override: Override fact type for all facts ('world', 'experience', 'opinion')
|
||||
confidence_score: Confidence score for opinions (0.0 to 1.0)
|
||||
fact_type_override: Override fact type for all facts ('world', 'experience')
|
||||
confidence_score: Confidence score (0.0 to 1.0)
|
||||
return_usage: If True, returns tuple of (unit_ids, TokenUsage). Default False for backward compatibility.
|
||||
|
||||
Returns:
|
||||
@@ -1554,16 +1637,16 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
if fact_type is None:
|
||||
fact_type = list(VALID_RECALL_FACT_TYPES)
|
||||
|
||||
# Validate fact types early
|
||||
# Filter out 'opinion' early (deprecated, silently ignore)
|
||||
fact_type = [ft for ft in fact_type if ft != "opinion"]
|
||||
|
||||
# Validate fact types
|
||||
invalid_types = set(fact_type) - VALID_RECALL_FACT_TYPES
|
||||
if invalid_types:
|
||||
raise ValueError(
|
||||
f"Invalid fact type(s): {', '.join(sorted(invalid_types))}. "
|
||||
f"Must be one of: {', '.join(sorted(VALID_RECALL_FACT_TYPES))}"
|
||||
)
|
||||
|
||||
# Filter out 'opinion' - opinions are no longer returned from recall
|
||||
fact_type = [ft for ft in fact_type if ft != "opinion"]
|
||||
if not fact_type:
|
||||
# All requested types were opinions - return empty result
|
||||
return RecallResultModel(results=[], entities={}, chunks={})
|
||||
@@ -2219,44 +2302,15 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
)
|
||||
top_results_dicts.append(result_dict)
|
||||
|
||||
# Get entities for each fact if include_entities is requested
|
||||
fact_entity_map = {} # unit_id -> list of (entity_id, entity_name)
|
||||
if include_entities and top_scored:
|
||||
unit_ids = [uuid.UUID(sr.id) for sr in top_scored]
|
||||
if unit_ids:
|
||||
async with acquire_with_retry(pool) as entity_conn:
|
||||
entity_rows = await entity_conn.fetch(
|
||||
f"""
|
||||
SELECT ue.unit_id, e.id as entity_id, e.canonical_name
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("entities")} e ON ue.entity_id = e.id
|
||||
WHERE ue.unit_id = ANY($1::uuid[])
|
||||
""",
|
||||
unit_ids,
|
||||
)
|
||||
for row in entity_rows:
|
||||
unit_id = str(row["unit_id"])
|
||||
if unit_id not in fact_entity_map:
|
||||
fact_entity_map[unit_id] = []
|
||||
fact_entity_map[unit_id].append(
|
||||
{"entity_id": str(row["entity_id"]), "canonical_name": row["canonical_name"]}
|
||||
)
|
||||
|
||||
# Convert results to MemoryFact objects
|
||||
memory_facts = []
|
||||
for result_dict in top_results_dicts:
|
||||
result_id = str(result_dict.get("id"))
|
||||
# Get entity names for this fact
|
||||
entity_names = None
|
||||
if include_entities and result_id in fact_entity_map:
|
||||
entity_names = [e["canonical_name"] for e in fact_entity_map[result_id]]
|
||||
|
||||
memory_facts.append(
|
||||
MemoryFact(
|
||||
id=result_id,
|
||||
id=str(result_dict.get("id")),
|
||||
text=result_dict.get("text"),
|
||||
fact_type=result_dict.get("fact_type", "world"),
|
||||
entities=entity_names,
|
||||
entities=None, # Entity observations removed
|
||||
context=result_dict.get("context"),
|
||||
occurred_start=result_dict.get("occurred_start"),
|
||||
occurred_end=result_dict.get("occurred_end"),
|
||||
@@ -2267,38 +2321,12 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
)
|
||||
)
|
||||
|
||||
# Fetch entity observations if requested
|
||||
# Entity observations removed - always set to None
|
||||
entities_dict = None
|
||||
total_entity_tokens = 0
|
||||
total_chunk_tokens = 0
|
||||
if include_entities and fact_entity_map:
|
||||
# Collect unique entities in order of fact relevance (preserving order from top_scored)
|
||||
# Use a list to maintain order, but track seen entities to avoid duplicates
|
||||
entities_ordered = [] # list of (entity_id, entity_name) tuples
|
||||
seen_entity_ids = set()
|
||||
|
||||
# Iterate through facts in relevance order
|
||||
for sr in top_scored:
|
||||
unit_id = sr.id
|
||||
if unit_id in fact_entity_map:
|
||||
for entity in fact_entity_map[unit_id]:
|
||||
entity_id = entity["entity_id"]
|
||||
entity_name = entity["canonical_name"]
|
||||
if entity_id not in seen_entity_ids:
|
||||
entities_ordered.append((entity_id, entity_name))
|
||||
seen_entity_ids.add(entity_id)
|
||||
|
||||
# Return entities with empty observations (summaries now live in mental models)
|
||||
entities_dict = {}
|
||||
for entity_id, entity_name in entities_ordered:
|
||||
entities_dict[entity_name] = EntityState(
|
||||
entity_id=entity_id,
|
||||
canonical_name=entity_name,
|
||||
observations=[], # Mental models provide this now
|
||||
)
|
||||
|
||||
# Fetch chunks if requested
|
||||
chunks_dict = None
|
||||
total_chunk_tokens = 0
|
||||
if include_chunks and top_scored:
|
||||
from .response_models import ChunkInfo
|
||||
|
||||
@@ -2367,7 +2395,6 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
# Log final recall stats
|
||||
total_time = time.time() - recall_start
|
||||
num_chunks = len(chunks_dict) if chunks_dict else 0
|
||||
num_entities = len(entities_dict) if entities_dict else 0
|
||||
# Include wait times in log if significant
|
||||
wait_parts = []
|
||||
if semaphore_wait > 0.01:
|
||||
@@ -2376,7 +2403,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
wait_parts.append(f"conn={max_conn_wait:.3f}s")
|
||||
wait_info = f" | waits: {', '.join(wait_parts)}" if wait_parts else ""
|
||||
log_buffer.append(
|
||||
f"[RECALL {recall_id}] Complete: {len(top_scored)} facts ({total_tokens} tok), {num_chunks} chunks ({total_chunk_tokens} tok), {num_entities} entities ({total_entity_tokens} tok) | {fact_type_summary} | {total_time:.3f}s{wait_info}"
|
||||
f"[RECALL {recall_id}] Complete: {len(top_scored)} facts ({total_tokens} tok), {num_chunks} chunks ({total_chunk_tokens} tok) | {fact_type_summary} | {total_time:.3f}s{wait_info}"
|
||||
)
|
||||
if not quiet:
|
||||
logger.info("\n" + "\n".join(log_buffer))
|
||||
@@ -3285,7 +3312,8 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
created_at,
|
||||
updated_at,
|
||||
LENGTH(original_text) as text_length,
|
||||
retain_params
|
||||
retain_params,
|
||||
tags
|
||||
FROM {fq_table("documents")}
|
||||
{where_clause}
|
||||
ORDER BY created_at DESC
|
||||
@@ -3341,6 +3369,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
"text_length": row["text_length"] or 0,
|
||||
"memory_unit_count": unit_count,
|
||||
"retain_params": row["retain_params"] if row["retain_params"] else None,
|
||||
"tags": row["tags"] if row["tags"] else [],
|
||||
}
|
||||
)
|
||||
|
||||
@@ -3550,7 +3579,6 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
ReflectResult containing:
|
||||
- text: Plain text answer
|
||||
- based_on: Empty dict (agent retrieves facts dynamically)
|
||||
- new_opinions: Empty list
|
||||
- structured_output: None (not yet supported for agentic reflect)
|
||||
"""
|
||||
# Use cached LLM config
|
||||
@@ -3645,12 +3673,14 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
# Load directives from the dedicated directives table
|
||||
# Directives are hard rules that must be followed in all responses
|
||||
# Use isolation_mode=True to prevent tag-scoped directives from leaking into untagged operations
|
||||
directives_raw = await self.list_directives(
|
||||
bank_id=bank_id,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
active_only=True,
|
||||
request_context=request_context,
|
||||
isolation_mode=True,
|
||||
)
|
||||
# Convert directive format to the expected format for reflect agent
|
||||
# The agent expects: name, description (optional), observations (list of {title, content})
|
||||
@@ -3719,7 +3749,16 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
# Extract memories from recall tool outputs - only include memories the agent actually used
|
||||
# agent_result.used_memory_ids contains validated IDs from the done action
|
||||
used_memory_ids_set = set(agent_result.used_memory_ids) if agent_result.used_memory_ids else set()
|
||||
based_on: dict[str, list[MemoryFact]] = {"world": [], "experience": [], "opinion": [], "observation": []}
|
||||
# based_on stores facts, mental models, and directives
|
||||
# Note: directives list stores raw directive dicts (not MemoryFact), which will be converted to Directive objects
|
||||
based_on: dict[str, list[MemoryFact] | list[dict[str, Any]]] = {
|
||||
"world": [],
|
||||
"experience": [],
|
||||
"opinion": [],
|
||||
"observation": [],
|
||||
"mental-models": [],
|
||||
"directives": [],
|
||||
}
|
||||
seen_memory_ids: set[str] = set()
|
||||
for tc in agent_result.tool_trace:
|
||||
if tc.tool == "recall" and "memories" in tc.output:
|
||||
@@ -3821,38 +3860,15 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
)
|
||||
# List all models lookup - don't add to based_on (too verbose, just a listing)
|
||||
|
||||
# Add directives to based_on["mental-models"] (they are mental models with subtype='directive')
|
||||
for directive in directives:
|
||||
# Extract summary from observations
|
||||
summary_parts: list[str] = []
|
||||
for obs in directive.get("observations", []):
|
||||
# Support both Pydantic Observation objects and dicts
|
||||
if hasattr(obs, "content"):
|
||||
content = obs.content
|
||||
title = obs.title
|
||||
else:
|
||||
content = obs.get("content", "")
|
||||
title = obs.get("title", "")
|
||||
if title and content:
|
||||
summary_parts.append(f"{title}: {content}")
|
||||
elif content:
|
||||
summary_parts.append(content)
|
||||
|
||||
# Fallback to description if no observations
|
||||
if not summary_parts and directive.get("description"):
|
||||
summary_parts.append(directive["description"])
|
||||
|
||||
directive_name = directive.get("name", "")
|
||||
directive_summary = "; ".join(summary_parts) if summary_parts else ""
|
||||
based_on["mental-models"].append(
|
||||
MemoryFact(
|
||||
id=directive.get("id", ""),
|
||||
text=f"{directive_name}: {directive_summary}",
|
||||
fact_type="mental-models",
|
||||
context="directive (directive)",
|
||||
occurred_start=None,
|
||||
occurred_end=None,
|
||||
)
|
||||
# Add directives to based_on["directives"]
|
||||
# Store raw directive dicts (with id, name, content) for http.py to convert to ReflectDirective
|
||||
for directive_raw in directives_raw:
|
||||
based_on["directives"].append(
|
||||
{
|
||||
"id": directive_raw["id"],
|
||||
"name": directive_raw["name"],
|
||||
"content": directive_raw["content"],
|
||||
}
|
||||
)
|
||||
|
||||
# Build directives_applied from agent result
|
||||
@@ -3875,7 +3891,6 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
result = ReflectResult(
|
||||
text=agent_result.text,
|
||||
based_on=based_on,
|
||||
new_opinions=[], # Learnings stored as mental models
|
||||
structured_output=agent_result.structured_output,
|
||||
usage=usage,
|
||||
tool_trace=tool_trace_result,
|
||||
@@ -3904,32 +3919,6 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
return result
|
||||
|
||||
async def get_entity_observations(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_id: str,
|
||||
*,
|
||||
limit: int = 10,
|
||||
request_context: "RequestContext",
|
||||
) -> list[Any]:
|
||||
"""
|
||||
Get observations for an entity.
|
||||
|
||||
NOTE: Entity observations/summaries have been moved to mental models.
|
||||
This method returns an empty list. Use mental models for entity summaries.
|
||||
|
||||
Args:
|
||||
bank_id: bank IDentifier
|
||||
entity_id: Entity UUID to get observations for
|
||||
limit: Ignored (kept for backwards compatibility)
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Empty list (observations now in mental models)
|
||||
"""
|
||||
await self._authenticate_tenant(request_context)
|
||||
return []
|
||||
|
||||
async def list_entities(
|
||||
self,
|
||||
bank_id: str,
|
||||
@@ -4116,36 +4105,6 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
await self._authenticate_tenant(request_context)
|
||||
return EntityState(entity_id=entity_id, canonical_name=entity_name, observations=[])
|
||||
|
||||
async def regenerate_entity_observations(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_id: str,
|
||||
entity_name: str,
|
||||
*,
|
||||
version: str | None = None,
|
||||
conn=None,
|
||||
request_context: "RequestContext",
|
||||
) -> list[str]:
|
||||
"""
|
||||
Regenerate observations for an entity.
|
||||
|
||||
NOTE: Entity observations/summaries have been moved to mental models.
|
||||
This method is now a no-op and returns an empty list.
|
||||
|
||||
Args:
|
||||
bank_id: bank IDentifier
|
||||
entity_id: Entity UUID
|
||||
entity_name: Canonical name of the entity
|
||||
version: Entity's last_seen timestamp when task was created (for deduplication)
|
||||
conn: Optional database connection (ignored)
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Empty list (observations now in mental models)
|
||||
"""
|
||||
await self._authenticate_tenant(request_context)
|
||||
return []
|
||||
|
||||
# =========================================================================
|
||||
# Statistics & Operations (for HTTP API layer)
|
||||
# =========================================================================
|
||||
@@ -4256,9 +4215,6 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
if not entity_row:
|
||||
return None
|
||||
|
||||
# Get observations for the entity
|
||||
observations = await self.get_entity_observations(bank_id, entity_id, limit=20, request_context=request_context)
|
||||
|
||||
return {
|
||||
"id": str(entity_row["id"]),
|
||||
"canonical_name": entity_row["canonical_name"],
|
||||
@@ -4266,7 +4222,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
"first_seen": entity_row["first_seen"].isoformat() if entity_row["first_seen"] else None,
|
||||
"last_seen": entity_row["last_seen"].isoformat() if entity_row["last_seen"] else None,
|
||||
"metadata": entity_row["metadata"] or {},
|
||||
"observations": observations,
|
||||
"observations": [],
|
||||
}
|
||||
|
||||
def _parse_observations(self, observations_raw: list):
|
||||
@@ -4761,11 +4717,19 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
if not mental_model:
|
||||
return None
|
||||
|
||||
# SECURITY: If the mental model has tags, pass them to reflect with "all_strict" matching
|
||||
# to ensure it can only access other mental models/memories with the SAME tags.
|
||||
# This prevents cross-tenant/cross-user information leakage by excluding untagged content.
|
||||
tags = mental_model.get("tags")
|
||||
tags_match = "all_strict" if tags else "any"
|
||||
|
||||
# Run reflect with the source query, excluding the mental model being refreshed
|
||||
reflect_result = await self.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query=mental_model["source_query"],
|
||||
request_context=request_context,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
exclude_mental_model_ids=[mental_model_id],
|
||||
)
|
||||
|
||||
@@ -4778,6 +4742,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
"id": str(fact.id),
|
||||
"text": fact.text,
|
||||
"type": fact_type,
|
||||
"context": fact.context, # Include context to distinguish directives from mental models in UI
|
||||
}
|
||||
for fact in facts
|
||||
]
|
||||
@@ -4966,6 +4931,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: "RequestContext",
|
||||
isolation_mode: bool = False,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""List directives for a bank.
|
||||
|
||||
@@ -4977,6 +4943,9 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
limit: Maximum number of results
|
||||
offset: Offset for pagination
|
||||
request_context: Request context for authentication
|
||||
isolation_mode: When True and tags=None, only return directives with no tags.
|
||||
This prevents tag-scoped directives from leaking into untagged operations.
|
||||
Default False (normal API behavior - returns all directives when tags=None)
|
||||
|
||||
Returns:
|
||||
List of directive dicts
|
||||
@@ -4986,6 +4955,8 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Build filters
|
||||
from .search.tags import build_tags_where_clause
|
||||
|
||||
filters = ["bank_id = $1"]
|
||||
params: list[Any] = [bank_id]
|
||||
param_idx = 2
|
||||
@@ -4993,15 +4964,23 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
if active_only:
|
||||
filters.append("is_active = TRUE")
|
||||
|
||||
# Apply tags filter:
|
||||
# - If tags provided: use standard filtering (with strict modes support)
|
||||
# - If tags=None and isolation_mode=True: only include directives with NO tags
|
||||
# (prevents tag-scoped directives from leaking into untagged reflect/refresh)
|
||||
# - If tags=None and isolation_mode=False: no filtering (normal API behavior)
|
||||
if tags:
|
||||
if tags_match == "all":
|
||||
filters.append(f"tags @> ${param_idx}::varchar[]")
|
||||
elif tags_match == "exact":
|
||||
filters.append(f"tags = ${param_idx}::varchar[]")
|
||||
else: # any
|
||||
filters.append(f"tags && ${param_idx}::varchar[]")
|
||||
params.append(tags)
|
||||
param_idx += 1
|
||||
tags_clause, tags_params, param_idx = build_tags_where_clause(
|
||||
tags=tags, param_offset=param_idx, table_alias="", match=tags_match
|
||||
)
|
||||
if tags_clause:
|
||||
# Remove leading "AND " from clause since we're building filters list
|
||||
filters.append(tags_clause.replace("AND ", "", 1))
|
||||
params.extend(tags_params)
|
||||
elif isolation_mode:
|
||||
# Isolation mode: only include directives with empty/null tags
|
||||
# This ensures tag-scoped directives don't apply to untagged operations
|
||||
filters.append("(tags IS NULL OR tags = '{}')")
|
||||
|
||||
params.extend([limit, offset])
|
||||
|
||||
@@ -5517,6 +5496,13 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
task_payload: dict[str, Any] = {"contents": contents}
|
||||
if document_tags:
|
||||
task_payload["document_tags"] = document_tags
|
||||
# Pass tenant_id and api_key_id through task payload so the worker
|
||||
# can propagate request context to downstream operations (e.g.,
|
||||
# consolidation and mental model refreshes triggered after retain).
|
||||
if request_context.tenant_id:
|
||||
task_payload["_tenant_id"] = request_context.tenant_id
|
||||
if request_context.api_key_id:
|
||||
task_payload["_api_key_id"] = request_context.api_key_id
|
||||
|
||||
result = await self._submit_async_operation(
|
||||
bank_id=bank_id,
|
||||
@@ -5549,11 +5535,21 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
Dict with operation_id
|
||||
"""
|
||||
await self._authenticate_tenant(request_context)
|
||||
|
||||
# Pass tenant_id and api_key_id through task payload so the worker
|
||||
# can provide request context to extension hooks (e.g., usage metering
|
||||
# for mental model refreshes triggered by consolidation).
|
||||
task_payload: dict[str, Any] = {}
|
||||
if request_context.tenant_id:
|
||||
task_payload["_tenant_id"] = request_context.tenant_id
|
||||
if request_context.api_key_id:
|
||||
task_payload["_api_key_id"] = request_context.api_key_id
|
||||
|
||||
return await self._submit_async_operation(
|
||||
bank_id=bank_id,
|
||||
operation_type="consolidation",
|
||||
task_type="consolidation",
|
||||
task_payload={},
|
||||
task_payload=task_payload,
|
||||
dedupe_by_bank=True,
|
||||
)
|
||||
|
||||
@@ -5583,13 +5579,21 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
if not mental_model:
|
||||
raise ValueError(f"Mental model {mental_model_id} not found in bank {bank_id}")
|
||||
|
||||
# Pass tenant_id and api_key_id through task payload so the worker
|
||||
# can provide request context to extension hooks.
|
||||
task_payload: dict[str, Any] = {
|
||||
"mental_model_id": mental_model_id,
|
||||
}
|
||||
if request_context.tenant_id:
|
||||
task_payload["_tenant_id"] = request_context.tenant_id
|
||||
if request_context.api_key_id:
|
||||
task_payload["_api_key_id"] = request_context.api_key_id
|
||||
|
||||
return await self._submit_async_operation(
|
||||
bank_id=bank_id,
|
||||
operation_type="refresh_mental_model",
|
||||
task_type="refresh_mental_model",
|
||||
task_payload={
|
||||
"mental_model_id": mental_model_id,
|
||||
},
|
||||
task_payload=task_payload,
|
||||
result_metadata={"mental_model_id": mental_model_id, "name": mental_model["name"]},
|
||||
dedupe_by_bank=False,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
"""
|
||||
LLM provider implementations.
|
||||
|
||||
This package contains concrete implementations of the LLMInterface for various providers.
|
||||
"""
|
||||
|
||||
from .anthropic_llm import AnthropicLLM
|
||||
from .claude_code_llm import ClaudeCodeLLM
|
||||
from .codex_llm import CodexLLM
|
||||
from .gemini_llm import GeminiLLM
|
||||
from .mock_llm import MockLLM
|
||||
from .openai_compatible_llm import OpenAICompatibleLLM
|
||||
|
||||
__all__ = ["AnthropicLLM", "ClaudeCodeLLM", "CodexLLM", "GeminiLLM", "MockLLM", "OpenAICompatibleLLM"]
|
||||
@@ -0,0 +1,434 @@
|
||||
"""
|
||||
Anthropic LLM provider using the Anthropic Python SDK.
|
||||
|
||||
This provider enables using Claude models from Anthropic with support for:
|
||||
- Structured JSON output
|
||||
- Tool/function calling with proper format conversion
|
||||
- Extended thinking mode
|
||||
- Retry logic with exponential backoff
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AnthropicLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using Anthropic's Claude models.
|
||||
|
||||
Supports structured output, tool calling, and extended thinking mode.
|
||||
Handles format conversion between OpenAI-style messages and Anthropic's format.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float = 300.0,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
Initialize Anthropic LLM provider.
|
||||
|
||||
Args:
|
||||
provider: Provider name (should be "anthropic").
|
||||
api_key: Anthropic API key.
|
||||
base_url: Base URL for the API (optional, uses Anthropic default if empty).
|
||||
model: Model name (e.g., "claude-sonnet-4-20250514").
|
||||
reasoning_effort: Reasoning effort level (not used by Anthropic).
|
||||
timeout: Request timeout in seconds.
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
"""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
|
||||
if not self.api_key:
|
||||
raise ValueError("API key is required for Anthropic provider")
|
||||
|
||||
# Import and initialize Anthropic client
|
||||
try:
|
||||
from anthropic import AsyncAnthropic
|
||||
|
||||
client_kwargs: dict[str, Any] = {"api_key": self.api_key}
|
||||
if self.base_url:
|
||||
client_kwargs["base_url"] = self.base_url
|
||||
if timeout:
|
||||
client_kwargs["timeout"] = timeout
|
||||
|
||||
self._client = AsyncAnthropic(**client_kwargs)
|
||||
logger.info(f"Anthropic client initialized for model: {self.model}")
|
||||
except ImportError as e:
|
||||
raise RuntimeError("Anthropic SDK not installed. Run: uv add anthropic or pip install anthropic") from e
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""
|
||||
Verify that the Anthropic provider is configured correctly by making a simple test call.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the connection test fails.
|
||||
"""
|
||||
try:
|
||||
test_messages = [{"role": "user", "content": "test"}]
|
||||
await self.call(
|
||||
messages=test_messages,
|
||||
max_completion_tokens=10,
|
||||
temperature=0.0,
|
||||
scope="test",
|
||||
max_retries=0,
|
||||
)
|
||||
logger.info("Anthropic connection verified successfully")
|
||||
except Exception as e:
|
||||
logger.error(f"Anthropic connection verification failed: {e}")
|
||||
raise RuntimeError(f"Failed to verify Anthropic connection: {e}") from e
|
||||
|
||||
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:
|
||||
"""
|
||||
Make an LLM API call with retry logic.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts with 'role' and 'content'.
|
||||
response_format: Optional Pydantic model for structured output.
|
||||
max_completion_tokens: Maximum tokens in response.
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Use strict JSON schema enforcement (not supported by Anthropic).
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
|
||||
Returns:
|
||||
If return_usage=False: Parsed response if response_format is provided, otherwise text content.
|
||||
If return_usage=True: Tuple of (result, TokenUsage) with token counts.
|
||||
|
||||
Raises:
|
||||
OutputTooLongError: If output exceeds token limits.
|
||||
Exception: Re-raises API errors after retries exhausted.
|
||||
"""
|
||||
from anthropic import APIConnectionError, APIStatusError, RateLimitError
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Convert OpenAI-style messages to Anthropic format
|
||||
system_prompt = None
|
||||
anthropic_messages = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
if system_prompt:
|
||||
system_prompt += "\n\n" + content
|
||||
else:
|
||||
system_prompt = content
|
||||
else:
|
||||
anthropic_messages.append({"role": role, "content": content})
|
||||
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
if system_prompt:
|
||||
system_prompt += schema_msg
|
||||
else:
|
||||
system_prompt = schema_msg
|
||||
|
||||
# Prepare parameters
|
||||
call_params: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"messages": anthropic_messages,
|
||||
"max_tokens": max_completion_tokens if max_completion_tokens is not None else 4096,
|
||||
}
|
||||
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if temperature is not None:
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await self._client.messages.create(**call_params)
|
||||
|
||||
# Anthropic response content is a list of blocks
|
||||
content = ""
|
||||
for block in response.content:
|
||||
if block.type == "text":
|
||||
content += block.text
|
||||
|
||||
if response_format is not None:
|
||||
# Models may wrap JSON in markdown code blocks
|
||||
clean_content = content
|
||||
if "```json" in content:
|
||||
clean_content = content.split("```json")[1].split("```")[0].strip()
|
||||
elif "```" in content:
|
||||
clean_content = content.split("```")[1].split("```")[0].strip()
|
||||
|
||||
try:
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError:
|
||||
# Fallback to parsing raw content if markdown stripping failed
|
||||
json_data = json.loads(content)
|
||||
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
else:
|
||||
result = response_format.model_validate(json_data)
|
||||
else:
|
||||
result = content
|
||||
|
||||
# Record metrics and log slow calls
|
||||
duration = time.time() - start_time
|
||||
input_tokens = response.usage.input_tokens or 0 if response.usage else 0
|
||||
output_tokens = response.usage.output_tokens or 0 if response.usage else 0
|
||||
total_tokens = input_tokens + output_tokens
|
||||
|
||||
# Record LLM metrics
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0:
|
||||
logger.info(
|
||||
f"slow llm call: scope={scope}, model={self.provider}/{self.model}, "
|
||||
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
|
||||
f"time={duration:.3f}s"
|
||||
)
|
||||
|
||||
if return_usage:
|
||||
token_usage = TokenUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=total_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning("Anthropic returned invalid JSON, retrying...")
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"Anthropic returned invalid JSON after {max_retries + 1} attempts")
|
||||
raise
|
||||
|
||||
except (APIConnectionError, RateLimitError, APIStatusError) as e:
|
||||
# Fast fail on 401/403
|
||||
if isinstance(e, APIStatusError) and e.status_code in (401, 403):
|
||||
logger.error(f"Anthropic auth error (HTTP {e.status_code}), not retrying: {str(e)}")
|
||||
raise
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
# Check if it's a rate limit or server error
|
||||
should_retry = isinstance(e, (APIConnectionError, RateLimitError)) or (
|
||||
isinstance(e, APIStatusError) and e.status_code >= 500
|
||||
)
|
||||
|
||||
if should_retry:
|
||||
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"Anthropic API error after {max_retries + 1} attempts: {str(e)}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error during Anthropic call: {type(e).__name__}: {str(e)}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("Anthropic 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:
|
||||
"""
|
||||
Make an LLM API call with tool/function calling support.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts. Can include tool results with role='tool'.
|
||||
tools: List of tool definitions in OpenAI format.
|
||||
max_completion_tokens: Maximum tokens in response.
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or specific function.
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
"""
|
||||
from anthropic import APIConnectionError, APIStatusError
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Convert OpenAI tool format to Anthropic format
|
||||
anthropic_tools = []
|
||||
for tool in tools:
|
||||
func = tool.get("function", {})
|
||||
anthropic_tools.append(
|
||||
{
|
||||
"name": func.get("name", ""),
|
||||
"description": func.get("description", ""),
|
||||
"input_schema": func.get("parameters", {"type": "object", "properties": {}}),
|
||||
}
|
||||
)
|
||||
|
||||
# Convert messages - handle tool results
|
||||
system_prompt = None
|
||||
anthropic_messages = []
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
system_prompt = (system_prompt + "\n\n" + content) if system_prompt else content
|
||||
elif role == "tool":
|
||||
# Anthropic uses tool_result blocks
|
||||
anthropic_messages.append(
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "tool_result", "tool_use_id": msg.get("tool_call_id", ""), "content": content}
|
||||
],
|
||||
}
|
||||
)
|
||||
elif role == "assistant" and msg.get("tool_calls"):
|
||||
# Convert assistant tool calls
|
||||
tool_use_blocks = []
|
||||
for tc in msg["tool_calls"]:
|
||||
tool_use_blocks.append(
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": tc.get("id", ""),
|
||||
"name": tc.get("function", {}).get("name", ""),
|
||||
"input": json.loads(tc.get("function", {}).get("arguments", "{}")),
|
||||
}
|
||||
)
|
||||
anthropic_messages.append({"role": "assistant", "content": tool_use_blocks})
|
||||
else:
|
||||
anthropic_messages.append({"role": role, "content": content})
|
||||
|
||||
call_params: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"messages": anthropic_messages,
|
||||
"tools": anthropic_tools,
|
||||
"max_tokens": max_completion_tokens or 4096,
|
||||
}
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if temperature is not None:
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await self._client.messages.create(**call_params)
|
||||
|
||||
# Extract content and tool calls
|
||||
content_parts = []
|
||||
tool_calls: list[LLMToolCall] = []
|
||||
|
||||
for block in response.content:
|
||||
if block.type == "text":
|
||||
content_parts.append(block.text)
|
||||
elif block.type == "tool_use":
|
||||
tool_calls.append(LLMToolCall(id=block.id, name=block.name, arguments=block.input or {}))
|
||||
|
||||
content = "".join(content_parts) if content_parts else None
|
||||
finish_reason = "tool_calls" if tool_calls else "stop"
|
||||
|
||||
# Extract token usage
|
||||
input_tokens = response.usage.input_tokens or 0
|
||||
output_tokens = response.usage.output_tokens or 0
|
||||
|
||||
# Record metrics
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=time.time() - start_time,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
|
||||
return LLMToolCallResult(
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
finish_reason=finish_reason,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
)
|
||||
|
||||
except (APIConnectionError, APIStatusError) as e:
|
||||
if isinstance(e, APIStatusError) and e.status_code in (401, 403):
|
||||
raise
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("Anthropic tool call failed")
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (close Anthropic client connections)."""
|
||||
if hasattr(self, "_client") and self._client:
|
||||
await self._client.close()
|
||||
@@ -0,0 +1,493 @@
|
||||
"""
|
||||
Claude Code LLM provider using Claude Agent SDK.
|
||||
|
||||
This provider enables using Claude Pro/Max subscriptions for API calls
|
||||
via the Claude CLI authentication. It uses the Claude Agent SDK which
|
||||
automatically handles authentication via `claude auth login` credentials.
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ClaudeCodeLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using Claude Code authentication.
|
||||
|
||||
Authenticates using Claude Pro/Max credentials via `claude auth login`
|
||||
and makes API calls through the Claude Agent SDK.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str, # Will be ignored, uses CLI auth
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Initialize Claude Code LLM provider."""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
|
||||
# Verify Claude Agent SDK is available
|
||||
try:
|
||||
self._verify_claude_code_available()
|
||||
logger.info("Claude Code: Using Claude Agent SDK (authentication via claude auth login)")
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to initialize Claude Code provider: {e}\n\n"
|
||||
"To set up Claude Code authentication:\n"
|
||||
"1. Install Claude Code CLI: npm install -g @anthropics/claude-code\n"
|
||||
"2. Login with your Pro/Max plan: claude auth login\n"
|
||||
"3. Verify authentication: claude --version\n\n"
|
||||
"Or use a different provider (anthropic, openai, gemini) with API keys."
|
||||
) from e
|
||||
|
||||
# Metrics collector is imported at module level
|
||||
|
||||
def _verify_claude_code_available(self) -> None:
|
||||
"""
|
||||
Verify that Claude Agent SDK can be imported and is properly configured.
|
||||
|
||||
Raises:
|
||||
ImportError: If Claude Agent SDK is not installed.
|
||||
RuntimeError: If Claude Code is not authenticated.
|
||||
"""
|
||||
try:
|
||||
# Import Claude Agent SDK
|
||||
# Reduce Claude Agent SDK logging verbosity
|
||||
import logging as sdk_logging
|
||||
|
||||
from claude_agent_sdk import query # noqa: F401
|
||||
|
||||
sdk_logging.getLogger("claude_agent_sdk").setLevel(sdk_logging.WARNING)
|
||||
sdk_logging.getLogger("claude_agent_sdk._internal").setLevel(sdk_logging.WARNING)
|
||||
|
||||
logger.debug("Claude Agent SDK imported successfully")
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"Claude Agent SDK not installed. Run: uv add claude-agent-sdk or pip install claude-agent-sdk"
|
||||
) from e
|
||||
|
||||
# SDK will automatically check for authentication when first used
|
||||
# No need to verify here - let it fail gracefully on first call with helpful error
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""
|
||||
Verify that the Claude Code provider is configured correctly by making a simple test call.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the connection test fails.
|
||||
"""
|
||||
try:
|
||||
test_messages = [{"role": "user", "content": "test"}]
|
||||
await self.call(
|
||||
messages=test_messages,
|
||||
max_completion_tokens=10,
|
||||
temperature=0.0,
|
||||
scope="test",
|
||||
max_retries=0,
|
||||
)
|
||||
logger.info("Claude Code connection verified successfully")
|
||||
except Exception as e:
|
||||
logger.error(f"Claude Code connection verification failed: {e}")
|
||||
raise RuntimeError(f"Failed to verify Claude Code connection: {e}") from e
|
||||
|
||||
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:
|
||||
"""
|
||||
Make an LLM API call with retry logic.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts with 'role' and 'content'.
|
||||
response_format: Optional Pydantic model for structured output.
|
||||
max_completion_tokens: Maximum tokens in response (ignored by Claude Agent SDK).
|
||||
temperature: Sampling temperature (ignored by Claude Agent SDK).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Use strict JSON schema enforcement (not supported).
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
|
||||
Returns:
|
||||
If return_usage=False: Parsed response if response_format is provided, otherwise text content.
|
||||
If return_usage=True: Tuple of (result, TokenUsage) with estimated token counts.
|
||||
|
||||
Raises:
|
||||
OutputTooLongError: If output exceeds token limits (not supported by Claude Agent SDK).
|
||||
Exception: Re-raises API errors after retries exhausted.
|
||||
"""
|
||||
from claude_agent_sdk import AssistantMessage, ClaudeAgentOptions, TextBlock, query
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Build system prompt
|
||||
system_prompt = ""
|
||||
user_content = ""
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
system_prompt += ("\n\n" + content) if system_prompt else content
|
||||
elif role == "user":
|
||||
user_content += ("\n\n" + content) if user_content else content
|
||||
elif role == "assistant":
|
||||
# Claude Agent SDK doesn't support multi-turn easily in query()
|
||||
# For now, prepend assistant messages to user content
|
||||
user_content += f"\n\n[Previous assistant response: {content}]"
|
||||
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_instruction = (
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}\n\n"
|
||||
"Respond with ONLY the JSON, no markdown formatting."
|
||||
)
|
||||
user_content += schema_instruction
|
||||
|
||||
# Configure SDK options
|
||||
options = ClaudeAgentOptions(
|
||||
system_prompt=system_prompt if system_prompt else None,
|
||||
max_turns=1, # Single-turn for API-style interactions
|
||||
allowed_tools=[], # Disable tools for standard LLM calls
|
||||
)
|
||||
|
||||
# Call Claude Agent SDK
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
# Collect streaming response
|
||||
full_text = ""
|
||||
|
||||
async for message in query(prompt=user_content, options=options):
|
||||
if isinstance(message, AssistantMessage):
|
||||
for block in message.content:
|
||||
if isinstance(block, TextBlock):
|
||||
full_text += block.text
|
||||
|
||||
# Handle structured output
|
||||
if response_format is not None:
|
||||
# Models may wrap JSON in markdown
|
||||
clean_text = full_text
|
||||
if "```json" in full_text:
|
||||
clean_text = full_text.split("```json")[1].split("```")[0].strip()
|
||||
elif "```" in full_text:
|
||||
clean_text = full_text.split("```")[1].split("```")[0].strip()
|
||||
|
||||
try:
|
||||
json_data = json.loads(clean_text)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"Claude Code JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {e}")
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
last_exception = e
|
||||
continue
|
||||
raise
|
||||
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
else:
|
||||
result = response_format.model_validate(json_data)
|
||||
else:
|
||||
result = full_text
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
# Estimate token usage (Claude Agent SDK doesn't report exact counts)
|
||||
# Use character count / 4 as rough estimate (1 token ≈ 4 characters)
|
||||
estimated_input = sum(len(m.get("content", "")) for m in messages) // 4
|
||||
estimated_output = len(full_text) // 4
|
||||
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0:
|
||||
logger.info(
|
||||
f"slow llm call: scope={scope}, model={self.provider}/{self.model}, time={duration:.3f}s"
|
||||
)
|
||||
|
||||
if return_usage:
|
||||
token_usage = TokenUsage(
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
total_tokens=estimated_input + estimated_output,
|
||||
)
|
||||
return result, token_usage
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
last_exception = e
|
||||
|
||||
# Check for authentication errors
|
||||
error_str = str(e).lower()
|
||||
if "auth" in error_str or "login" in error_str or "credential" in error_str:
|
||||
logger.error(f"Claude Code authentication error: {e}")
|
||||
raise RuntimeError(
|
||||
f"Claude Code authentication failed: {e}\n\n"
|
||||
"Run 'claude auth login' to authenticate with Claude Pro/Max."
|
||||
) from e
|
||||
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
logger.warning(f"Claude Code error (attempt {attempt + 1}/{max_retries + 1}): {e}")
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"Claude Code error after {max_retries + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("Claude Code 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:
|
||||
"""
|
||||
Make an LLM API call with tool/function calling support using Claude Agent SDK.
|
||||
|
||||
This implementation uses ClaudeSDKClient (not query()) because custom tools via
|
||||
SDK MCP servers are only supported with the client. Tools are converted from OpenAI
|
||||
format to SDK MCP tools, and tool names are formatted as mcp__hindsight_tools__{name}.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts. Can include tool results with role='tool'.
|
||||
tools: List of tool definitions in OpenAI format.
|
||||
max_completion_tokens: Maximum tokens in response (not used by Claude Agent SDK).
|
||||
temperature: Sampling temperature (not used by Claude Agent SDK).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools (not used by Claude Agent SDK).
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
"""
|
||||
from claude_agent_sdk import (
|
||||
AssistantMessage,
|
||||
ClaudeAgentOptions,
|
||||
ClaudeSDKClient,
|
||||
SdkMcpTool,
|
||||
TextBlock,
|
||||
ToolUseBlock,
|
||||
create_sdk_mcp_server,
|
||||
)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Convert OpenAI tool format to Claude Agent SDK SdkMcpTool format
|
||||
sdk_tools: list[SdkMcpTool] = []
|
||||
tool_names: list[str] = []
|
||||
|
||||
for tool in tools:
|
||||
func = tool.get("function", {})
|
||||
tool_name = func.get("name", "")
|
||||
tool_description = func.get("description", "")
|
||||
parameters = func.get("parameters", {})
|
||||
|
||||
# Create a handler with proper closure to avoid transport issues
|
||||
def make_handler(name: str):
|
||||
async def handler(args: dict[str, Any]) -> dict[str, Any]:
|
||||
# Return immediately with success - tool execution happens externally
|
||||
return {
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"[Tool {name} called successfully]",
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
return handler
|
||||
|
||||
sdk_tools.append(
|
||||
SdkMcpTool(
|
||||
name=tool_name,
|
||||
description=tool_description,
|
||||
input_schema=parameters,
|
||||
handler=make_handler(tool_name),
|
||||
)
|
||||
)
|
||||
tool_names.append(tool_name)
|
||||
|
||||
# Create an MCP server with the tools
|
||||
mcp_server = create_sdk_mcp_server(
|
||||
name="hindsight_tools",
|
||||
version="1.0.0",
|
||||
tools=sdk_tools if sdk_tools else None,
|
||||
)
|
||||
|
||||
# Build system prompt and user content from messages
|
||||
system_prompt = ""
|
||||
user_content = ""
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
system_prompt += ("\n\n" + content) if system_prompt else content
|
||||
elif role == "user":
|
||||
user_content += ("\n\n" + content) if user_content else content
|
||||
elif role == "assistant":
|
||||
# Include previous assistant messages as context
|
||||
user_content += f"\n\n[Previous assistant response: {content}]"
|
||||
elif role == "tool":
|
||||
# Tool results are already in tool_results_map, append to user context
|
||||
tool_call_id = msg.get("tool_call_id", "")
|
||||
user_content += f"\n\n[Tool result for {tool_call_id}: {content}]"
|
||||
|
||||
# Format tool names for SDK MCP servers: mcp__{server_name}__{tool_name}
|
||||
# This is required by the Claude Agent SDK for MCP server tools
|
||||
allowed_tool_names = [f"mcp__hindsight_tools__{name}" for name in tool_names]
|
||||
|
||||
# Configure SDK options with MCP server
|
||||
options = ClaudeAgentOptions(
|
||||
system_prompt=system_prompt if system_prompt else None,
|
||||
max_turns=1, # Single-turn for API-style interactions
|
||||
mcp_servers={"hindsight_tools": mcp_server} if sdk_tools else {},
|
||||
allowed_tools=allowed_tool_names if allowed_tool_names else [],
|
||||
)
|
||||
|
||||
# Call Claude Agent SDK with retry logic
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
full_text = ""
|
||||
tool_calls: list[LLMToolCall] = []
|
||||
|
||||
# Use ClaudeSDKClient for tool calling support
|
||||
# Note: query() does NOT support custom tools, only ClaudeSDKClient does
|
||||
async with ClaudeSDKClient(options=options) as client:
|
||||
# Send the query
|
||||
await client.query(user_content)
|
||||
|
||||
# Receive response
|
||||
async for message in client.receive_response():
|
||||
if isinstance(message, AssistantMessage):
|
||||
for block in message.content:
|
||||
if isinstance(block, TextBlock):
|
||||
full_text += block.text
|
||||
elif isinstance(block, ToolUseBlock):
|
||||
# SDK returns tool names with MCP prefix (mcp__hindsight_tools__{name})
|
||||
# Strip the prefix to return original tool name expected by caller
|
||||
tool_name = block.name
|
||||
if tool_name.startswith("mcp__hindsight_tools__"):
|
||||
tool_name = tool_name.replace("mcp__hindsight_tools__", "", 1)
|
||||
|
||||
tool_calls.append(
|
||||
LLMToolCall(
|
||||
id=block.id,
|
||||
name=tool_name,
|
||||
arguments=block.input,
|
||||
)
|
||||
)
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
# Estimate token usage (Claude Agent SDK doesn't report exact counts)
|
||||
estimated_input = sum(len(m.get("content", "")) for m in messages) // 4
|
||||
estimated_output = len(full_text) // 4
|
||||
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0:
|
||||
logger.info(
|
||||
f"slow llm call: scope={scope}, model={self.provider}/{self.model}, time={duration:.3f}s"
|
||||
)
|
||||
|
||||
return LLMToolCallResult(
|
||||
content=full_text if full_text else None,
|
||||
tool_calls=tool_calls,
|
||||
finish_reason="tool_calls" if tool_calls else "stop",
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
last_exception = e
|
||||
|
||||
# Check for authentication errors
|
||||
error_str = str(e).lower()
|
||||
if "auth" in error_str or "login" in error_str or "credential" in error_str:
|
||||
logger.error(f"Claude Code authentication error: {e}")
|
||||
raise RuntimeError(
|
||||
f"Claude Code authentication failed: {e}\n\n"
|
||||
"Run 'claude auth login' to authenticate with Claude Pro/Max."
|
||||
) from e
|
||||
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
logger.warning(f"Claude Code tool call error (attempt {attempt + 1}/{max_retries + 1}): {e}")
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"Claude Code tool call error after {max_retries + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("Claude Code tool call failed after all retries")
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (no HTTP client to close for Claude Agent SDK)."""
|
||||
pass
|
||||
@@ -0,0 +1,578 @@
|
||||
"""
|
||||
OpenAI Codex LLM provider using ChatGPT Plus/Pro OAuth authentication.
|
||||
|
||||
This provider enables using ChatGPT Plus/Pro subscriptions for API calls
|
||||
without separate OpenAI Platform API credits. It uses OAuth tokens from
|
||||
~/.codex/auth.json and communicates with the ChatGPT backend API.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CodexLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using OpenAI Codex OAuth authentication.
|
||||
|
||||
Authenticates using ChatGPT Plus/Pro credentials stored in ~/.codex/auth.json
|
||||
and makes API calls to chatgpt.com/backend-api/codex/responses.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str, # Will be ignored, reads from ~/.codex/auth.json
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Initialize Codex LLM provider."""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
|
||||
# Load Codex OAuth credentials
|
||||
try:
|
||||
self.access_token, self.account_id = self._load_codex_auth()
|
||||
logger.info(f"Loaded Codex OAuth credentials for account: {self.account_id}")
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to load Codex OAuth credentials from ~/.codex/auth.json: {e}\n\n"
|
||||
"To set up Codex authentication:\n"
|
||||
"1. Install Codex CLI: npm install -g @openai/codex\n"
|
||||
"2. Login: codex auth login\n"
|
||||
"3. Verify: ls ~/.codex/auth.json\n\n"
|
||||
"Or use a different provider (openai, anthropic, gemini) with API keys."
|
||||
) from e
|
||||
|
||||
# Use ChatGPT backend API endpoint
|
||||
if not self.base_url:
|
||||
self.base_url = "https://chatgpt.com/backend-api"
|
||||
|
||||
# Normalize model name (strip openai/ prefix if present)
|
||||
if self.model.startswith("openai/"):
|
||||
self.model = self.model[len("openai/") :]
|
||||
|
||||
# Map reasoning effort to Codex reasoning summary format
|
||||
# Codex supports: "auto", "concise", "detailed"
|
||||
self.reasoning_summary = self._map_reasoning_effort(reasoning_effort)
|
||||
|
||||
# HTTP client for SSE streaming
|
||||
self._client = httpx.AsyncClient(timeout=120.0)
|
||||
|
||||
def _load_codex_auth(self) -> tuple[str, str]:
|
||||
"""
|
||||
Load OAuth credentials from ~/.codex/auth.json.
|
||||
|
||||
Returns:
|
||||
Tuple of (access_token, account_id).
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If auth file doesn't exist.
|
||||
ValueError: If auth file is invalid.
|
||||
"""
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(
|
||||
f"Codex auth file not found: {auth_file}\nRun 'codex auth login' to authenticate with ChatGPT Plus/Pro."
|
||||
)
|
||||
|
||||
with open(auth_file) as f:
|
||||
data = json.load(f)
|
||||
|
||||
# Validate auth structure
|
||||
auth_mode = data.get("auth_mode")
|
||||
if auth_mode != "chatgpt":
|
||||
raise ValueError(f"Expected auth_mode='chatgpt', got: {auth_mode}")
|
||||
|
||||
tokens = data.get("tokens", {})
|
||||
access_token = tokens.get("access_token")
|
||||
account_id = tokens.get("account_id")
|
||||
|
||||
if not access_token:
|
||||
raise ValueError("No access_token found in Codex auth file. Run 'codex auth login' again.")
|
||||
|
||||
return access_token, account_id
|
||||
|
||||
def _map_reasoning_effort(self, effort: str) -> str:
|
||||
"""
|
||||
Map standard reasoning effort to Codex reasoning summary format.
|
||||
|
||||
Args:
|
||||
effort: Standard effort level ("low", "medium", "high", "xhigh").
|
||||
|
||||
Returns:
|
||||
Codex reasoning summary: "concise", "detailed", or "auto".
|
||||
"""
|
||||
mapping = {
|
||||
"low": "concise",
|
||||
"medium": "auto",
|
||||
"high": "detailed",
|
||||
"xhigh": "detailed",
|
||||
}
|
||||
return mapping.get(effort.lower(), "auto")
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""Verify Codex connection by making a simple test call."""
|
||||
try:
|
||||
logger.info(f"Verifying Codex LLM: model={self.model}, account={self.account_id}...")
|
||||
await self.call(
|
||||
messages=[{"role": "user", "content": "Say 'ok'"}],
|
||||
max_completion_tokens=10,
|
||||
max_retries=2,
|
||||
initial_backoff=0.5,
|
||||
max_backoff=2.0,
|
||||
)
|
||||
logger.info(f"Codex LLM verified: {self.model}")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Codex LLM connection verification failed for {self.model}: {e}") from e
|
||||
|
||||
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:
|
||||
"""Make API call to Codex backend with SSE streaming."""
|
||||
start_time = time.time()
|
||||
|
||||
# Prepare system instructions
|
||||
system_instruction = ""
|
||||
user_messages = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
system_instruction += ("\n\n" + content) if system_instruction else content
|
||||
else:
|
||||
user_messages.append(msg)
|
||||
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
system_instruction += schema_msg
|
||||
|
||||
# gpt-5.2-codex only supports "detailed" reasoning summary
|
||||
reasoning_summary = "detailed" if "5.2" in self.model else self.reasoning_summary
|
||||
|
||||
# Build Codex request payload
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"instructions": system_instruction,
|
||||
"input": [
|
||||
{
|
||||
"type": "message",
|
||||
"role": msg.get("role", "user"),
|
||||
"content": msg.get("content", ""),
|
||||
}
|
||||
for msg in user_messages
|
||||
],
|
||||
"tools": [],
|
||||
"tool_choice": "auto",
|
||||
"parallel_tool_calls": True,
|
||||
"reasoning": {"summary": reasoning_summary},
|
||||
"store": False, # Codex uses stateless mode
|
||||
"stream": True, # SSE streaming
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"prompt_cache_key": str(uuid.uuid4()),
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.access_token}",
|
||||
"Content-Type": "application/json",
|
||||
"OpenAI-Account-ID": self.account_id,
|
||||
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)",
|
||||
"Origin": "https://chatgpt.com",
|
||||
}
|
||||
|
||||
url = f"{self.base_url}/codex/responses"
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
|
||||
response.raise_for_status()
|
||||
|
||||
# Parse SSE stream
|
||||
content = await self._parse_sse_stream(response)
|
||||
|
||||
# Handle structured output
|
||||
if response_format is not None:
|
||||
# Models may wrap JSON in markdown
|
||||
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 as e:
|
||||
logger.warning(f"Codex JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {e}")
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
last_exception = e
|
||||
continue
|
||||
raise
|
||||
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
else:
|
||||
result = response_format.model_validate(json_data)
|
||||
else:
|
||||
result = content
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=0, # Codex doesn't report token counts in SSE
|
||||
output_tokens=0,
|
||||
success=True,
|
||||
)
|
||||
|
||||
if return_usage:
|
||||
# Codex doesn't provide token counts, estimate based on content
|
||||
estimated_input = sum(len(m.get("content", "")) for m in messages) // 4
|
||||
estimated_output = len(content) // 4
|
||||
token_usage = TokenUsage(
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
total_tokens=estimated_input + estimated_output,
|
||||
)
|
||||
return result, token_usage
|
||||
|
||||
return result
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
last_exception = e
|
||||
status_code = e.response.status_code
|
||||
|
||||
# Fast fail on auth errors
|
||||
if status_code in (401, 403):
|
||||
logger.error(f"Codex auth error (HTTP {status_code}): {e.response.text[:200]}")
|
||||
raise RuntimeError(
|
||||
"Codex authentication failed. Your OAuth token may have expired.\n"
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
) from e
|
||||
|
||||
# Log the actual error message from the API
|
||||
error_detail = e.response.text[:500] if hasattr(e.response, "text") else str(e)
|
||||
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
logger.warning(
|
||||
f"Codex HTTP error {status_code} (attempt {attempt + 1}/{max_retries + 1}): {error_detail}"
|
||||
)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(
|
||||
f"Codex HTTP error after {max_retries + 1} attempts: Status {status_code}, Detail: {error_detail}"
|
||||
)
|
||||
raise
|
||||
|
||||
except httpx.RequestError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
logger.warning(f"Codex connection error (attempt {attempt + 1}/{max_retries + 1}): {e}")
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"Codex connection error after {max_retries + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected Codex error: {type(e).__name__}: {e}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("Codex call failed after all retries")
|
||||
|
||||
async def _parse_sse_stream(self, response: httpx.Response) -> str:
|
||||
"""
|
||||
Parse Server-Sent Events (SSE) stream from Codex API.
|
||||
|
||||
Args:
|
||||
response: HTTP response with SSE stream.
|
||||
|
||||
Returns:
|
||||
Extracted text content from stream.
|
||||
"""
|
||||
full_text = ""
|
||||
event_type = None
|
||||
|
||||
async for line in response.aiter_lines():
|
||||
if not line:
|
||||
continue
|
||||
|
||||
# Track event type
|
||||
if line.startswith("event: "):
|
||||
event_type = line[7:]
|
||||
|
||||
# Parse data
|
||||
elif line.startswith("data: "):
|
||||
data_str = line[6:]
|
||||
if data_str == "[DONE]":
|
||||
break
|
||||
|
||||
try:
|
||||
data = json.loads(data_str)
|
||||
|
||||
# Extract content based on event type
|
||||
if event_type == "response.text.delta" and "delta" in data:
|
||||
full_text += data["delta"]
|
||||
elif event_type == "response.content_part.delta" and "delta" in data:
|
||||
full_text += data["delta"]
|
||||
# Check for item content
|
||||
elif "item" in data:
|
||||
item = data["item"]
|
||||
if "content" in item:
|
||||
content = item["content"]
|
||||
if isinstance(content, list):
|
||||
for part in content:
|
||||
if isinstance(part, dict) and "text" in part:
|
||||
full_text += part["text"]
|
||||
elif isinstance(content, str):
|
||||
full_text += content
|
||||
|
||||
except json.JSONDecodeError:
|
||||
# Skip malformed JSON events
|
||||
pass
|
||||
|
||||
return full_text
|
||||
|
||||
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:
|
||||
"""
|
||||
Make API call with tool calling support.
|
||||
|
||||
Parses Codex SSE stream to extract tool calls from response.output_item.done events.
|
||||
Tools are converted from OpenAI format to Codex format (flat structure at top level).
|
||||
|
||||
Args:
|
||||
messages: List of message dicts. Can include tool results with role='tool'.
|
||||
tools: List of tool definitions in OpenAI format.
|
||||
max_completion_tokens: Maximum tokens in response.
|
||||
temperature: Sampling temperature.
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or specific function.
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Prepare system instructions
|
||||
system_instruction = ""
|
||||
user_messages = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
system_instruction += ("\n\n" + content) if system_instruction else content
|
||||
elif role == "tool":
|
||||
# Handle tool results
|
||||
user_messages.append(
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": f"Tool result: {content}",
|
||||
}
|
||||
)
|
||||
else:
|
||||
user_messages.append(
|
||||
{
|
||||
"type": "message",
|
||||
"role": role,
|
||||
"content": content,
|
||||
}
|
||||
)
|
||||
|
||||
# Convert tools to Codex format
|
||||
# Codex expects tools with type and name/description/parameters at top level
|
||||
codex_tools = []
|
||||
for tool in tools:
|
||||
func = tool.get("function", {})
|
||||
codex_tools.append(
|
||||
{
|
||||
"type": "function",
|
||||
"name": func.get("name", ""),
|
||||
"description": func.get("description", ""),
|
||||
"parameters": func.get("parameters", {}),
|
||||
}
|
||||
)
|
||||
|
||||
# gpt-5.2-codex only supports "detailed" reasoning summary
|
||||
reasoning_summary = "detailed" if "5.2" in self.model else self.reasoning_summary
|
||||
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"instructions": system_instruction,
|
||||
"input": user_messages,
|
||||
"tools": codex_tools,
|
||||
"tool_choice": tool_choice,
|
||||
"parallel_tool_calls": True,
|
||||
"reasoning": {"summary": reasoning_summary},
|
||||
"store": False,
|
||||
"stream": True,
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"prompt_cache_key": str(uuid.uuid4()),
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.access_token}",
|
||||
"Content-Type": "application/json",
|
||||
"OpenAI-Account-ID": self.account_id,
|
||||
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)",
|
||||
"Origin": "https://chatgpt.com",
|
||||
}
|
||||
|
||||
url = f"{self.base_url}/codex/responses"
|
||||
|
||||
# Debug logging for troubleshooting
|
||||
logger.debug(f"Codex tool call request: url={url}, model={payload['model']}, tools={len(codex_tools)}")
|
||||
|
||||
try:
|
||||
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
|
||||
|
||||
# Log response details on error
|
||||
if response.status_code != 200:
|
||||
logger.error(f"Codex API error {response.status_code}: {response.text[:500]}")
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
# Parse SSE for tool calls and content
|
||||
content, tool_calls = await self._parse_sse_tool_stream(response)
|
||||
|
||||
duration = time.time() - start_time
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
success=True,
|
||||
)
|
||||
|
||||
return LLMToolCallResult(
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
finish_reason="tool_calls" if tool_calls else "stop",
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Codex tool call error: {e}")
|
||||
raise
|
||||
|
||||
async def _parse_sse_tool_stream(self, response: httpx.Response) -> tuple[str | None, list[LLMToolCall]]:
|
||||
"""
|
||||
Parse SSE stream for tool calls and content.
|
||||
|
||||
Returns:
|
||||
Tuple of (content, tool_calls).
|
||||
"""
|
||||
content = ""
|
||||
tool_calls: list[LLMToolCall] = []
|
||||
event_type = None
|
||||
|
||||
async for line in response.aiter_lines():
|
||||
if not line:
|
||||
continue
|
||||
|
||||
if line.startswith("event: "):
|
||||
event_type = line[7:]
|
||||
|
||||
elif line.startswith("data: "):
|
||||
data_str = line[6:]
|
||||
if data_str == "[DONE]":
|
||||
break
|
||||
|
||||
try:
|
||||
data = json.loads(data_str)
|
||||
|
||||
# Extract text content
|
||||
if event_type == "response.text.delta" and "delta" in data:
|
||||
content += data["delta"]
|
||||
|
||||
# Extract completed tool calls from response.output_item.done
|
||||
elif event_type == "response.output_item.done":
|
||||
item = data.get("item", {})
|
||||
if item.get("type") == "function_call" and item.get("status") == "completed":
|
||||
tool_name = item.get("name", "")
|
||||
arguments_str = item.get("arguments", "{}")
|
||||
call_id = item.get("call_id", "")
|
||||
|
||||
try:
|
||||
arguments = json.loads(arguments_str)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to parse tool arguments: {arguments_str}")
|
||||
arguments = {}
|
||||
|
||||
tool_calls.append(
|
||||
LLMToolCall(
|
||||
id=call_id,
|
||||
name=tool_name,
|
||||
arguments=arguments,
|
||||
)
|
||||
)
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"Failed to parse SSE data: {e}, data_str: {data_str[:200]}")
|
||||
|
||||
return content if content else None, tool_calls
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up HTTP client."""
|
||||
await self._client.aclose()
|
||||
@@ -0,0 +1,502 @@
|
||||
"""
|
||||
Google Gemini/VertexAI LLM provider.
|
||||
|
||||
This provider supports both:
|
||||
1. Gemini API (api.generativeai.google.com) with API key authentication
|
||||
2. Vertex AI with service account or Application Default Credentials (ADC)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from google import genai
|
||||
from google.genai import errors as genai_errors
|
||||
from google.genai import types as genai_types
|
||||
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Vertex AI imports (optional)
|
||||
try:
|
||||
import google.auth
|
||||
from google.oauth2 import service_account
|
||||
|
||||
VERTEXAI_AVAILABLE = True
|
||||
except ImportError:
|
||||
VERTEXAI_AVAILABLE = False
|
||||
|
||||
|
||||
class GeminiLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider for Google Gemini and Vertex AI.
|
||||
|
||||
Supports:
|
||||
- Gemini API: provider="gemini", requires api_key
|
||||
- Vertex AI: provider="vertexai", requires project_id and region, uses ADC or service account
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Initialize Gemini/VertexAI LLM provider."""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
|
||||
self._client = None
|
||||
self._is_vertexai = self.provider == "vertexai"
|
||||
|
||||
if self._is_vertexai:
|
||||
self._init_vertexai(**kwargs)
|
||||
else:
|
||||
self._init_gemini()
|
||||
|
||||
def _init_gemini(self) -> None:
|
||||
"""Initialize Gemini API client."""
|
||||
if not self.api_key:
|
||||
raise ValueError("Gemini provider requires api_key")
|
||||
|
||||
self._client = genai.Client(api_key=self.api_key)
|
||||
logger.info(f"Gemini API: model={self.model}")
|
||||
|
||||
def _init_vertexai(self, **kwargs: Any) -> None:
|
||||
"""Initialize Vertex AI client with project, region, and credentials."""
|
||||
# Extract Vertex AI config from kwargs
|
||||
project_id = kwargs.get("vertexai_project_id")
|
||||
region = kwargs.get("vertexai_region", "us-central1")
|
||||
service_account_key = kwargs.get("vertexai_service_account_key")
|
||||
credentials = kwargs.get("vertexai_credentials") # Pre-loaded credentials object
|
||||
|
||||
if not project_id:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID is required for Vertex AI provider. "
|
||||
"Set it to your GCP project ID."
|
||||
)
|
||||
|
||||
auth_method = "ADC"
|
||||
|
||||
# Use pre-loaded credentials if provided (passed from LLMProvider)
|
||||
if credentials is not None:
|
||||
auth_method = "service_account"
|
||||
# Otherwise, load explicit service account credentials if path provided
|
||||
elif service_account_key:
|
||||
if not VERTEXAI_AVAILABLE:
|
||||
raise ValueError(
|
||||
"Vertex AI service account auth requires 'google-auth' package. "
|
||||
"Install with: pip install google-auth"
|
||||
)
|
||||
credentials = service_account.Credentials.from_service_account_file(
|
||||
service_account_key,
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
)
|
||||
auth_method = "service_account"
|
||||
logger.info(f"Vertex AI: Using service account key: {service_account_key}")
|
||||
|
||||
# Strip google/ prefix from model name — native SDK uses bare names
|
||||
# e.g. "google/gemini-2.0-flash-lite-001" -> "gemini-2.0-flash-lite-001"
|
||||
if self.model.startswith("google/"):
|
||||
self.model = self.model[len("google/") :]
|
||||
|
||||
# Create Vertex AI client
|
||||
client_kwargs: dict[str, Any] = {
|
||||
"vertexai": True,
|
||||
"project": project_id,
|
||||
"location": region,
|
||||
}
|
||||
if credentials is not None:
|
||||
client_kwargs["credentials"] = credentials
|
||||
|
||||
self._client = genai.Client(**client_kwargs)
|
||||
|
||||
logger.info(f"Vertex AI: project={project_id}, region={region}, model={self.model}, auth={auth_method}")
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""
|
||||
Verify that the Gemini/VertexAI provider is configured correctly.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the connection test fails.
|
||||
"""
|
||||
try:
|
||||
logger.info(f"Verifying {self.provider.upper()}: model={self.model}...")
|
||||
await self.call(
|
||||
messages=[{"role": "user", "content": "Say 'ok'"}],
|
||||
max_completion_tokens=100,
|
||||
max_retries=2,
|
||||
initial_backoff=0.5,
|
||||
max_backoff=2.0,
|
||||
)
|
||||
logger.info(f"{self.provider.upper()} connection verified successfully")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to verify {self.provider.upper()} connection: {e}") from e
|
||||
|
||||
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:
|
||||
"""
|
||||
Make a Gemini/VertexAI API call with retry logic.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts with 'role' and 'content'.
|
||||
response_format: Optional Pydantic model for structured output.
|
||||
max_completion_tokens: Maximum tokens in response (not supported by Gemini).
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Use strict JSON schema enforcement (not supported by Gemini).
|
||||
return_usage: If True, return tuple (result, TokenUsage).
|
||||
|
||||
Returns:
|
||||
If return_usage=False: Parsed response if response_format provided, else text.
|
||||
If return_usage=True: Tuple of (result, TokenUsage).
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Convert OpenAI-style messages to Gemini format
|
||||
system_instruction = None
|
||||
gemini_contents = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
if system_instruction:
|
||||
system_instruction += "\n\n" + content
|
||||
else:
|
||||
system_instruction = content
|
||||
elif role == "assistant":
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=[genai_types.Part(text=content)]))
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
|
||||
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
if system_instruction:
|
||||
system_instruction += schema_msg
|
||||
else:
|
||||
system_instruction = schema_msg
|
||||
|
||||
# Build generation config
|
||||
config_kwargs: dict[str, Any] = {}
|
||||
if system_instruction:
|
||||
config_kwargs["system_instruction"] = system_instruction
|
||||
if response_format is not None:
|
||||
config_kwargs["response_mime_type"] = "application/json"
|
||||
config_kwargs["response_schema"] = response_format
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
|
||||
generation_config = genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await self._client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=gemini_contents,
|
||||
config=generation_config,
|
||||
)
|
||||
|
||||
content = response.text
|
||||
|
||||
# Handle empty response
|
||||
if content is None:
|
||||
block_reason = None
|
||||
if hasattr(response, "candidates") and response.candidates:
|
||||
candidate = response.candidates[0]
|
||||
if hasattr(candidate, "finish_reason"):
|
||||
block_reason = candidate.finish_reason
|
||||
|
||||
if attempt < max_retries:
|
||||
logger.warning(f"Gemini returned empty response (reason: {block_reason}), retrying...")
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
raise RuntimeError(f"Gemini returned empty response after {max_retries + 1} attempts")
|
||||
|
||||
# Parse structured output if requested
|
||||
if response_format is not None:
|
||||
json_data = json.loads(content)
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
else:
|
||||
result = response_format.model_validate(json_data)
|
||||
else:
|
||||
result = content
|
||||
|
||||
# Extract token usage
|
||||
input_tokens = 0
|
||||
output_tokens = 0
|
||||
if hasattr(response, "usage_metadata") and response.usage_metadata:
|
||||
usage = response.usage_metadata
|
||||
input_tokens = usage.prompt_token_count or 0
|
||||
output_tokens = usage.candidates_token_count or 0
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0 and input_tokens > 0:
|
||||
logger.info(
|
||||
f"slow llm call: scope={scope}, model={self.provider}/{self.model}, "
|
||||
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
|
||||
f"time={duration:.3f}s"
|
||||
)
|
||||
|
||||
if return_usage:
|
||||
token_usage = TokenUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=input_tokens + output_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning("Gemini returned invalid JSON, retrying...")
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"Gemini returned invalid JSON after {max_retries + 1} attempts")
|
||||
raise
|
||||
|
||||
except genai_errors.APIError as e:
|
||||
# Fast fail on auth errors - these won't recover with retries
|
||||
if e.code in (401, 403):
|
||||
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
|
||||
raise
|
||||
|
||||
# Retry on retryable errors (rate limits, server errors, client errors)
|
||||
if e.code in (400, 429, 500, 502, 503, 504) or (e.code and e.code >= 500):
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
|
||||
await asyncio.sleep(backoff + jitter)
|
||||
else:
|
||||
logger.error(f"Gemini API error after {max_retries + 1} attempts: {str(e)}")
|
||||
raise
|
||||
else:
|
||||
logger.error(f"Gemini API error: {type(e).__name__}: {str(e)}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error during Gemini call: {type(e).__name__}: {str(e)}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("Gemini 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:
|
||||
"""
|
||||
Make a Gemini/VertexAI API call with tool/function calling support.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts. Can include tool results with role='tool'.
|
||||
tools: List of tool definitions in OpenAI format.
|
||||
max_completion_tokens: Maximum tokens (not supported by Gemini).
|
||||
temperature: Sampling temperature.
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools (Gemini uses "auto" only).
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Convert tools to Gemini format
|
||||
gemini_tools = []
|
||||
for tool in tools:
|
||||
func = tool.get("function", {})
|
||||
gemini_tools.append(
|
||||
genai_types.Tool(
|
||||
function_declarations=[
|
||||
genai_types.FunctionDeclaration(
|
||||
name=func.get("name", ""),
|
||||
description=func.get("description", ""),
|
||||
parameters=func.get("parameters"),
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
# Convert messages
|
||||
system_instruction = None
|
||||
gemini_contents = []
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
system_instruction = (system_instruction + "\n\n" + content) if system_instruction else content
|
||||
elif role == "tool":
|
||||
# Gemini uses function_response
|
||||
gemini_contents.append(
|
||||
genai_types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
genai_types.Part(
|
||||
function_response=genai_types.FunctionResponse(
|
||||
name=msg.get("name", ""),
|
||||
response={"result": content},
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
)
|
||||
elif role == "assistant":
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=[genai_types.Part(text=content)]))
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
|
||||
|
||||
config_kwargs: dict[str, Any] = {"tools": gemini_tools}
|
||||
if system_instruction:
|
||||
config_kwargs["system_instruction"] = system_instruction
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
|
||||
config = genai_types.GenerateContentConfig(**config_kwargs)
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await self._client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=gemini_contents,
|
||||
config=config,
|
||||
)
|
||||
|
||||
# Extract content and tool calls
|
||||
content = None
|
||||
tool_calls: list[LLMToolCall] = []
|
||||
|
||||
if response.candidates and response.candidates[0].content:
|
||||
parts = response.candidates[0].content.parts
|
||||
if parts:
|
||||
for part in parts:
|
||||
if hasattr(part, "text") and part.text:
|
||||
content = part.text
|
||||
if hasattr(part, "function_call") and part.function_call:
|
||||
fc = part.function_call
|
||||
tool_calls.append(
|
||||
LLMToolCall(
|
||||
id=f"gemini_{len(tool_calls)}",
|
||||
name=fc.name,
|
||||
arguments=dict(fc.args) if fc.args else {},
|
||||
)
|
||||
)
|
||||
|
||||
finish_reason = "tool_calls" if tool_calls else "stop"
|
||||
|
||||
# Extract token usage
|
||||
input_tokens = 0
|
||||
output_tokens = 0
|
||||
if response.usage_metadata:
|
||||
input_tokens = response.usage_metadata.prompt_token_count or 0
|
||||
output_tokens = response.usage_metadata.candidates_token_count or 0
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
|
||||
return LLMToolCallResult(
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
finish_reason=finish_reason,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
)
|
||||
|
||||
except genai_errors.APIError as e:
|
||||
# Fast fail on auth errors
|
||||
if e.code in (401, 403):
|
||||
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
|
||||
raise
|
||||
|
||||
# Retry on retryable errors
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error during Gemini tool call: {type(e).__name__}: {str(e)}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("Gemini tool call failed")
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (close connections, etc.)."""
|
||||
# Gemini client doesn't require explicit cleanup
|
||||
pass
|
||||
@@ -0,0 +1,234 @@
|
||||
"""
|
||||
Mock LLM provider for testing.
|
||||
|
||||
This provider allows tests to record LLM calls and return configurable mock responses
|
||||
without making actual API calls to external LLM services.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from ..llm_interface import LLMInterface
|
||||
from ..response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MockLLM(LLMInterface):
|
||||
"""
|
||||
Mock LLM provider for testing.
|
||||
|
||||
This provider records all calls and returns configurable mock responses,
|
||||
enabling tests to verify LLM interactions without making real API calls.
|
||||
|
||||
Example:
|
||||
# Create mock provider
|
||||
mock_llm = MockLLM(provider="mock", api_key="", base_url="", model="mock-model")
|
||||
|
||||
# Set mock response
|
||||
mock_llm.set_mock_response({"answer": "test"})
|
||||
|
||||
# Make calls
|
||||
result = await mock_llm.call(
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
response_format=MyResponseModel
|
||||
)
|
||||
|
||||
# Verify calls
|
||||
calls = mock_llm.get_mock_calls()
|
||||
assert len(calls) == 1
|
||||
assert calls[0]["scope"] == "memory"
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
Initialize mock LLM provider.
|
||||
|
||||
Args:
|
||||
provider: Provider name (should be "mock").
|
||||
api_key: Not used for mock provider.
|
||||
base_url: Not used for mock provider.
|
||||
model: Model name for tracking.
|
||||
reasoning_effort: Not used for mock provider.
|
||||
**kwargs: Additional parameters (not used).
|
||||
"""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
|
||||
# Storage for test verification
|
||||
self._mock_calls: list[dict] = []
|
||||
self._mock_response: Any = None
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""
|
||||
Verify mock provider (always succeeds).
|
||||
|
||||
Mock provider doesn't need connection verification since it doesn't
|
||||
make real API calls.
|
||||
"""
|
||||
logger.debug("Mock LLM: connection verification (always succeeds)")
|
||||
|
||||
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:
|
||||
"""
|
||||
Make a mock LLM API call.
|
||||
|
||||
Records the call for test verification and returns the configured mock response.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts with 'role' and 'content'.
|
||||
response_format: Optional Pydantic model for structured output.
|
||||
max_completion_tokens: Not used in mock.
|
||||
temperature: Not used in mock.
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Not used in mock.
|
||||
initial_backoff: Not used in mock.
|
||||
max_backoff: Not used in mock.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Not used in mock.
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
|
||||
Returns:
|
||||
If return_usage=False: Parsed response if response_format is provided, otherwise text content.
|
||||
If return_usage=True: Tuple of (result, TokenUsage) with mock token counts.
|
||||
"""
|
||||
# Record the call for test verification
|
||||
call_record = {
|
||||
"provider": self.provider,
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"response_format": response_format.__name__
|
||||
if response_format and hasattr(response_format, "__name__")
|
||||
else str(response_format),
|
||||
"scope": scope,
|
||||
}
|
||||
self._mock_calls.append(call_record)
|
||||
logger.debug(f"Mock LLM call recorded: scope={scope}, model={self.model}")
|
||||
|
||||
# Return mock response
|
||||
if self._mock_response is not None:
|
||||
result = self._mock_response
|
||||
elif response_format is not None:
|
||||
# Try to create a minimal valid instance of the response format
|
||||
try:
|
||||
# For Pydantic models, try to create with minimal valid data
|
||||
result = {"mock": True}
|
||||
except Exception:
|
||||
result = {"mock": True}
|
||||
else:
|
||||
result = "mock response"
|
||||
|
||||
if return_usage:
|
||||
token_usage = TokenUsage(input_tokens=10, output_tokens=5, total_tokens=15)
|
||||
return result, token_usage
|
||||
return result
|
||||
|
||||
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:
|
||||
"""
|
||||
Make a mock LLM API call with tool/function calling support.
|
||||
|
||||
Records the call for test verification and returns the configured mock response.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts. Can include tool results with role='tool'.
|
||||
tools: List of tool definitions in OpenAI format.
|
||||
max_completion_tokens: Not used in mock.
|
||||
temperature: Not used in mock.
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Not used in mock.
|
||||
initial_backoff: Not used in mock.
|
||||
max_backoff: Not used in mock.
|
||||
tool_choice: Not used in mock.
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
"""
|
||||
# Record the call for test verification
|
||||
call_record = {
|
||||
"provider": self.provider,
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"tools": [t.get("function", {}).get("name") for t in tools],
|
||||
"scope": scope,
|
||||
}
|
||||
self._mock_calls.append(call_record)
|
||||
|
||||
if self._mock_response is not None:
|
||||
if isinstance(self._mock_response, LLMToolCallResult):
|
||||
return self._mock_response
|
||||
# Allow setting just tool calls as a list
|
||||
if isinstance(self._mock_response, list):
|
||||
return LLMToolCallResult(
|
||||
tool_calls=[
|
||||
LLMToolCall(id=f"mock_{i}", name=tc["name"], arguments=tc.get("arguments", {}))
|
||||
for i, tc in enumerate(self._mock_response)
|
||||
],
|
||||
finish_reason="tool_calls",
|
||||
)
|
||||
|
||||
return LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (no-op for mock provider)."""
|
||||
pass
|
||||
|
||||
def set_mock_response(self, response: Any) -> None:
|
||||
"""
|
||||
Set the response to return from mock calls.
|
||||
|
||||
Args:
|
||||
response: The response to return. Can be:
|
||||
- A dict/Pydantic model for regular calls
|
||||
- An LLMToolCallResult for tool calls
|
||||
- A list of tool call dicts for tool calls
|
||||
- Any other value to return as-is
|
||||
"""
|
||||
self._mock_response = response
|
||||
|
||||
def get_mock_calls(self) -> list[dict]:
|
||||
"""
|
||||
Get the list of recorded mock calls.
|
||||
|
||||
Returns:
|
||||
List of call records, each containing:
|
||||
- provider: Provider name
|
||||
- model: Model name
|
||||
- messages: Messages sent
|
||||
- response_format/tools: Format or tools used
|
||||
- scope: Call scope
|
||||
"""
|
||||
return self._mock_calls
|
||||
|
||||
def clear_mock_calls(self) -> None:
|
||||
"""Clear the recorded mock calls."""
|
||||
self._mock_calls = []
|
||||
@@ -0,0 +1,745 @@
|
||||
"""
|
||||
OpenAI-compatible LLM provider supporting OpenAI, Groq, Ollama, and LMStudio.
|
||||
|
||||
This provider handles all OpenAI API-compatible models including:
|
||||
- OpenAI: GPT-4, GPT-4o, GPT-5, o1, o3 (reasoning models)
|
||||
- Groq: Fast inference with seed control and service tiers
|
||||
- Ollama: Local models with native streaming API support
|
||||
- LMStudio: Local models with OpenAI-compatible API
|
||||
|
||||
Features:
|
||||
- Reasoning models with extended thinking (o1, o3, GPT-5 families)
|
||||
- Strict JSON schema enforcement (OpenAI)
|
||||
- Provider-specific parameters (Groq seed, service tier)
|
||||
- Native Ollama streaming for better structured output
|
||||
- Automatic token limit handling per model family
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from openai import APIConnectionError, APIStatusError, AsyncOpenAI, LengthFinishReasonError
|
||||
|
||||
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Seed applied to every Groq request for deterministic behavior
|
||||
DEFAULT_LLM_SEED = 4242
|
||||
|
||||
|
||||
class OpenAICompatibleLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider for OpenAI-compatible APIs.
|
||||
|
||||
Supports:
|
||||
- OpenAI: Standard models (GPT-4, GPT-4o) and reasoning models (o1, o3, GPT-5)
|
||||
- Groq: Fast inference with seed control and service tiers
|
||||
- Ollama: Local models with native streaming API for better structured output
|
||||
- LMStudio: Local models with OpenAI-compatible API
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float | None = None,
|
||||
groq_service_tier: str | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
Initialize OpenAI-compatible LLM provider.
|
||||
|
||||
Args:
|
||||
provider: Provider name ("openai", "groq", "ollama", "lmstudio").
|
||||
api_key: API key (optional for ollama/lmstudio).
|
||||
base_url: Base URL for the API (uses defaults for groq/ollama/lmstudio if empty).
|
||||
model: Model name.
|
||||
reasoning_effort: Reasoning effort level for supported models ("low", "medium", "high").
|
||||
timeout: Request timeout in seconds (uses env var or 300s default).
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto").
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
"""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
|
||||
# Validate provider
|
||||
valid_providers = ["openai", "groq", "ollama", "lmstudio"]
|
||||
if self.provider not in valid_providers:
|
||||
raise ValueError(f"OpenAICompatibleLLM only supports: {', '.join(valid_providers)}. Got: {self.provider}")
|
||||
|
||||
# Set default base URLs
|
||||
if not self.base_url:
|
||||
if self.provider == "groq":
|
||||
self.base_url = "https://api.groq.com/openai/v1"
|
||||
elif self.provider == "ollama":
|
||||
self.base_url = "http://localhost:11434/v1"
|
||||
elif self.provider == "lmstudio":
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
|
||||
# For ollama/lmstudio, use dummy key if not provided
|
||||
if self.provider in ("ollama", "lmstudio") and not self.api_key:
|
||||
self.api_key = "local"
|
||||
|
||||
# Validate API key for cloud providers
|
||||
if self.provider in ("openai", "groq") and not self.api_key:
|
||||
raise ValueError(f"API key is required for {self.provider}")
|
||||
|
||||
# Groq service tier configuration
|
||||
self.groq_service_tier = groq_service_tier or os.getenv("HINDSIGHT_API_LLM_GROQ_SERVICE_TIER", "auto")
|
||||
|
||||
# Get timeout config
|
||||
self.timeout = timeout or float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT)))
|
||||
|
||||
# Create OpenAI client
|
||||
client_kwargs: dict[str, Any] = {"api_key": self.api_key, "max_retries": 0}
|
||||
if self.base_url:
|
||||
client_kwargs["base_url"] = self.base_url
|
||||
if self.timeout:
|
||||
client_kwargs["timeout"] = self.timeout
|
||||
|
||||
self._client = AsyncOpenAI(**client_kwargs)
|
||||
logger.info(
|
||||
f"OpenAI-compatible client initialized: provider={self.provider}, model={self.model}, "
|
||||
f"base_url={self.base_url or 'default'}"
|
||||
)
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""
|
||||
Verify that the provider is configured correctly by making a simple test call.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the connection test fails.
|
||||
"""
|
||||
try:
|
||||
logger.info(f"Verifying connection: {self.provider}/{self.model}")
|
||||
await self.call(
|
||||
messages=[{"role": "user", "content": "Say 'ok'"}],
|
||||
max_completion_tokens=100,
|
||||
max_retries=2,
|
||||
initial_backoff=0.5,
|
||||
max_backoff=2.0,
|
||||
)
|
||||
logger.info(f"Connection verified: {self.provider}/{self.model}")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Connection verification failed for {self.provider}/{self.model}: {e}") from e
|
||||
|
||||
def _supports_reasoning_model(self) -> bool:
|
||||
"""Check if the current model is a reasoning model (o1, o3, GPT-5, DeepSeek)."""
|
||||
model_lower = self.model.lower()
|
||||
return any(x in model_lower for x in ["gpt-5", "o1", "o3", "deepseek"])
|
||||
|
||||
def _get_max_reasoning_tokens(self) -> int | None:
|
||||
"""Get max reasoning tokens for reasoning models."""
|
||||
model_lower = self.model.lower()
|
||||
|
||||
# GPT-4 and GPT-4.1 models have different caps
|
||||
if any(x in model_lower for x in ["gpt-4.1", "gpt-4-"]):
|
||||
return 32000
|
||||
elif "gpt-4o" in model_lower:
|
||||
return 16384
|
||||
|
||||
return 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:
|
||||
"""
|
||||
Make an LLM API call with retry logic.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts with 'role' and 'content'.
|
||||
response_format: Optional Pydantic model for structured output.
|
||||
max_completion_tokens: Maximum tokens in response.
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Use strict JSON schema enforcement (OpenAI only).
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
|
||||
Returns:
|
||||
If return_usage=False: Parsed response if response_format is provided, otherwise text content.
|
||||
If return_usage=True: Tuple of (result, TokenUsage) with token counts.
|
||||
|
||||
Raises:
|
||||
OutputTooLongError: If output exceeds token limits.
|
||||
Exception: Re-raises API errors after retries exhausted.
|
||||
"""
|
||||
# Handle Ollama with native API for structured output (better schema enforcement)
|
||||
if self.provider == "ollama" and response_format is not None:
|
||||
return await self._call_ollama_native(
|
||||
messages=messages,
|
||||
response_format=response_format,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
temperature=temperature,
|
||||
max_retries=max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
skip_validation=skip_validation,
|
||||
scope=scope,
|
||||
return_usage=return_usage,
|
||||
)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Build call parameters
|
||||
call_params: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
}
|
||||
|
||||
# Check if model supports reasoning parameter
|
||||
is_reasoning_model = self._supports_reasoning_model()
|
||||
|
||||
# Apply model-specific token limits
|
||||
if max_completion_tokens is not None:
|
||||
max_tokens_cap = self._get_max_reasoning_tokens()
|
||||
if max_tokens_cap and max_completion_tokens > max_tokens_cap:
|
||||
max_completion_tokens = max_tokens_cap
|
||||
# For reasoning models, enforce minimum to ensure space for reasoning + output
|
||||
if is_reasoning_model and max_completion_tokens < 16000:
|
||||
max_completion_tokens = 16000
|
||||
call_params["max_completion_tokens"] = max_completion_tokens
|
||||
|
||||
# Temperature - reasoning models don't support custom temperature
|
||||
if temperature is not None and not is_reasoning_model:
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
# Set reasoning_effort for reasoning models
|
||||
if is_reasoning_model:
|
||||
call_params["reasoning_effort"] = self.reasoning_effort
|
||||
|
||||
# Provider-specific parameters
|
||||
if self.provider == "groq":
|
||||
call_params["seed"] = DEFAULT_LLM_SEED
|
||||
extra_body: dict[str, Any] = {}
|
||||
# Add service_tier if configured
|
||||
if self.groq_service_tier:
|
||||
extra_body["service_tier"] = self.groq_service_tier
|
||||
# Add reasoning parameters for reasoning models
|
||||
if is_reasoning_model:
|
||||
extra_body["include_reasoning"] = False
|
||||
if extra_body:
|
||||
call_params["extra_body"] = extra_body
|
||||
|
||||
# Prepare response format ONCE before retry loop
|
||||
if response_format is not None:
|
||||
schema = None
|
||||
if hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
|
||||
if strict_schema and schema is not None:
|
||||
# Use OpenAI's strict JSON schema enforcement
|
||||
call_params["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "response",
|
||||
"strict": True,
|
||||
"schema": schema,
|
||||
},
|
||||
}
|
||||
else:
|
||||
# Soft enforcement: add schema to prompt and use json_object mode
|
||||
if schema is not None:
|
||||
schema_msg = (
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
)
|
||||
|
||||
if call_params["messages"] and call_params["messages"][0].get("role") == "system":
|
||||
first_msg = call_params["messages"][0]
|
||||
if isinstance(first_msg, dict) and isinstance(first_msg.get("content"), str):
|
||||
first_msg["content"] += schema_msg
|
||||
elif call_params["messages"]:
|
||||
first_msg = call_params["messages"][0]
|
||||
if isinstance(first_msg, dict) and isinstance(first_msg.get("content"), str):
|
||||
first_msg["content"] = schema_msg + "\n\n" + first_msg["content"]
|
||||
if self.provider not in ("lmstudio", "ollama"):
|
||||
# LM Studio and Ollama don't support json_object response format reliably
|
||||
call_params["response_format"] = {"type": "json_object"}
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
if response_format is not None:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
|
||||
content = response.choices[0].message.content
|
||||
|
||||
# Strip reasoning model thinking tags
|
||||
# Supports: <think>, <thinking>, <reasoning>, |startthink|/|endthink|
|
||||
if content:
|
||||
original_len = len(content)
|
||||
content = re.sub(r"<think>.*?</think>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"<thinking>.*?</thinking>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"<reasoning>.*?</reasoning>", "", content, flags=re.DOTALL)
|
||||
content = re.sub(r"\|startthink\|.*?\|endthink\|", "", content, flags=re.DOTALL)
|
||||
content = content.strip()
|
||||
if len(content) < original_len:
|
||||
logger.debug(f"Stripped {original_len - len(content)} chars of reasoning tokens")
|
||||
|
||||
# For local models, they may wrap JSON in markdown code blocks
|
||||
if self.provider in ("lmstudio", "ollama"):
|
||||
clean_content = content
|
||||
if "```json" in content:
|
||||
clean_content = content.split("```json")[1].split("```")[0].strip()
|
||||
elif "```" in content:
|
||||
clean_content = content.split("```")[1].split("```")[0].strip()
|
||||
try:
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError:
|
||||
# Fallback to parsing raw content
|
||||
json_data = json.loads(content)
|
||||
else:
|
||||
# Log raw LLM response for debugging JSON parse issues
|
||||
try:
|
||||
json_data = json.loads(content)
|
||||
except json.JSONDecodeError as json_err:
|
||||
# Truncate content for logging
|
||||
content_preview = content[:500] if content else "<empty>"
|
||||
if content and len(content) > 700:
|
||||
content_preview = f"{content[:500]}...TRUNCATED...{content[-200:]}"
|
||||
logger.warning(
|
||||
f"JSON parse error from LLM response (attempt {attempt + 1}/{max_retries + 1}): {json_err}\n"
|
||||
f" Model: {self.provider}/{self.model}\n"
|
||||
f" Content length: {len(content) if content else 0} chars\n"
|
||||
f" Content preview: {content_preview!r}\n"
|
||||
f" Finish reason: {response.choices[0].finish_reason if response.choices else 'unknown'}"
|
||||
)
|
||||
# Retry on JSON parse errors
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
last_exception = json_err
|
||||
continue
|
||||
else:
|
||||
logger.error(f"JSON parse error after {max_retries + 1} attempts, giving up")
|
||||
raise
|
||||
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
else:
|
||||
result = response_format.model_validate(json_data)
|
||||
else:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
result = response.choices[0].message.content
|
||||
|
||||
# Record token usage metrics
|
||||
duration = time.time() - start_time
|
||||
usage = response.usage
|
||||
input_tokens = usage.prompt_tokens or 0 if usage else 0
|
||||
output_tokens = usage.completion_tokens or 0 if usage else 0
|
||||
total_tokens = usage.total_tokens or 0 if usage else 0
|
||||
|
||||
# Record LLM metrics
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0 and usage:
|
||||
ratio = max(1, output_tokens) / max(1, input_tokens)
|
||||
cached_tokens = 0
|
||||
if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
|
||||
cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
|
||||
cache_info = f", cached_tokens={cached_tokens}" if cached_tokens > 0 else ""
|
||||
logger.info(
|
||||
f"slow llm call: scope={scope}, model={self.provider}/{self.model}, "
|
||||
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
|
||||
f"total_tokens={total_tokens}{cache_info}, time={duration:.3f}s, ratio out/in={ratio:.2f}"
|
||||
)
|
||||
|
||||
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 LengthFinishReasonError as e:
|
||||
logger.warning(f"LLM output exceeded token limits: {str(e)}")
|
||||
raise OutputTooLongError(
|
||||
"LLM output exceeded token limits. Input may need to be split into smaller chunks."
|
||||
) from e
|
||||
|
||||
except APIConnectionError as e:
|
||||
last_exception = e
|
||||
status_code = getattr(e, "status_code", None) or getattr(
|
||||
getattr(e, "response", None), "status_code", None
|
||||
)
|
||||
logger.warning(f"APIConnectionError (HTTP {status_code}), attempt {attempt + 1}: {str(e)[:200]}")
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"Connection error after {max_retries + 1} attempts: {str(e)}")
|
||||
raise
|
||||
|
||||
except APIStatusError as e:
|
||||
# Fast fail only on 401 (unauthorized) and 403 (forbidden)
|
||||
if e.status_code in (401, 403):
|
||||
logger.error(f"Auth error (HTTP {e.status_code}), not retrying: {str(e)}")
|
||||
raise
|
||||
|
||||
# Handle tool_use_failed error - model outputted in tool call format
|
||||
if e.status_code == 400 and response_format is not None:
|
||||
try:
|
||||
error_body = e.body if hasattr(e, "body") else {}
|
||||
if isinstance(error_body, dict):
|
||||
error_info: dict[str, Any] = error_body.get("error") or {}
|
||||
if error_info.get("code") == "tool_use_failed":
|
||||
failed_gen = error_info.get("failed_generation", "")
|
||||
if failed_gen:
|
||||
# Parse tool call format and convert to expected format
|
||||
tool_call = json.loads(failed_gen)
|
||||
tool_name = tool_call.get("name", "")
|
||||
tool_args = tool_call.get("arguments", {})
|
||||
converted = {"actions": [{"tool": tool_name, **tool_args}]}
|
||||
if skip_validation:
|
||||
result = converted
|
||||
else:
|
||||
result = response_format.model_validate(converted)
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
success=True,
|
||||
)
|
||||
if return_usage:
|
||||
return result, TokenUsage(input_tokens=0, output_tokens=0, total_tokens=0)
|
||||
return result
|
||||
except (json.JSONDecodeError, KeyError, TypeError):
|
||||
pass # Failed to parse tool_use_failed, continue with normal retry
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
|
||||
sleep_time = backoff + jitter
|
||||
await asyncio.sleep(sleep_time)
|
||||
else:
|
||||
logger.error(f"API error after {max_retries + 1} attempts: {str(e)}")
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("LLM call failed after all retries with no exception captured")
|
||||
|
||||
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:
|
||||
"""
|
||||
Make an LLM API call with tool/function calling support.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts. Can include tool results with role='tool'.
|
||||
tools: List of tool definitions in OpenAI format.
|
||||
max_completion_tokens: Maximum tokens in response.
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or specific function.
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Build call parameters
|
||||
call_params: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"tools": tools,
|
||||
"tool_choice": tool_choice,
|
||||
}
|
||||
|
||||
if max_completion_tokens is not None:
|
||||
call_params["max_completion_tokens"] = max_completion_tokens
|
||||
if temperature is not None:
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
# Provider-specific parameters
|
||||
if self.provider == "groq":
|
||||
call_params["seed"] = DEFAULT_LLM_SEED
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
|
||||
message = response.choices[0].message
|
||||
finish_reason = response.choices[0].finish_reason
|
||||
|
||||
# Extract tool calls if present
|
||||
tool_calls: list[LLMToolCall] = []
|
||||
if message.tool_calls:
|
||||
for tc in message.tool_calls:
|
||||
try:
|
||||
args = json.loads(tc.function.arguments) if tc.function.arguments else {}
|
||||
except json.JSONDecodeError:
|
||||
args = {"_raw": tc.function.arguments}
|
||||
tool_calls.append(LLMToolCall(id=tc.id, name=tc.function.name, arguments=args))
|
||||
|
||||
content = message.content
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
usage = response.usage
|
||||
input_tokens = usage.prompt_tokens or 0 if usage else 0
|
||||
output_tokens = usage.completion_tokens or 0 if usage else 0
|
||||
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
|
||||
return LLMToolCallResult(
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
finish_reason=finish_reason,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
)
|
||||
|
||||
except APIConnectionError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
raise
|
||||
|
||||
except APIStatusError as e:
|
||||
if e.status_code in (401, 403):
|
||||
raise
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("Tool call failed after all retries")
|
||||
|
||||
async def _call_ollama_native(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
response_format: Any,
|
||||
max_completion_tokens: int | None,
|
||||
temperature: float | None,
|
||||
max_retries: int,
|
||||
initial_backoff: float,
|
||||
max_backoff: float,
|
||||
skip_validation: bool,
|
||||
scope: str = "memory",
|
||||
return_usage: bool = False,
|
||||
) -> Any:
|
||||
"""
|
||||
Call Ollama using native API with JSON schema enforcement.
|
||||
|
||||
Ollama's native API supports passing a full JSON schema in the 'format' parameter,
|
||||
which provides better structured output control than the OpenAI-compatible API.
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Get the JSON schema from the Pydantic model
|
||||
schema = response_format.model_json_schema() if hasattr(response_format, "model_json_schema") else None
|
||||
|
||||
# Build the base URL for Ollama's native API
|
||||
# Default OpenAI-compatible URL is http://localhost:11434/v1
|
||||
# Native API is at http://localhost:11434/api/chat
|
||||
base_url = self.base_url or "http://localhost:11434/v1"
|
||||
if base_url.endswith("/v1"):
|
||||
native_url = base_url[:-3] + "/api/chat"
|
||||
else:
|
||||
native_url = base_url.rstrip("/") + "/api/chat"
|
||||
|
||||
# Build request payload
|
||||
payload: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"stream": False,
|
||||
}
|
||||
|
||||
# Add schema as format parameter for structured output
|
||||
if schema:
|
||||
payload["format"] = schema
|
||||
|
||||
# Add optional parameters with optimized defaults for Ollama
|
||||
options: dict[str, Any] = {
|
||||
"num_ctx": 16384, # 16k context window for larger prompts
|
||||
"num_batch": 512, # Optimal batch size for prompt processing
|
||||
}
|
||||
if max_completion_tokens:
|
||||
options["num_predict"] = max_completion_tokens
|
||||
if temperature is not None:
|
||||
options["temperature"] = temperature
|
||||
payload["options"] = options
|
||||
|
||||
last_exception = None
|
||||
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await client.post(native_url, json=payload)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
content = result.get("message", {}).get("content", "")
|
||||
|
||||
# Parse JSON response
|
||||
try:
|
||||
json_data = json.loads(content)
|
||||
except json.JSONDecodeError as json_err:
|
||||
content_preview = content[:500] if content else "<empty>"
|
||||
if content and len(content) > 700:
|
||||
content_preview = f"{content[:500]}...TRUNCATED...{content[-200:]}"
|
||||
logger.warning(
|
||||
f"Ollama JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {json_err}\n"
|
||||
f" Model: ollama/{self.model}\n"
|
||||
f" Content length: {len(content) if content else 0} chars\n"
|
||||
f" Content preview: {content_preview!r}"
|
||||
)
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
last_exception = json_err
|
||||
continue
|
||||
else:
|
||||
raise
|
||||
|
||||
# Extract token usage from Ollama response
|
||||
duration = time.time() - start_time
|
||||
input_tokens = result.get("prompt_eval_count", 0) or 0
|
||||
output_tokens = result.get("eval_count", 0) or 0
|
||||
total_tokens = input_tokens + output_tokens
|
||||
|
||||
# Record LLM metrics
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Validate against Pydantic model or return raw JSON
|
||||
if skip_validation:
|
||||
validated_result = json_data
|
||||
else:
|
||||
validated_result = response_format.model_validate(json_data)
|
||||
|
||||
if return_usage:
|
||||
token_usage = TokenUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=total_tokens,
|
||||
)
|
||||
return validated_result, token_usage
|
||||
return validated_result
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"Ollama HTTP error (attempt {attempt + 1}/{max_retries + 1}): {e.response.status_code}"
|
||||
)
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"Ollama HTTP error after {max_retries + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
except httpx.RequestError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(f"Ollama connection error (attempt {attempt + 1}/{max_retries + 1}): {e}")
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"Ollama connection error after {max_retries + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error during Ollama call: {type(e).__name__}: {e}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("Ollama call failed after all retries")
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (close OpenAI client connections)."""
|
||||
if hasattr(self, "_client") and self._client:
|
||||
await self._client.close()
|
||||
@@ -871,21 +871,21 @@ async def _execute_tool(
|
||||
query = args.get("query")
|
||||
if not query:
|
||||
return {"error": "search_mental_models requires a query parameter"}
|
||||
max_results = args.get("max_results") or 5
|
||||
max_results = int(args.get("max_results") or 5)
|
||||
return await search_mental_models_fn(query, max_results)
|
||||
|
||||
elif tool_name == "search_observations":
|
||||
query = args.get("query")
|
||||
if not query:
|
||||
return {"error": "search_observations requires a query parameter"}
|
||||
max_tokens = max(args.get("max_tokens") or 5000, 1000) # Default 5000, min 1000
|
||||
max_tokens = max(int(args.get("max_tokens") or 5000), 1000) # Default 5000, min 1000
|
||||
return await search_observations_fn(query, max_tokens)
|
||||
|
||||
elif tool_name == "recall":
|
||||
query = args.get("query")
|
||||
if not query:
|
||||
return {"error": "recall requires a query parameter"}
|
||||
max_tokens = max(args.get("max_tokens") or 2048, 1000) # Default 2048, min 1000
|
||||
max_tokens = max(int(args.get("max_tokens") or 2048), 1000) # Default 2048, min 1000
|
||||
return await recall_fn(query, max_tokens)
|
||||
|
||||
elif tool_name == "expand":
|
||||
@@ -904,18 +904,18 @@ def _summarize_input(tool_name: str, args: dict[str, Any]) -> str:
|
||||
if tool_name == "search_mental_models":
|
||||
query = args.get("query", "")
|
||||
query_preview = f"'{query[:30]}...'" if len(query) > 30 else f"'{query}'"
|
||||
max_results = args.get("max_results") or 5
|
||||
max_results = int(args.get("max_results") or 5)
|
||||
return f"(query={query_preview}, max_results={max_results})"
|
||||
elif tool_name == "search_observations":
|
||||
query = args.get("query", "")
|
||||
query_preview = f"'{query[:30]}...'" if len(query) > 30 else f"'{query}'"
|
||||
max_tokens = max(args.get("max_tokens") or 5000, 1000)
|
||||
max_tokens = max(int(args.get("max_tokens") or 5000), 1000)
|
||||
return f"(query={query_preview}, max_tokens={max_tokens})"
|
||||
elif tool_name == "recall":
|
||||
query = args.get("query", "")
|
||||
query_preview = f"'{query[:30]}...'" if len(query) > 30 else f"'{query}'"
|
||||
# Show actual value used (default 2048, min 1000)
|
||||
max_tokens = max(args.get("max_tokens") or 2048, 1000)
|
||||
max_tokens = max(int(args.get("max_tokens") or 2048), 1000)
|
||||
return f"(query={query_preview}, max_tokens={max_tokens})"
|
||||
elif tool_name == "expand":
|
||||
memory_ids = args.get("memory_ids", [])
|
||||
|
||||
@@ -463,9 +463,12 @@ def build_final_prompt(
|
||||
parts.append(
|
||||
"\n## Instructions\n"
|
||||
"Provide a thoughtful answer by synthesizing and reasoning from the retrieved data above. "
|
||||
"You can make reasonable inferences from the memories, but don't completely fabricate information."
|
||||
"You can make reasonable inferences from the memories, but don't completely fabricate information. "
|
||||
"If the exact answer isn't stated, use what IS stated to give the best possible answer. "
|
||||
"Only say 'I don't have information' if the retrieved data is truly unrelated to the question."
|
||||
"Only say 'I don't have information' if the retrieved data is truly unrelated to the question.\n\n"
|
||||
"IMPORTANT: Output ONLY the final answer. Do NOT include meta-commentary like "
|
||||
'"I\'ll search..." or "Let me analyze...". Do NOT explain your reasoning process. '
|
||||
"Just provide the direct synthesized answer."
|
||||
)
|
||||
|
||||
return "\n".join(parts)
|
||||
@@ -480,4 +483,10 @@ Your approach:
|
||||
- Be helpful - if you have related information, use it to give the best possible answer
|
||||
|
||||
Only say "I don't have information" if the retrieved data is truly unrelated to the question.
|
||||
Do NOT fabricate information that has no basis in the retrieved data."""
|
||||
Do NOT fabricate information that has no basis in the retrieved data.
|
||||
|
||||
CRITICAL: Output ONLY the final synthesized answer. Do NOT include:
|
||||
- Meta-commentary about what you're doing ("I'll search...", "Let me analyze...")
|
||||
- Explanations of your reasoning process
|
||||
- Descriptions of your approach
|
||||
Just provide the direct answer."""
|
||||
|
||||
@@ -54,19 +54,18 @@ async def tool_search_mental_models(
|
||||
Dict with matching mental models including content and freshness info
|
||||
"""
|
||||
from ..memory_engine import fq_table
|
||||
from ..search.tags import build_tags_where_clause
|
||||
|
||||
# Build filters dynamically
|
||||
filters = ""
|
||||
params: list[Any] = [bank_id, str(query_embedding), max_results]
|
||||
next_param = 4
|
||||
|
||||
# Use the centralized tag filtering logic
|
||||
if tags:
|
||||
if tags_match == "all":
|
||||
filters += f" AND tags @> ${next_param}::varchar[]"
|
||||
else:
|
||||
filters += f" AND (tags && ${next_param}::varchar[] OR tags IS NULL OR tags = '{{}}')"
|
||||
params.append(tags)
|
||||
next_param += 1
|
||||
tag_clause, tag_params, next_param = build_tags_where_clause(tags, param_offset=next_param, match=tags_match)
|
||||
filters += f" {tag_clause}"
|
||||
params.extend(tag_params)
|
||||
|
||||
if exclude_ids:
|
||||
filters += f" AND id != ALL(${next_param}::text[])"
|
||||
|
||||
@@ -263,7 +263,6 @@ class ReflectResult(BaseModel):
|
||||
}
|
||||
],
|
||||
},
|
||||
"new_opinions": ["Machine learning has great potential in healthcare"],
|
||||
"structured_output": {"summary": "ML in healthcare", "confidence": 0.9},
|
||||
"usage": {"input_tokens": 1500, "output_tokens": 500, "total_tokens": 2000},
|
||||
}
|
||||
@@ -272,9 +271,8 @@ class ReflectResult(BaseModel):
|
||||
|
||||
text: str = Field(description="The formulated answer text")
|
||||
based_on: dict[str, Any] = Field(
|
||||
description="Facts used to formulate the answer, organized by type (world, experience, opinion, mental_models, directives)"
|
||||
description="Facts used to formulate the answer, organized by type (world, experience, mental_models, directives)"
|
||||
)
|
||||
new_opinions: list[str] = Field(default_factory=list, description="List of newly formed opinions during reflection")
|
||||
structured_output: dict[str, Any] | None = Field(
|
||||
default=None,
|
||||
description="Structured output parsed according to the provided response schema. Only present when response_schema was provided.",
|
||||
@@ -297,24 +295,6 @@ class ReflectResult(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class Opinion(BaseModel):
|
||||
"""
|
||||
An opinion with confidence score.
|
||||
|
||||
Opinions represent the bank's formed perspectives on topics,
|
||||
with a confidence level indicating strength of belief.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"example": {"text": "Machine learning has great potential in healthcare", "confidence": 0.85}
|
||||
}
|
||||
)
|
||||
|
||||
text: str = Field(description="The opinion text")
|
||||
confidence: float = Field(description="Confidence score between 0.0 and 1.0")
|
||||
|
||||
|
||||
class EntityObservation(BaseModel):
|
||||
"""
|
||||
An observation about an entity.
|
||||
|
||||
@@ -57,21 +57,25 @@ def _infer_temporal_date(fact_text: str, event_date: datetime) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def _sanitize_text(text: str) -> str:
|
||||
def _sanitize_text(text: str | None) -> str | None:
|
||||
"""
|
||||
Sanitize text by removing invalid Unicode surrogate characters.
|
||||
Sanitize text by removing characters that break downstream systems.
|
||||
|
||||
Surrogate characters (U+D800 to U+DFFF) are used in UTF-16 encoding
|
||||
but cannot be encoded in UTF-8. They can appear in Python strings
|
||||
from improperly decoded data (e.g., from JavaScript or broken files).
|
||||
Removes:
|
||||
- Null bytes (\\x00): Invalid in PostgreSQL UTF-8 encoding
|
||||
- Unicode surrogates (U+D800-U+DFFF): Invalid in UTF-8, break LLM APIs
|
||||
|
||||
This function removes unpaired surrogates to prevent UnicodeEncodeError
|
||||
when the text is sent to the LLM API.
|
||||
Surrogate characters are used in UTF-16 encoding but cannot be encoded
|
||||
in UTF-8. They can appear in Python strings from improperly decoded data
|
||||
(e.g., from JavaScript or broken files). Null bytes commonly appear in
|
||||
OCR output, PDF extraction, or copy-paste from binary sources.
|
||||
"""
|
||||
if text is None:
|
||||
return None
|
||||
if not text:
|
||||
return text
|
||||
# Remove surrogate characters (U+D800 to U+DFFF) using regex
|
||||
# These are invalid in UTF-8 and cause encoding errors
|
||||
# Remove null bytes and surrogate characters
|
||||
text = text.replace("\x00", "")
|
||||
return re.sub(r"[\ud800-\udfff]", "", text)
|
||||
|
||||
|
||||
@@ -693,7 +697,6 @@ async def _extract_facts_from_chunk(
|
||||
context: str,
|
||||
llm_config: "LLMConfig",
|
||||
agent_name: str = None,
|
||||
extract_opinions: bool = False,
|
||||
) -> tuple[list[dict[str, str]], TokenUsage]:
|
||||
"""
|
||||
Extract facts from a single chunk (internal helper for parallel processing).
|
||||
@@ -707,17 +710,9 @@ async def _extract_facts_from_chunk(
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
memory_bank_context = f"\n- Your name: {agent_name}" if agent_name and extract_opinions else ""
|
||||
|
||||
# Determine which fact types to extract based on the flag
|
||||
# Determine which fact types to extract
|
||||
# Note: We use "assistant" in the prompt but convert to "bank" for storage
|
||||
if extract_opinions:
|
||||
# Opinion extraction uses a separate prompt (not this one)
|
||||
fact_types_instruction = "Extract ONLY 'opinion' type facts (formed opinions, beliefs, and perspectives). DO NOT extract 'world' or 'assistant' facts."
|
||||
else:
|
||||
fact_types_instruction = (
|
||||
"Extract ONLY 'world' and 'assistant' type facts. DO NOT extract opinions - those are extracted separately."
|
||||
)
|
||||
fact_types_instruction = "Extract ONLY 'world' and 'assistant' type facts."
|
||||
|
||||
# Check config for extraction mode and causal link extraction
|
||||
config = get_config()
|
||||
@@ -768,9 +763,12 @@ async def _extract_facts_from_chunk(
|
||||
|
||||
# Build user message with metadata and chunk content in a clear format
|
||||
# Format event_date with day of week for better temporal reasoning
|
||||
# Handle both datetime objects and ISO string formats (from deserialized async tasks)
|
||||
from .orchestrator import parse_datetime_flexible
|
||||
|
||||
event_date = parse_datetime_flexible(event_date)
|
||||
event_date_formatted = event_date.strftime("%A, %B %d, %Y") # e.g., "Monday, June 10, 2024"
|
||||
user_message = f"""Extract facts from the following text chunk.
|
||||
{memory_bank_context}
|
||||
|
||||
Chunk: {chunk_index + 1}/{total_chunks}
|
||||
Event Date: {event_date_formatted} ({event_date.isoformat()})
|
||||
@@ -782,12 +780,28 @@ Text:
|
||||
usage = TokenUsage() # Track cumulative usage across retries
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
# Use retain-specific overrides if set, otherwise fall back to global LLM config
|
||||
max_retries = (
|
||||
config.retain_llm_max_retries if config.retain_llm_max_retries is not None else config.llm_max_retries
|
||||
)
|
||||
initial_backoff = (
|
||||
config.retain_llm_initial_backoff
|
||||
if config.retain_llm_initial_backoff is not None
|
||||
else config.llm_initial_backoff
|
||||
)
|
||||
max_backoff = (
|
||||
config.retain_llm_max_backoff if config.retain_llm_max_backoff is not None else config.llm_max_backoff
|
||||
)
|
||||
|
||||
extraction_response_json, call_usage = await llm_config.call(
|
||||
messages=[{"role": "system", "content": prompt}, {"role": "user", "content": user_message}],
|
||||
response_format=response_schema,
|
||||
scope="memory_extract_facts",
|
||||
temperature=0.1,
|
||||
max_completion_tokens=config.retain_max_completion_tokens,
|
||||
max_retries=max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
skip_validation=True, # Get raw JSON, we'll validate leniently
|
||||
return_usage=True,
|
||||
)
|
||||
@@ -1013,7 +1027,6 @@ async def _extract_facts_with_auto_split(
|
||||
context: str,
|
||||
llm_config: LLMConfig,
|
||||
agent_name: str = None,
|
||||
extract_opinions: bool = False,
|
||||
) -> tuple[list[dict[str, str]], TokenUsage]:
|
||||
"""
|
||||
Extract facts from a chunk with automatic splitting if output exceeds token limits.
|
||||
@@ -1029,7 +1042,6 @@ async def _extract_facts_with_auto_split(
|
||||
context: Context about the conversation/document
|
||||
llm_config: LLM configuration to use
|
||||
agent_name: Optional agent name (memory owner)
|
||||
extract_opinions: If True, extract ONLY opinions. If False, extract world and agent facts (no opinions)
|
||||
|
||||
Returns:
|
||||
Tuple of (facts list, token usage) extracted from the chunk (possibly from sub-chunks)
|
||||
@@ -1048,7 +1060,6 @@ async def _extract_facts_with_auto_split(
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name=agent_name,
|
||||
extract_opinions=extract_opinions,
|
||||
)
|
||||
except OutputTooLongError:
|
||||
# Output exceeded token limits - split the chunk in half and retry
|
||||
@@ -1093,7 +1104,6 @@ async def _extract_facts_with_auto_split(
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name=agent_name,
|
||||
extract_opinions=extract_opinions,
|
||||
),
|
||||
_extract_facts_with_auto_split(
|
||||
chunk=second_half,
|
||||
@@ -1103,7 +1113,6 @@ async def _extract_facts_with_auto_split(
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name=agent_name,
|
||||
extract_opinions=extract_opinions,
|
||||
),
|
||||
]
|
||||
|
||||
@@ -1127,7 +1136,6 @@ async def extract_facts_from_text(
|
||||
llm_config: LLMConfig,
|
||||
agent_name: str,
|
||||
context: str = "",
|
||||
extract_opinions: bool = False,
|
||||
) -> tuple[list[Fact], list[tuple[str, int]], TokenUsage]:
|
||||
"""
|
||||
Extract semantic facts from conversational or narrative text using LLM.
|
||||
@@ -1144,7 +1152,6 @@ async def extract_facts_from_text(
|
||||
context: Context about the conversation/document
|
||||
llm_config: LLM configuration to use
|
||||
agent_name: Agent name (memory owner)
|
||||
extract_opinions: If True, extract ONLY opinions. If False, extract world and bank facts (no opinions)
|
||||
|
||||
Returns:
|
||||
Tuple of (facts, chunks, usage) where:
|
||||
@@ -1172,7 +1179,6 @@ async def extract_facts_from_text(
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name=agent_name,
|
||||
extract_opinions=extract_opinions,
|
||||
)
|
||||
for i, chunk in enumerate(chunks)
|
||||
]
|
||||
@@ -1204,7 +1210,7 @@ SECONDS_PER_FACT = 10
|
||||
|
||||
|
||||
async def extract_facts_from_contents(
|
||||
contents: list[RetainContent], llm_config, agent_name: str, extract_opinions: bool = False
|
||||
contents: list[RetainContent], llm_config, agent_name: str
|
||||
) -> tuple[list[ExtractedFactType], list[ChunkMetadata], TokenUsage]:
|
||||
"""
|
||||
Extract facts from multiple content items in parallel.
|
||||
@@ -1219,7 +1225,6 @@ async def extract_facts_from_contents(
|
||||
contents: List of RetainContent objects to process
|
||||
llm_config: LLM configuration for fact extraction
|
||||
agent_name: Name of the agent (for agent-related fact detection)
|
||||
extract_opinions: If True, extract only opinions; otherwise world/bank facts
|
||||
|
||||
Returns:
|
||||
Tuple of (extracted_facts, chunks_metadata, usage)
|
||||
@@ -1238,7 +1243,6 @@ async def extract_facts_from_contents(
|
||||
context=item.context,
|
||||
llm_config=llm_config,
|
||||
agent_name=agent_name,
|
||||
extract_opinions=extract_opinions,
|
||||
)
|
||||
fact_extraction_tasks.append(task)
|
||||
|
||||
@@ -1350,6 +1354,8 @@ def _add_temporal_offsets(facts: list[ExtractedFactType], contents: list[RetainC
|
||||
|
||||
Modifies facts in place.
|
||||
"""
|
||||
from .orchestrator import parse_datetime_flexible
|
||||
|
||||
# Group facts by content_index
|
||||
current_content_idx = 0
|
||||
content_fact_start = 0
|
||||
@@ -1364,10 +1370,10 @@ def _add_temporal_offsets(facts: list[ExtractedFactType], contents: list[RetainC
|
||||
fact_position = i - content_fact_start
|
||||
offset = timedelta(seconds=fact_position * SECONDS_PER_FACT)
|
||||
|
||||
# Apply offset to all temporal fields
|
||||
# Apply offset to all temporal fields (handle both datetime objects and ISO strings)
|
||||
if fact.occurred_start:
|
||||
fact.occurred_start = fact.occurred_start + offset
|
||||
fact.occurred_start = parse_datetime_flexible(fact.occurred_start) + offset
|
||||
if fact.occurred_end:
|
||||
fact.occurred_end = fact.occurred_end + offset
|
||||
fact.occurred_end = parse_datetime_flexible(fact.occurred_end) + offset
|
||||
if fact.mentioned_at:
|
||||
fact.mentioned_at = fact.mentioned_at + offset
|
||||
fact.mentioned_at = parse_datetime_flexible(fact.mentioned_at) + offset
|
||||
|
||||
@@ -8,6 +8,7 @@ import json
|
||||
import logging
|
||||
|
||||
from ..memory_engine import fq_table
|
||||
from .fact_extraction import _sanitize_text
|
||||
from .types import ProcessedFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -47,7 +48,7 @@ async def insert_facts_batch(
|
||||
tags_list = []
|
||||
|
||||
for fact in facts:
|
||||
fact_texts.append(fact.fact_text)
|
||||
fact_texts.append(_sanitize_text(fact.fact_text))
|
||||
# Convert embedding to string for asyncpg vector type
|
||||
embeddings.append(str(fact.embedding))
|
||||
# event_date: Use occurred_start if available, otherwise use mentioned_at
|
||||
@@ -56,7 +57,7 @@ async def insert_facts_batch(
|
||||
occurred_starts.append(fact.occurred_start)
|
||||
occurred_ends.append(fact.occurred_end)
|
||||
mentioned_ats.append(fact.mentioned_at)
|
||||
contexts.append(fact.context)
|
||||
contexts.append(_sanitize_text(fact.context))
|
||||
fact_types.append(fact.fact_type)
|
||||
# confidence_score is only for opinion facts
|
||||
confidence_scores.append(1.0 if fact.fact_type == "opinion" else None)
|
||||
@@ -157,7 +158,8 @@ async def handle_document_tracking(
|
||||
"""
|
||||
import hashlib
|
||||
|
||||
# Calculate content hash
|
||||
# Sanitize and calculate content hash
|
||||
combined_content = _sanitize_text(combined_content) or ""
|
||||
content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
|
||||
|
||||
# Always delete old document first if it exists (cascades to units and links)
|
||||
|
||||
@@ -8,6 +8,7 @@ import logging
|
||||
import time
|
||||
import uuid
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from . import bank_utils
|
||||
@@ -18,6 +19,39 @@ def utcnow():
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
def parse_datetime_flexible(value: Any) -> datetime:
|
||||
"""
|
||||
Parse a datetime value that could be either a datetime object or an ISO string.
|
||||
|
||||
This handles datetime values from both direct Python calls and deserialized JSON
|
||||
(where datetime objects are serialized as ISO strings).
|
||||
|
||||
Args:
|
||||
value: Either a datetime object or an ISO format string
|
||||
|
||||
Returns:
|
||||
datetime object (timezone-aware)
|
||||
|
||||
Raises:
|
||||
TypeError: If value is neither datetime nor string
|
||||
ValueError: If string is not a valid ISO datetime
|
||||
"""
|
||||
if isinstance(value, datetime):
|
||||
# Ensure timezone-aware
|
||||
if value.tzinfo is None:
|
||||
return value.replace(tzinfo=UTC)
|
||||
return value
|
||||
elif isinstance(value, str):
|
||||
# Parse ISO format string (handles both 'Z' and '+00:00' timezone formats)
|
||||
dt = datetime.fromisoformat(value.replace("Z", "+00:00"))
|
||||
# Ensure timezone-aware
|
||||
if dt.tzinfo is None:
|
||||
return dt.replace(tzinfo=UTC)
|
||||
return dt
|
||||
else:
|
||||
raise TypeError(f"Expected datetime or string, got {type(value).__name__}")
|
||||
|
||||
|
||||
from ..response_models import TokenUsage
|
||||
from . import (
|
||||
chunk_storage,
|
||||
@@ -89,10 +123,18 @@ async def retain_batch(
|
||||
# Merge item-level tags with document-level tags
|
||||
item_tags = item.get("tags", []) or []
|
||||
merged_tags = list(set(item_tags + (document_tags or [])))
|
||||
|
||||
# Handle event_date: parse flexibly (handles both datetime objects and ISO strings)
|
||||
event_date_value = item.get("event_date")
|
||||
if event_date_value:
|
||||
event_date_value = parse_datetime_flexible(event_date_value)
|
||||
else:
|
||||
event_date_value = utcnow()
|
||||
|
||||
content = RetainContent(
|
||||
content=item["content"],
|
||||
context=item.get("context", ""),
|
||||
event_date=item.get("event_date") or utcnow(),
|
||||
event_date=event_date_value,
|
||||
metadata=item.get("metadata", {}),
|
||||
entities=item.get("entities", []),
|
||||
tags=merged_tags,
|
||||
@@ -101,11 +143,8 @@ async def retain_batch(
|
||||
|
||||
# Step 1: Extract facts from all contents
|
||||
step_start = time.time()
|
||||
extract_opinions = fact_type_override == "opinion"
|
||||
|
||||
extracted_facts, chunks, usage = await fact_extraction.extract_facts_from_contents(
|
||||
contents, llm_config, agent_name, extract_opinions
|
||||
)
|
||||
extracted_facts, chunks, usage = await fact_extraction.extract_facts_from_contents(contents, llm_config, agent_name)
|
||||
log_buffer.append(
|
||||
f"[1] Extract facts: {len(extracted_facts)} facts, {len(chunks)} chunks from {len(contents)} contents in {time.time() - step_start:.3f}s"
|
||||
)
|
||||
@@ -119,6 +158,13 @@ async def retain_batch(
|
||||
# Handle document tracking even with no facts
|
||||
if document_id:
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
# Collect tags from all content items and merge with document_tags
|
||||
all_tags = set(document_tags or [])
|
||||
for item in contents_dicts:
|
||||
item_tags = item.get("tags", []) or []
|
||||
all_tags.update(item_tags)
|
||||
merged_tags = list(all_tags)
|
||||
|
||||
retain_params = {}
|
||||
if contents_dicts:
|
||||
first_item = contents_dicts[0]
|
||||
@@ -133,7 +179,7 @@ async def retain_batch(
|
||||
if first_item.get("metadata"):
|
||||
retain_params["metadata"] = first_item["metadata"]
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn, bank_id, document_id, combined_content, is_first_batch, retain_params, document_tags
|
||||
conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags
|
||||
)
|
||||
else:
|
||||
# Check for per-item document_ids
|
||||
@@ -147,6 +193,13 @@ async def retain_batch(
|
||||
|
||||
for doc_id, doc_contents in contents_by_doc.items():
|
||||
combined_content = "\n".join([c.get("content", "") for _, c in doc_contents])
|
||||
# Collect tags from all content items for this document and merge with document_tags
|
||||
all_tags = set(document_tags or [])
|
||||
for _, item in doc_contents:
|
||||
item_tags = item.get("tags", []) or []
|
||||
all_tags.update(item_tags)
|
||||
merged_tags = list(all_tags)
|
||||
|
||||
retain_params = {}
|
||||
if doc_contents:
|
||||
first_item = doc_contents[0][1]
|
||||
@@ -161,7 +214,7 @@ async def retain_batch(
|
||||
if first_item.get("metadata"):
|
||||
retain_params["metadata"] = first_item["metadata"]
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn, bank_id, doc_id, combined_content, is_first_batch, retain_params, document_tags
|
||||
conn, bank_id, doc_id, combined_content, is_first_batch, retain_params, merged_tags
|
||||
)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
@@ -213,6 +266,13 @@ async def retain_batch(
|
||||
# Legacy: single document_id parameter
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
retain_params = {}
|
||||
# Collect tags from all content items and merge with document_tags
|
||||
all_tags = set(document_tags or [])
|
||||
for item in contents_dicts:
|
||||
item_tags = item.get("tags", []) or []
|
||||
all_tags.update(item_tags)
|
||||
merged_tags = list(all_tags)
|
||||
|
||||
if contents_dicts:
|
||||
first_item = contents_dicts[0]
|
||||
if first_item.get("context"):
|
||||
@@ -227,7 +287,7 @@ async def retain_batch(
|
||||
retain_params["metadata"] = first_item["metadata"]
|
||||
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn, bank_id, document_id, combined_content, is_first_batch, retain_params, document_tags
|
||||
conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags
|
||||
)
|
||||
document_ids_added.append(document_id)
|
||||
doc_id_mapping[None] = document_id # For backwards compatibility
|
||||
@@ -255,6 +315,13 @@ async def retain_batch(
|
||||
# Combine content for this document
|
||||
combined_content = "\n".join([c.get("content", "") for _, c in doc_contents])
|
||||
|
||||
# Collect tags from all content items for this document and merge with document_tags
|
||||
all_tags = set(document_tags or [])
|
||||
for _, item in doc_contents:
|
||||
item_tags = item.get("tags", []) or []
|
||||
all_tags.update(item_tags)
|
||||
merged_tags = list(all_tags)
|
||||
|
||||
# Extract retain params from first content item
|
||||
retain_params = {}
|
||||
if doc_contents:
|
||||
@@ -277,7 +344,7 @@ async def retain_batch(
|
||||
combined_content,
|
||||
is_first_batch,
|
||||
retain_params,
|
||||
document_tags,
|
||||
merged_tags,
|
||||
)
|
||||
document_ids_added.append(actual_doc_id)
|
||||
|
||||
|
||||
@@ -182,7 +182,16 @@ class BrokerTaskBackend(TaskBackend):
|
||||
operation_id = task_dict.get("operation_id")
|
||||
task_type = task_dict.get("type", "unknown")
|
||||
bank_id = task_dict.get("bank_id")
|
||||
payload_json = json.dumps(task_dict)
|
||||
|
||||
# Custom encoder to handle datetime objects
|
||||
from datetime import datetime
|
||||
|
||||
def datetime_encoder(obj):
|
||||
if isinstance(obj, datetime):
|
||||
return obj.isoformat()
|
||||
raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
|
||||
|
||||
payload_json = json.dumps(task_dict, default=datetime_encoder)
|
||||
|
||||
schema = self._schema_getter() if self._schema_getter else self._schema
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
@@ -19,7 +19,6 @@ async def extract_facts(
|
||||
context: str = "",
|
||||
llm_config: "LLMConfig" = None,
|
||||
agent_name: str = None,
|
||||
extract_opinions: bool = False,
|
||||
) -> tuple[list["Fact"], list[tuple[str, int]]]:
|
||||
"""
|
||||
Extract semantic facts from text using LLM.
|
||||
@@ -36,7 +35,6 @@ async def extract_facts(
|
||||
context: Context about the conversation/document
|
||||
llm_config: LLM configuration to use
|
||||
agent_name: Optional agent name to help identify agent-related facts
|
||||
extract_opinions: If True, extract ONLY opinions. If False, extract world and agent facts (no opinions)
|
||||
|
||||
Returns:
|
||||
Tuple of (facts, chunks) where:
|
||||
@@ -55,7 +53,6 @@ async def extract_facts(
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name=agent_name,
|
||||
extract_opinions=extract_opinions,
|
||||
)
|
||||
|
||||
if not facts:
|
||||
|
||||
@@ -24,6 +24,11 @@ from hindsight_api.extensions.operation_validator import (
|
||||
# Consolidation operation
|
||||
ConsolidateContext,
|
||||
ConsolidateResult,
|
||||
# Mental Model operations
|
||||
MentalModelGetContext,
|
||||
MentalModelGetResult,
|
||||
MentalModelRefreshContext,
|
||||
MentalModelRefreshResult,
|
||||
# Core operations
|
||||
OperationValidationError,
|
||||
OperationValidatorExtension,
|
||||
@@ -65,6 +70,11 @@ __all__ = [
|
||||
# Operation Validator - Consolidation
|
||||
"ConsolidateContext",
|
||||
"ConsolidateResult",
|
||||
# Operation Validator - Mental Model
|
||||
"MentalModelGetContext",
|
||||
"MentalModelGetResult",
|
||||
"MentalModelRefreshContext",
|
||||
"MentalModelRefreshResult",
|
||||
# Tenant/Auth
|
||||
"ApiKeyTenantExtension",
|
||||
"AuthenticationError",
|
||||
|
||||
@@ -5,6 +5,42 @@ from hindsight_api.extensions.tenant import AuthenticationError, Tenant, TenantC
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
class DefaultTenantExtension(TenantExtension):
|
||||
"""
|
||||
Default single-tenant extension with no authentication.
|
||||
|
||||
This is the default extension used when no tenant extension is configured.
|
||||
It provides single-tenant behavior using the configured schema from
|
||||
HINDSIGHT_API_DATABASE_SCHEMA (defaults to 'public').
|
||||
|
||||
Features:
|
||||
- No authentication required (passes all requests)
|
||||
- Uses configured schema from environment
|
||||
- Perfect for single-tenant deployments without auth
|
||||
|
||||
Configuration:
|
||||
HINDSIGHT_API_DATABASE_SCHEMA=your-schema (optional, defaults to 'public')
|
||||
|
||||
This is automatically enabled by default. To use custom authentication,
|
||||
configure a different tenant extension:
|
||||
HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict[str, str]):
|
||||
super().__init__(config)
|
||||
# Cache the schema at initialization for consistency
|
||||
# Support explicit schema override via config, otherwise use environment
|
||||
self._schema = config.get("schema", get_config().database_schema)
|
||||
|
||||
async def authenticate(self, context: RequestContext) -> TenantContext:
|
||||
"""Return configured schema without any authentication."""
|
||||
return TenantContext(schema_name=self._schema)
|
||||
|
||||
async def list_tenants(self) -> list[Tenant]:
|
||||
"""Return configured schema for single-tenant setup."""
|
||||
return [Tenant(schema=self._schema)]
|
||||
|
||||
|
||||
class ApiKeyTenantExtension(TenantExtension):
|
||||
"""
|
||||
Built-in tenant extension that validates API key against an environment variable.
|
||||
|
||||
@@ -196,6 +196,57 @@ class ConsolidateResult:
|
||||
error: str | None = None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Mental Model Contexts
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class MentalModelGetContext:
|
||||
"""Context for a mental model GET operation validation (pre-operation)."""
|
||||
|
||||
bank_id: str
|
||||
mental_model_id: str
|
||||
request_context: "RequestContext"
|
||||
|
||||
|
||||
@dataclass
|
||||
class MentalModelRefreshContext:
|
||||
"""Context for a mental model refresh/create operation validation (pre-operation)."""
|
||||
|
||||
bank_id: str
|
||||
mental_model_id: str | None # None for create (not yet assigned)
|
||||
request_context: "RequestContext"
|
||||
|
||||
|
||||
@dataclass
|
||||
class MentalModelGetResult:
|
||||
"""Result context for post-mental-model-GET hook."""
|
||||
|
||||
bank_id: str
|
||||
mental_model_id: str
|
||||
request_context: "RequestContext"
|
||||
output_tokens: int # tokens in the returned content
|
||||
success: bool = True
|
||||
error: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class MentalModelRefreshResult:
|
||||
"""Result context for post-mental-model-refresh hook."""
|
||||
|
||||
bank_id: str
|
||||
mental_model_id: str
|
||||
request_context: "RequestContext"
|
||||
query_tokens: int # tokens in source_query
|
||||
output_tokens: int # tokens in generated content
|
||||
context_tokens: int # tokens in context (if any)
|
||||
facts_used: int # facts referenced in based_on
|
||||
mental_models_used: int # mental models referenced in based_on
|
||||
success: bool = True
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class OperationValidatorExtension(Extension, ABC):
|
||||
"""
|
||||
Validates and hooks into retain/recall/reflect/consolidate operations.
|
||||
@@ -402,3 +453,81 @@ class OperationValidatorExtension(Extension, ABC):
|
||||
- error: Error message (if failed)
|
||||
"""
|
||||
pass
|
||||
|
||||
# =========================================================================
|
||||
# Mental Model - Pre-operation validation hook (optional - override to implement)
|
||||
# =========================================================================
|
||||
|
||||
async def validate_mental_model_get(self, ctx: MentalModelGetContext) -> ValidationResult:
|
||||
"""
|
||||
Validate a mental model GET operation before execution.
|
||||
|
||||
Override to implement custom validation logic for mental model retrieval.
|
||||
|
||||
Args:
|
||||
ctx: Context containing:
|
||||
- bank_id: Bank identifier
|
||||
- mental_model_id: Mental model identifier
|
||||
- request_context: Request context with auth info
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the operation is allowed.
|
||||
"""
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_mental_model_refresh(self, ctx: MentalModelRefreshContext) -> ValidationResult:
|
||||
"""
|
||||
Validate a mental model refresh/create operation before execution.
|
||||
|
||||
Override to implement custom validation logic for mental model refresh.
|
||||
|
||||
Args:
|
||||
ctx: Context containing:
|
||||
- bank_id: Bank identifier
|
||||
- mental_model_id: Mental model identifier (None for create)
|
||||
- request_context: Request context with auth info
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the operation is allowed.
|
||||
"""
|
||||
return ValidationResult.accept()
|
||||
|
||||
# =========================================================================
|
||||
# Mental Model - Post-operation hooks (optional - override to implement)
|
||||
# =========================================================================
|
||||
|
||||
async def on_mental_model_get_complete(self, result: MentalModelGetResult) -> None:
|
||||
"""
|
||||
Called after a mental model GET operation completes (success or failure).
|
||||
|
||||
Override to implement post-operation logic such as tracking or audit logging.
|
||||
|
||||
Args:
|
||||
result: Result context containing:
|
||||
- bank_id: Bank identifier
|
||||
- mental_model_id: Mental model identifier
|
||||
- output_tokens: Token count of the returned content
|
||||
- success: Whether the operation succeeded
|
||||
- error: Error message (if failed)
|
||||
"""
|
||||
pass
|
||||
|
||||
async def on_mental_model_refresh_complete(self, result: MentalModelRefreshResult) -> None:
|
||||
"""
|
||||
Called after a mental model refresh operation completes (success or failure).
|
||||
|
||||
Override to implement post-operation logic such as tracking or audit logging.
|
||||
|
||||
Args:
|
||||
result: Result context containing:
|
||||
- bank_id: Bank identifier
|
||||
- mental_model_id: Mental model identifier
|
||||
- query_tokens: Tokens in source_query
|
||||
- output_tokens: Tokens in generated content
|
||||
- context_tokens: Tokens in context
|
||||
- facts_used: Number of facts referenced
|
||||
- mental_models_used: Number of mental models referenced
|
||||
- success: Whether the operation succeeded
|
||||
- error: Error message (if failed)
|
||||
"""
|
||||
pass
|
||||
|
||||
@@ -20,14 +20,13 @@ import warnings
|
||||
|
||||
import uvicorn
|
||||
|
||||
from . import MemoryEngine
|
||||
from . import MemoryEngine, __version__
|
||||
from .api import create_app
|
||||
from .banner import print_banner
|
||||
from .config import DEFAULT_WORKERS, ENV_WORKERS, HindsightConfig, get_config
|
||||
from .daemon import (
|
||||
DEFAULT_DAEMON_PORT,
|
||||
DEFAULT_IDLE_TIMEOUT,
|
||||
DaemonLock,
|
||||
IdleTimeoutMiddleware,
|
||||
daemonize,
|
||||
)
|
||||
@@ -136,30 +135,15 @@ def main():
|
||||
|
||||
# Daemon mode handling
|
||||
if args.daemon:
|
||||
# Use fixed daemon port
|
||||
args.port = DEFAULT_DAEMON_PORT
|
||||
# Use port from args (may be custom for profiles)
|
||||
if args.port == config.port: # No custom port specified
|
||||
args.port = DEFAULT_DAEMON_PORT
|
||||
args.host = "127.0.0.1" # Only bind to localhost for security
|
||||
|
||||
# Check if another daemon is already running
|
||||
daemon_lock = DaemonLock()
|
||||
if not daemon_lock.acquire():
|
||||
print(f"Daemon already running (PID: {daemon_lock.get_pid()})", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# Fork into background
|
||||
# No lockfile needed - port binding prevents duplicate daemons
|
||||
daemonize()
|
||||
|
||||
# Re-acquire lock in child process
|
||||
daemon_lock = DaemonLock()
|
||||
if not daemon_lock.acquire():
|
||||
sys.exit(1)
|
||||
|
||||
# Register cleanup to release lock
|
||||
def release_lock():
|
||||
daemon_lock.release()
|
||||
|
||||
atexit.register(release_lock)
|
||||
|
||||
# Print banner (not in daemon mode)
|
||||
if not args.daemon:
|
||||
print()
|
||||
@@ -176,19 +160,40 @@ def main():
|
||||
llm_model=config.llm_model,
|
||||
llm_base_url=config.llm_base_url,
|
||||
llm_max_concurrent=config.llm_max_concurrent,
|
||||
llm_max_retries=config.llm_max_retries,
|
||||
llm_initial_backoff=config.llm_initial_backoff,
|
||||
llm_max_backoff=config.llm_max_backoff,
|
||||
llm_timeout=config.llm_timeout,
|
||||
llm_vertexai_project_id=config.llm_vertexai_project_id,
|
||||
llm_vertexai_region=config.llm_vertexai_region,
|
||||
llm_vertexai_service_account_key=config.llm_vertexai_service_account_key,
|
||||
retain_llm_provider=config.retain_llm_provider,
|
||||
retain_llm_api_key=config.retain_llm_api_key,
|
||||
retain_llm_model=config.retain_llm_model,
|
||||
retain_llm_base_url=config.retain_llm_base_url,
|
||||
retain_llm_max_concurrent=config.retain_llm_max_concurrent,
|
||||
retain_llm_max_retries=config.retain_llm_max_retries,
|
||||
retain_llm_initial_backoff=config.retain_llm_initial_backoff,
|
||||
retain_llm_max_backoff=config.retain_llm_max_backoff,
|
||||
retain_llm_timeout=config.retain_llm_timeout,
|
||||
reflect_llm_provider=config.reflect_llm_provider,
|
||||
reflect_llm_api_key=config.reflect_llm_api_key,
|
||||
reflect_llm_model=config.reflect_llm_model,
|
||||
reflect_llm_base_url=config.reflect_llm_base_url,
|
||||
reflect_llm_max_concurrent=config.reflect_llm_max_concurrent,
|
||||
reflect_llm_max_retries=config.reflect_llm_max_retries,
|
||||
reflect_llm_initial_backoff=config.reflect_llm_initial_backoff,
|
||||
reflect_llm_max_backoff=config.reflect_llm_max_backoff,
|
||||
reflect_llm_timeout=config.reflect_llm_timeout,
|
||||
consolidation_llm_provider=config.consolidation_llm_provider,
|
||||
consolidation_llm_api_key=config.consolidation_llm_api_key,
|
||||
consolidation_llm_model=config.consolidation_llm_model,
|
||||
consolidation_llm_base_url=config.consolidation_llm_base_url,
|
||||
consolidation_llm_max_concurrent=config.consolidation_llm_max_concurrent,
|
||||
consolidation_llm_max_retries=config.consolidation_llm_max_retries,
|
||||
consolidation_llm_initial_backoff=config.consolidation_llm_initial_backoff,
|
||||
consolidation_llm_max_backoff=config.consolidation_llm_max_backoff,
|
||||
consolidation_llm_timeout=config.consolidation_llm_timeout,
|
||||
embeddings_provider=config.embeddings_provider,
|
||||
embeddings_local_model=config.embeddings_local_model,
|
||||
embeddings_local_force_cpu=config.embeddings_local_force_cpu,
|
||||
@@ -218,7 +223,6 @@ def main():
|
||||
retain_extract_causal_links=config.retain_extract_causal_links,
|
||||
retain_extraction_mode=config.retain_extraction_mode,
|
||||
retain_custom_instructions=config.retain_custom_instructions,
|
||||
retain_observations_async=config.retain_observations_async,
|
||||
enable_observations=config.enable_observations,
|
||||
consolidation_batch_size=config.consolidation_batch_size,
|
||||
consolidation_max_tokens=config.consolidation_max_tokens,
|
||||
@@ -233,8 +237,9 @@ def main():
|
||||
worker_id=config.worker_id,
|
||||
worker_poll_interval_ms=config.worker_poll_interval_ms,
|
||||
worker_max_retries=config.worker_max_retries,
|
||||
worker_batch_size=config.worker_batch_size,
|
||||
worker_http_port=config.worker_http_port,
|
||||
worker_max_slots=config.worker_max_slots,
|
||||
worker_consolidation_max_slots=config.worker_consolidation_max_slots,
|
||||
reflect_max_iterations=config.reflect_max_iterations,
|
||||
mental_model_refresh_concurrency=config.mental_model_refresh_concurrency,
|
||||
)
|
||||
@@ -341,6 +346,7 @@ def main():
|
||||
embeddings_provider=config.embeddings_provider,
|
||||
reranker_provider=config.reranker_provider,
|
||||
mcp_enabled=config.mcp_enabled,
|
||||
version=__version__,
|
||||
)
|
||||
|
||||
# Start idle checker in daemon mode
|
||||
@@ -362,7 +368,7 @@ def main():
|
||||
|
||||
threading.Thread(target=run_idle_checker, daemon=True).start()
|
||||
|
||||
uvicorn.run(**uvicorn_config) # type: ignore[invalid-argument-type] - dict kwargs
|
||||
uvicorn.run(**uvicorn_config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -32,6 +32,9 @@ class MCPToolsConfig:
|
||||
# How to resolve bank_id for operations
|
||||
bank_id_resolver: Callable[[], str | None]
|
||||
|
||||
# How to resolve API key for tenant auth (optional)
|
||||
api_key_resolver: Callable[[], str | None] | None = None
|
||||
|
||||
# Whether to include bank_id as a parameter on tools (for multi-bank support)
|
||||
include_bank_id_param: bool = False
|
||||
|
||||
@@ -46,6 +49,16 @@ class MCPToolsConfig:
|
||||
retain_fire_and_forget: bool = False # If True, use asyncio.create_task pattern
|
||||
|
||||
|
||||
def _get_request_context(config: MCPToolsConfig) -> RequestContext:
|
||||
"""Create RequestContext with API key from resolver if available.
|
||||
|
||||
This enables tenant auth to work with MCP tools by propagating
|
||||
the Bearer token from the MCP middleware to the memory engine.
|
||||
"""
|
||||
api_key = config.api_key_resolver() if config.api_key_resolver else None
|
||||
return RequestContext(api_key=api_key)
|
||||
|
||||
|
||||
def parse_timestamp(timestamp: str) -> datetime | None:
|
||||
"""Parse an ISO format timestamp string.
|
||||
|
||||
@@ -155,12 +168,14 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
if error:
|
||||
return {"status": "error", "message": error}
|
||||
|
||||
request_context = _get_request_context(config)
|
||||
|
||||
async def _retain():
|
||||
try:
|
||||
await memory.retain_batch_async(
|
||||
bank_id=target_bank,
|
||||
contents=[content_dict],
|
||||
request_context=RequestContext(),
|
||||
request_context=request_context,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing memory: {e}", exc_info=True)
|
||||
@@ -196,16 +211,17 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
return f"Error: {error}"
|
||||
|
||||
contents = [content_dict]
|
||||
request_context = _get_request_context(config)
|
||||
if async_processing:
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=target_bank, contents=contents, request_context=RequestContext()
|
||||
bank_id=target_bank, contents=contents, request_context=request_context
|
||||
)
|
||||
return f"Memory queued for background processing (operation_id: {result.get('operation_id', 'N/A')})"
|
||||
else:
|
||||
await memory.retain_batch_async(
|
||||
bank_id=target_bank,
|
||||
contents=contents,
|
||||
request_context=RequestContext(),
|
||||
request_context=request_context,
|
||||
)
|
||||
return f"Memory stored successfully in bank '{target_bank}'"
|
||||
except Exception as e:
|
||||
@@ -237,12 +253,14 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
if error:
|
||||
return {"status": "error", "message": error}
|
||||
|
||||
request_context = _get_request_context(config)
|
||||
|
||||
async def _retain():
|
||||
try:
|
||||
await memory.retain_batch_async(
|
||||
bank_id=target_bank,
|
||||
contents=[content_dict],
|
||||
request_context=RequestContext(),
|
||||
request_context=request_context,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing memory: {e}", exc_info=True)
|
||||
@@ -280,7 +298,7 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
fact_type=list(VALID_RECALL_FACT_TYPES),
|
||||
budget=Budget.HIGH,
|
||||
max_tokens=max_tokens,
|
||||
request_context=RequestContext(),
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
|
||||
return recall_result.model_dump_json(indent=2)
|
||||
@@ -311,7 +329,7 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
fact_type=list(VALID_RECALL_FACT_TYPES),
|
||||
budget=Budget.HIGH,
|
||||
max_tokens=max_tokens,
|
||||
request_context=RequestContext(),
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
|
||||
return recall_result.model_dump()
|
||||
@@ -370,7 +388,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
query=query,
|
||||
budget=budget_enum,
|
||||
context=context,
|
||||
request_context=RequestContext(),
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
|
||||
return reflect_result.model_dump_json(indent=2)
|
||||
@@ -423,7 +441,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
query=query,
|
||||
budget=budget_enum,
|
||||
context=context,
|
||||
request_context=RequestContext(),
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
|
||||
return reflect_result.model_dump()
|
||||
@@ -447,7 +465,7 @@ def _register_list_banks(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
JSON list of banks with their IDs, names, dispositions, and missions.
|
||||
"""
|
||||
try:
|
||||
banks = await memory.list_banks(request_context=RequestContext())
|
||||
banks = await memory.list_banks(request_context=_get_request_context(config))
|
||||
return json.dumps({"banks": banks}, indent=2)
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing banks: {e}", exc_info=True)
|
||||
@@ -471,8 +489,9 @@ def _register_create_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
mission: Optional mission describing who the agent is and what they're trying to accomplish
|
||||
"""
|
||||
try:
|
||||
request_context = _get_request_context(config)
|
||||
# get_bank_profile auto-creates bank if it doesn't exist
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=RequestContext())
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Update name/mission if provided
|
||||
if name is not None or mission is not None:
|
||||
@@ -480,10 +499,10 @@ def _register_create_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
bank_id,
|
||||
name=name,
|
||||
mission=mission,
|
||||
request_context=RequestContext(),
|
||||
request_context=request_context,
|
||||
)
|
||||
# Fetch updated profile
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=RequestContext())
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Serialize disposition if it's a Pydantic model
|
||||
if "disposition" in profile and hasattr(profile["disposition"], "model_dump"):
|
||||
|
||||
@@ -189,7 +189,7 @@ class MetricsCollectorBase:
|
||||
Args:
|
||||
provider: LLM provider name (openai, anthropic, gemini, groq, ollama, lmstudio)
|
||||
model: Model name
|
||||
scope: Scope identifier (e.g., "memory", "reflect", "entity_observation")
|
||||
scope: Scope identifier (e.g., "memory", "reflect", "consolidation")
|
||||
duration: Call duration in seconds
|
||||
input_tokens: Number of input/prompt tokens
|
||||
output_tokens: Number of output/completion tokens
|
||||
@@ -321,7 +321,7 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
pass
|
||||
|
||||
Args:
|
||||
operation: Operation name (retain, recall, reflect, entity_observation)
|
||||
operation: Operation name (retain, recall, reflect, consolidation)
|
||||
bank_id: Memory bank ID
|
||||
source: Source of the operation (api, reflect, internal)
|
||||
budget: Optional budget level (low, mid, high)
|
||||
@@ -371,7 +371,7 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
Args:
|
||||
provider: LLM provider name (openai, anthropic, gemini, groq, ollama, lmstudio)
|
||||
model: Model name
|
||||
scope: Scope identifier (e.g., "memory", "reflect", "entity_observation")
|
||||
scope: Scope identifier (e.g., "memory", "reflect", "consolidation")
|
||||
duration: Call duration in seconds
|
||||
input_tokens: Number of input/prompt tokens
|
||||
output_tokens: Number of output/completion tokens
|
||||
|
||||
@@ -165,6 +165,81 @@ def run_migrations(
|
||||
logger.debug("Migration advisory lock acquired")
|
||||
|
||||
try:
|
||||
# Ensure pgvector extension is installed globally BEFORE schema migrations
|
||||
# This is critical: the extension must exist database-wide before any schema
|
||||
# migrations run, otherwise custom schemas won't have access to vector types
|
||||
logger.debug("Checking pgvector extension availability...")
|
||||
|
||||
# First, check if extension already exists
|
||||
ext_check = conn.execute(
|
||||
text(
|
||||
"SELECT extname, nspname FROM pg_extension e "
|
||||
"JOIN pg_namespace n ON e.extnamespace = n.oid "
|
||||
"WHERE extname = 'vector'"
|
||||
)
|
||||
).fetchone()
|
||||
|
||||
if ext_check:
|
||||
# Extension exists - check if in correct schema
|
||||
ext_schema = ext_check[1]
|
||||
if ext_schema == "public":
|
||||
logger.info("pgvector extension found in public schema - ready to use")
|
||||
else:
|
||||
# Extension in wrong schema - try to fix if we have permissions
|
||||
logger.warning(
|
||||
f"pgvector extension found in schema '{ext_schema}' instead of 'public'. "
|
||||
f"Attempting to relocate..."
|
||||
)
|
||||
try:
|
||||
conn.execute(text("DROP EXTENSION vector CASCADE"))
|
||||
conn.execute(text("SET search_path TO public"))
|
||||
conn.execute(text("CREATE EXTENSION vector"))
|
||||
conn.commit()
|
||||
logger.info("pgvector extension relocated to public schema")
|
||||
except Exception as e:
|
||||
# Failed to relocate - log but don't fail if extension exists somewhere
|
||||
logger.warning(
|
||||
f"Could not relocate pgvector extension to public schema: {e}. "
|
||||
f"Continuing with extension in '{ext_schema}' schema."
|
||||
)
|
||||
conn.rollback()
|
||||
else:
|
||||
# Extension doesn't exist - try to install
|
||||
logger.info("pgvector extension not found, attempting to install...")
|
||||
try:
|
||||
conn.execute(text("SET search_path TO public"))
|
||||
conn.execute(text("CREATE EXTENSION vector"))
|
||||
conn.commit()
|
||||
logger.info("pgvector extension installed in public schema")
|
||||
except Exception as e:
|
||||
# Installation failed - this is only fatal if extension truly doesn't exist
|
||||
# Check one more time in case another process installed it
|
||||
conn.rollback()
|
||||
ext_recheck = conn.execute(
|
||||
text(
|
||||
"SELECT nspname FROM pg_extension e "
|
||||
"JOIN pg_namespace n ON e.extnamespace = n.oid "
|
||||
"WHERE extname = 'vector'"
|
||||
)
|
||||
).fetchone()
|
||||
|
||||
if ext_recheck:
|
||||
logger.warning(
|
||||
f"Could not install pgvector extension (permission denied?), "
|
||||
f"but extension exists in '{ext_recheck[0]}' schema. Continuing..."
|
||||
)
|
||||
else:
|
||||
# Extension truly doesn't exist and we can't install it
|
||||
logger.error(
|
||||
f"pgvector extension is not installed and cannot be installed: {e}. "
|
||||
f"Please ensure pgvector is installed by a database administrator. "
|
||||
f"See: https://github.com/pgvector/pgvector#installation"
|
||||
)
|
||||
raise RuntimeError(
|
||||
"pgvector extension is required but not installed. "
|
||||
"Please install it with: CREATE EXTENSION vector;"
|
||||
) from e
|
||||
|
||||
# Run migrations while holding the lock
|
||||
_run_migrations_internal(database_url, script_location, schema=schema)
|
||||
finally:
|
||||
|
||||
@@ -40,7 +40,7 @@ class EmbeddedPostgres:
|
||||
# Only set port if explicitly specified
|
||||
if self.port is not None:
|
||||
kwargs["port"] = self.port
|
||||
self._pg0 = Pg0(**kwargs) # type: ignore[invalid-argument-type] - dict kwargs
|
||||
self._pg0 = Pg0(**kwargs)
|
||||
return self._pg0
|
||||
|
||||
async def start(self, max_retries: int = 5, retry_delay: float = 4.0) -> str:
|
||||
|
||||
@@ -124,12 +124,6 @@ def main():
|
||||
default=config.worker_poll_interval_ms,
|
||||
help=f"Poll interval in milliseconds (default: {config.worker_poll_interval_ms}, env: HINDSIGHT_API_WORKER_POLL_INTERVAL_MS)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-size",
|
||||
type=int,
|
||||
default=config.worker_batch_size,
|
||||
help=f"Tasks to claim per poll (default: {config.worker_batch_size}, env: HINDSIGHT_API_WORKER_BATCH_SIZE)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-retries",
|
||||
type=int,
|
||||
@@ -168,8 +162,9 @@ def main():
|
||||
|
||||
print(f"Starting Hindsight Worker: {args.worker_id}")
|
||||
print(f" Poll interval: {args.poll_interval}ms")
|
||||
print(f" Batch size: {args.batch_size}")
|
||||
print(f" Max retries: {args.max_retries}")
|
||||
print(f" Max slots: {config.worker_max_slots}")
|
||||
print(f" Consolidation max slots: {config.worker_consolidation_max_slots}")
|
||||
print(f" HTTP server: {args.http_host}:{args.http_port}")
|
||||
print()
|
||||
|
||||
@@ -181,7 +176,19 @@ def main():
|
||||
nonlocal memory, poller
|
||||
import uvicorn
|
||||
|
||||
from ..extensions import TenantExtension, load_extension
|
||||
from ..extensions import OperationValidatorExtension, TenantExtension, load_extension
|
||||
|
||||
# Load tenant extension BEFORE creating MemoryEngine so it can
|
||||
# set correct schema context during task execution. Without this,
|
||||
# _authenticate_tenant sees no extension and resets schema to "public",
|
||||
# causing worker writes to land in the wrong schema.
|
||||
tenant_extension = load_extension("TENANT", TenantExtension)
|
||||
|
||||
# Load operation validator so workers can record usage metering
|
||||
# for async operations (e.g. refresh_mental_model after consolidation)
|
||||
operation_validator = load_extension("OPERATION_VALIDATOR", OperationValidatorExtension)
|
||||
if operation_validator:
|
||||
logger.info(f"Loaded operation validator: {operation_validator.__class__.__name__}")
|
||||
|
||||
# Initialize MemoryEngine
|
||||
# Workers use SyncTaskBackend because they execute tasks directly,
|
||||
@@ -189,43 +196,63 @@ def main():
|
||||
memory = MemoryEngine(
|
||||
run_migrations=False, # Workers don't run migrations
|
||||
task_backend=SyncTaskBackend(),
|
||||
tenant_extension=tenant_extension,
|
||||
operation_validator=operation_validator,
|
||||
)
|
||||
|
||||
await memory.initialize()
|
||||
|
||||
print(f"Database connected: {config.database_url}")
|
||||
|
||||
# Load tenant extension for dynamic schema discovery
|
||||
tenant_extension = load_extension("TENANT", TenantExtension)
|
||||
|
||||
if tenant_extension:
|
||||
print("Tenant extension loaded - schemas will be discovered dynamically on each poll")
|
||||
else:
|
||||
print("No tenant extension configured, using public schema only")
|
||||
print(f"No tenant extension configured, using schema: {config.database_schema}")
|
||||
|
||||
# Create a single poller that handles all schemas dynamically
|
||||
# Convert default schema to None for SQL compatibility (no schema prefix)
|
||||
from hindsight_api.config import DEFAULT_DATABASE_SCHEMA
|
||||
|
||||
schema = None if config.database_schema == DEFAULT_DATABASE_SCHEMA else config.database_schema
|
||||
poller = WorkerPoller(
|
||||
pool=memory._pool,
|
||||
worker_id=args.worker_id,
|
||||
executor=memory.execute_task,
|
||||
poll_interval_ms=args.poll_interval,
|
||||
batch_size=args.batch_size,
|
||||
max_retries=args.max_retries,
|
||||
schema=schema,
|
||||
tenant_extension=tenant_extension,
|
||||
max_slots=config.worker_max_slots,
|
||||
consolidation_max_slots=config.worker_consolidation_max_slots,
|
||||
)
|
||||
|
||||
# Create the HTTP app for metrics/health
|
||||
app = create_worker_app(poller, memory)
|
||||
|
||||
# Setup signal handlers for graceful shutdown
|
||||
# Setup signal handlers for graceful shutdown using asyncio
|
||||
shutdown_requested = asyncio.Event()
|
||||
force_exit = False
|
||||
|
||||
def signal_handler(signum, frame):
|
||||
print(f"\nReceived signal {signum}, initiating graceful shutdown...")
|
||||
shutdown_requested.set()
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
signal.signal(signal.SIGTERM, signal_handler)
|
||||
def signal_handler():
|
||||
nonlocal force_exit
|
||||
if shutdown_requested.is_set():
|
||||
# Second signal = force exit
|
||||
print("\nReceived second signal, forcing immediate exit...")
|
||||
force_exit = True
|
||||
# Restore default handler so third signal kills process
|
||||
loop.remove_signal_handler(signal.SIGINT)
|
||||
loop.remove_signal_handler(signal.SIGTERM)
|
||||
sys.exit(1)
|
||||
else:
|
||||
print("\nReceived shutdown signal, initiating graceful shutdown...")
|
||||
print("(Press Ctrl+C again to force immediate exit)")
|
||||
shutdown_requested.set()
|
||||
|
||||
# Use asyncio's signal handlers which work properly with the event loop
|
||||
loop.add_signal_handler(signal.SIGINT, signal_handler)
|
||||
loop.add_signal_handler(signal.SIGTERM, signal_handler)
|
||||
|
||||
# Create uvicorn config and server
|
||||
uvicorn_config = uvicorn.Config(
|
||||
@@ -244,7 +271,10 @@ def main():
|
||||
print(f"Worker started. Metrics available at http://{args.http_host}:{args.http_port}/metrics")
|
||||
|
||||
# Wait for shutdown signal
|
||||
await shutdown_requested.wait()
|
||||
try:
|
||||
await shutdown_requested.wait()
|
||||
except KeyboardInterrupt:
|
||||
print("\nReceived interrupt, initiating graceful shutdown...")
|
||||
|
||||
# Graceful shutdown
|
||||
print("Shutting down HTTP server...")
|
||||
|
||||
@@ -57,10 +57,11 @@ class WorkerPoller:
|
||||
worker_id: str,
|
||||
executor: Callable[[dict[str, Any]], Awaitable[None]],
|
||||
poll_interval_ms: int = 500,
|
||||
batch_size: int = 10,
|
||||
max_retries: int = 3,
|
||||
schema: str | None = None,
|
||||
tenant_extension: "TenantExtension | None" = None,
|
||||
max_slots: int = 10,
|
||||
consolidation_max_slots: int = 2,
|
||||
):
|
||||
"""
|
||||
Initialize the worker poller.
|
||||
@@ -70,91 +71,158 @@ class WorkerPoller:
|
||||
worker_id: Unique identifier for this worker
|
||||
executor: Async function to execute tasks (typically MemoryEngine.execute_task)
|
||||
poll_interval_ms: Interval between polls when no tasks found (milliseconds)
|
||||
batch_size: Maximum number of tasks to claim per poll cycle
|
||||
max_retries: Maximum retry attempts before marking task as failed
|
||||
schema: Database schema for single-tenant support (ignored if tenant_extension is set)
|
||||
tenant_extension: Extension for dynamic multi-tenant discovery. If set, list_tenants()
|
||||
is called on each poll cycle to discover schemas dynamically.
|
||||
schema: Database schema for single-tenant support (deprecated, use tenant_extension)
|
||||
tenant_extension: Extension for dynamic multi-tenant discovery. If None, creates a
|
||||
DefaultTenantExtension with the configured schema.
|
||||
max_slots: Maximum concurrent tasks per worker
|
||||
consolidation_max_slots: Maximum concurrent consolidation tasks per worker
|
||||
"""
|
||||
self._pool = pool
|
||||
self._worker_id = worker_id
|
||||
self._executor = executor
|
||||
self._poll_interval_ms = poll_interval_ms
|
||||
self._batch_size = batch_size
|
||||
self._max_retries = max_retries
|
||||
self._schema = schema
|
||||
# Always set tenant extension (use DefaultTenantExtension if none provided)
|
||||
if tenant_extension is None:
|
||||
from ..extensions.builtin.tenant import DefaultTenantExtension
|
||||
|
||||
# Pass schema parameter to DefaultTenantExtension if explicitly provided
|
||||
config = {"schema": schema} if schema else {}
|
||||
tenant_extension = DefaultTenantExtension(config=config)
|
||||
self._tenant_extension = tenant_extension
|
||||
self._max_slots = max_slots
|
||||
self._consolidation_max_slots = consolidation_max_slots
|
||||
self._shutdown = asyncio.Event()
|
||||
self._current_tasks: set[asyncio.Task] = set()
|
||||
self._in_flight_count = 0
|
||||
self._in_flight_lock = asyncio.Lock()
|
||||
self._last_progress_log = 0.0
|
||||
self._tasks_completed_since_log = 0
|
||||
# Track active tasks locally: operation_id -> (op_type, bank_id, schema)
|
||||
self._active_tasks: dict[str, tuple[str, str, str | None]] = {}
|
||||
# Track active tasks locally: operation_id -> (op_type, bank_id, schema, asyncio.Task)
|
||||
self._active_tasks: dict[str, tuple[str, str, str | None, asyncio.Task]] = {}
|
||||
# Track in-flight tasks by operation type
|
||||
self._in_flight_by_type: dict[str, int] = {}
|
||||
|
||||
async def _get_schemas(self) -> list[str | None]:
|
||||
"""Get list of schemas to poll. Returns [None] for public schema."""
|
||||
if self._tenant_extension is not None:
|
||||
tenants = await self._tenant_extension.list_tenants()
|
||||
# Convert "public" to None for SQL compatibility, keep others as-is
|
||||
return [t.schema if t.schema != "public" else None for t in tenants]
|
||||
# Single schema mode
|
||||
return [self._schema]
|
||||
"""Get list of schemas to poll. Returns [None] for default schema (no prefix)."""
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA
|
||||
|
||||
tenants = await self._tenant_extension.list_tenants()
|
||||
# Convert default schema to None for SQL compatibility (no prefix), keep others as-is
|
||||
return [t.schema if t.schema != DEFAULT_DATABASE_SCHEMA else None for t in tenants]
|
||||
|
||||
async def _get_available_slots(self) -> tuple[int, int]:
|
||||
"""
|
||||
Calculate available slots for claiming tasks.
|
||||
|
||||
Returns:
|
||||
(total_available, consolidation_available) tuple
|
||||
"""
|
||||
async with self._in_flight_lock:
|
||||
total_in_flight = self._in_flight_count
|
||||
consolidation_in_flight = self._in_flight_by_type.get("consolidation", 0)
|
||||
|
||||
total_available = max(0, self._max_slots - total_in_flight)
|
||||
consolidation_available = max(0, self._consolidation_max_slots - consolidation_in_flight)
|
||||
|
||||
return total_available, consolidation_available
|
||||
|
||||
async def wait_for_active_tasks(self, timeout: float = 10.0) -> bool:
|
||||
"""
|
||||
Wait for all active background tasks to complete (test helper).
|
||||
|
||||
This is a test-only utility that allows tests to synchronize with
|
||||
fire-and-forget background tasks without using sleep().
|
||||
|
||||
Args:
|
||||
timeout: Maximum time to wait in seconds
|
||||
|
||||
Returns:
|
||||
True if all tasks completed, False if timeout was reached
|
||||
"""
|
||||
start_time = asyncio.get_event_loop().time()
|
||||
while True:
|
||||
async with self._in_flight_lock:
|
||||
if self._in_flight_count == 0:
|
||||
return True
|
||||
|
||||
elapsed = asyncio.get_event_loop().time() - start_time
|
||||
if elapsed >= timeout:
|
||||
return False
|
||||
|
||||
# Short sleep to avoid busy-waiting
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
async def claim_batch(self) -> list[ClaimedTask]:
|
||||
"""
|
||||
Claim up to batch_size pending tasks atomically across all tenant schemas.
|
||||
Claim pending tasks atomically across all tenant schemas,
|
||||
respecting slot limits (total and consolidation).
|
||||
|
||||
Uses FOR UPDATE SKIP LOCKED to ensure no conflicts with other workers.
|
||||
|
||||
For consolidation tasks specifically, skips pending tasks if there's already
|
||||
a processing consolidation for the same bank (to avoid duplicate work).
|
||||
|
||||
If tenant_extension is configured, dynamically discovers schemas on each call.
|
||||
|
||||
Returns:
|
||||
List of ClaimedTask objects containing operation_id, task_dict, and schema
|
||||
"""
|
||||
# Calculate available slots
|
||||
total_available, consolidation_available = await self._get_available_slots()
|
||||
|
||||
if total_available <= 0:
|
||||
return []
|
||||
|
||||
schemas = await self._get_schemas()
|
||||
all_tasks: list[ClaimedTask] = []
|
||||
remaining_batch = self._batch_size
|
||||
remaining_total = total_available
|
||||
remaining_consolidation = consolidation_available
|
||||
|
||||
for schema in schemas:
|
||||
if remaining_batch <= 0:
|
||||
if remaining_total <= 0:
|
||||
break
|
||||
|
||||
tasks = await self._claim_batch_for_schema(schema, remaining_batch)
|
||||
tasks = await self._claim_batch_for_schema(schema, remaining_total, remaining_consolidation)
|
||||
|
||||
# Update remaining slots based on what was claimed
|
||||
for task in tasks:
|
||||
op_type = task.task_dict.get("operation_type", "unknown")
|
||||
if op_type == "consolidation":
|
||||
remaining_consolidation -= 1
|
||||
|
||||
all_tasks.extend(tasks)
|
||||
remaining_batch -= len(tasks)
|
||||
remaining_total -= len(tasks)
|
||||
|
||||
return all_tasks
|
||||
|
||||
async def _claim_batch_for_schema(self, schema: str | None, limit: int) -> list[ClaimedTask]:
|
||||
"""Claim tasks from a specific schema."""
|
||||
async def _claim_batch_for_schema(
|
||||
self, schema: str | None, limit: int, consolidation_limit: int
|
||||
) -> list[ClaimedTask]:
|
||||
"""Claim tasks from a specific schema respecting slot limits."""
|
||||
try:
|
||||
return await self._claim_batch_for_schema_inner(schema, limit, consolidation_limit)
|
||||
except Exception as e:
|
||||
# Format schema for logging: custom schemas in quotes, None as-is
|
||||
schema_display = f'"{schema}"' if schema else str(schema)
|
||||
logger.warning(f"Worker {self._worker_id} failed to claim tasks for schema {schema_display}: {e}")
|
||||
return []
|
||||
|
||||
async def _claim_batch_for_schema_inner(
|
||||
self, schema: str | None, limit: int, consolidation_limit: int
|
||||
) -> list[ClaimedTask]:
|
||||
"""Inner implementation for claiming tasks from a specific schema with slot limits."""
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
async with self._pool.acquire() as conn:
|
||||
async with conn.transaction():
|
||||
# Select and lock pending tasks
|
||||
# For consolidation: skip if same bank already has one processing
|
||||
rows = await conn.fetch(
|
||||
# Strategy: Claim non-consolidation tasks first, then consolidation up to limit
|
||||
|
||||
# 1. Claim non-consolidation tasks (up to limit)
|
||||
non_consolidation_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, task_payload
|
||||
FROM {table} AS pending
|
||||
WHERE status = 'pending' AND task_payload IS NOT NULL
|
||||
AND (
|
||||
-- Non-consolidation tasks: always claimable
|
||||
operation_type != 'consolidation'
|
||||
OR
|
||||
-- Consolidation: only if no other consolidation processing for same bank
|
||||
NOT EXISTS (
|
||||
SELECT 1 FROM {table} AS processing
|
||||
WHERE processing.bank_id = pending.bank_id
|
||||
AND processing.operation_type = 'consolidation'
|
||||
AND processing.status = 'processing'
|
||||
)
|
||||
)
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type != 'consolidation'
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
@@ -162,11 +230,39 @@ class WorkerPoller:
|
||||
limit,
|
||||
)
|
||||
|
||||
if not rows:
|
||||
claimed_count = len(non_consolidation_rows)
|
||||
remaining_limit = limit - claimed_count
|
||||
|
||||
# 2. Claim consolidation tasks (up to consolidation_limit and remaining_limit)
|
||||
consolidation_rows = []
|
||||
if consolidation_limit > 0 and remaining_limit > 0:
|
||||
consolidation_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, task_payload
|
||||
FROM {table} AS pending
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {table} AS processing
|
||||
WHERE processing.bank_id = pending.bank_id
|
||||
AND processing.operation_type = 'consolidation'
|
||||
AND processing.status = 'processing'
|
||||
)
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
min(consolidation_limit, remaining_limit),
|
||||
)
|
||||
|
||||
all_rows = non_consolidation_rows + consolidation_rows
|
||||
|
||||
if not all_rows:
|
||||
return []
|
||||
|
||||
# Claim the tasks by updating status and worker_id
|
||||
operation_ids = [row["operation_id"] for row in rows]
|
||||
operation_ids = [row["operation_id"] for row in all_rows]
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
@@ -184,7 +280,7 @@ class WorkerPoller:
|
||||
task_dict=json.loads(row["task_payload"]),
|
||||
schema=schema,
|
||||
)
|
||||
for row in rows
|
||||
for row in all_rows
|
||||
]
|
||||
|
||||
async def _mark_completed(self, operation_id: str, schema: str | None):
|
||||
@@ -250,18 +346,43 @@ class WorkerPoller:
|
||||
logger.warning(f"Task {operation_id} failed, will retry (attempt {retry_count + 1}/{self._max_retries})")
|
||||
|
||||
async def execute_task(self, task: ClaimedTask):
|
||||
"""Execute a single task and update its status."""
|
||||
"""Execute a single task as a background job (fire-and-forget)."""
|
||||
task_type = task.task_dict.get("type", "unknown")
|
||||
operation_type = task.task_dict.get("operation_type", "unknown")
|
||||
bank_id = task.task_dict.get("bank_id", "unknown")
|
||||
|
||||
# Create background task
|
||||
bg_task = asyncio.create_task(self._execute_task_inner(task))
|
||||
|
||||
# Track this task as active
|
||||
async with self._in_flight_lock:
|
||||
self._active_tasks[task.operation_id] = (task_type, bank_id, task.schema)
|
||||
self._active_tasks[task.operation_id] = (task_type, bank_id, task.schema, bg_task)
|
||||
self._in_flight_count += 1
|
||||
self._in_flight_by_type[operation_type] = self._in_flight_by_type.get(operation_type, 0) + 1
|
||||
|
||||
# Add cleanup callback
|
||||
bg_task.add_done_callback(lambda _: asyncio.create_task(self._cleanup_task(task.operation_id, operation_type)))
|
||||
|
||||
async def _cleanup_task(self, operation_id: str, operation_type: str):
|
||||
"""Remove task from tracking after completion."""
|
||||
async with self._in_flight_lock:
|
||||
if operation_id in self._active_tasks:
|
||||
self._active_tasks.pop(operation_id, None)
|
||||
self._in_flight_count -= 1
|
||||
count = self._in_flight_by_type.get(operation_type, 0)
|
||||
if count > 0:
|
||||
self._in_flight_by_type[operation_type] = count - 1
|
||||
if self._in_flight_by_type[operation_type] == 0:
|
||||
del self._in_flight_by_type[operation_type]
|
||||
|
||||
async def _execute_task_inner(self, task: ClaimedTask):
|
||||
"""Inner task execution with error handling."""
|
||||
task_type = task.task_dict.get("type", "unknown")
|
||||
bank_id = task.task_dict.get("bank_id", "unknown")
|
||||
|
||||
try:
|
||||
schema_info = f", schema={task.schema}" if task.schema else ""
|
||||
logger.debug(f"Executing task {task.operation_id} (type={task_type}, bank={bank_id}{schema_info})")
|
||||
# Pass schema to executor so it can set the correct context
|
||||
if task.schema:
|
||||
task.task_dict["_schema"] = task.schema
|
||||
await self._executor(task.task_dict)
|
||||
@@ -271,10 +392,6 @@ class WorkerPoller:
|
||||
error_msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
|
||||
logger.error(f"Task {task.operation_id} failed: {e}")
|
||||
await self._retry_or_fail(task.operation_id, error_msg, task.schema)
|
||||
finally:
|
||||
# Remove from active tasks
|
||||
async with self._in_flight_lock:
|
||||
self._active_tasks.pop(task.operation_id, None)
|
||||
|
||||
async def recover_own_tasks(self) -> int:
|
||||
"""
|
||||
@@ -293,20 +410,25 @@ class WorkerPoller:
|
||||
total_count = 0
|
||||
|
||||
for schema in schemas:
|
||||
table = fq_table("async_operations", schema)
|
||||
try:
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
result = await self._pool.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
|
||||
WHERE status = 'processing' AND worker_id = $1
|
||||
""",
|
||||
self._worker_id,
|
||||
)
|
||||
result = await self._pool.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
|
||||
WHERE status = 'processing' AND worker_id = $1
|
||||
""",
|
||||
self._worker_id,
|
||||
)
|
||||
|
||||
# Parse "UPDATE N" to get count
|
||||
count = int(result.split()[-1]) if result else 0
|
||||
total_count += count
|
||||
# Parse "UPDATE N" to get count
|
||||
count = int(result.split()[-1]) if result else 0
|
||||
total_count += count
|
||||
except Exception as e:
|
||||
# Format schema for logging: custom schemas in quotes, None as-is
|
||||
schema_display = f'"{schema}"' if schema else str(schema)
|
||||
logger.warning(f"Worker {self._worker_id} failed to recover tasks for schema {schema_display}: {e}")
|
||||
|
||||
if total_count > 0:
|
||||
logger.info(f"Worker {self._worker_id} recovered {total_count} stale tasks from previous run")
|
||||
@@ -314,59 +436,60 @@ class WorkerPoller:
|
||||
|
||||
async def run(self):
|
||||
"""
|
||||
Main polling loop.
|
||||
Main polling loop with fire-and-forget task execution.
|
||||
|
||||
Continuously polls for pending tasks, claims them, and executes them
|
||||
until shutdown is signaled.
|
||||
|
||||
If tenant_extension is configured, dynamically discovers schemas on each poll.
|
||||
Continuously polls for pending tasks, spawns them as background tasks,
|
||||
and immediately continues polling (up to slot limits).
|
||||
"""
|
||||
# Recover any tasks from a previous crash before starting
|
||||
await self.recover_own_tasks()
|
||||
|
||||
logger.info(f"Worker {self._worker_id} starting polling loop")
|
||||
logger.info(
|
||||
f"Worker {self._worker_id} starting polling loop "
|
||||
f"(max_slots={self._max_slots}, consolidation_max_slots={self._consolidation_max_slots})"
|
||||
)
|
||||
|
||||
while not self._shutdown.is_set():
|
||||
try:
|
||||
# Claim a batch of tasks (across all tenant schemas if configured)
|
||||
# Claim a batch of tasks (respecting slot limits)
|
||||
tasks = await self.claim_batch()
|
||||
|
||||
if tasks:
|
||||
# Log batch info
|
||||
task_types: dict[str, int] = {}
|
||||
schemas_seen: set[str | None] = set()
|
||||
consolidation_count = 0
|
||||
for task in tasks:
|
||||
t = task.task_dict.get("type", "unknown")
|
||||
op_type = task.task_dict.get("operation_type", "unknown")
|
||||
task_types[t] = task_types.get(t, 0) + 1
|
||||
schemas_seen.add(task.schema)
|
||||
if op_type == "consolidation":
|
||||
consolidation_count += 1
|
||||
|
||||
types_str = ", ".join(f"{k}:{v}" for k, v in task_types.items())
|
||||
schemas_str = ", ".join(s or "public" for s in schemas_seen)
|
||||
# Display None as "default" in logs
|
||||
schemas_str = ", ".join(s if s else "default" for s in schemas_seen)
|
||||
logger.info(
|
||||
f"Worker {self._worker_id} claimed {len(tasks)} tasks: {types_str} (schemas: {schemas_str})"
|
||||
f"Worker {self._worker_id} claimed {len(tasks)} tasks "
|
||||
f"({consolidation_count} consolidation): {types_str} (schemas: {schemas_str})"
|
||||
)
|
||||
|
||||
# Track in-flight tasks
|
||||
async with self._in_flight_lock:
|
||||
self._in_flight_count += len(tasks)
|
||||
# Spawn tasks as background jobs (fire-and-forget)
|
||||
for task in tasks:
|
||||
await self.execute_task(task)
|
||||
|
||||
# Execute tasks concurrently
|
||||
try:
|
||||
await asyncio.gather(
|
||||
*[self.execute_task(task) for task in tasks],
|
||||
return_exceptions=True,
|
||||
)
|
||||
finally:
|
||||
async with self._in_flight_lock:
|
||||
self._in_flight_count -= len(tasks)
|
||||
else:
|
||||
# No tasks found, wait before polling again
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._shutdown.wait(),
|
||||
timeout=self._poll_interval_ms / 1000,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
pass # Normal timeout, continue polling
|
||||
# Continue immediately to claim more tasks (if slots available)
|
||||
continue
|
||||
|
||||
# No tasks claimed (either no pending tasks or slots full)
|
||||
# Wait before polling again
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._shutdown.wait(),
|
||||
timeout=self._poll_interval_ms / 1000,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
pass # Normal timeout, continue polling
|
||||
|
||||
# Log progress stats periodically
|
||||
await self._log_progress_if_due()
|
||||
@@ -397,15 +520,27 @@ class WorkerPoller:
|
||||
while asyncio.get_event_loop().time() - start_time < timeout:
|
||||
async with self._in_flight_lock:
|
||||
in_flight = self._in_flight_count
|
||||
active_task_objects = [task_info[3] for task_info in self._active_tasks.values()]
|
||||
|
||||
if in_flight == 0:
|
||||
logger.info(f"Worker {self._worker_id} graceful shutdown complete")
|
||||
return
|
||||
|
||||
logger.info(f"Worker {self._worker_id} waiting for {in_flight} in-flight tasks")
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
logger.warning(f"Worker {self._worker_id} shutdown timeout after {timeout}s")
|
||||
# Wait for at least one task to complete
|
||||
if active_task_objects:
|
||||
done, _ = await asyncio.wait(active_task_objects, timeout=0.5, return_when=asyncio.FIRST_COMPLETED)
|
||||
else:
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
logger.warning(f"Worker {self._worker_id} shutdown timeout after {timeout}s, cancelling remaining tasks")
|
||||
|
||||
# Cancel remaining tasks
|
||||
async with self._in_flight_lock:
|
||||
for operation_id, (_, _, _, bg_task) in list(self._active_tasks.items()):
|
||||
if not bg_task.done():
|
||||
bg_task.cancel()
|
||||
|
||||
async def _log_progress_if_due(self):
|
||||
"""Log progress stats every PROGRESS_LOG_INTERVAL seconds."""
|
||||
@@ -416,14 +551,19 @@ class WorkerPoller:
|
||||
self._last_progress_log = now
|
||||
|
||||
try:
|
||||
# Get local active tasks (this worker only)
|
||||
# Get local active tasks
|
||||
async with self._in_flight_lock:
|
||||
in_flight = self._in_flight_count
|
||||
active_tasks = dict(self._active_tasks) # Copy to avoid holding lock
|
||||
in_flight_by_type = dict(self._in_flight_by_type)
|
||||
active_tasks = dict(self._active_tasks)
|
||||
|
||||
# Build local processing breakdown grouped by (op_type, bank_id)
|
||||
consolidation_count = in_flight_by_type.get("consolidation", 0)
|
||||
available_slots = self._max_slots - in_flight
|
||||
available_consolidation_slots = self._consolidation_max_slots - consolidation_count
|
||||
|
||||
# Build local processing breakdown
|
||||
task_groups: dict[tuple[str, str], int] = {}
|
||||
for op_type, bank_id, _ in active_tasks.values():
|
||||
for op_type, bank_id, _, _ in active_tasks.values():
|
||||
key = (op_type, bank_id)
|
||||
task_groups[key] = task_groups.get(key, 0) + 1
|
||||
|
||||
@@ -432,7 +572,7 @@ class WorkerPoller:
|
||||
if len(processing_info) > 10:
|
||||
processing_str += f" +{len(processing_info) - 10} more"
|
||||
|
||||
# Get global stats from DB across all schemas
|
||||
# Get global stats from DB
|
||||
schemas = await self._get_schemas()
|
||||
global_pending = 0
|
||||
all_worker_counts: dict[str, int] = {}
|
||||
@@ -444,7 +584,6 @@ class WorkerPoller:
|
||||
row = await conn.fetchrow(f"SELECT COUNT(*) as count FROM {table} WHERE status = 'pending'")
|
||||
global_pending += row["count"] if row else 0
|
||||
|
||||
# Get processing breakdown by worker
|
||||
worker_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT worker_id, COUNT(*) as count
|
||||
@@ -457,16 +596,18 @@ class WorkerPoller:
|
||||
wid = wr["worker_id"] or "unknown"
|
||||
all_worker_counts[wid] = all_worker_counts.get(wid, 0) + wr["count"]
|
||||
|
||||
# Format other workers' processing counts
|
||||
other_workers = []
|
||||
for wid, cnt in all_worker_counts.items():
|
||||
if wid != self._worker_id:
|
||||
other_workers.append(f"{wid}:{cnt}")
|
||||
others_str = ", ".join(other_workers) if other_workers else "none"
|
||||
|
||||
schemas_str = ", ".join(s or "public" for s in schemas)
|
||||
# Display None as "default" in logs
|
||||
schemas_str = ", ".join(s if s else "default" for s in schemas)
|
||||
logger.info(
|
||||
f"[WORKER_STATS] worker={self._worker_id} in_flight={in_flight} | "
|
||||
f"[WORKER_STATS] worker={self._worker_id} "
|
||||
f"slots={in_flight}/{self._max_slots} (consolidation={consolidation_count}/{self._consolidation_max_slots}) | "
|
||||
f"available={available_slots} (consolidation={available_consolidation_slots}) | "
|
||||
f"global: pending={global_pending} (schemas: {schemas_str}) | "
|
||||
f"others: {others_str} | "
|
||||
f"my_active: {processing_str}"
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api"
|
||||
version = "0.4.1"
|
||||
version = "0.4.9"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -25,7 +25,7 @@ dependencies = [
|
||||
"psycopg2-binary>=2.9.11",
|
||||
"tiktoken>=0.12.0",
|
||||
"httpx>=0.27.0",
|
||||
"fastmcp>=2.14.0", # CVE-2025-66416
|
||||
"fastmcp>=2.14.0", # CVE-2025-66416
|
||||
"pg0-embedded>=0.11.0",
|
||||
"python-dateutil>=2.8.0",
|
||||
"opentelemetry-api>=1.20.0",
|
||||
@@ -34,22 +34,24 @@ dependencies = [
|
||||
"opentelemetry-exporter-prometheus>=0.41b0",
|
||||
"dateparser>=1.2.2",
|
||||
"google-genai>=1.0.0",
|
||||
"google-auth>=2.0.0",
|
||||
"anthropic>=0.40.0",
|
||||
"typer>=0.9.0",
|
||||
"cohere>=5.0.0",
|
||||
"flashrank>=0.2.0",
|
||||
# Local ML models for embeddings/reranking - can be excluded in Docker with INCLUDE_LOCAL_MODELS=false
|
||||
"sentence-transformers>=3.3.0",
|
||||
"transformers>=4.53.0", # Security fixes for ReDoS vulnerabilities
|
||||
"torch>=2.6.0", # CVE fix for remote code execution
|
||||
"transformers>=4.53.0", # Security fixes for ReDoS vulnerabilities
|
||||
"torch>=2.6.0", # CVE fix for remote code execution
|
||||
"uvloop>=0.22.1",
|
||||
# Transitive dependency security fixes
|
||||
"pyasn1>=0.6.2", # DoS vulnerability fix
|
||||
"urllib3>=2.6.3", # Decompression-bomb safeguards bypass fix
|
||||
"langchain-core>=1.2.5", # Serialization injection vulnerability fix
|
||||
"filelock>=3.20.1", # TOCTOU race condition fix
|
||||
"authlib>=1.6.6", # Account takeover vulnerability fix
|
||||
"aiohttp>=3.13.3", # Multiple DoS vulnerabilities
|
||||
"pyasn1>=0.6.2", # DoS vulnerability fix
|
||||
"urllib3>=2.6.3", # Decompression-bomb safeguards bypass fix
|
||||
"langchain-core>=1.2.5", # Serialization injection vulnerability fix
|
||||
"filelock>=3.20.1", # TOCTOU race condition fix
|
||||
"authlib>=1.6.6", # Account takeover vulnerability fix
|
||||
"aiohttp>=3.13.3", # Multiple DoS vulnerabilities
|
||||
"claude-agent-sdk>=0.1.27",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
@@ -141,6 +143,11 @@ known-third-party = ["alembic"]
|
||||
quote-style = "double"
|
||||
indent-style = "space"
|
||||
|
||||
[tool.uv]
|
||||
# Allow uv to search all configured indexes for packages, not just the first one
|
||||
# This prevents dependency resolution failures when using pytorch index + PyPI
|
||||
index-strategy = "unsafe-best-match"
|
||||
|
||||
[tool.ty]
|
||||
# Type checking configuration
|
||||
# ty is an extremely fast Python type checker from Astral (same team as ruff/uv)
|
||||
|
||||
@@ -220,3 +220,34 @@ async def memory(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
await mem.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope="function")
|
||||
async def memory_no_llm_verify(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
"""
|
||||
Provide a MemoryEngine instance that skips LLM connection verification.
|
||||
|
||||
This fixture is useful for tests that override the LLM configuration
|
||||
after initialization (e.g., to test specific providers).
|
||||
"""
|
||||
mem = MemoryEngine(
|
||||
db_url=pg0_db_url,
|
||||
memory_llm_provider="mock", # Use mock provider as placeholder
|
||||
memory_llm_api_key="",
|
||||
memory_llm_model="mock",
|
||||
embeddings=embeddings,
|
||||
cross_encoder=cross_encoder,
|
||||
query_analyzer=query_analyzer,
|
||||
pool_min_size=1,
|
||||
pool_max_size=5,
|
||||
run_migrations=False,
|
||||
task_backend=SyncTaskBackend(),
|
||||
skip_llm_verification=True, # Skip verification - will be overridden by test
|
||||
)
|
||||
await mem.initialize()
|
||||
yield mem
|
||||
try:
|
||||
if mem._pool and not mem._pool._closing:
|
||||
await mem.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -346,11 +346,11 @@ class TestConsolidationIntegration:
|
||||
or when one directly updates another (e.g., location change).
|
||||
|
||||
Given:
|
||||
- "Nicolò lives in Italy"
|
||||
- "Nicolò moved to the US recently" (updates the living location)
|
||||
- "Alex lives in Italy"
|
||||
- "Alex moved to the US recently" (updates the living location)
|
||||
|
||||
The second fact should UPDATE the first, not create a separate observation.
|
||||
But unrelated facts like "Nicolò works at Vectorize" should stay separate.
|
||||
But unrelated facts like "Alex works at Vectorize" should stay separate.
|
||||
"""
|
||||
bank_id = f"test-consolidation-merge-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
@@ -360,14 +360,14 @@ class TestConsolidationIntegration:
|
||||
# Retain a memory about living location
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Nicolò lives in Italy.",
|
||||
content="Alex lives in Italy.",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Retain an unrelated memory (different topic - should NOT merge)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Nicolò works at Vectorize as an engineer.",
|
||||
content="Alex works at Vectorize as an engineer.",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
@@ -384,7 +384,7 @@ class TestConsolidationIntegration:
|
||||
# Add a memory that UPDATES the living location (should merge with first)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Nicolò recently moved to the United States.",
|
||||
content="Alex recently moved to the United States.",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
@@ -485,9 +485,9 @@ class TestConsolidationIntegration:
|
||||
they should be merged into ONE observation that captures the change.
|
||||
|
||||
Example:
|
||||
- "Nicolò loves pizza"
|
||||
- "Nicolò hates pizza"
|
||||
→ Should become: "Nicolò used to love pizza but now hates it" (or similar)
|
||||
- "Alex loves pizza"
|
||||
- "Alex hates pizza"
|
||||
→ Should become: "Alex used to love pizza but now hates it" (or similar)
|
||||
"""
|
||||
bank_id = f"test-consolidation-contradict-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
@@ -497,7 +497,7 @@ class TestConsolidationIntegration:
|
||||
# Add initial fact
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Nicolò loves pizza.",
|
||||
content="Alex loves pizza.",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
@@ -515,7 +515,7 @@ class TestConsolidationIntegration:
|
||||
# Add contradicting fact (same person, same topic, opposite sentiment)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Nicolò hates pizza.",
|
||||
content="Alex hates pizza.",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
|
||||
@@ -58,7 +58,6 @@ async def test_fact_extraction_basic_analysis(llm_config):
|
||||
llm_config=llm_config,
|
||||
agent_name="test-agent",
|
||||
context="Friday Standup meeting",
|
||||
extract_opinions=False,
|
||||
)
|
||||
|
||||
duration = time.time() - start_time
|
||||
|
||||
@@ -1098,3 +1098,134 @@ async def test_version_endpoint_returns_correct_version(api_client):
|
||||
assert isinstance(features["worker"], bool)
|
||||
|
||||
print(f"Version endpoint returned: api_version={result['api_version']}, features={features}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_with_timestamp_async(api_client, test_bank_id):
|
||||
"""Test that async retain accepts timestamp field and serializes correctly."""
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "Test memory with timestamp",
|
||||
"context": "test",
|
||||
"timestamp": "2026-01-30T11:45:00Z"
|
||||
}
|
||||
],
|
||||
"async": True
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200, f"Expected 200, got {response.status_code}: {response.text}"
|
||||
data = response.json()
|
||||
assert data["success"] is True
|
||||
assert data["async"] is True
|
||||
assert "operation_id" in data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_with_timestamp_sync(api_client, test_bank_id):
|
||||
"""Test that sync retain accepts timestamp field."""
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "Test memory with timestamp sync",
|
||||
"context": "test",
|
||||
"timestamp": "2026-01-30T11:45:00Z"
|
||||
}
|
||||
],
|
||||
"async": False
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200, f"Expected 200, got {response.status_code}: {response.text}"
|
||||
data = response.json()
|
||||
assert data["success"] is True
|
||||
assert data["async"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_with_multiple_timestamps(api_client, test_bank_id):
|
||||
"""Test that multiple items with different timestamp formats work."""
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "Event 1",
|
||||
"timestamp": "2026-01-30T11:45:00Z" # With Z
|
||||
},
|
||||
{
|
||||
"content": "Event 2",
|
||||
"timestamp": "2026-01-30T12:00:00+00:00" # With timezone
|
||||
},
|
||||
{
|
||||
"content": "Event 3" # No timestamp
|
||||
}
|
||||
],
|
||||
"async": True
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200, f"Expected 200, got {response.status_code}: {response.text}"
|
||||
data = response.json()
|
||||
assert data["success"] is True
|
||||
assert data["items_count"] == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_with_timestamp_async_complete_processing(api_client, test_bank_id):
|
||||
"""Test that async retain with timestamp completes full processing including fact extraction."""
|
||||
# Submit async retain with timestamp
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "The quarterly meeting was held on January 30th 2026",
|
||||
"context": "meetings",
|
||||
"timestamp": "2026-01-30T11:45:00Z"
|
||||
}
|
||||
],
|
||||
"async": True
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200, f"Expected 200, got {response.status_code}: {response.text}"
|
||||
data = response.json()
|
||||
assert data["success"] is True
|
||||
assert data["async"] is True
|
||||
operation_id = data["operation_id"]
|
||||
|
||||
# Wait for async processing to complete (poll operation status)
|
||||
max_wait_seconds = 30
|
||||
poll_interval = 0.5
|
||||
elapsed = 0
|
||||
operation_completed = False
|
||||
|
||||
while elapsed < max_wait_seconds:
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/operations/{operation_id}")
|
||||
if response.status_code == 200:
|
||||
op_status = response.json()
|
||||
if op_status.get("status") == "completed":
|
||||
operation_completed = True
|
||||
break
|
||||
elif op_status.get("status") == "failed":
|
||||
raise AssertionError(f"Operation failed: {op_status.get('error_message')}")
|
||||
|
||||
await asyncio.sleep(poll_interval)
|
||||
elapsed += poll_interval
|
||||
|
||||
assert operation_completed, f"Async operation did not complete within {max_wait_seconds} seconds"
|
||||
|
||||
# Verify memories were actually stored
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/memories/list",
|
||||
params={"limit": 10}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
items = response.json()["items"]
|
||||
assert len(items) > 0, "Should have stored memories after async processing"
|
||||
|
||||
@@ -1,5 +1,10 @@
|
||||
"""
|
||||
Test LLM provider with different models using actual memory operations.
|
||||
Test LLM provider with different models using actual Hindsight memory operations.
|
||||
|
||||
Tests validate that providers work correctly with:
|
||||
1. Retain (memory ingestion with fact extraction)
|
||||
2. Reflect (memory retrieval with tool calling)
|
||||
3. Mental models (consolidated knowledge generation)
|
||||
"""
|
||||
import os
|
||||
from datetime import datetime
|
||||
@@ -19,6 +24,10 @@ MODEL_MATRIX = [
|
||||
("openai", "gpt-5-nano"),
|
||||
("openai", "gpt-5"),
|
||||
("openai", "gpt-5.2"),
|
||||
# Anthropic models
|
||||
("anthropic", "claude-sonnet-4-20250514"),
|
||||
("anthropic", "claude-opus-4-5-20251101"),
|
||||
("anthropic", "claude-haiku-4-20250514"),
|
||||
# Groq models
|
||||
("groq", "openai/gpt-oss-120b"),
|
||||
("groq", "openai/gpt-oss-20b"),
|
||||
@@ -29,6 +38,12 @@ MODEL_MATRIX = [
|
||||
# Ollama models (local)
|
||||
("ollama", "gemma3:12b"),
|
||||
("ollama", "gemma3:1b"),
|
||||
# Claude Code (uses Claude Agent SDK with Claude models)
|
||||
("claude-code", "claude-sonnet-4-20250514"),
|
||||
# OpenAI Codex (uses MCP with Codex-specific models)
|
||||
("openai-codex", "gpt-5.2-codex"),
|
||||
# Mock provider (for testing)
|
||||
("mock", "mock"),
|
||||
]
|
||||
|
||||
|
||||
@@ -36,6 +51,7 @@ def get_api_key_for_provider(provider: str) -> str | None:
|
||||
"""Get API key for provider from environment variables."""
|
||||
provider_key_map = {
|
||||
"openai": "OPENAI_API_KEY",
|
||||
"anthropic": "ANTHROPIC_API_KEY",
|
||||
"groq": "GROQ_API_KEY",
|
||||
"gemini": "GEMINI_API_KEY",
|
||||
}
|
||||
@@ -43,6 +59,165 @@ def get_api_key_for_provider(provider: str) -> str | None:
|
||||
return os.getenv(env_var) if env_var else None
|
||||
|
||||
|
||||
def should_skip_provider(provider: str, model: str = "") -> tuple[bool, str]:
|
||||
"""Check if provider should be skipped and return reason."""
|
||||
# Never skip mock provider
|
||||
if provider == "mock":
|
||||
return False, ""
|
||||
|
||||
# Skip claude-code and openai-codex in CI (require local auth)
|
||||
if os.getenv("CI") and provider in ("claude-code", "openai-codex"):
|
||||
return True, f"{provider} not available in CI (requires local authentication)"
|
||||
|
||||
# Skip Ollama in CI (no models available)
|
||||
if provider == "ollama" and os.getenv("CI"):
|
||||
return True, "Ollama not available in CI"
|
||||
|
||||
# Skip Ollama gemma models (don't support tool calling)
|
||||
if provider == "ollama" and "gemma" in model.lower():
|
||||
return True, f"Ollama {model} does not support tool calling"
|
||||
|
||||
# Other providers need an API key
|
||||
if provider not in ("ollama", "claude-code", "openai-codex", "mock"):
|
||||
api_key = get_api_key_for_provider(provider)
|
||||
if not api_key:
|
||||
return True, f"No API key available (set {provider.upper()}_API_KEY)"
|
||||
|
||||
return False, ""
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider,model", MODEL_MATRIX)
|
||||
@pytest.mark.asyncio
|
||||
async def test_llm_provider_api_methods(provider: str, model: str):
|
||||
"""
|
||||
Test all LLM API methods used by Hindsight at runtime.
|
||||
This validates that the provider correctly implements the LLMInterface.
|
||||
|
||||
Tests:
|
||||
1. verify_connection() - Connection verification
|
||||
2. call() with plain text - Basic LLM call
|
||||
3. call() with response_format - Structured output (used in fact extraction)
|
||||
4. call_with_tools() - Tool calling (used in reflect agent)
|
||||
"""
|
||||
# Skip mock provider - it's a test stub, not a real LLM implementation
|
||||
if provider == "mock":
|
||||
pytest.skip("Mock provider is a test stub, not a real LLM")
|
||||
|
||||
should_skip, reason = should_skip_provider(provider, model)
|
||||
if should_skip:
|
||||
pytest.skip(f"Skipping {provider}/{model}: {reason}")
|
||||
|
||||
api_key = get_api_key_for_provider(provider)
|
||||
|
||||
llm = LLMProvider(
|
||||
provider=provider,
|
||||
api_key=api_key or "",
|
||||
base_url="",
|
||||
model=model,
|
||||
)
|
||||
|
||||
print(f"\n{provider}/{model} - API methods test:")
|
||||
|
||||
# Test 1: verify_connection()
|
||||
try:
|
||||
await llm.verify_connection()
|
||||
print(" ✓ verify_connection()")
|
||||
except Exception as e:
|
||||
pytest.fail(f"{provider}/{model} verify_connection() failed: {e}")
|
||||
|
||||
# Test 2: call() with plain text
|
||||
try:
|
||||
response = await llm.call(
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "What is 2+2? Answer in one word."},
|
||||
],
|
||||
max_completion_tokens=50,
|
||||
)
|
||||
assert response is not None, "call() returned None"
|
||||
assert len(response) > 0, "call() returned empty string"
|
||||
print(f" ✓ call() plain text: {response[:50]}")
|
||||
except Exception as e:
|
||||
pytest.fail(f"{provider}/{model} call() plain text failed: {e}")
|
||||
|
||||
# Test 3: call() with response_format (structured output)
|
||||
try:
|
||||
from pydantic import BaseModel
|
||||
|
||||
class TestResponse(BaseModel):
|
||||
answer: str
|
||||
confidence: str
|
||||
|
||||
response = await llm.call(
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a math assistant."},
|
||||
{"role": "user", "content": "What is the capital of France?"},
|
||||
],
|
||||
response_format=TestResponse,
|
||||
max_completion_tokens=100,
|
||||
)
|
||||
assert isinstance(response, TestResponse), f"Expected TestResponse, got {type(response)}"
|
||||
assert hasattr(response, "answer"), "Structured output missing 'answer' field"
|
||||
assert hasattr(response, "confidence"), "Structured output missing 'confidence' field"
|
||||
print(f" ✓ call() structured output: answer={response.answer}, confidence={response.confidence}")
|
||||
except Exception as e:
|
||||
pytest.fail(f"{provider}/{model} call() structured output failed: {e}")
|
||||
|
||||
# Test 4: call_with_tools() (tool calling)
|
||||
try:
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather for a location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {"type": "string", "description": "City name"},
|
||||
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
result = await llm.call_with_tools(
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant with access to tools."},
|
||||
{"role": "user", "content": "What's the weather like in Paris?"},
|
||||
],
|
||||
tools=tools,
|
||||
max_completion_tokens=200,
|
||||
)
|
||||
|
||||
assert result is not None, "call_with_tools() returned None"
|
||||
assert hasattr(result, "tool_calls"), "Result missing 'tool_calls' attribute"
|
||||
|
||||
# Nano models may hit token limits before making tool calls - that's acceptable
|
||||
is_nano_model = "nano" in model.lower()
|
||||
if is_nano_model and len(result.tool_calls) == 0:
|
||||
# Check if it hit length limit (expected for nano models)
|
||||
if hasattr(result, "finish_reason") and result.finish_reason == "length":
|
||||
print(f" ✓ call_with_tools(): nano model hit token limit (expected)")
|
||||
else:
|
||||
pytest.fail(f"Nano model made 0 tool calls but didn't hit length limit (finish_reason={getattr(result, 'finish_reason', 'unknown')})")
|
||||
else:
|
||||
assert len(result.tool_calls) > 0, f"Expected at least 1 tool call, got {len(result.tool_calls)}"
|
||||
|
||||
# Verify tool call structure
|
||||
tool_call = result.tool_calls[0]
|
||||
assert hasattr(tool_call, "name"), "Tool call missing 'name'"
|
||||
assert hasattr(tool_call, "arguments"), "Tool call missing 'arguments'"
|
||||
assert tool_call.name == "get_weather", f"Expected 'get_weather', got '{tool_call.name}'"
|
||||
assert "location" in tool_call.arguments, "Tool call arguments missing 'location'"
|
||||
|
||||
print(f" ✓ call_with_tools(): {tool_call.name}({tool_call.arguments})")
|
||||
except Exception as e:
|
||||
pytest.fail(f"{provider}/{model} call_with_tools() failed: {e}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider,model", MODEL_MATRIX)
|
||||
@pytest.mark.asyncio
|
||||
async def test_llm_provider_memory_operations(provider: str, model: str):
|
||||
@@ -50,16 +225,16 @@ async def test_llm_provider_memory_operations(provider: str, model: str):
|
||||
Test LLM provider with actual memory operations: fact extraction and reflect.
|
||||
All models must pass this test.
|
||||
"""
|
||||
# Skip mock provider - it's a test stub, not designed for real operations
|
||||
if provider == "mock":
|
||||
pytest.skip("Mock provider is a test stub, not designed for real operations")
|
||||
|
||||
should_skip, reason = should_skip_provider(provider, model)
|
||||
if should_skip:
|
||||
pytest.skip(f"Skipping {provider}/{model}: {reason}")
|
||||
|
||||
api_key = get_api_key_for_provider(provider)
|
||||
|
||||
# Skip Ollama tests in CI (no models available)
|
||||
if provider == "ollama" and os.getenv("CI"):
|
||||
pytest.skip(f"Skipping {provider}/{model}: Ollama not available in CI")
|
||||
|
||||
# Other providers need an API key
|
||||
if provider != "ollama" and not api_key:
|
||||
pytest.skip(f"Skipping {provider}/{model}: no API key available")
|
||||
|
||||
llm = LLMProvider(
|
||||
provider=provider,
|
||||
api_key=api_key or "",
|
||||
@@ -117,3 +292,115 @@ async def test_llm_provider_memory_operations(provider: str, model: str):
|
||||
|
||||
assert response is not None, f"{provider}/{model} reflect returned None"
|
||||
assert len(response) > 10, f"{provider}/{model} reflect response too short"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider,model", [
|
||||
("claude-code", "claude-sonnet-4-20250514"),
|
||||
("openai-codex", "gpt-5.2-codex"),
|
||||
])
|
||||
@pytest.mark.asyncio
|
||||
async def test_llm_provider_consolidation(memory_no_llm_verify, request_context, provider: str, model: str):
|
||||
"""
|
||||
Test LLM provider with consolidation (automatic mental model generation from observations).
|
||||
This validates that the provider can generate synthesized knowledge from raw memories.
|
||||
|
||||
This test is limited to claude-code and codex since they're the critical providers
|
||||
that needed tool calling fixes for reflect and consolidation operations.
|
||||
"""
|
||||
should_skip, reason = should_skip_provider(provider, model)
|
||||
if should_skip:
|
||||
pytest.skip(f"Skipping {provider}/{model}: {reason}")
|
||||
|
||||
# Use provider-specific LLM for this test
|
||||
api_key = get_api_key_for_provider(provider)
|
||||
memory_no_llm_verify._consolidation_llm = LLMProvider(
|
||||
provider=provider,
|
||||
api_key=api_key or "",
|
||||
base_url="",
|
||||
model=model,
|
||||
)
|
||||
# Also need retain LLM for ingesting data
|
||||
memory_no_llm_verify._retain_llm = memory_no_llm_verify._consolidation_llm
|
||||
|
||||
test_bank_id = f"llm_test_consolidation_{provider}_{model}_{datetime.now().timestamp()}"
|
||||
|
||||
# Enable observations for this bank
|
||||
from hindsight_api.config import get_config
|
||||
config = get_config()
|
||||
original_value = config.enable_observations
|
||||
config.enable_observations = True
|
||||
|
||||
try:
|
||||
# Retain memories to consolidate
|
||||
test_content = """
|
||||
Bob prefers functional programming with Rust and Haskell.
|
||||
He emphasizes immutability and pure functions in code reviews.
|
||||
Bob advocates for type safety and compile-time guarantees.
|
||||
He avoids mutable state and prefers declarative code patterns.
|
||||
"""
|
||||
|
||||
await memory_no_llm_verify.retain_async(
|
||||
bank_id=test_bank_id,
|
||||
content=test_content,
|
||||
context="Team coding preferences",
|
||||
event_date=datetime(2024, 12, 1),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n{provider}/{model} - Consolidation test:")
|
||||
|
||||
# Run consolidation to generate observations (mental models)
|
||||
from hindsight_api.engine.consolidation.consolidator import run_consolidation_job
|
||||
|
||||
result = await run_consolidation_job(
|
||||
memory_engine=memory_no_llm_verify,
|
||||
bank_id=test_bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f" Processed: {result.get('memories_processed', 0)} memories")
|
||||
print(f" Created: {result.get('observations_created', 0)} observations")
|
||||
print(f" Updated: {result.get('observations_updated', 0)} observations")
|
||||
|
||||
# Verify consolidation ran successfully
|
||||
assert result["status"] in ["success", "no_new_memories"], f"{provider}/{model} consolidation failed"
|
||||
|
||||
# If observations were created, verify they contain relevant content
|
||||
if result.get("observations_created", 0) > 0:
|
||||
observations = await memory_no_llm_verify.list_mental_models_consolidated(
|
||||
bank_id=test_bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(observations) > 0, f"{provider}/{model} consolidation created 0 observations"
|
||||
|
||||
# Check first observation contains relevant information
|
||||
obs_content = observations[0].get("content", "").lower()
|
||||
relevant_terms = ["bob", "functional", "rust", "immutab", "type"]
|
||||
matches = [term for term in relevant_terms if term in obs_content]
|
||||
|
||||
print(f" Observation preview: {observations[0].get('content', '')[:200]}...")
|
||||
print(f" Found {len(matches)} relevant terms: {matches}")
|
||||
|
||||
assert len(matches) >= 2, (
|
||||
f"{provider}/{model} consolidated observation doesn't contain relevant info. "
|
||||
f"Expected at least 2 of {relevant_terms}, found {len(matches)}: {matches}"
|
||||
)
|
||||
|
||||
finally:
|
||||
# Restore original config
|
||||
config.enable_observations = original_value
|
||||
|
||||
|
||||
# NOTE: The tests above validate the critical Hindsight operations:
|
||||
#
|
||||
# test_llm_provider_memory_operations (ALL providers):
|
||||
# - Fact extraction (retain): tests structured output generation
|
||||
# - Reflect: tests memory retrieval and reasoning (uses tool calling for claude-code/codex)
|
||||
#
|
||||
# test_llm_provider_consolidation (claude-code and codex only):
|
||||
# - Consolidation: tests automatic mental model generation from observations
|
||||
# - Requires MemoryEngine fixture with working LLM (from .env or env vars)
|
||||
# - Run your local LLM server OR set HINDSIGHT_API_LLM_PROVIDER/API_KEY/MODEL env vars
|
||||
#
|
||||
# For full end-to-end integration tests using the HTTP API, see tests/test_http_api_integration.py
|
||||
|
||||
@@ -30,7 +30,8 @@ async def test_llm_metrics_recorded_for_groq():
|
||||
# Create a mock metrics collector to track record_llm_call calls
|
||||
mock_collector = MagicMock(spec=MetricsCollector)
|
||||
|
||||
with patch("hindsight_api.engine.llm_wrapper.get_metrics_collector", return_value=mock_collector):
|
||||
# Patch the provider module where get_metrics_collector is actually called
|
||||
with patch("hindsight_api.engine.providers.openai_compatible_llm.get_metrics_collector", return_value=mock_collector):
|
||||
llm = LLMProvider(
|
||||
provider="groq",
|
||||
api_key=api_key,
|
||||
@@ -90,7 +91,8 @@ async def test_llm_metrics_recorded_for_structured_output():
|
||||
|
||||
mock_collector = MagicMock(spec=MetricsCollector)
|
||||
|
||||
with patch("hindsight_api.engine.llm_wrapper.get_metrics_collector", return_value=mock_collector):
|
||||
# Patch the provider module where get_metrics_collector is actually called
|
||||
with patch("hindsight_api.engine.providers.openai_compatible_llm.get_metrics_collector", return_value=mock_collector):
|
||||
llm = LLMProvider(
|
||||
provider="groq",
|
||||
api_key=api_key,
|
||||
|
||||
@@ -97,3 +97,47 @@ def test_path_parsing_logic():
|
||||
bank_id, remaining = parse_path("/my-bank/some/path")
|
||||
assert bank_id == "my-bank"
|
||||
assert remaining == "/some/path"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_key_context_variable():
|
||||
"""Test that API key context variable works correctly."""
|
||||
from hindsight_api.api.mcp import get_current_api_key, _current_api_key
|
||||
|
||||
# Initially None
|
||||
assert get_current_api_key() is None
|
||||
|
||||
# Set and verify
|
||||
token = _current_api_key.set("test-api-key-123")
|
||||
try:
|
||||
assert get_current_api_key() == "test-api-key-123"
|
||||
finally:
|
||||
_current_api_key.reset(token)
|
||||
|
||||
# Back to None after reset
|
||||
assert get_current_api_key() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_tools_propagate_api_key(mock_memory):
|
||||
"""Test that MCP tools propagate API key to RequestContext."""
|
||||
from hindsight_api.api.mcp import create_mcp_server, _current_bank_id, _current_api_key
|
||||
|
||||
mcp_server = create_mcp_server(mock_memory)
|
||||
tools = mcp_server._tool_manager._tools
|
||||
|
||||
# Set both bank_id and api_key context
|
||||
bank_token = _current_bank_id.set("test-bank")
|
||||
api_key_token = _current_api_key.set("test-bearer-token")
|
||||
try:
|
||||
retain_tool = tools["retain"]
|
||||
result = await retain_tool.fn(content="test content", context="test_context", async_processing=False)
|
||||
assert "successfully" in result.lower()
|
||||
|
||||
# Verify the memory was called with request_context containing api_key
|
||||
mock_memory.retain_batch_async.assert_called_once()
|
||||
call_kwargs = mock_memory.retain_batch_async.call_args.kwargs
|
||||
assert call_kwargs["request_context"].api_key == "test-bearer-token"
|
||||
finally:
|
||||
_current_bank_id.reset(bank_token)
|
||||
_current_api_key.reset(api_key_token)
|
||||
|
||||
@@ -0,0 +1,206 @@
|
||||
"""Unit tests for mental model operation validator hooks.
|
||||
|
||||
Tests that the operation validator hooks are called correctly for
|
||||
mental model GET and refresh operations.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.extensions.operation_validator import (
|
||||
MentalModelGetContext,
|
||||
MentalModelGetResult,
|
||||
MentalModelRefreshResult,
|
||||
OperationValidatorExtension,
|
||||
ValidationResult,
|
||||
)
|
||||
|
||||
|
||||
class TestMentalModelGetContextDataclass:
|
||||
"""Tests for MentalModelGetContext dataclass."""
|
||||
|
||||
def test_create_context(self):
|
||||
"""Test creating a MentalModelGetContext."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
request_context = MagicMock()
|
||||
ctx = MentalModelGetContext(
|
||||
bank_id="bank-1",
|
||||
mental_model_id="mm-1",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert ctx.bank_id == "bank-1"
|
||||
assert ctx.mental_model_id == "mm-1"
|
||||
assert ctx.request_context is request_context
|
||||
|
||||
|
||||
class TestMentalModelGetResultDataclass:
|
||||
"""Tests for MentalModelGetResult dataclass."""
|
||||
|
||||
def test_create_result_success(self):
|
||||
"""Test creating a successful MentalModelGetResult."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
request_context = MagicMock()
|
||||
result = MentalModelGetResult(
|
||||
bank_id="bank-1",
|
||||
mental_model_id="mm-1",
|
||||
request_context=request_context,
|
||||
output_tokens=250,
|
||||
)
|
||||
|
||||
assert result.bank_id == "bank-1"
|
||||
assert result.mental_model_id == "mm-1"
|
||||
assert result.output_tokens == 250
|
||||
assert result.success is True
|
||||
assert result.error is None
|
||||
|
||||
def test_create_result_failure(self):
|
||||
"""Test creating a failed MentalModelGetResult."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
result = MentalModelGetResult(
|
||||
bank_id="bank-1",
|
||||
mental_model_id="mm-1",
|
||||
request_context=MagicMock(),
|
||||
output_tokens=0,
|
||||
success=False,
|
||||
error="Not found",
|
||||
)
|
||||
|
||||
assert result.success is False
|
||||
assert result.error == "Not found"
|
||||
|
||||
|
||||
class TestMentalModelRefreshResultDataclass:
|
||||
"""Tests for MentalModelRefreshResult dataclass."""
|
||||
|
||||
def test_create_result_with_all_fields(self):
|
||||
"""Test creating a MentalModelRefreshResult with all fields."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
result = MentalModelRefreshResult(
|
||||
bank_id="bank-1",
|
||||
mental_model_id="mm-1",
|
||||
request_context=MagicMock(),
|
||||
query_tokens=50,
|
||||
output_tokens=500,
|
||||
context_tokens=0,
|
||||
facts_used=10,
|
||||
mental_models_used=2,
|
||||
)
|
||||
|
||||
assert result.query_tokens == 50
|
||||
assert result.output_tokens == 500
|
||||
assert result.context_tokens == 0
|
||||
assert result.facts_used == 10
|
||||
assert result.mental_models_used == 2
|
||||
assert result.success is True
|
||||
assert result.error is None
|
||||
|
||||
def test_create_result_failure(self):
|
||||
"""Test creating a failed MentalModelRefreshResult."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
result = MentalModelRefreshResult(
|
||||
bank_id="bank-1",
|
||||
mental_model_id="mm-1",
|
||||
request_context=MagicMock(),
|
||||
query_tokens=50,
|
||||
output_tokens=0,
|
||||
context_tokens=0,
|
||||
facts_used=0,
|
||||
mental_models_used=0,
|
||||
success=False,
|
||||
error="Reflect failed",
|
||||
)
|
||||
|
||||
assert result.success is False
|
||||
assert result.error == "Reflect failed"
|
||||
|
||||
|
||||
class TestDefaultHookBehavior:
|
||||
"""Tests for default (no-op) behavior of mental model hooks on OperationValidatorExtension."""
|
||||
|
||||
@pytest.fixture
|
||||
def validator(self):
|
||||
"""Create a concrete subclass for testing default behavior."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Create a concrete subclass that implements the abstract methods
|
||||
class TestValidator(OperationValidatorExtension):
|
||||
async def validate_retain(self, ctx):
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_recall(self, ctx):
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_reflect(self, ctx):
|
||||
return ValidationResult.accept()
|
||||
|
||||
return TestValidator(config={})
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_mental_model_get_default_accepts(self, validator):
|
||||
"""Test that default validate_mental_model_get accepts."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
ctx = MentalModelGetContext(
|
||||
bank_id="bank-1",
|
||||
mental_model_id="mm-1",
|
||||
request_context=MagicMock(),
|
||||
)
|
||||
|
||||
result = await validator.validate_mental_model_get(ctx)
|
||||
|
||||
assert result.allowed is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_mental_model_get_complete_default_noop(self, validator):
|
||||
"""Test that default on_mental_model_get_complete is a no-op."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
result = MentalModelGetResult(
|
||||
bank_id="bank-1",
|
||||
mental_model_id="mm-1",
|
||||
request_context=MagicMock(),
|
||||
output_tokens=100,
|
||||
)
|
||||
|
||||
# Should not raise
|
||||
await validator.on_mental_model_get_complete(result)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_mental_model_refresh_complete_default_noop(self, validator):
|
||||
"""Test that default on_mental_model_refresh_complete is a no-op."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
result = MentalModelRefreshResult(
|
||||
bank_id="bank-1",
|
||||
mental_model_id="mm-1",
|
||||
request_context=MagicMock(),
|
||||
query_tokens=50,
|
||||
output_tokens=500,
|
||||
context_tokens=0,
|
||||
facts_used=5,
|
||||
mental_models_used=1,
|
||||
)
|
||||
|
||||
# Should not raise
|
||||
await validator.on_mental_model_refresh_complete(result)
|
||||
|
||||
|
||||
class TestExportsAvailable:
|
||||
"""Test that mental model hooks are properly exported."""
|
||||
|
||||
def test_imports_from_extensions_package(self):
|
||||
"""Test that all mental model types can be imported from hindsight_api.extensions."""
|
||||
from hindsight_api.extensions import (
|
||||
MentalModelGetContext,
|
||||
MentalModelGetResult,
|
||||
MentalModelRefreshResult,
|
||||
)
|
||||
|
||||
assert MentalModelGetContext is not None
|
||||
assert MentalModelGetResult is not None
|
||||
assert MentalModelRefreshResult is not None
|
||||
@@ -312,6 +312,49 @@ class TestDirectiveTags:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_list_all_directives_without_filter(self, memory: MemoryEngine, request_context):
|
||||
"""Test that listing directives without tags returns ALL directives (both tagged and untagged)."""
|
||||
bank_id = f"test-directive-list-all-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Ensure bank exists
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Create untagged directive
|
||||
await memory.create_directive(
|
||||
bank_id=bank_id,
|
||||
name="Untagged Directive",
|
||||
content="This has no tags",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Create tagged directive
|
||||
await memory.create_directive(
|
||||
bank_id=bank_id,
|
||||
name="Tagged Directive",
|
||||
content="This has tags",
|
||||
tags=["project-x"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# List ALL directives (no tag filter, isolation_mode defaults to False)
|
||||
all_directives = await memory.list_directives(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should return BOTH tagged and untagged directives
|
||||
assert len(all_directives) == 2
|
||||
directive_names = {d["name"] for d in all_directives}
|
||||
assert "Untagged Directive" in directive_names
|
||||
assert "Tagged Directive" in directive_names
|
||||
|
||||
# Verify the tagged directive has its tags
|
||||
tagged = next(d for d in all_directives if d["name"] == "Tagged Directive")
|
||||
assert tagged["tags"] == ["project-x"]
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
class TestReflect:
|
||||
"""Test reflect endpoint."""
|
||||
@@ -399,6 +442,161 @@ class TestDirectivesInReflect:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_tagged_directive_not_applied_without_tags(self, memory: MemoryEngine, request_context):
|
||||
"""Test that directives with tags are NOT applied to untagged reflect operations."""
|
||||
bank_id = f"test-directive-isolation-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Ensure bank exists
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Add some untagged content
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{"content": "The sky is blue."},
|
||||
{"content": "Water is wet."},
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Add some tagged content for the project-x context
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{"content": "The sky is blue according to project X standards.", "tags": ["project-x"]},
|
||||
{"content": "Project X color guidelines specify sky is blue.", "tags": ["project-x"]},
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Create an untagged directive (should be applied)
|
||||
await memory.create_directive(
|
||||
bank_id=bank_id,
|
||||
name="General Policy",
|
||||
content="Always be polite and start responses with 'Hello!'",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Create a tagged directive (should NOT be applied to untagged reflect)
|
||||
await memory.create_directive(
|
||||
bank_id=bank_id,
|
||||
name="Tagged Policy",
|
||||
content="ALWAYS respond in ALL CAPS and end with 'PROJECT-X ONLY'",
|
||||
tags=["project-x"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Run reflect without tags - should only apply the untagged directive
|
||||
result = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="What color is the sky?",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
response_lower = result.text.lower()
|
||||
|
||||
# Should follow the untagged directive (polite greeting)
|
||||
assert "hello" in response_lower, f"Expected 'Hello' from untagged directive, but got: {result.text}"
|
||||
|
||||
# Should NOT follow the tagged directive (all caps and PROJECT-X)
|
||||
# If it did follow, the entire response would be in caps
|
||||
all_caps = result.text.replace(" ", "").replace("!", "").replace(".", "").isupper()
|
||||
assert not all_caps, f"Tagged directive was incorrectly applied to untagged operation: {result.text}"
|
||||
assert "project-x only" not in response_lower, f"Tagged directive was incorrectly applied: {result.text}"
|
||||
|
||||
# Now run reflect WITH the tag - should apply BOTH directives
|
||||
result_tagged = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="What color is the sky?",
|
||||
tags=["project-x"],
|
||||
tags_match="all_strict",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
response_tagged_lower = result_tagged.text.lower()
|
||||
|
||||
# With strict matching and tags, should apply the tagged directive
|
||||
assert "project-x only" in response_tagged_lower, f"Tagged directive should be applied with tags: {result_tagged.text}"
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_reflect_based_on_structure(self, memory: MemoryEngine, request_context):
|
||||
"""Test that reflect returns correct based_on structure with directives and memories separated."""
|
||||
bank_id = f"test-reflect-based-on-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Ensure bank exists
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Add some memories
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{"content": "Alice works at Google as a software engineer."},
|
||||
{"content": "Bob is a product manager at Microsoft."},
|
||||
{"content": "The team meets every Monday at 9am."},
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Create a directive
|
||||
directive = await memory.create_directive(
|
||||
bank_id=bank_id,
|
||||
name="Professional Tone",
|
||||
content="Always maintain a professional and formal tone in responses.",
|
||||
request_context=request_context,
|
||||
)
|
||||
directive_id = directive["id"]
|
||||
|
||||
# Run reflect which returns the core result
|
||||
result = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="Who works at Google?",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify based_on structure exists
|
||||
assert result.based_on is not None
|
||||
|
||||
# Verify directives key exists and contains our directive
|
||||
assert "directives" in result.based_on
|
||||
directives_list = result.based_on.get("directives", [])
|
||||
|
||||
# Verify directives are dicts with id, name, content (not MemoryFact objects)
|
||||
assert len(directives_list) > 0, "Should have at least one directive"
|
||||
directive_found = False
|
||||
for d in directives_list:
|
||||
assert isinstance(d, dict), f"Directive should be dict, got {type(d)}"
|
||||
assert "id" in d, "Directive dict should have 'id'"
|
||||
assert "name" in d, "Directive dict should have 'name'"
|
||||
assert "content" in d, "Directive dict should have 'content'"
|
||||
# Check if this is our directive
|
||||
if d["id"] == directive_id:
|
||||
directive_found = True
|
||||
assert d["name"] == "Professional Tone"
|
||||
assert "professional" in d["content"].lower()
|
||||
|
||||
assert directive_found, f"Our directive {directive_id} should be in based_on.directives"
|
||||
|
||||
# Verify memories (world/experience) are separate from directives
|
||||
has_memories = "world" in result.based_on or "experience" in result.based_on
|
||||
assert has_memories, "Should have world or experience memories"
|
||||
|
||||
# Verify that if mental-models key exists, it's separate from directives
|
||||
if "mental-models" in result.based_on:
|
||||
mental_models = result.based_on.get("mental-models", [])
|
||||
# Verify mental models are MemoryFact objects, not dicts like directives
|
||||
for mm in mental_models:
|
||||
assert hasattr(mm, "fact_type"), "Mental model should be MemoryFact with fact_type"
|
||||
assert mm.fact_type == "mental-models"
|
||||
assert hasattr(mm, "context")
|
||||
assert "mental model" in mm.context.lower()
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
class TestDirectivesPromptInjection:
|
||||
"""Test that directives are properly injected into the system prompt."""
|
||||
@@ -451,3 +649,198 @@ class TestDirectivesPromptInjection:
|
||||
directives_pos = prompt.find("## DIRECTIVES")
|
||||
critical_rules_pos = prompt.find("## CRITICAL RULES")
|
||||
assert directives_pos < critical_rules_pos
|
||||
|
||||
|
||||
class TestMentalModelRefreshTagSecurity:
|
||||
"""Test that mental model refresh respects tag-based security boundaries."""
|
||||
|
||||
async def test_refresh_with_tags_only_accesses_same_tagged_models(
|
||||
self, memory: MemoryEngine, request_context
|
||||
):
|
||||
"""Test that refreshing a mental model with tags can only access other models with the same tags.
|
||||
|
||||
This is a security test to ensure that mental models with tags (e.g., user:alice)
|
||||
cannot access mental models from other scopes (e.g., user:bob or no tags) during refresh.
|
||||
"""
|
||||
bank_id = f"test-refresh-tags-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Ensure bank exists
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Add some facts with different tags
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{"content": "Alice works on the frontend React project. Alice's favorite color is blue.", "tags": ["user:alice"]},
|
||||
{"content": "Alice prefers working in the morning. Alice drinks coffee every day.", "tags": ["user:alice"]},
|
||||
{"content": "Bob works on the backend API services. Bob's favorite language is Python.", "tags": ["user:bob"]},
|
||||
{"content": "Bob prefers working at night. Bob drinks tea every day.", "tags": ["user:bob"]},
|
||||
{"content": "The company has 100 employees and is growing fast.", "tags": []}, # No tags
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Wait for background processing
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Create mental model for user:alice with sensitive data
|
||||
mm_alice = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Alice's Work Profile",
|
||||
source_query="What does Alice work on?",
|
||||
content="Alice is a frontend engineer specializing in React",
|
||||
tags=["user:alice"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Create mental model for user:bob with sensitive data
|
||||
mm_bob = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Bob's Work Profile",
|
||||
source_query="What does Bob work on?",
|
||||
content="Bob is a backend engineer specializing in Python",
|
||||
tags=["user:bob"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Create mental model with no tags (should not be accessible from tagged models)
|
||||
mm_untagged = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Company Info",
|
||||
source_query="What is the company info?",
|
||||
content="The company has 100 employees",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Create a mental model for user:alice that will be refreshed
|
||||
mm_alice_refresh = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Alice's Summary",
|
||||
source_query="What are all the facts about work and preferences?", # Broad query that should match all facts
|
||||
content="Initial content",
|
||||
tags=["user:alice"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Refresh Alice's mental model
|
||||
refreshed = await memory.refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mm_alice_refresh["id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# SECURITY CHECK: The refreshed content should ONLY include information from
|
||||
# memories/models tagged with user:alice, NOT from user:bob or untagged
|
||||
refreshed_content = refreshed["content"].lower()
|
||||
|
||||
# Should include Alice's content (either from facts or mental models)
|
||||
assert "alice" in refreshed_content, \
|
||||
"Refreshed model should access memories/models with matching tags (user:alice)"
|
||||
|
||||
# MUST NOT include Bob's content (security violation)
|
||||
assert "bob" not in refreshed_content and "python" not in refreshed_content and "tea" not in refreshed_content, \
|
||||
f"SECURITY VIOLATION: Refreshed model accessed memories/models with different tags (user:bob). Content: {refreshed_content}"
|
||||
|
||||
# MUST NOT include untagged content (security violation)
|
||||
assert "100 employees" not in refreshed_content and "growing fast" not in refreshed_content, \
|
||||
f"SECURITY VIOLATION: Refreshed model accessed untagged memories/models. Content: {refreshed_content}"
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_consolidation_only_refreshes_matching_tagged_models(
|
||||
self, memory: MemoryEngine, request_context
|
||||
):
|
||||
"""Test that consolidation only triggers refresh for mental models with matching tags.
|
||||
|
||||
This is a security test to ensure that when tagged memories are consolidated,
|
||||
only mental models with overlapping tags get refreshed, not all mental models.
|
||||
"""
|
||||
bank_id = f"test-consolidation-refresh-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Ensure bank exists
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Create mental models with different tags, all with refresh_after_consolidation=true
|
||||
mm_alice = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Alice's Model",
|
||||
source_query="What about Alice?",
|
||||
content="Initial Alice content",
|
||||
tags=["user:alice"],
|
||||
trigger={"refresh_after_consolidation": True},
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
mm_bob = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Bob's Model",
|
||||
source_query="What about Bob?",
|
||||
content="Initial Bob content",
|
||||
tags=["user:bob"],
|
||||
trigger={"refresh_after_consolidation": True},
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
mm_untagged = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Untagged Model",
|
||||
source_query="What about general stuff?",
|
||||
content="Initial untagged content",
|
||||
trigger={"refresh_after_consolidation": True},
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Record initial last_refreshed_at timestamps
|
||||
alice_initial = mm_alice["last_refreshed_at"]
|
||||
bob_initial = mm_bob["last_refreshed_at"]
|
||||
untagged_initial = mm_untagged["last_refreshed_at"]
|
||||
|
||||
# Add memories with user:alice tags
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{"content": "Alice likes React", "tags": ["user:alice"]},
|
||||
{"content": "Alice drinks coffee", "tags": ["user:alice"]},
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Trigger consolidation manually (this should only refresh Alice's mental model)
|
||||
from hindsight_api.engine.consolidation.consolidator import run_consolidation_job
|
||||
|
||||
result = await run_consolidation_job(
|
||||
memory_engine=memory,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Wait for background refresh tasks to complete
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check that mental models were refreshed appropriately
|
||||
mm_alice_after = await memory.get_mental_model(
|
||||
bank_id, mm_alice["id"], request_context=request_context
|
||||
)
|
||||
mm_bob_after = await memory.get_mental_model(
|
||||
bank_id, mm_bob["id"], request_context=request_context
|
||||
)
|
||||
mm_untagged_after = await memory.get_mental_model(
|
||||
bank_id, mm_untagged["id"], request_context=request_context
|
||||
)
|
||||
|
||||
# SECURITY CHECK: Only Alice's mental model and untagged model should be refreshed
|
||||
# Alice's model should be refreshed (tags match)
|
||||
assert mm_alice_after["last_refreshed_at"] != alice_initial or mm_alice_after["content"] != mm_alice["content"], \
|
||||
"Alice's mental model should be refreshed when user:alice memories are consolidated"
|
||||
|
||||
# Bob's model should NOT be refreshed (tags don't match)
|
||||
assert mm_bob_after["last_refreshed_at"] == bob_initial, \
|
||||
"SECURITY VIOLATION: Bob's mental model was refreshed even though user:bob memories were not consolidated"
|
||||
|
||||
# Untagged model should be refreshed (untagged models are always refreshed)
|
||||
assert mm_untagged_after["last_refreshed_at"] != untagged_initial or mm_untagged_after["content"] != mm_untagged["content"], \
|
||||
"Untagged mental model should be refreshed after any consolidation"
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -358,7 +358,7 @@ class TestLLMMetrics:
|
||||
collector.record_llm_call(
|
||||
provider="gemini",
|
||||
model="gemini-pro",
|
||||
scope="entity_observation",
|
||||
scope="memory",
|
||||
duration=2.0,
|
||||
success=True,
|
||||
)
|
||||
@@ -369,11 +369,11 @@ class TestLLMMetrics:
|
||||
assert call_args[0][0] == 1
|
||||
assert call_args[0][1]["provider"] == "gemini"
|
||||
assert call_args[0][1]["model"] == "gemini-pro"
|
||||
assert call_args[0][1]["scope"] == "entity_observation"
|
||||
assert call_args[0][1]["scope"] == "memory"
|
||||
|
||||
def test_record_llm_call_different_scopes(self, collector):
|
||||
"""Test recording LLM calls with different scopes."""
|
||||
scopes = ["memory", "reflect", "entity_observation", "answer"]
|
||||
scopes = ["memory", "reflect", "consolidation", "answer"]
|
||||
|
||||
for scope in scopes:
|
||||
collector.llm_duration.record.reset_mock()
|
||||
|
||||
@@ -469,7 +469,6 @@ async def test_mixed_language_entities(memory, request_context):
|
||||
budget=Budget.MID,
|
||||
max_tokens=1000,
|
||||
fact_type=["world"],
|
||||
include_entities=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
|
||||
@@ -91,156 +91,13 @@ async def test_entity_extraction_on_retain(memory, request_context):
|
||||
await conn.execute("DELETE FROM entities WHERE bank_id = $1", bank_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regenerate_entity_observations(memory, request_context):
|
||||
"""
|
||||
Test explicit regeneration of summary for an entity.
|
||||
"""
|
||||
bank_id = f"test_regen_obs_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Store facts about an entity
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Sarah is a product manager who loves user research and data analysis.",
|
||||
context="work info",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Find the Sarah entity
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
entity_row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT id, canonical_name
|
||||
FROM entities
|
||||
WHERE bank_id = $1 AND LOWER(canonical_name) LIKE '%sarah%'
|
||||
LIMIT 1
|
||||
""",
|
||||
bank_id
|
||||
)
|
||||
|
||||
if entity_row:
|
||||
entity_id = str(entity_row['id'])
|
||||
entity_name = entity_row['canonical_name']
|
||||
|
||||
# Manually regenerate summary (via observations API for backwards compat)
|
||||
created_ids = await memory.regenerate_entity_observations(
|
||||
bank_id=bank_id,
|
||||
entity_id=entity_id,
|
||||
entity_name=entity_name,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Regenerated Summary ===")
|
||||
print(f"Created {len(created_ids)} summary for {entity_name}")
|
||||
|
||||
# Get entity state
|
||||
state = await memory.get_entity_state(
|
||||
bank_id, entity_id, entity_name, request_context=request_context
|
||||
)
|
||||
for obs in state.observations:
|
||||
print(f" - {obs.text}")
|
||||
|
||||
# Verify summary was created
|
||||
if len(created_ids) > 0:
|
||||
assert len(state.observations) == 1, "Should have exactly 1 observation (the summary)"
|
||||
print(f"Summary regenerated successfully")
|
||||
else:
|
||||
print(f"Note: No summary was regenerated")
|
||||
|
||||
else:
|
||||
print(f"Note: No 'Sarah' entity was extracted")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM memory_units WHERE bank_id = $1", bank_id)
|
||||
await conn.execute("DELETE FROM entities WHERE bank_id = $1", bank_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_entity_state_retrieval(memory, request_context):
|
||||
"""
|
||||
Test retrieving entity state with facts.
|
||||
"""
|
||||
bank_id = f"test_entity_state_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Store facts
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google as a senior software engineer.",
|
||||
context="work info",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice loves hiking and outdoor photography.",
|
||||
context="hobbies",
|
||||
event_date=datetime(2024, 1, 16, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Find the Alice entity
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
entity_row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT id, canonical_name
|
||||
FROM entities
|
||||
WHERE bank_id = $1 AND LOWER(canonical_name) LIKE '%alice%'
|
||||
LIMIT 1
|
||||
""",
|
||||
bank_id
|
||||
)
|
||||
|
||||
assert entity_row is not None, "Alice entity should have been extracted"
|
||||
|
||||
entity_id = str(entity_row['id'])
|
||||
entity_name = entity_row['canonical_name']
|
||||
|
||||
# Check fact count
|
||||
async with pool.acquire() as conn:
|
||||
fact_count = await conn.fetchval(
|
||||
"SELECT COUNT(*) FROM unit_entities WHERE entity_id = $1",
|
||||
entity_row['id']
|
||||
)
|
||||
|
||||
print(f"\n=== Entity State Test ===")
|
||||
print(f"Entity: {entity_name} (id: {entity_id})")
|
||||
print(f"Linked facts: {fact_count}")
|
||||
|
||||
# Get entity state
|
||||
state = await memory.get_entity_state(
|
||||
bank_id, entity_id, entity_name, request_context=request_context
|
||||
)
|
||||
|
||||
assert state.entity_id == entity_id
|
||||
assert state.canonical_name == entity_name
|
||||
print(f"Entity state retrieved successfully")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM memory_units WHERE bank_id = $1", bank_id)
|
||||
await conn.execute("DELETE FROM entities WHERE bank_id = $1", bank_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_with_include_entities(memory, request_context):
|
||||
"""
|
||||
Test that search with include_entities=True returns entity information.
|
||||
Test that recall accepts include_entities parameter for backwards compatibility.
|
||||
|
||||
This test verifies that:
|
||||
1. Entities are extracted after retain
|
||||
2. Entity info is returned in recall results with include_entities=True
|
||||
Note: Entity observations have been deprecated. This test verifies the parameter
|
||||
is still accepted without errors.
|
||||
"""
|
||||
bank_id = f"test_search_ent_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -249,10 +106,6 @@ async def test_search_with_include_entities(memory, request_context):
|
||||
contents = [
|
||||
"Alice is a data scientist who works on recommendation systems at Netflix.",
|
||||
"Alice presented her research at the ML conference last month.",
|
||||
"Alice is an expert in deep learning and neural networks.",
|
||||
"Alice graduated from Stanford with a PhD in Computer Science.",
|
||||
"Alice leads a team of 5 data scientists at Netflix.",
|
||||
"Alice published a paper on collaborative filtering algorithms.",
|
||||
]
|
||||
|
||||
for i, content in enumerate(contents):
|
||||
@@ -267,7 +120,7 @@ async def test_search_with_include_entities(memory, request_context):
|
||||
# Wait for background tasks
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Search with include_entities=True
|
||||
# Search with include_entities=True (should be accepted for backwards compatibility)
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="What does Alice do?",
|
||||
@@ -279,98 +132,9 @@ async def test_search_with_include_entities(memory, request_context):
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Search Results ===")
|
||||
print(f"Found {len(result.results)} facts")
|
||||
for fact in result.results:
|
||||
print(f" - {fact.text}")
|
||||
if fact.entities:
|
||||
print(f" Entities: {', '.join(fact.entities)}")
|
||||
|
||||
# Verify results
|
||||
# Verify recall works
|
||||
assert len(result.results) > 0, "Should find some facts"
|
||||
|
||||
# Check if entities are included in facts
|
||||
facts_with_entities = [f for f in result.results if f.entities]
|
||||
assert len(facts_with_entities) > 0, "Some facts should have entity information"
|
||||
print(f"{len(facts_with_entities)} facts have entity information")
|
||||
|
||||
# Check if entity info is returned
|
||||
if result.entities:
|
||||
print(f"Entity info included for {len(result.entities)} entities")
|
||||
|
||||
# Verify Alice entity is in results
|
||||
alice_found = False
|
||||
for name, state in result.entities.items():
|
||||
assert state.canonical_name == name, "Entity canonical_name should match key"
|
||||
assert state.entity_id, "Entity should have an ID"
|
||||
if "alice" in name.lower():
|
||||
alice_found = True
|
||||
print(f"Alice entity found: {name}")
|
||||
|
||||
assert alice_found, "Alice entity should be in recall results"
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM memory_units WHERE bank_id = $1", bank_id)
|
||||
await conn.execute("DELETE FROM entities WHERE bank_id = $1", bank_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_entity_state(memory, request_context):
|
||||
"""
|
||||
Test getting the full state of an entity.
|
||||
"""
|
||||
bank_id = f"test_entity_state_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Store facts
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob is a frontend developer who specializes in React and TypeScript.",
|
||||
context="work info",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Find entity
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
entity_row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT id, canonical_name
|
||||
FROM entities
|
||||
WHERE bank_id = $1 AND LOWER(canonical_name) LIKE '%bob%'
|
||||
LIMIT 1
|
||||
""",
|
||||
bank_id
|
||||
)
|
||||
|
||||
if entity_row:
|
||||
entity_id = str(entity_row['id'])
|
||||
entity_name = entity_row['canonical_name']
|
||||
|
||||
# Get entity state
|
||||
state = await memory.get_entity_state(
|
||||
bank_id=bank_id,
|
||||
entity_id=entity_id,
|
||||
entity_name=entity_name,
|
||||
limit=10,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Entity State for {entity_name} ===")
|
||||
print(f"Entity ID: {state.entity_id}")
|
||||
print(f"Canonical Name: {state.canonical_name}")
|
||||
print(f"Observations: {len(state.observations)}")
|
||||
for obs in state.observations:
|
||||
print(f" - {obs.text}")
|
||||
|
||||
assert state.entity_id == entity_id, "Entity ID should match"
|
||||
assert state.canonical_name == entity_name, "Canonical name should match"
|
||||
print(f"Found {len(result.results)} facts")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
|
||||
@@ -275,3 +275,88 @@ class TestReflectUsesReflectLLMConfig:
|
||||
|
||||
# Verify it's different from the retain config
|
||||
assert engine._reflect_llm_config.model != engine._retain_llm_config.model
|
||||
|
||||
|
||||
class TestRetryAndBackoffConfiguration:
|
||||
"""Test retry and backoff configuration options."""
|
||||
|
||||
def test_global_retry_backoff_config_defaults(self):
|
||||
"""Test that global retry/backoff settings have correct defaults."""
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
config = get_config()
|
||||
|
||||
# Verify global defaults
|
||||
assert config.llm_max_retries == 10
|
||||
assert config.llm_initial_backoff == 1.0
|
||||
assert config.llm_max_backoff == 60.0
|
||||
|
||||
def test_per_operation_retry_backoff_config_from_env(self):
|
||||
"""Test that per-operation retry/backoff settings are loaded from environment."""
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
# Set per-operation overrides
|
||||
os.environ["HINDSIGHT_API_RETAIN_LLM_MAX_RETRIES"] = "3"
|
||||
os.environ["HINDSIGHT_API_RETAIN_LLM_INITIAL_BACKOFF"] = "2.0"
|
||||
os.environ["HINDSIGHT_API_RETAIN_LLM_MAX_BACKOFF"] = "120.0"
|
||||
os.environ["HINDSIGHT_API_REFLECT_LLM_MAX_RETRIES"] = "5"
|
||||
os.environ["HINDSIGHT_API_REFLECT_LLM_INITIAL_BACKOFF"] = "1.5"
|
||||
os.environ["HINDSIGHT_API_REFLECT_LLM_MAX_BACKOFF"] = "90.0"
|
||||
|
||||
try:
|
||||
clear_config_cache()
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
config = get_config()
|
||||
|
||||
# Verify retain overrides
|
||||
assert config.retain_llm_max_retries == 3
|
||||
assert config.retain_llm_initial_backoff == 2.0
|
||||
assert config.retain_llm_max_backoff == 120.0
|
||||
|
||||
# Verify reflect overrides
|
||||
assert config.reflect_llm_max_retries == 5
|
||||
assert config.reflect_llm_initial_backoff == 1.5
|
||||
assert config.reflect_llm_max_backoff == 90.0
|
||||
|
||||
# Verify global defaults remain unchanged
|
||||
assert config.llm_max_retries == 10
|
||||
assert config.llm_initial_backoff == 1.0
|
||||
assert config.llm_max_backoff == 60.0
|
||||
finally:
|
||||
# Clean up
|
||||
os.environ.pop("HINDSIGHT_API_RETAIN_LLM_MAX_RETRIES", None)
|
||||
os.environ.pop("HINDSIGHT_API_RETAIN_LLM_INITIAL_BACKOFF", None)
|
||||
os.environ.pop("HINDSIGHT_API_RETAIN_LLM_MAX_BACKOFF", None)
|
||||
os.environ.pop("HINDSIGHT_API_REFLECT_LLM_MAX_RETRIES", None)
|
||||
os.environ.pop("HINDSIGHT_API_REFLECT_LLM_INITIAL_BACKOFF", None)
|
||||
os.environ.pop("HINDSIGHT_API_REFLECT_LLM_MAX_BACKOFF", None)
|
||||
clear_config_cache()
|
||||
|
||||
def test_per_operation_retry_backoff_fallback_to_global(self):
|
||||
"""Test that per-operation settings fall back to global when not set."""
|
||||
from hindsight_api.config import clear_config_cache, get_config
|
||||
|
||||
# Set only global values
|
||||
os.environ["HINDSIGHT_API_LLM_MAX_RETRIES"] = "7"
|
||||
os.environ["HINDSIGHT_API_LLM_INITIAL_BACKOFF"] = "3.0"
|
||||
os.environ["HINDSIGHT_API_LLM_MAX_BACKOFF"] = "180.0"
|
||||
|
||||
try:
|
||||
clear_config_cache()
|
||||
config = get_config()
|
||||
|
||||
# Per-operation should be None (will fall back to global at runtime)
|
||||
assert config.retain_llm_max_retries is None
|
||||
assert config.retain_llm_initial_backoff is None
|
||||
assert config.retain_llm_max_backoff is None
|
||||
|
||||
# Global values should be set
|
||||
assert config.llm_max_retries == 7
|
||||
assert config.llm_initial_backoff == 3.0
|
||||
assert config.llm_max_backoff == 180.0
|
||||
finally:
|
||||
os.environ.pop("HINDSIGHT_API_LLM_MAX_RETRIES", None)
|
||||
os.environ.pop("HINDSIGHT_API_LLM_INITIAL_BACKOFF", None)
|
||||
os.environ.pop("HINDSIGHT_API_LLM_MAX_BACKOFF", None)
|
||||
clear_config_cache()
|
||||
|
||||
@@ -0,0 +1,123 @@
|
||||
"""Test provider-specific default models in config."""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_provider_default_models():
|
||||
"""Test that each provider has a default model and it's used when model is not explicitly set."""
|
||||
from hindsight_api.config import PROVIDER_DEFAULT_MODELS, HindsightConfig, clear_config_cache
|
||||
|
||||
# Save original env vars
|
||||
original_provider = os.environ.get("HINDSIGHT_API_LLM_PROVIDER")
|
||||
original_model = os.environ.get("HINDSIGHT_API_LLM_MODEL")
|
||||
|
||||
try:
|
||||
# Test each provider has a default
|
||||
for provider, expected_model in PROVIDER_DEFAULT_MODELS.items():
|
||||
clear_config_cache()
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = provider
|
||||
# Remove explicit model setting to test default
|
||||
if "HINDSIGHT_API_LLM_MODEL" in os.environ:
|
||||
del os.environ["HINDSIGHT_API_LLM_MODEL"]
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_provider == provider, f"Provider mismatch for {provider}"
|
||||
assert config.llm_model == expected_model, f"Expected {expected_model} for {provider}, got {config.llm_model}"
|
||||
|
||||
finally:
|
||||
# Restore original env vars
|
||||
clear_config_cache()
|
||||
if original_provider:
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = original_provider
|
||||
elif "HINDSIGHT_API_LLM_PROVIDER" in os.environ:
|
||||
del os.environ["HINDSIGHT_API_LLM_PROVIDER"]
|
||||
|
||||
if original_model:
|
||||
os.environ["HINDSIGHT_API_LLM_MODEL"] = original_model
|
||||
elif "HINDSIGHT_API_LLM_MODEL" in os.environ:
|
||||
del os.environ["HINDSIGHT_API_LLM_MODEL"]
|
||||
|
||||
|
||||
def test_explicit_model_overrides_provider_default():
|
||||
"""Test that explicit model setting overrides provider default."""
|
||||
from hindsight_api.config import HindsightConfig, clear_config_cache
|
||||
|
||||
original_provider = os.environ.get("HINDSIGHT_API_LLM_PROVIDER")
|
||||
original_model = os.environ.get("HINDSIGHT_API_LLM_MODEL")
|
||||
|
||||
try:
|
||||
clear_config_cache()
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "anthropic"
|
||||
os.environ["HINDSIGHT_API_LLM_MODEL"] = "claude-sonnet-4-5-20250929"
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_provider == "anthropic"
|
||||
assert config.llm_model == "claude-sonnet-4-5-20250929", "Explicit model should override default"
|
||||
|
||||
finally:
|
||||
clear_config_cache()
|
||||
if original_provider:
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = original_provider
|
||||
elif "HINDSIGHT_API_LLM_PROVIDER" in os.environ:
|
||||
del os.environ["HINDSIGHT_API_LLM_PROVIDER"]
|
||||
|
||||
if original_model:
|
||||
os.environ["HINDSIGHT_API_LLM_MODEL"] = original_model
|
||||
elif "HINDSIGHT_API_LLM_MODEL" in os.environ:
|
||||
del os.environ["HINDSIGHT_API_LLM_MODEL"]
|
||||
|
||||
|
||||
def test_per_operation_provider_default_model():
|
||||
"""Test that per-operation providers use their own default models."""
|
||||
from hindsight_api.config import HindsightConfig, clear_config_cache
|
||||
|
||||
original_provider = os.environ.get("HINDSIGHT_API_LLM_PROVIDER")
|
||||
original_model = os.environ.get("HINDSIGHT_API_LLM_MODEL")
|
||||
original_retain_provider = os.environ.get("HINDSIGHT_API_RETAIN_LLM_PROVIDER")
|
||||
original_retain_model = os.environ.get("HINDSIGHT_API_RETAIN_LLM_MODEL")
|
||||
|
||||
try:
|
||||
clear_config_cache()
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "openai"
|
||||
# Remove explicit model to use provider default
|
||||
if "HINDSIGHT_API_LLM_MODEL" in os.environ:
|
||||
del os.environ["HINDSIGHT_API_LLM_MODEL"]
|
||||
|
||||
# Set retain-specific provider but not model
|
||||
os.environ["HINDSIGHT_API_RETAIN_LLM_PROVIDER"] = "anthropic"
|
||||
if "HINDSIGHT_API_RETAIN_LLM_MODEL" in os.environ:
|
||||
del os.environ["HINDSIGHT_API_RETAIN_LLM_MODEL"]
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
# Global LLM should use OpenAI default
|
||||
assert config.llm_model == "o3-mini", f"Expected o3-mini, got {config.llm_model}"
|
||||
|
||||
# Retain should use Anthropic default
|
||||
assert (
|
||||
config.retain_llm_model == "claude-haiku-4-5-20251001"
|
||||
), f"Expected claude-haiku-4-5-20251001, got {config.retain_llm_model}"
|
||||
|
||||
finally:
|
||||
clear_config_cache()
|
||||
if original_provider:
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = original_provider
|
||||
elif "HINDSIGHT_API_LLM_PROVIDER" in os.environ:
|
||||
del os.environ["HINDSIGHT_API_LLM_PROVIDER"]
|
||||
|
||||
if original_model:
|
||||
os.environ["HINDSIGHT_API_LLM_MODEL"] = original_model
|
||||
elif "HINDSIGHT_API_LLM_MODEL" in os.environ:
|
||||
del os.environ["HINDSIGHT_API_LLM_MODEL"]
|
||||
|
||||
if original_retain_provider:
|
||||
os.environ["HINDSIGHT_API_RETAIN_LLM_PROVIDER"] = original_retain_provider
|
||||
elif "HINDSIGHT_API_RETAIN_LLM_PROVIDER" in os.environ:
|
||||
del os.environ["HINDSIGHT_API_RETAIN_LLM_PROVIDER"]
|
||||
|
||||
if original_retain_model:
|
||||
os.environ["HINDSIGHT_API_RETAIN_LLM_MODEL"] = original_retain_model
|
||||
elif "HINDSIGHT_API_RETAIN_LLM_MODEL" in os.environ:
|
||||
del os.environ["HINDSIGHT_API_RETAIN_LLM_MODEL"]
|
||||
@@ -175,6 +175,45 @@ class TestMentalModelsCRUD:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_mental_model_with_custom_id(self, memory: MemoryEngine, request_context):
|
||||
"""Test creating a mental model with a custom ID."""
|
||||
bank_id = f"test-mental-model-custom-id-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Create the bank first
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
# Create a mental model with a custom ID
|
||||
custom_id = "team-communication-preferences"
|
||||
mental_model = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=custom_id,
|
||||
name="Team Communication Preferences",
|
||||
source_query="How does the team prefer to communicate?",
|
||||
content="The team prefers async communication via Slack",
|
||||
tags=["team", "communication"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify the custom ID was used
|
||||
assert mental_model["id"] == custom_id
|
||||
assert mental_model["name"] == "Team Communication Preferences"
|
||||
assert mental_model["tags"] == ["team", "communication"]
|
||||
|
||||
# Verify we can retrieve it with the custom ID
|
||||
fetched = await memory.get_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=custom_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert fetched is not None
|
||||
assert fetched["id"] == custom_id
|
||||
assert fetched["name"] == "Team Communication Preferences"
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
class TestObservationsAPI:
|
||||
"""Test observations API endpoints.
|
||||
|
||||
@@ -16,7 +16,6 @@ async def test_retain_with_chunks(memory, request_context):
|
||||
Test that retain function:
|
||||
1. Stores facts with associated chunks
|
||||
2. Recall returns chunk_id for each fact
|
||||
3. Recall with include_entities=True also works (for compatibility)
|
||||
"""
|
||||
bank_id = f"test_chunks_{datetime.now(timezone.utc).timestamp()}"
|
||||
document_id = "test_doc_123"
|
||||
@@ -56,7 +55,6 @@ async def test_retain_with_chunks(memory, request_context):
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world"], # Search for world facts
|
||||
include_entities=False, # Disable entities for simpler test
|
||||
include_chunks=True, # Enable chunks
|
||||
max_chunk_tokens=8192,
|
||||
request_context=request_context,
|
||||
@@ -146,7 +144,6 @@ async def test_chunks_and_entities_follow_fact_order(memory, request_context):
|
||||
budget=Budget.MID,
|
||||
max_tokens=1000,
|
||||
fact_type=["world"],
|
||||
include_entities=True,
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=8192,
|
||||
request_context=request_context,
|
||||
@@ -2196,3 +2193,65 @@ If the text contains both Italian and English content, extract ONLY the Italian
|
||||
|
||||
# Clear cache again to restore original config
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_batch_with_per_item_tags_on_document(memory, request_context):
|
||||
"""
|
||||
Test that per-item tags are correctly stored on documents.
|
||||
|
||||
This test verifies the fix for a bug where per-item tags in content dictionaries
|
||||
were not being merged and passed to document tracking, causing tags to be lost
|
||||
even though they were correctly sent through the API.
|
||||
|
||||
Without the fix, this test would fail because:
|
||||
- Tags are correctly passed in the content dict
|
||||
- Tags are correctly stored on memory_units (facts)
|
||||
- BUT tags were NOT stored on the document record itself
|
||||
"""
|
||||
bank_id = f"test_doc_tags_{datetime.now(timezone.utc).timestamp()}"
|
||||
document_id = "app-state-testuser"
|
||||
|
||||
try:
|
||||
# Retain content with per-item tags (simulating the TasteAI use case)
|
||||
contents = [
|
||||
{
|
||||
"content": '{"username":"testuser","meals":[],"preferences":{"nickname":"testuser"}}',
|
||||
"document_id": document_id,
|
||||
"tags": ["user:testuser", "app-type:taste-ai"],
|
||||
}
|
||||
]
|
||||
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result) > 0, "Should have retained content"
|
||||
print(f"\n=== Retained content with tags ===")
|
||||
|
||||
# Retrieve the document
|
||||
doc = await memory.get_document(
|
||||
document_id=document_id,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert doc is not None, "Document should exist"
|
||||
assert "tags" in doc, "Document should have tags field"
|
||||
|
||||
# This is the critical assertion - tags should be stored on the document
|
||||
doc_tags = doc["tags"] or []
|
||||
print(f"Document tags: {doc_tags}")
|
||||
|
||||
assert "user:testuser" in doc_tags, \
|
||||
f"Document should have 'user:testuser' tag, but got: {doc_tags}"
|
||||
assert "app-type:taste-ai" in doc_tags, \
|
||||
f"Document should have 'app-type:taste-ai' tag, but got: {doc_tags}"
|
||||
|
||||
print("✓ Per-item tags correctly stored on document")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
print(f"\n=== Cleaned up bank: {bank_id} ===")
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
"""
|
||||
Test think function for opinion generation and consistency.
|
||||
Test reflect (think) function.
|
||||
"""
|
||||
import pytest
|
||||
from datetime import datetime, timezone
|
||||
@@ -7,131 +7,6 @@ from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_think_opinion_consistency(memory, request_context):
|
||||
"""
|
||||
Test that think function:
|
||||
1. Generates an opinion
|
||||
2. Stores the opinion in the database
|
||||
3. Returns consistent response on subsequent calls with the same query
|
||||
"""
|
||||
bank_id = f"test_think_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
|
||||
# Store some initial facts to give context for opinion formation
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice is a software engineer who has worked on 5 major projects. She always delivers on time and writes clean, well-documented code.",
|
||||
context="performance review",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob recently joined the team. He missed his first deadline and his code had many bugs.",
|
||||
context="performance review",
|
||||
event_date=datetime(2024, 2, 1, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# First think call - should generate opinions
|
||||
query = "Who is a more reliable engineer?"
|
||||
result1 = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
budget=Budget.LOW,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== First Think Call ===")
|
||||
print(f"Answer: {result1.text}")
|
||||
|
||||
# Verify we got an answer
|
||||
assert result1.text, "First think call should return an answer"
|
||||
assert result1.based_on, "Should return based_on facts"
|
||||
|
||||
# Wait for background opinion processing tasks to complete
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Search for stored opinions to verify they were actually saved
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
stored_opinions = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, confidence_score, fact_type
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'opinion'
|
||||
ORDER BY created_at DESC
|
||||
""",
|
||||
bank_id
|
||||
)
|
||||
|
||||
print(f"\n=== Stored Opinions in Database ===")
|
||||
print(f"Total opinions stored: {len(stored_opinions)}")
|
||||
for op in stored_opinions:
|
||||
print(f" - {op['text']} (confidence: {op['confidence_score']:.2f})")
|
||||
|
||||
# Verify opinions were actually written to database
|
||||
# NOTE: Opinion extraction may not always detect opinions depending on the LLM response format
|
||||
if len(stored_opinions) > 0:
|
||||
assert all(op['fact_type'] == 'opinion' for op in stored_opinions), "All stored items should have fact_type='opinion'"
|
||||
print(f"✓ Opinions were successfully stored in database")
|
||||
else:
|
||||
print(f"⚠ Note: No opinions were extracted/stored (this can happen if the LLM response format doesn't trigger opinion extraction)")
|
||||
|
||||
# Second think call - should use the stored opinions
|
||||
result2 = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
budget=Budget.LOW,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Second Think Call ===")
|
||||
print(f"Answer: {result2.text}")
|
||||
print(f"Existing opinions used: {len(result2.based_on.get('opinion', []))}")
|
||||
for opinion in result2.based_on.get('opinion', []):
|
||||
print(f" - {opinion.text}")
|
||||
|
||||
# Verify second call also got an answer
|
||||
assert result2.text, "Second think call should return an answer"
|
||||
|
||||
# Verify second call used the stored opinions (if any were stored)
|
||||
if len(stored_opinions) > 0:
|
||||
assert len(result2.based_on.get('opinion', [])) > 0, "Second call should retrieve stored opinions"
|
||||
|
||||
# The responses should be consistent (both should mention the same person as more reliable)
|
||||
# We'll do a basic check that they're not contradictory
|
||||
text1_lower = result1.text.lower()
|
||||
text2_lower = result2.text.lower()
|
||||
|
||||
print(f"\n=== Consistency Check ===")
|
||||
|
||||
# Check if Alice is mentioned as more reliable in first response
|
||||
if 'alice' in text1_lower and ('reliable' in text1_lower or 'better' in text1_lower):
|
||||
print("First response favors Alice")
|
||||
# Second response should also favor Alice (consistency)
|
||||
assert 'alice' in text2_lower, "Second response should also mention Alice"
|
||||
print("Second response also mentions Alice - CONSISTENT ✓")
|
||||
|
||||
# Check if Bob is mentioned
|
||||
if 'bob' in text1_lower:
|
||||
print("First response mentions Bob")
|
||||
if 'bob' in text2_lower:
|
||||
print("Second response also mentions Bob - CONSISTENT ✓")
|
||||
|
||||
print(f"\n✅ Test passed - opinions were formed, stored, and used consistently")
|
||||
|
||||
finally:
|
||||
# Clean up agent data
|
||||
try:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception as e:
|
||||
print(f"Warning: Error during cleanup: {e}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_think_without_prior_context(memory, request_context):
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,244 @@
|
||||
"""
|
||||
Test Vertex AI provider integration using native genai SDK.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# Skip all tests if google-auth not available
|
||||
pytest.importorskip("google.auth")
|
||||
|
||||
|
||||
def test_llm_wrapper_vertexai_missing_dependency():
|
||||
"""Test error when google-auth is not available and service account key is set."""
|
||||
from hindsight_api.engine import llm_wrapper
|
||||
|
||||
# VERTEXAI_AVAILABLE only matters when a service account key is provided
|
||||
original_available = llm_wrapper.VERTEXAI_AVAILABLE
|
||||
try:
|
||||
llm_wrapper.VERTEXAI_AVAILABLE = False
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": "test-project",
|
||||
"HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY": "/path/to/key.json",
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
with pytest.raises(ValueError, match="google-auth"):
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="google/gemini-2.0-flash-001",
|
||||
)
|
||||
|
||||
clear_config_cache()
|
||||
finally:
|
||||
llm_wrapper.VERTEXAI_AVAILABLE = original_available
|
||||
|
||||
|
||||
def test_llm_wrapper_vertexai_missing_project_id():
|
||||
"""Test error when project ID is not configured."""
|
||||
with patch.dict(os.environ, {"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": ""}, clear=False):
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
with pytest.raises(ValueError, match="HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID"):
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="google/gemini-2.0-flash-001",
|
||||
)
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_llm_wrapper_vertexai_adc_auth():
|
||||
"""Test Vertex AI with ADC authentication creates native genai client."""
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": "test-project"},
|
||||
clear=False,
|
||||
):
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
# genai.Client handles ADC internally — just verify it creates the client
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
|
||||
provider = LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="google/gemini-2.0-flash-001",
|
||||
)
|
||||
|
||||
assert provider.provider == "vertexai"
|
||||
assert provider.model == "gemini-2.0-flash-001" # google/ prefix stripped
|
||||
assert provider._gemini_client is not None
|
||||
|
||||
# Verify genai.Client was called with vertexai=True
|
||||
mock_client_cls.assert_called_once_with(
|
||||
vertexai=True,
|
||||
project="test-project",
|
||||
location="us-central1",
|
||||
)
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_llm_wrapper_vertexai_sa_auth():
|
||||
"""Test Vertex AI with service account authentication passes credentials to genai client."""
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
mock_credentials = MagicMock()
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": "test-project",
|
||||
"HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY": "/path/to/key.json",
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
with patch(
|
||||
"google.oauth2.service_account.Credentials.from_service_account_file",
|
||||
return_value=mock_credentials,
|
||||
):
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
|
||||
provider = LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="google/gemini-2.0-flash-001",
|
||||
)
|
||||
|
||||
assert provider.provider == "vertexai"
|
||||
assert provider._gemini_client is not None
|
||||
|
||||
# Verify credentials were passed to genai.Client
|
||||
mock_client_cls.assert_called_once_with(
|
||||
vertexai=True,
|
||||
project="test-project",
|
||||
location="us-central1",
|
||||
credentials=mock_credentials,
|
||||
)
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_llm_wrapper_vertexai_strips_google_prefix():
|
||||
"""Test that google/ prefix is stripped from model name for native SDK."""
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": "test-project"},
|
||||
clear=False,
|
||||
):
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
|
||||
provider = LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="google/gemini-2.0-flash-lite-001",
|
||||
)
|
||||
|
||||
assert provider.model == "gemini-2.0-flash-lite-001"
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_llm_wrapper_vertexai_no_prefix_model():
|
||||
"""Test that model without google/ prefix is unchanged."""
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": "test-project"},
|
||||
clear=False,
|
||||
):
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
|
||||
provider = LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="gemini-2.0-flash-001",
|
||||
)
|
||||
|
||||
assert provider.model == "gemini-2.0-flash-001"
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(
|
||||
not os.getenv("HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID"),
|
||||
reason="Vertex AI integration tests require HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID",
|
||||
)
|
||||
async def test_vertexai_integration_actual_api():
|
||||
"""
|
||||
Integration test with actual Vertex AI API.
|
||||
|
||||
Requires:
|
||||
- HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID
|
||||
- ADC or HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY
|
||||
"""
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
provider = LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="google/gemini-2.0-flash-001",
|
||||
)
|
||||
|
||||
try:
|
||||
# Simple test call
|
||||
response = await provider.call(
|
||||
messages=[{"role": "user", "content": "Say 'ok' and nothing else"}],
|
||||
max_completion_tokens=10,
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, str)
|
||||
assert len(response) > 0
|
||||
|
||||
finally:
|
||||
await provider.cleanup()
|
||||
@@ -156,7 +156,6 @@ class TestWorkerPoller:
|
||||
pool=pool,
|
||||
worker_id="test-worker-1",
|
||||
executor=mock_executor,
|
||||
batch_size=10,
|
||||
)
|
||||
|
||||
claimed = await poller.claim_batch()
|
||||
@@ -177,8 +176,8 @@ class TestWorkerPoller:
|
||||
assert row["worker_id"] == "test-worker-1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claim_batch_respects_batch_size(self, pool, clean_operations):
|
||||
"""Test that claim_batch respects the batch_size limit."""
|
||||
async def test_claim_batch_respects_max_slots(self, pool, clean_operations):
|
||||
"""Test that claim_batch respects the max_slots limit."""
|
||||
from hindsight_api.worker import WorkerPoller
|
||||
|
||||
# Create 10 pending tasks
|
||||
@@ -196,12 +195,11 @@ class TestWorkerPoller:
|
||||
payload,
|
||||
)
|
||||
|
||||
# Claim with batch_size=3
|
||||
poller = WorkerPoller(
|
||||
pool=pool,
|
||||
worker_id="test-worker-1",
|
||||
executor=lambda x: None,
|
||||
batch_size=3,
|
||||
max_slots=3, # Limit to 3 concurrent tasks
|
||||
)
|
||||
|
||||
claimed = await poller.claim_batch()
|
||||
@@ -238,11 +236,14 @@ class TestWorkerPoller:
|
||||
executor=mock_executor,
|
||||
)
|
||||
|
||||
# Execute the task
|
||||
# Execute the task (fire-and-forget)
|
||||
task_dict = json.loads(payload)
|
||||
claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
|
||||
await poller.execute_task(claimed_task)
|
||||
|
||||
# Wait for background task to complete
|
||||
completed = await poller.wait_for_active_tasks(timeout=5.0)
|
||||
assert completed, "Task did not complete within timeout"
|
||||
assert len(executed) == 1
|
||||
|
||||
# Verify task is marked as completed
|
||||
@@ -283,11 +284,15 @@ class TestWorkerPoller:
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
# Execute (should fail and retry)
|
||||
# Execute (should fail and retry) - fire-and-forget
|
||||
task_dict = json.loads(payload)
|
||||
claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
|
||||
await poller.execute_task(claimed_task)
|
||||
|
||||
# Wait for background task to complete
|
||||
completed = await poller.wait_for_active_tasks(timeout=5.0)
|
||||
assert completed, "Task did not complete within timeout"
|
||||
|
||||
# Verify task is back to pending with incremented retry_count
|
||||
row = await pool.fetchrow(
|
||||
"SELECT status, retry_count, worker_id FROM async_operations WHERE operation_id = $1",
|
||||
@@ -327,11 +332,15 @@ class TestWorkerPoller:
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
# Execute (should fail permanently)
|
||||
# Execute (should fail permanently) - fire-and-forget
|
||||
task_dict = json.loads(payload)
|
||||
claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
|
||||
await poller.execute_task(claimed_task)
|
||||
|
||||
# Wait for background task to complete
|
||||
completed = await poller.wait_for_active_tasks(timeout=5.0)
|
||||
assert completed, "Task did not complete within timeout"
|
||||
|
||||
# Verify task is marked as failed
|
||||
row = await pool.fetchrow(
|
||||
"SELECT status, error_message FROM async_operations WHERE operation_id = $1",
|
||||
@@ -388,7 +397,6 @@ class TestWorkerPoller:
|
||||
pool=pool,
|
||||
worker_id="test-worker-1",
|
||||
executor=lambda x: None,
|
||||
batch_size=10,
|
||||
)
|
||||
|
||||
claimed = await poller.claim_batch()
|
||||
@@ -440,7 +448,6 @@ class TestWorkerPoller:
|
||||
pool=pool,
|
||||
worker_id="test-worker-1",
|
||||
executor=lambda x: None,
|
||||
batch_size=10,
|
||||
)
|
||||
|
||||
claimed = await poller.claim_batch()
|
||||
@@ -607,7 +614,6 @@ class TestConcurrentWorkers:
|
||||
pool=pool,
|
||||
worker_id=worker_id,
|
||||
executor=lambda x: None,
|
||||
batch_size=5, # Each worker tries to claim 5
|
||||
)
|
||||
claimed = await poller.claim_batch()
|
||||
workers_claimed[worker_id] = [task.operation_id for task in claimed]
|
||||
@@ -680,7 +686,6 @@ class TestConcurrentWorkers:
|
||||
pool=pool,
|
||||
worker_id="new-worker",
|
||||
executor=lambda x: None,
|
||||
batch_size=10,
|
||||
)
|
||||
|
||||
claimed = await poller.claim_batch()
|
||||
@@ -879,7 +884,6 @@ class TestDynamicTenantDiscovery:
|
||||
pool=pool,
|
||||
worker_id="test-worker-1",
|
||||
executor=lambda x: None,
|
||||
batch_size=10,
|
||||
tenant_extension=mock_extension,
|
||||
)
|
||||
|
||||
@@ -946,7 +950,6 @@ class TestDynamicTenantDiscovery:
|
||||
pool=pool,
|
||||
worker_id="test-worker-1",
|
||||
executor=lambda x: None,
|
||||
batch_size=10,
|
||||
tenant_extension=dynamic_extension,
|
||||
)
|
||||
|
||||
@@ -1008,7 +1011,6 @@ class TestDynamicTenantDiscovery:
|
||||
pool=pool,
|
||||
worker_id="test-worker-1",
|
||||
executor=lambda x: None,
|
||||
batch_size=10,
|
||||
)
|
||||
|
||||
claimed = await poller.claim_batch()
|
||||
@@ -1017,3 +1019,273 @@ class TestDynamicTenantDiscovery:
|
||||
# All tasks should have schema=None (public)
|
||||
for task in claimed:
|
||||
assert task.schema is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poller_with_custom_schema(self, pool):
|
||||
"""Test that poller uses custom schema when schema parameter is provided."""
|
||||
from hindsight_api.worker import WorkerPoller
|
||||
|
||||
# Create a custom schema for testing
|
||||
test_schema = "test_custom_schema"
|
||||
|
||||
try:
|
||||
# Create schema and copy table structure
|
||||
await pool.execute(f'CREATE SCHEMA IF NOT EXISTS "{test_schema}"')
|
||||
await pool.execute(
|
||||
f"""
|
||||
CREATE TABLE "{test_schema}".async_operations (
|
||||
LIKE public.async_operations INCLUDING ALL
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Create pending tasks in the custom schema
|
||||
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
|
||||
task_ids = []
|
||||
for i in range(3):
|
||||
op_id = uuid.uuid4()
|
||||
task_ids.append(str(op_id))
|
||||
payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
|
||||
await pool.execute(
|
||||
f"""
|
||||
INSERT INTO "{test_schema}".async_operations (operation_id, bank_id, operation_type, status, task_payload)
|
||||
VALUES ($1, $2, 'test', 'pending', $3::jsonb)
|
||||
""",
|
||||
op_id,
|
||||
bank_id,
|
||||
payload,
|
||||
)
|
||||
|
||||
# Create poller with custom schema
|
||||
poller = WorkerPoller(
|
||||
pool=pool,
|
||||
worker_id="test-worker-custom-schema",
|
||||
executor=lambda x: None,
|
||||
schema=test_schema,
|
||||
)
|
||||
|
||||
# Claim tasks
|
||||
claimed = await poller.claim_batch()
|
||||
assert len(claimed) == 3, f"Expected 3 tasks, got {len(claimed)}"
|
||||
|
||||
# All tasks should have schema=test_schema
|
||||
claimed_ids = []
|
||||
for task in claimed:
|
||||
assert task.schema == test_schema, f"Expected schema '{test_schema}', got '{task.schema}'"
|
||||
claimed_ids.append(task.operation_id)
|
||||
|
||||
# Verify claimed tasks match what we inserted
|
||||
assert set(claimed_ids) == set(task_ids)
|
||||
|
||||
# Verify tasks are marked as processing in the custom schema
|
||||
rows = await pool.fetch(
|
||||
f"""
|
||||
SELECT operation_id, status, worker_id
|
||||
FROM "{test_schema}".async_operations
|
||||
WHERE operation_id = ANY($1)
|
||||
""",
|
||||
[uuid.UUID(tid) for tid in task_ids],
|
||||
)
|
||||
assert len(rows) == 3
|
||||
for row in rows:
|
||||
assert row["status"] == "processing"
|
||||
assert row["worker_id"] == "test-worker-custom-schema"
|
||||
|
||||
finally:
|
||||
# Clean up: drop the custom schema
|
||||
await pool.execute(f'DROP SCHEMA IF EXISTS "{test_schema}" CASCADE')
|
||||
|
||||
|
||||
async def test_worker_fire_and_forget_nonblocking(pool, clean_operations):
|
||||
"""
|
||||
Test that worker continues polling while tasks run (fire-and-forget pattern).
|
||||
|
||||
This test verifies the FIX: With the old blocking behavior, the worker would
|
||||
wait for all tasks in a batch to complete before claiming more. This test
|
||||
would FAIL with the old code because tasks 3-4 wouldn't be claimed until
|
||||
tasks 1-2 complete. With fire-and-forget, tasks 3-4 are claimed immediately.
|
||||
"""
|
||||
from hindsight_api.worker.poller import WorkerPoller
|
||||
|
||||
task_started = {} # operation_id -> Event (set when task starts)
|
||||
task_canfinish = {} # operation_id -> Event (wait before finishing)
|
||||
|
||||
async def blocking_executor(task_dict: dict):
|
||||
op_id = task_dict["operation_id"]
|
||||
# Signal that this task has started
|
||||
started = asyncio.Event()
|
||||
task_started[op_id] = started
|
||||
started.set()
|
||||
|
||||
# Block until we're told to finish
|
||||
finish = asyncio.Event()
|
||||
task_canfinish[op_id] = finish
|
||||
await finish.wait()
|
||||
|
||||
poller = WorkerPoller(
|
||||
pool=pool,
|
||||
worker_id="test-worker",
|
||||
executor=blocking_executor,
|
||||
poll_interval_ms=50, # Fast polling
|
||||
max_slots=10,
|
||||
consolidation_max_slots=2,
|
||||
)
|
||||
|
||||
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Submit initial 2 tasks
|
||||
task_ids = []
|
||||
for i in range(2):
|
||||
op_id = uuid.uuid4()
|
||||
task_ids.append(str(op_id))
|
||||
payload = json.dumps({"type": "test", "operation_type": "retain", "operation_id": str(op_id), "bank_id": bank_id})
|
||||
await pool.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
|
||||
VALUES ($1, $2, 'retain', 'pending', $3::jsonb)
|
||||
""",
|
||||
op_id,
|
||||
bank_id,
|
||||
payload,
|
||||
)
|
||||
|
||||
poll_task = asyncio.create_task(poller.run())
|
||||
|
||||
try:
|
||||
# Wait for first 2 tasks to start executing (but not finish)
|
||||
for i in range(100): # Try for up to 1 second
|
||||
if len(task_started) >= 2:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert len(task_started) == 2, f"Expected 2 tasks started, got {len(task_started)}"
|
||||
|
||||
# Verify tasks are in_flight
|
||||
async with poller._in_flight_lock:
|
||||
assert poller._in_flight_count == 2
|
||||
|
||||
# NOW submit 2 more tasks WHILE the first 2 are still running
|
||||
for i in range(2):
|
||||
op_id = uuid.uuid4()
|
||||
task_ids.append(str(op_id))
|
||||
payload = json.dumps({"type": "test", "operation_type": "retain", "operation_id": str(op_id), "bank_id": bank_id})
|
||||
await pool.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
|
||||
VALUES ($1, $2, 'retain', 'pending', $3::jsonb)
|
||||
""",
|
||||
op_id,
|
||||
bank_id,
|
||||
payload,
|
||||
)
|
||||
|
||||
# KEY ASSERTION: Worker should claim tasks 3-4 WITHOUT waiting for 1-2 to finish
|
||||
# This would FAIL with the old blocking behavior
|
||||
for i in range(100): # Try for up to 1 second
|
||||
if len(task_started) >= 4:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert len(task_started) == 4, (
|
||||
f"Fire-and-forget FAILED: Expected 4 tasks started, got {len(task_started)}. "
|
||||
"This means the worker blocked waiting for the first batch to complete."
|
||||
)
|
||||
|
||||
# Verify all 4 tasks are in-flight
|
||||
async with poller._in_flight_lock:
|
||||
assert poller._in_flight_count == 4
|
||||
|
||||
# Clean up: allow all tasks to finish
|
||||
for event in task_canfinish.values():
|
||||
event.set()
|
||||
|
||||
finally:
|
||||
# Ensure cleanup
|
||||
for event in task_canfinish.values():
|
||||
event.set()
|
||||
await poller.shutdown_graceful(timeout=2.0)
|
||||
try:
|
||||
await asyncio.wait_for(poll_task, timeout=1.0)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
|
||||
async def test_worker_slot_limits_enforced(pool, clean_operations):
|
||||
"""Test that worker respects max_slots and won't exceed the limit."""
|
||||
from hindsight_api.worker.poller import WorkerPoller
|
||||
|
||||
tasks_started = set()
|
||||
task_events = {}
|
||||
|
||||
async def controlled_executor(task_dict: dict):
|
||||
op_id = task_dict["operation_id"]
|
||||
tasks_started.add(op_id)
|
||||
event = asyncio.Event()
|
||||
task_events[op_id] = event
|
||||
await event.wait()
|
||||
|
||||
poller = WorkerPoller(
|
||||
pool=pool,
|
||||
worker_id="test-worker",
|
||||
executor=controlled_executor,
|
||||
poll_interval_ms=50,
|
||||
max_slots=3, # Only allow 3 concurrent tasks
|
||||
consolidation_max_slots=1,
|
||||
)
|
||||
|
||||
# Submit 10 tasks
|
||||
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
|
||||
for i in range(10):
|
||||
op_id = uuid.uuid4()
|
||||
payload = json.dumps({"type": "test", "operation_type": "retain", "operation_id": str(op_id), "bank_id": bank_id})
|
||||
await pool.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
|
||||
VALUES ($1, $2, 'retain', 'pending', $3::jsonb)
|
||||
""",
|
||||
op_id,
|
||||
bank_id,
|
||||
payload,
|
||||
)
|
||||
|
||||
poll_task = asyncio.create_task(poller.run())
|
||||
|
||||
try:
|
||||
# Wait for slots to fill
|
||||
for i in range(100):
|
||||
if len(tasks_started) >= 3:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
# Should have claimed exactly 3 tasks (slot limit)
|
||||
assert len(tasks_started) == 3
|
||||
|
||||
# Wait to ensure no additional tasks are claimed
|
||||
for i in range(30):
|
||||
await asyncio.sleep(0.01)
|
||||
assert len(tasks_started) == 3, "Worker exceeded slot limit!"
|
||||
|
||||
# Release tasks one by one and verify remaining are claimed
|
||||
completed = 0
|
||||
while completed < 10 and len(tasks_started) < 10:
|
||||
# Release the next batch
|
||||
events_to_release = list(task_events.values())[completed:completed+3]
|
||||
for event in events_to_release:
|
||||
event.set()
|
||||
completed += len(events_to_release)
|
||||
|
||||
# Wait for new tasks to be claimed
|
||||
for i in range(100):
|
||||
if len(tasks_started) >= min(completed + 3, 10):
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert len(tasks_started) == 10
|
||||
|
||||
finally:
|
||||
for event in task_events.values():
|
||||
event.set()
|
||||
await poller.shutdown_graceful(timeout=2.0)
|
||||
try:
|
||||
await asyncio.wait_for(poll_task, timeout=1.0)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "hindsight-cli"
|
||||
version = "0.4.1"
|
||||
version = "0.4.9"
|
||||
edition = "2021"
|
||||
authors = ["Hindsight Team"]
|
||||
description = "A beautiful CLI for Hindsight - semantic memory system"
|
||||
|
||||
@@ -500,6 +500,8 @@ pub fn delete(
|
||||
pub fn consolidate(
|
||||
client: &ApiClient,
|
||||
bank_id: &str,
|
||||
wait: bool,
|
||||
poll_interval: u64,
|
||||
verbose: bool,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
@@ -517,17 +519,82 @@ pub fn consolidate(
|
||||
|
||||
match response {
|
||||
Ok(result) => {
|
||||
let operation_id = result.operation_id.clone();
|
||||
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_success("Consolidation triggered");
|
||||
println!(" {} {}", ui::dim("Operation ID:"), result.operation_id);
|
||||
println!(" {} {}", ui::dim("Operation ID:"), operation_id);
|
||||
if result.deduplicated {
|
||||
println!(" {} {}", ui::dim("Note:"), "Reusing existing pending consolidation task");
|
||||
}
|
||||
println!();
|
||||
println!("{}", ui::dim("Use 'hindsight operation get' to check the operation status."));
|
||||
} else {
|
||||
output::print_output(&result, output_format)?;
|
||||
}
|
||||
|
||||
if !wait {
|
||||
if output_format == OutputFormat::Pretty {
|
||||
println!();
|
||||
println!("{}", ui::dim("Use --wait to poll for completion, or 'hindsight operation get' to check status."));
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Poll for completion
|
||||
if output_format == OutputFormat::Pretty {
|
||||
println!();
|
||||
println!("{}", ui::dim(&format!("Polling every {}s for completion...", poll_interval)));
|
||||
}
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
loop {
|
||||
std::thread::sleep(std::time::Duration::from_secs(poll_interval));
|
||||
let elapsed = start.elapsed().as_secs();
|
||||
|
||||
let ops_result = client.list_operations(bank_id, verbose);
|
||||
match ops_result {
|
||||
Ok(ops) => {
|
||||
// Find the operation by ID
|
||||
let op = ops.operations.iter().find(|o| o.id == operation_id);
|
||||
|
||||
match op.map(|o| o.status.as_str()) {
|
||||
Some("completed") => {
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_success(&format!("Consolidation completed ({}s)", elapsed));
|
||||
}
|
||||
break;
|
||||
}
|
||||
Some("failed") => {
|
||||
let error_msg = op
|
||||
.and_then(|o| o.error_message.as_ref())
|
||||
.map(|s| s.as_str())
|
||||
.unwrap_or("Unknown error");
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_error(&format!("Consolidation failed: {}", error_msg));
|
||||
}
|
||||
std::process::exit(1);
|
||||
}
|
||||
Some(status) => {
|
||||
if output_format == OutputFormat::Pretty {
|
||||
println!(" ⏳ {} ({}s elapsed)", status, elapsed);
|
||||
}
|
||||
}
|
||||
None => {
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_warning(&format!("Operation {} not found in list", operation_id));
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_error(&format!("Failed to check operation status: {}", e));
|
||||
}
|
||||
return Err(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
use anyhow::Result;
|
||||
use chrono::{Duration as ChronoDuration, NaiveDate, Utc};
|
||||
use std::collections::BTreeMap;
|
||||
use crate::api::ApiClient;
|
||||
use crate::output::{self, OutputFormat};
|
||||
use crate::ui;
|
||||
@@ -7,11 +9,17 @@ pub fn list(
|
||||
client: &ApiClient,
|
||||
agent_id: &str,
|
||||
query: Option<String>,
|
||||
date: Option<String>,
|
||||
limit: i32,
|
||||
offset: i32,
|
||||
verbose: bool,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
// If date filter is provided, use the date-aware listing
|
||||
if date.is_some() {
|
||||
return list_with_date(client, agent_id, date.as_deref(), verbose, output_format);
|
||||
}
|
||||
|
||||
let spinner = if output_format == OutputFormat::Pretty {
|
||||
Some(ui::create_spinner("Fetching documents..."))
|
||||
} else {
|
||||
@@ -50,6 +58,139 @@ pub fn list(
|
||||
}
|
||||
}
|
||||
|
||||
/// List documents with date filtering
|
||||
fn list_with_date(
|
||||
client: &ApiClient,
|
||||
bank_id: &str,
|
||||
date_filter: Option<&str>,
|
||||
verbose: bool,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
let spinner = if output_format == OutputFormat::Pretty {
|
||||
Some(ui::create_spinner("Fetching all documents..."))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Fetch all documents with pagination
|
||||
let all_docs = fetch_all_documents(client, bank_id, verbose)?;
|
||||
|
||||
if let Some(mut sp) = spinner {
|
||||
sp.finish();
|
||||
}
|
||||
|
||||
// Parse the date filter
|
||||
let target_date = parse_date_filter(date_filter)?;
|
||||
|
||||
// Filter and group documents by date
|
||||
let mut by_date: BTreeMap<String, Vec<serde_json::Value>> = BTreeMap::new();
|
||||
let mut filtered_count = 0;
|
||||
|
||||
for doc in all_docs {
|
||||
let created_at = doc.get("created_at")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("");
|
||||
|
||||
// Parse the date part (YYYY-MM-DD) from created_at
|
||||
let doc_date = created_at.split('T').next().unwrap_or("");
|
||||
|
||||
// Apply date filter if specified
|
||||
if let Some(ref target) = target_date {
|
||||
let target_str = target.format("%Y-%m-%d").to_string();
|
||||
if doc_date != target_str {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
filtered_count += 1;
|
||||
by_date.entry(doc_date.to_string()).or_default().push(doc);
|
||||
}
|
||||
|
||||
// Output
|
||||
if output_format == OutputFormat::Pretty {
|
||||
let filter_desc = match date_filter {
|
||||
None | Some("yesterday") => "yesterday".to_string(),
|
||||
Some("today") => "today".to_string(),
|
||||
Some("all") => "all dates".to_string(),
|
||||
Some(d) => d.to_string(),
|
||||
};
|
||||
|
||||
ui::print_info(&format!(
|
||||
"Documents for bank '{}' (filter: {}, showing: {})",
|
||||
bank_id, filter_desc, filtered_count
|
||||
));
|
||||
println!();
|
||||
|
||||
// Show documents grouped by date (reverse order - newest first)
|
||||
for (date_str, docs) in by_date.iter().rev() {
|
||||
println!(" {} ({} documents)", date_str, docs.len());
|
||||
for doc in docs {
|
||||
let id = doc.get("id").and_then(|v| v.as_str()).unwrap_or("unknown");
|
||||
let mem_count = doc.get("memory_unit_count").and_then(|v| v.as_i64()).unwrap_or(0);
|
||||
println!(" - {} ({} memories)", id, mem_count);
|
||||
}
|
||||
println!();
|
||||
}
|
||||
} else {
|
||||
// JSON/YAML output - convert to a list structure
|
||||
let output: Vec<serde_json::Value> = by_date.values().flatten().cloned().collect();
|
||||
output::print_output(&output, output_format)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Fetch all documents with pagination
|
||||
fn fetch_all_documents(
|
||||
client: &ApiClient,
|
||||
bank_id: &str,
|
||||
verbose: bool,
|
||||
) -> Result<Vec<serde_json::Value>> {
|
||||
let mut all_docs = Vec::new();
|
||||
let mut offset = 0;
|
||||
let limit = 500;
|
||||
|
||||
loop {
|
||||
let response = client.list_documents(bank_id, None, Some(limit), Some(offset), verbose)?;
|
||||
|
||||
if response.items.is_empty() {
|
||||
break;
|
||||
}
|
||||
|
||||
// Convert Map<String, Value> to Value for each item
|
||||
for item in response.items {
|
||||
all_docs.push(serde_json::Value::Object(item));
|
||||
}
|
||||
|
||||
offset += limit;
|
||||
|
||||
// Check if we've fetched everything
|
||||
if all_docs.len() >= response.total as usize {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(all_docs)
|
||||
}
|
||||
|
||||
/// Parse date filter string into a NaiveDate
|
||||
fn parse_date_filter(filter: Option<&str>) -> Result<Option<NaiveDate>> {
|
||||
match filter {
|
||||
None | Some("yesterday") => {
|
||||
// Default to yesterday
|
||||
Ok(Some(Utc::now().date_naive() - ChronoDuration::days(1)))
|
||||
}
|
||||
Some("today") => Ok(Some(Utc::now().date_naive())),
|
||||
Some("all") => Ok(None), // No filtering
|
||||
Some(date_str) => {
|
||||
// Try to parse as YYYY-MM-DD
|
||||
NaiveDate::parse_from_str(date_str, "%Y-%m-%d")
|
||||
.map(Some)
|
||||
.map_err(|e| anyhow::anyhow!("Invalid date format '{}': {}. Use YYYY-MM-DD, 'yesterday', 'today', or 'all'", date_str, e))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get(
|
||||
client: &ApiClient,
|
||||
agent_id: &str,
|
||||
|
||||
@@ -98,6 +98,7 @@ pub fn create(
|
||||
bank_id: &str,
|
||||
name: &str,
|
||||
source_query: &str,
|
||||
id: Option<&str>,
|
||||
verbose: bool,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
@@ -108,6 +109,7 @@ pub fn create(
|
||||
};
|
||||
|
||||
let request = types::CreateMentalModelRequest {
|
||||
id: id.map(|s| s.to_string()),
|
||||
name: name.to_string(),
|
||||
source_query: source_query.to_string(),
|
||||
max_tokens: 2048,
|
||||
|
||||
@@ -260,6 +260,14 @@ enum BankCommands {
|
||||
Consolidate {
|
||||
/// Bank ID
|
||||
bank_id: String,
|
||||
|
||||
/// Wait for consolidation to complete (poll for status)
|
||||
#[arg(long)]
|
||||
wait: bool,
|
||||
|
||||
/// Poll interval in seconds (only used with --wait)
|
||||
#[arg(long, default_value = "10")]
|
||||
poll_interval: u64,
|
||||
},
|
||||
|
||||
/// Clear all observations for a bank
|
||||
@@ -441,6 +449,10 @@ enum DocumentCommands {
|
||||
#[arg(short = 'q', long)]
|
||||
query: Option<String>,
|
||||
|
||||
/// Filter by date (yesterday, today, YYYY-MM-DD, or all)
|
||||
#[arg(short = 'd', long)]
|
||||
date: Option<String>,
|
||||
|
||||
/// Maximum number of results
|
||||
#[arg(short = 'l', long, default_value = "100")]
|
||||
limit: i32,
|
||||
@@ -584,6 +596,10 @@ enum MentalModelCommands {
|
||||
|
||||
/// Source query to generate the mental model from
|
||||
source_query: String,
|
||||
|
||||
/// Optional custom ID for the mental model (alphanumeric lowercase with hyphens)
|
||||
#[arg(long)]
|
||||
id: Option<String>,
|
||||
},
|
||||
|
||||
/// Update a mental model
|
||||
@@ -754,8 +770,8 @@ fn run() -> Result<()> {
|
||||
BankCommands::Delete { bank_id, yes } => {
|
||||
commands::bank::delete(&client, &bank_id, yes, verbose, output_format)
|
||||
}
|
||||
BankCommands::Consolidate { bank_id } => {
|
||||
commands::bank::consolidate(&client, &bank_id, verbose, output_format)
|
||||
BankCommands::Consolidate { bank_id, wait, poll_interval } => {
|
||||
commands::bank::consolidate(&client, &bank_id, wait, poll_interval, verbose, output_format)
|
||||
}
|
||||
BankCommands::ClearObservations { bank_id, yes } => {
|
||||
commands::bank::clear_observations(&client, &bank_id, yes, verbose, output_format)
|
||||
@@ -792,8 +808,8 @@ fn run() -> Result<()> {
|
||||
|
||||
// Document commands
|
||||
Commands::Document(doc_cmd) => match doc_cmd {
|
||||
DocumentCommands::List { bank_id, query, limit, offset } => {
|
||||
commands::document::list(&client, &bank_id, query, limit, offset, verbose, output_format)
|
||||
DocumentCommands::List { bank_id, query, date, limit, offset } => {
|
||||
commands::document::list(&client, &bank_id, query, date, limit, offset, verbose, output_format)
|
||||
}
|
||||
DocumentCommands::Get { bank_id, document_id } => {
|
||||
commands::document::get(&client, &bank_id, &document_id, verbose, output_format)
|
||||
@@ -851,8 +867,8 @@ fn run() -> Result<()> {
|
||||
MentalModelCommands::Get { bank_id, mental_model_id } => {
|
||||
commands::mental_model::get(&client, &bank_id, &mental_model_id, verbose, output_format)
|
||||
}
|
||||
MentalModelCommands::Create { bank_id, name, source_query } => {
|
||||
commands::mental_model::create(&client, &bank_id, &name, &source_query, verbose, output_format)
|
||||
MentalModelCommands::Create { bank_id, name, source_query, id } => {
|
||||
commands::mental_model::create(&client, &bank_id, &name, &source_query, id.as_deref(), verbose, output_format)
|
||||
}
|
||||
MentalModelCommands::Update { bank_id, mental_model_id, name } => {
|
||||
commands::mental_model::update(&client, &bank_id, &mental_model_id, name, verbose, output_format)
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
@@ -1644,7 +1644,7 @@ class MemoryApi:
|
||||
) -> RecallResponse:
|
||||
"""Recall memory
|
||||
|
||||
Recall memory using semantic similarity and spreading activation. The type parameter is optional and must be one of: - `world`: General knowledge about people, places, events, and things that happen - `experience`: Memories about experience, conversations, actions taken, and tasks performed - `opinion`: The bank's formed beliefs, perspectives, and viewpoints Set `include_entities=true` to get entity observations alongside recall results.
|
||||
Recall memory using semantic similarity and spreading activation. The type parameter is optional and must be one of: - `world`: General knowledge about people, places, events, and things that happen - `experience`: Memories about experience, conversations, actions taken, and tasks performed
|
||||
|
||||
:param bank_id: (required)
|
||||
:type bank_id: str
|
||||
@@ -1720,7 +1720,7 @@ class MemoryApi:
|
||||
) -> ApiResponse[RecallResponse]:
|
||||
"""Recall memory
|
||||
|
||||
Recall memory using semantic similarity and spreading activation. The type parameter is optional and must be one of: - `world`: General knowledge about people, places, events, and things that happen - `experience`: Memories about experience, conversations, actions taken, and tasks performed - `opinion`: The bank's formed beliefs, perspectives, and viewpoints Set `include_entities=true` to get entity observations alongside recall results.
|
||||
Recall memory using semantic similarity and spreading activation. The type parameter is optional and must be one of: - `world`: General knowledge about people, places, events, and things that happen - `experience`: Memories about experience, conversations, actions taken, and tasks performed
|
||||
|
||||
:param bank_id: (required)
|
||||
:type bank_id: str
|
||||
@@ -1796,7 +1796,7 @@ class MemoryApi:
|
||||
) -> RESTResponseType:
|
||||
"""Recall memory
|
||||
|
||||
Recall memory using semantic similarity and spreading activation. The type parameter is optional and must be one of: - `world`: General knowledge about people, places, events, and things that happen - `experience`: Memories about experience, conversations, actions taken, and tasks performed - `opinion`: The bank's formed beliefs, perspectives, and viewpoints Set `include_entities=true` to get entity observations alongside recall results.
|
||||
Recall memory using semantic similarity and spreading activation. The type parameter is optional and must be one of: - `world`: General knowledge about people, places, events, and things that happen - `experience`: Memories about experience, conversations, actions taken, and tasks performed
|
||||
|
||||
:param bank_id: (required)
|
||||
:type bank_id: str
|
||||
@@ -1950,7 +1950,7 @@ class MemoryApi:
|
||||
) -> ReflectResponse:
|
||||
"""Reflect and generate answer
|
||||
|
||||
Reflect and formulate an answer using bank identity, world facts, and opinions. This endpoint: 1. Retrieves experience (conversations and events) 2. Retrieves world facts relevant to the query 3. Retrieves existing opinions (bank's perspectives) 4. Uses LLM to formulate a contextual answer 5. Extracts and stores any new opinions formed 6. Returns plain text answer, the facts used, and new opinions
|
||||
Reflect and formulate an answer using bank identity, world facts, and opinions. This endpoint: 1. Retrieves experience (conversations and events) 2. Retrieves world facts relevant to the query 3. Retrieves existing opinions (bank's perspectives) 4. Uses LLM to formulate a contextual answer 5. Returns plain text answer and the facts used
|
||||
|
||||
:param bank_id: (required)
|
||||
:type bank_id: str
|
||||
@@ -2026,7 +2026,7 @@ class MemoryApi:
|
||||
) -> ApiResponse[ReflectResponse]:
|
||||
"""Reflect and generate answer
|
||||
|
||||
Reflect and formulate an answer using bank identity, world facts, and opinions. This endpoint: 1. Retrieves experience (conversations and events) 2. Retrieves world facts relevant to the query 3. Retrieves existing opinions (bank's perspectives) 4. Uses LLM to formulate a contextual answer 5. Extracts and stores any new opinions formed 6. Returns plain text answer, the facts used, and new opinions
|
||||
Reflect and formulate an answer using bank identity, world facts, and opinions. This endpoint: 1. Retrieves experience (conversations and events) 2. Retrieves world facts relevant to the query 3. Retrieves existing opinions (bank's perspectives) 4. Uses LLM to formulate a contextual answer 5. Returns plain text answer and the facts used
|
||||
|
||||
:param bank_id: (required)
|
||||
:type bank_id: str
|
||||
@@ -2102,7 +2102,7 @@ class MemoryApi:
|
||||
) -> RESTResponseType:
|
||||
"""Reflect and generate answer
|
||||
|
||||
Reflect and formulate an answer using bank identity, world facts, and opinions. This endpoint: 1. Retrieves experience (conversations and events) 2. Retrieves world facts relevant to the query 3. Retrieves existing opinions (bank's perspectives) 4. Uses LLM to formulate a contextual answer 5. Extracts and stores any new opinions formed 6. Returns plain text answer, the facts used, and new opinions
|
||||
Reflect and formulate an answer using bank identity, world facts, and opinions. This endpoint: 1. Retrieves experience (conversations and events) 2. Retrieves world facts relevant to the query 3. Retrieves existing opinions (bank's perspectives) 4. Uses LLM to formulate a contextual answer 5. Returns plain text answer and the facts used
|
||||
|
||||
:param bank_id: (required)
|
||||
:type bank_id: str
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
@@ -489,7 +489,7 @@ class Configuration:
|
||||
return "Python SDK Debug Report:\n"\
|
||||
"OS: {env}\n"\
|
||||
"Python Version: {pyversion}\n"\
|
||||
"Version of the API: 0.4.1\n"\
|
||||
"Version of the API: 0.4.9\n"\
|
||||
"SDK Package Version: 0.0.7".\
|
||||
format(env=sys.platform, pyversion=sys.version)
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
+1
-1
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.1
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user