Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
10bff9a1da | ||
|
|
3fa39f3385 |
@@ -31,22 +31,10 @@ HINDSIGHT_API_HOST=0.0.0.0
|
||||
HINDSIGHT_API_PORT=8888
|
||||
HINDSIGHT_API_LOG_LEVEL=info
|
||||
|
||||
# Base Path / Reverse Proxy Support (Optional)
|
||||
# Set these when deploying behind a reverse proxy with path-based routing
|
||||
# Example: To deploy at example.com/hindsight/, set both to "/hindsight"
|
||||
# HINDSIGHT_API_BASE_PATH=/hindsight
|
||||
# NEXT_PUBLIC_BASE_PATH=/hindsight
|
||||
|
||||
# Database (Optional - uses embedded pg0 by default)
|
||||
# HINDSIGHT_API_DATABASE_URL=postgresql://user:pass@host:5432/db
|
||||
# HINDSIGHT_API_DATABASE_SCHEMA=public # PostgreSQL schema name (default: public)
|
||||
|
||||
# Vector Extension (Optional - uses pgvector by default)
|
||||
# Options: "pgvector" (default), "vchord", "pgvectorscale" (DiskANN)
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvector
|
||||
# For Azure PostgreSQL with DiskANN:
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale # Auto-detects pg_diskann on Azure
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default) or "tei" (HuggingFace Text Embeddings Inference)
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
|
||||
+21
-76
@@ -648,80 +648,6 @@ jobs:
|
||||
echo "=== API Server Logs ==="
|
||||
cat /tmp/api-server.log || echo "No API server log found"
|
||||
|
||||
test-go-client:
|
||||
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
|
||||
HINDSIGHT_API_URL: http://localhost:8888
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
# Prefer CPU-only PyTorch in CI (but keep PyPI for everything else)
|
||||
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: Set up Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: '1.23'
|
||||
cache-dependency-path: hindsight-clients/go/go.sum
|
||||
|
||||
- name: Build API
|
||||
working-directory: ./hindsight-api
|
||||
run: uv build
|
||||
|
||||
- name: Install API dependencies
|
||||
working-directory: ./hindsight-api
|
||||
run: uv sync --frozen --no-install-project --index-strategy unsafe-best-match
|
||||
|
||||
- name: Create .env file
|
||||
run: |
|
||||
cat > .env << EOF
|
||||
HINDSIGHT_API_LLM_PROVIDER=${{ env.HINDSIGHT_API_LLM_PROVIDER }}
|
||||
HINDSIGHT_API_LLM_API_KEY=${{ env.HINDSIGHT_API_LLM_API_KEY }}
|
||||
HINDSIGHT_API_LLM_MODEL=${{ env.HINDSIGHT_API_LLM_MODEL }}
|
||||
EOF
|
||||
|
||||
- name: Start API server
|
||||
run: |
|
||||
./scripts/dev/start-api.sh > /tmp/api-server.log 2>&1 &
|
||||
echo "Waiting for API server to be ready..."
|
||||
for i in {1..60}; do
|
||||
if curl -sf http://localhost:8888/health > /dev/null 2>&1; then
|
||||
echo "API server is ready after ${i}s"
|
||||
break
|
||||
fi
|
||||
if [ $i -eq 60 ]; then
|
||||
echo "API server failed to start after 60s"
|
||||
cat /tmp/api-server.log
|
||||
exit 1
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
|
||||
- name: Run Go client tests
|
||||
working-directory: ./hindsight-clients/go
|
||||
run: go test -v -tags=integration
|
||||
|
||||
- name: Show API server logs
|
||||
if: always()
|
||||
run: |
|
||||
echo "=== API Server Logs ==="
|
||||
cat /tmp/api-server.log || echo "No API server log found"
|
||||
|
||||
test-integration:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
@@ -1015,11 +941,30 @@ jobs:
|
||||
sleep 1
|
||||
done
|
||||
|
||||
- name: Run Python doc examples
|
||||
working-directory: ./hindsight-clients/python
|
||||
run: |
|
||||
for f in ../../hindsight-docs/examples/api/*.py; do
|
||||
echo "Running $f..."
|
||||
uv run python "$f"
|
||||
done
|
||||
|
||||
- name: Run Node.js doc examples
|
||||
run: |
|
||||
for f in hindsight-docs/examples/api/*.mjs; do
|
||||
echo "Running $f..."
|
||||
node "$f"
|
||||
done
|
||||
|
||||
- name: Configure CLI
|
||||
run: hindsight configure --api-url http://localhost:8888
|
||||
|
||||
- name: Run all doc examples
|
||||
run: ./scripts/test-doc-examples.sh
|
||||
- name: Run CLI doc examples
|
||||
run: |
|
||||
for f in hindsight-docs/examples/api/*.sh; do
|
||||
echo "Running $f..."
|
||||
bash "$f"
|
||||
done
|
||||
|
||||
- name: Show API server logs
|
||||
if: always()
|
||||
|
||||
@@ -46,7 +46,6 @@ hindsight-docs/static/llms-full.txt
|
||||
hindsight-dev/benchmarks/locomo/results/
|
||||
hindsight-dev/benchmarks/longmemeval/results/
|
||||
hindsight-dev/benchmarks/consolidation/results/
|
||||
hindsight-dev/benchmarks/perf/results/
|
||||
benchmarks/results/
|
||||
hindsight-cli/target
|
||||
hindsight-clients/rust/target
|
||||
|
||||
@@ -57,15 +57,8 @@ cd hindsight-control-plane && npm run dev
|
||||
|
||||
### Benchmarks
|
||||
```bash
|
||||
# Accuracy benchmarks
|
||||
./scripts/benchmarks/run-longmemeval.sh
|
||||
./scripts/benchmarks/run-locomo.sh
|
||||
|
||||
# Performance benchmarks
|
||||
./scripts/benchmarks/run-consolidation.sh
|
||||
./scripts/benchmarks/run-retain-perf.sh --document <path> # Requires API server running
|
||||
|
||||
# Results viewer
|
||||
./scripts/benchmarks/start-visualizer.sh # View results at localhost:8001
|
||||
```
|
||||
|
||||
@@ -245,61 +238,26 @@ def process(data: UserData) -> str:
|
||||
|
||||
### Adding New API Configuration Flags
|
||||
|
||||
Configuration follows a hierarchical system: **Global (env vars) → Tenant (via extension) → Bank (database)**.
|
||||
|
||||
Fields must be categorized as either **hierarchical** (can be overridden per-tenant/bank) or **static** (server-level only).
|
||||
|
||||
#### Adding a New Configuration Field
|
||||
When adding a new environment variable configuration:
|
||||
|
||||
1. **config.py** (`hindsight-api/hindsight_api/config.py`):
|
||||
- Add `ENV_*` constant for the environment variable name (e.g., `ENV_MY_SETTING = "HINDSIGHT_API_MY_SETTING"`)
|
||||
- Add `ENV_*` constant for the environment variable name
|
||||
- Add `DEFAULT_*` constant for the default value
|
||||
- Add field to `HindsightConfig` dataclass with type annotation
|
||||
- **Mark as hierarchical or static** by adding to `_HIERARCHICAL_FIELDS` set (hierarchical) or leaving it out (static)
|
||||
- Add field to `HindsightConfig` dataclass
|
||||
- Add initialization in `from_env()` method
|
||||
|
||||
```python
|
||||
# Hierarchical field (can be overridden per-bank)
|
||||
_HIERARCHICAL_FIELDS = {
|
||||
...,
|
||||
"my_setting", # Add here for hierarchical
|
||||
}
|
||||
|
||||
# Static field - just don't add to _HIERARCHICAL_FIELDS
|
||||
```
|
||||
|
||||
2. **main.py** (`hindsight-api/hindsight_api/main.py`):
|
||||
- Add field to the manual `HindsightConfig()` constructor call (search for "CLI override")
|
||||
|
||||
3. **Use hierarchical config in MemoryEngine**:
|
||||
```python
|
||||
# Config is resolved automatically per bank via ConfigResolver
|
||||
config_dict = await self._config_resolver.get_bank_config(bank_id, context)
|
||||
value = config_dict["my_setting"]
|
||||
```
|
||||
|
||||
4. **Use static config** (non-hierarchical):
|
||||
3. **Use the config** in code:
|
||||
```python
|
||||
from ...config import get_config
|
||||
config = get_config()
|
||||
value = config.my_static_field
|
||||
value = config.your_new_field
|
||||
```
|
||||
|
||||
5. **Documentation** (`hindsight-docs/docs/developer/configuration.md`):
|
||||
4. **Documentation** (`hindsight-docs/docs/developer/configuration.md`):
|
||||
- Add to appropriate section table with Variable, Description, Default
|
||||
- Mark if it's hierarchical (can be overridden per-bank)
|
||||
|
||||
#### Hierarchical vs Static Guidelines
|
||||
|
||||
**Hierarchical** (per-bank overridable):
|
||||
- LLM settings (provider, model, API key, base URL)
|
||||
- Operation-specific settings (retain mode, chunk size, etc.)
|
||||
- Feature flags that vary by customer/bank
|
||||
|
||||
**Static** (server-level only):
|
||||
- Infrastructure settings (database URL, port, host)
|
||||
- Global limits (max concurrent operations)
|
||||
- System-wide feature flags
|
||||
|
||||
## Environment Setup
|
||||
|
||||
@@ -323,4 +281,3 @@ Optional (uses local models by default):
|
||||
- `HINDSIGHT_API_EMBEDDINGS_PROVIDER`: local (default) or tei
|
||||
- `HINDSIGHT_API_RERANKER_PROVIDER`: local (default) or tei
|
||||
- `HINDSIGHT_API_DATABASE_URL`: External PostgreSQL (uses embedded pg0 by default)
|
||||
- `HINDSIGHT_API_ENABLE_BANK_CONFIG_API`: Enable per-bank config API (default: false, disabled for security)
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||

|
||||
|
||||
[Documentation](https://hindsight.vectorize.io) • [Paper](https://arxiv.org/abs/2512.12818) • [Cookbook](https://hindsight.vectorize.io/cookbook) • [Hindsight Cloud](https://ui.hindsight.vectorize.io/signup)
|
||||
[Documentation](https://hindsight.vectorize.io) • [Paper](https://arxiv.org/abs/2512.12818) • [Cookbook](https://hindsight.vectorize.io/cookbook) • [Hindsight Cloud](https://vectorize.io/hindsight/cloud)
|
||||
|
||||
[](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml)
|
||||
[](https://join.slack.com/t/hindsight-space/shared_invite/zt-3nhbm4w29-LeSJ5Ixi6j8PdiYOCPlOgg)
|
||||
|
||||
@@ -1,96 +0,0 @@
|
||||
# Nginx Reverse Proxy with Custom Base Path
|
||||
|
||||
Deploy Hindsight API under `/hindsight` (or any custom path) using Nginx reverse proxy.
|
||||
|
||||
## Quick Start (Published Image - API Only)
|
||||
|
||||
```bash
|
||||
docker-compose up
|
||||
```
|
||||
|
||||
- **API:** http://localhost:8080/hindsight/docs
|
||||
- **Control Plane:** http://localhost:9999 (direct access, not proxied)
|
||||
|
||||
## Full Stack with Custom Base Path (Requires Build)
|
||||
|
||||
**Important:** You cannot rebuild from the published image with build args. You must build from source.
|
||||
|
||||
### Build from Source with Custom Base Path
|
||||
|
||||
1. **Clone the repository** (if you haven't):
|
||||
```bash
|
||||
git clone https://github.com/vectorize-io/hindsight.git
|
||||
cd hindsight
|
||||
```
|
||||
|
||||
2. **Build with base path**:
|
||||
```bash
|
||||
docker build \
|
||||
--build-arg NEXT_PUBLIC_BASE_PATH=/hindsight \
|
||||
-f docker/standalone/Dockerfile \
|
||||
-t hindsight:custom \
|
||||
.
|
||||
```
|
||||
|
||||
3. **Update docker-compose.yml** to use your built image:
|
||||
```yaml
|
||||
services:
|
||||
hindsight:
|
||||
image: hindsight:custom # ← Change this
|
||||
environment:
|
||||
HINDSIGHT_API_BASE_PATH: /hindsight
|
||||
NEXT_PUBLIC_BASE_PATH: /hindsight
|
||||
```
|
||||
|
||||
4. **Update nginx.conf** to handle Control Plane routes (see below)
|
||||
|
||||
5. **Run**:
|
||||
```bash
|
||||
docker-compose up
|
||||
```
|
||||
|
||||
### Required nginx.conf for Full Stack
|
||||
|
||||
Replace the current `nginx.conf` with this to proxy both API and Control Plane:
|
||||
|
||||
```nginx
|
||||
events { worker_connections 1024; }
|
||||
|
||||
http {
|
||||
include /etc/nginx/mime.types;
|
||||
default_type application/octet-stream;
|
||||
|
||||
upstream hindsight_api { server hindsight:8888; }
|
||||
upstream hindsight_cp { server hindsight:9999; }
|
||||
|
||||
server {
|
||||
listen 80;
|
||||
|
||||
# API
|
||||
location ~ ^/hindsight/(docs|openapi\.json|health|metrics|v1|mcp) {
|
||||
proxy_pass http://hindsight_api;
|
||||
proxy_set_header Host $http_host;
|
||||
}
|
||||
|
||||
# Control Plane static files
|
||||
location ~ ^/hindsight/_next/ {
|
||||
proxy_pass http://hindsight_cp;
|
||||
proxy_set_header Host $http_host;
|
||||
}
|
||||
|
||||
# Control Plane UI
|
||||
location /hindsight {
|
||||
proxy_pass http://hindsight_cp;
|
||||
proxy_set_header Host $http_host;
|
||||
}
|
||||
|
||||
location = / { return 301 /hindsight; }
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Why Build is Required
|
||||
|
||||
Next.js requires `basePath` at **build time**. The published image was built without a custom base path, so you must rebuild from source with the `NEXT_PUBLIC_BASE_PATH` build arg to deploy the Control Plane under a subpath.
|
||||
|
||||
The API works without rebuild because `HINDSIGHT_API_BASE_PATH` is a runtime environment variable.
|
||||
@@ -1,88 +0,0 @@
|
||||
# Hindsight API deployment with Nginx reverse proxy (API-only)
|
||||
#
|
||||
# This example deploys Hindsight API under the path /hindsight with:
|
||||
# - Hindsight standalone image (API + Control Plane + embedded pg0)
|
||||
# - Nginx reverse proxy (API only)
|
||||
#
|
||||
# Quick Start:
|
||||
# docker-compose -f docker/docker-compose/nginx/docker-compose.yml up
|
||||
#
|
||||
# Access:
|
||||
# API (via nginx): http://localhost:8080/hindsight/docs
|
||||
# Control Plane (direct): http://localhost:9999
|
||||
#
|
||||
# For full stack deployment (API + Control Plane both under /hindsight):
|
||||
# See README.md in this directory for instructions on building with basePath.
|
||||
#
|
||||
# Note: This configuration uses the published image (no build required).
|
||||
# Control Plane is served directly because Next.js basePath requires
|
||||
# build-time configuration. See README.md for the full stack option.
|
||||
|
||||
services:
|
||||
# Hindsight (API + Control Plane + embedded pg0)
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:latest
|
||||
ports:
|
||||
- "9999:9999" # Control Plane (direct access, not proxied)
|
||||
environment:
|
||||
# API base path for reverse proxy
|
||||
HINDSIGHT_API_BASE_PATH: /hindsight
|
||||
|
||||
# LLM configuration
|
||||
# Using mock provider for testing (no API key needed)
|
||||
# For production, set OPENAI_API_KEY or ANTHROPIC_API_KEY and use a real provider
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-mock}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-not-needed-for-mock}
|
||||
HINDSIGHT_API_LLM_MODEL: ${HINDSIGHT_API_LLM_MODEL:-mock-model}
|
||||
|
||||
# Production examples (uncomment and set appropriate API key):
|
||||
# HINDSIGHT_API_LLM_PROVIDER: openai
|
||||
# HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY}
|
||||
# HINDSIGHT_API_LLM_MODEL: gpt-4o-mini
|
||||
|
||||
# HINDSIGHT_API_LLM_PROVIDER: anthropic
|
||||
# HINDSIGHT_API_LLM_API_KEY: ${ANTHROPIC_API_KEY}
|
||||
# HINDSIGHT_API_LLM_MODEL: claude-sonnet-4-20250514
|
||||
|
||||
# Server config
|
||||
HINDSIGHT_API_HOST: 0.0.0.0
|
||||
HINDSIGHT_API_PORT: 8888
|
||||
HINDSIGHT_API_LOG_LEVEL: info
|
||||
|
||||
# Control Plane config
|
||||
HINDSIGHT_CP_DATAPLANE_API_URL: http://localhost:8888
|
||||
volumes:
|
||||
# Persist embedded pg0 database
|
||||
- hindsight_data:/app/data
|
||||
# Note: Ports not exposed - access via Nginx at localhost:8080/hindsight/
|
||||
# To debug directly, uncomment these ports:
|
||||
# ports:
|
||||
# - "8888:8888" # API
|
||||
# - "9999:9999" # Control Plane
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-f", "http://localhost:8888/hindsight/health"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 3
|
||||
start_period: 30s
|
||||
networks:
|
||||
- hindsight
|
||||
|
||||
# Nginx reverse proxy
|
||||
nginx:
|
||||
image: nginx:alpine
|
||||
ports:
|
||||
- "8080:80"
|
||||
volumes:
|
||||
- ./nginx.conf:/etc/nginx/nginx.conf:ro
|
||||
depends_on:
|
||||
hindsight:
|
||||
condition: service_healthy
|
||||
networks:
|
||||
- hindsight
|
||||
|
||||
volumes:
|
||||
hindsight_data:
|
||||
|
||||
networks:
|
||||
hindsight:
|
||||
@@ -1,40 +0,0 @@
|
||||
# Nginx configuration for API-only reverse proxy
|
||||
# Control Plane accessed directly (not through nginx)
|
||||
|
||||
events {
|
||||
worker_connections 1024;
|
||||
}
|
||||
|
||||
http {
|
||||
include /etc/nginx/mime.types;
|
||||
default_type application/octet-stream;
|
||||
|
||||
# Logging
|
||||
access_log /var/log/nginx/access.log;
|
||||
error_log /var/log/nginx/error.log;
|
||||
|
||||
# Upstream - Hindsight API
|
||||
upstream hindsight_api {
|
||||
server hindsight:8888;
|
||||
}
|
||||
|
||||
server {
|
||||
listen 80;
|
||||
server_name _;
|
||||
|
||||
# API endpoints - forward with /hindsight prefix
|
||||
location /hindsight/ {
|
||||
proxy_pass http://hindsight_api;
|
||||
|
||||
proxy_set_header Host $http_host;
|
||||
proxy_set_header X-Real-IP $remote_addr;
|
||||
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
|
||||
proxy_set_header X-Forwarded-Proto $scheme;
|
||||
}
|
||||
|
||||
# Redirect root to API docs
|
||||
location = / {
|
||||
return 301 /hindsight/docs;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,32 +0,0 @@
|
||||
# PostgreSQL with pgvector and pg_textsearch extensions
|
||||
# Note: pg_textsearch requires PostgreSQL 17+
|
||||
FROM postgres:17
|
||||
|
||||
# Install build dependencies
|
||||
RUN apt-get update && apt-get install -y \
|
||||
build-essential \
|
||||
git \
|
||||
postgresql-server-dev-17 \
|
||||
libpq-dev \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install pgvector
|
||||
RUN cd /tmp && \
|
||||
git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git && \
|
||||
cd pgvector && \
|
||||
make && \
|
||||
make install
|
||||
|
||||
# Install pg_textsearch
|
||||
RUN cd /tmp && \
|
||||
git clone https://github.com/timescale/pg_textsearch.git && \
|
||||
cd pg_textsearch && \
|
||||
make && \
|
||||
make install
|
||||
|
||||
# Clean up source files and build dependencies
|
||||
RUN rm -rf /tmp/pgvector /tmp/pg_textsearch && \
|
||||
apt-get purge -y --auto-remove build-essential git postgresql-server-dev-17
|
||||
|
||||
# Ensure extensions are preloaded
|
||||
RUN echo "shared_preload_libraries = 'pg_textsearch'" >> /usr/share/postgresql/postgresql.conf.sample
|
||||
@@ -1,91 +0,0 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and Timescale pg_textsearch
|
||||
# docker compose -f docker/docker-compose/pg_textsearch/docker-compose.yaml down && sleep 2 && docker compose -f docker/docker-compose/pg_textsearch/docker-compose.yaml up -d
|
||||
# Make sure to set the required environment variables before running:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see below in the hindsight service)
|
||||
#
|
||||
# Usage:
|
||||
# docker compose up -d
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
|
||||
services:
|
||||
db:
|
||||
# Use custom PostgreSQL image with pgvector and pg_textsearch extensions
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
# Expose PostgreSQL port
|
||||
ports:
|
||||
- "5437:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
pg-textsearch-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pg_textsearch CASCADE;';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvector
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pg_textsearch
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -1,83 +0,0 @@
|
||||
# Docker Compose file for Hindsight with S3 file storage (SeaweedFS)
|
||||
#
|
||||
# SeaweedFS (Apache 2.0) provides an S3-compatible object storage backend
|
||||
# for storing uploaded files instead of PostgreSQL BYTEA storage.
|
||||
#
|
||||
# Make sure to set the required environment variables before running:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see below in the hindsight service)
|
||||
#
|
||||
# Usage:
|
||||
# docker compose up -d
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_DB_VERSION: PostgreSQL version (default: 18)
|
||||
# - SEAWEEDFS_S3_ACCESS_KEY: S3 access key (default: hindsight_s3_key)
|
||||
# - SEAWEEDFS_S3_SECRET_KEY: S3 secret key (default: hindsight_s3_secret)
|
||||
|
||||
services:
|
||||
db:
|
||||
image: pgvector/pgvector:pg${HINDSIGHT_DB_VERSION:-18}
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:?Please set the HINDSIGHT_DB_PASSWORD env variable}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/${HINDSIGHT_DB_VERSION:-18}/docker
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
seaweedfs:
|
||||
image: chrislusf/seaweedfs:latest
|
||||
container_name: hindsight-seaweedfs
|
||||
restart: always
|
||||
# Single-node mode: master + volume + filer + S3 gateway all in one process
|
||||
command: >
|
||||
server
|
||||
-s3
|
||||
-s3.port=8333
|
||||
-s3.config=/etc/seaweedfs/s3.json
|
||||
-ip.bind=0.0.0.0
|
||||
volumes:
|
||||
- seaweedfs_data:/data
|
||||
- ./s3.json:/etc/seaweedfs/s3.json:ro
|
||||
# Expose S3 API port (uncomment to access from host)
|
||||
# ports:
|
||||
# - "8333:8333"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
- HINDSIGHT_API_LLM_API_KEY=${OPENAI_API_KEY?Please set the OPENAI_API_KEY env variable}
|
||||
- HINDSIGHT_API_DATABASE_URL=postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:?Please set the HINDSIGHT_DB_PASSWORD env variable}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
# S3 file storage configuration (SeaweedFS)
|
||||
- HINDSIGHT_API_FILE_STORAGE_TYPE=s3
|
||||
- HINDSIGHT_API_FILE_STORAGE_S3_BUCKET=hindsight
|
||||
- HINDSIGHT_API_FILE_STORAGE_S3_ENDPOINT=http://seaweedfs:8333
|
||||
- HINDSIGHT_API_FILE_STORAGE_S3_REGION=us-east-1
|
||||
- HINDSIGHT_API_FILE_STORAGE_S3_ACCESS_KEY_ID=${SEAWEEDFS_S3_ACCESS_KEY:-hindsight_s3_key}
|
||||
- HINDSIGHT_API_FILE_STORAGE_S3_SECRET_ACCESS_KEY=${SEAWEEDFS_S3_SECRET_KEY:-hindsight_s3_secret}
|
||||
depends_on:
|
||||
- db
|
||||
- seaweedfs
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
seaweedfs_data:
|
||||
@@ -1,19 +0,0 @@
|
||||
{
|
||||
"identities": [
|
||||
{
|
||||
"name": "hindsight",
|
||||
"credentials": [
|
||||
{
|
||||
"accessKey": "hindsight_s3_key",
|
||||
"secretKey": "hindsight_s3_secret"
|
||||
}
|
||||
],
|
||||
"actions": [
|
||||
"Admin",
|
||||
"Read",
|
||||
"Write",
|
||||
"List"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -1,16 +0,0 @@
|
||||
# Git
|
||||
.git
|
||||
.gitignore
|
||||
.gitattributes
|
||||
|
||||
# Docker
|
||||
docker-compose.yaml
|
||||
.dockerignore
|
||||
|
||||
# Documentation
|
||||
README.md
|
||||
*.md
|
||||
|
||||
# Environment
|
||||
.env
|
||||
.env.example
|
||||
@@ -1,25 +0,0 @@
|
||||
# PostgreSQL Configuration
|
||||
HINDSIGHT_DB_USER=hindsight_user
|
||||
HINDSIGHT_DB_PASSWORD=change-me-to-secure-password
|
||||
HINDSIGHT_DB_NAME=hindsight_db
|
||||
|
||||
# Hindsight Version
|
||||
HINDSIGHT_VERSION=latest
|
||||
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
OPENAI_API_KEY=your-openai-api-key-here
|
||||
|
||||
# Alternative LLM providers (uncomment and configure as needed):
|
||||
# HINDSIGHT_API_LLM_PROVIDER=anthropic
|
||||
# ANTHROPIC_API_KEY=your-anthropic-api-key
|
||||
|
||||
# HINDSIGHT_API_LLM_PROVIDER=gemini
|
||||
# GEMINI_API_KEY=your-gemini-api-key
|
||||
|
||||
# HINDSIGHT_API_LLM_PROVIDER=groq
|
||||
# GROQ_API_KEY=your-groq-api-key
|
||||
|
||||
# Vector and Text Search (already configured in docker-compose.yaml)
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION=pg_textsearch
|
||||
@@ -1,55 +0,0 @@
|
||||
# PostgreSQL with pgvector, pgvectorscale, and pg_textsearch extensions
|
||||
# All three extensions from Timescale/pgvector for high-performance vector and text search
|
||||
# Note: Requires PostgreSQL 16+
|
||||
FROM postgres:17
|
||||
|
||||
# Install build dependencies and Rust toolchain
|
||||
RUN apt-get update && apt-get install -y \
|
||||
build-essential \
|
||||
git \
|
||||
postgresql-server-dev-17 \
|
||||
libpq-dev \
|
||||
cmake \
|
||||
curl \
|
||||
pkg-config \
|
||||
libssl-dev \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install Rust toolchain (required for pgvectorscale)
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
|
||||
ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
|
||||
# Install pgvector (required by pgvectorscale)
|
||||
RUN cd /tmp && \
|
||||
git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git && \
|
||||
cd pgvector && \
|
||||
make && \
|
||||
make install && \
|
||||
rm -rf /tmp/pgvector
|
||||
|
||||
# Install cargo-pgrx (PostgreSQL extension framework for Rust)
|
||||
RUN cargo install cargo-pgrx --version 0.12.5 --locked && \
|
||||
cargo pgrx init --pg17 /usr/bin/pg_config
|
||||
|
||||
# Install pgvectorscale (DiskANN index support)
|
||||
RUN cd /tmp && \
|
||||
git clone --branch 0.5.1 https://github.com/timescale/pgvectorscale.git && \
|
||||
cd pgvectorscale/pgvectorscale && \
|
||||
cargo pgrx install --release && \
|
||||
rm -rf /tmp/pgvectorscale
|
||||
|
||||
# Install pg_textsearch (BM25 text search)
|
||||
RUN cd /tmp && \
|
||||
git clone https://github.com/timescale/pg_textsearch.git && \
|
||||
cd pg_textsearch && \
|
||||
make && \
|
||||
make install && \
|
||||
rm -rf /tmp/pg_textsearch
|
||||
|
||||
# Clean up build dependencies (keep runtime dependencies)
|
||||
RUN apt-get purge -y --auto-remove git cmake curl && \
|
||||
rm -rf /root/.cargo/registry /root/.cargo/git
|
||||
|
||||
# Ensure extensions are preloaded (pg_textsearch requires preloading)
|
||||
RUN echo "shared_preload_libraries = 'pg_textsearch'" >> /usr/share/postgresql/postgresql.conf.sample
|
||||
|
||||
@@ -1,101 +0,0 @@
|
||||
# Hindsight with Timescale Extensions
|
||||
|
||||
This Docker Compose setup provides a complete Hindsight deployment with **Timescale extensions**:
|
||||
- **pgvectorscale** - DiskANN algorithm for disk-based scalable vector search
|
||||
- **pg_textsearch** - High-performance BM25 text search
|
||||
|
||||
Both extensions are from [Timescale](https://github.com/timescale) and provide production-grade performance.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Docker and Docker Compose installed
|
||||
- OpenAI API key (or another LLM provider)
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
# Set environment variables
|
||||
export HINDSIGHT_DB_PASSWORD="your-secure-password"
|
||||
export OPENAI_API_KEY="your-openai-api-key"
|
||||
|
||||
# Build and start
|
||||
docker compose -f docker/docker-compose/timescale/docker-compose.yaml up -d --build
|
||||
|
||||
# Check logs
|
||||
|
||||
docker compose -f docker/docker-compose/timescale/docker-compose.yaml logs -f
|
||||
```
|
||||
|
||||
**Access:**
|
||||
- API: http://localhost:8888
|
||||
- Control Plane: http://localhost:9999
|
||||
|
||||
## Stop and Clean Up
|
||||
|
||||
```bash
|
||||
# Stop services
|
||||
docker compose -f docker/docker-compose/timescale/docker-compose.yaml down
|
||||
|
||||
# Remove volumes (deletes all data)
|
||||
docker compose -f docker/docker-compose/timescale/docker-compose.yaml down -v
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
### Environment Variables
|
||||
|
||||
| Variable | Description | Default |
|
||||
|----------|-------------|---------|
|
||||
| `HINDSIGHT_DB_PASSWORD` | PostgreSQL password | `hindsight_password` |
|
||||
| `HINDSIGHT_DB_USER` | PostgreSQL username | `hindsight_user` |
|
||||
| `HINDSIGHT_DB_NAME` | Database name | `hindsight_db` |
|
||||
| `HINDSIGHT_VERSION` | Hindsight Docker image version | `latest` |
|
||||
| `OPENAI_API_KEY` | OpenAI API key | (required) |
|
||||
| `HINDSIGHT_API_LLM_PROVIDER` | LLM provider | `openai` |
|
||||
|
||||
### Why Timescale Extensions?
|
||||
|
||||
**pgvectorscale (DiskANN):**
|
||||
- 28x lower p95 latency vs dedicated vector databases
|
||||
- 16x higher query throughput at 99% recall
|
||||
- 60-75% cost reduction (disk is cheaper than RAM)
|
||||
- Best for large datasets (10M+ vectors)
|
||||
|
||||
**pg_textsearch (BM25):**
|
||||
- High-performance keyword retrieval
|
||||
- Native BM25 ranking algorithm
|
||||
- Optimized for full-text search
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Extensions not installed
|
||||
|
||||
Check if extensions are available:
|
||||
|
||||
```bash
|
||||
docker exec -it hindsight-db-timescale psql -U hindsight_user -d hindsight_db -c "\dx"
|
||||
```
|
||||
|
||||
You should see:
|
||||
- `vector` (pgvector)
|
||||
- `vectorscale` (pgvectorscale/DiskANN)
|
||||
- `pg_textsearch` (BM25 search)
|
||||
|
||||
### Build fails
|
||||
|
||||
If the Docker build fails during pgvectorscale compilation:
|
||||
|
||||
1. Ensure you have sufficient memory (recommended: 4GB+)
|
||||
2. Check Docker build logs for Rust compilation errors
|
||||
3. Try building with more resources: `docker compose build --no-cache --memory 4g`
|
||||
|
||||
### Port conflicts
|
||||
|
||||
If port 5438 is already in use, modify the `ports` section in docker-compose.yaml.
|
||||
|
||||
## Learn More
|
||||
|
||||
- [pgvectorscale GitHub](https://github.com/timescale/pgvectorscale)
|
||||
- [pg_textsearch GitHub](https://github.com/timescale/pg_textsearch)
|
||||
- [HNSW vs DiskANN](https://www.tigerdata.com/learn/hnsw-vs-diskann)
|
||||
- [Hindsight Documentation](https://hindsight.dev)
|
||||
@@ -1,108 +0,0 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with Timescale extensions
|
||||
# - pgvectorscale: DiskANN vector search (disk-based, scalable)
|
||||
# - pg_textsearch: BM25 text search (high-performance keyword retrieval)
|
||||
#
|
||||
# Quick start:
|
||||
# docker compose -f docker/docker-compose/timescale/docker-compose.yaml up -d --build
|
||||
#
|
||||
# Required environment variables:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - OPENAI_API_KEY (or configure another LLM provider)
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
|
||||
services:
|
||||
db:
|
||||
# Custom PostgreSQL image with Timescale extensions (pgvectorscale + pg_textsearch)
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db-timescale
|
||||
restart: always
|
||||
# Expose PostgreSQL port (using 5438 to avoid conflicts with other setups)
|
||||
ports:
|
||||
- "5438:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
# Health check to ensure database is ready
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U hindsight_user"]
|
||||
interval: 5s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
|
||||
timescale-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
db:
|
||||
condition: service_healthy
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db-timescale -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Installing Timescale extensions...';
|
||||
echo '1/3: Installing pgvector (required by pgvectorscale)...';
|
||||
psql -h hindsight-db-timescale -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
echo '2/3: Installing pgvectorscale (DiskANN vector search)...';
|
||||
psql -h hindsight-db-timescale -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vectorscale CASCADE;';
|
||||
echo '3/3: Installing pg_textsearch (BM25 text search)...';
|
||||
psql -h hindsight-db-timescale -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pg_textsearch CASCADE;';
|
||||
echo '';
|
||||
echo '✅ Timescale extensions installed successfully';
|
||||
echo '';
|
||||
echo 'Installed extensions:';
|
||||
psql -h hindsight-db-timescale -p 5432 -U hindsight_user -d hindsight_db -c \"\\dx\" | grep -E '(vector|vectorscale|pg_textsearch)';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app-timescale
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Timescale Extensions
|
||||
# pgvectorscale: DiskANN algorithm for disk-based scalable vector search
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvectorscale
|
||||
# pg_textsearch: High-performance BM25 text search
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pg_textsearch
|
||||
|
||||
depends_on:
|
||||
db:
|
||||
condition: service_healthy
|
||||
timescale-init:
|
||||
condition: service_completed_successfully
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -1,93 +0,0 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and vectorchord
|
||||
# docker compose -f docker/docker-compose/docker-compose.yaml down && sleep 2 && docker compose -f docker/docker-compose/docker-compose.yaml up -d
|
||||
# Make sure to set the required environment variables before running:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see below in the hindsight service)
|
||||
#
|
||||
# Usage:
|
||||
# docker compose up -d
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_DB_VERSION: PostgreSQL version (default: 18)
|
||||
|
||||
services:
|
||||
db:
|
||||
# Use a PostgreSQL-Image with vectorchord extension pre-installed
|
||||
image: tensorchord/vchord-suite:pg${HINDSIGHT_DB_VERSION:-18-latest}
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
# Expose PostgreSQL port
|
||||
ports:
|
||||
- "5436:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/${HINDSIGHT_DB_VERSION:-18}/docker
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
vectorchord-init:
|
||||
image: tensorchord/vchord-suite:pg18-latest
|
||||
#container_name: vectorchord-init
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pg_tokenizer CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vchord_bm25 CASCADE;';
|
||||
echo 'Creating llmlingua2 tokenizer';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c \"SELECT create_tokenizer('llmlingua2', \\$\\$ model = \\\"llmlingua2\\\" \\$\\$);\" 2>/dev/null || echo 'Tokenizer already exists or creation skipped';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration (uses OpenAI for testing vchord)
|
||||
# LLM configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: vchord
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: vchord
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -112,10 +112,6 @@ RUN rm -f package-lock.json && sed -i '/"@vectorize-io\/hindsight-client":/d' pa
|
||||
# Copy built SDK directly into node_modules (more reliable than npm link in Docker)
|
||||
COPY --from=sdk-builder /app/hindsight-clients/typescript ./node_modules/@vectorize-io/hindsight-client
|
||||
|
||||
# Accept base path as build argument for reverse proxy deployments
|
||||
# Usage: docker build --build-arg NEXT_PUBLIC_BASE_PATH=/hindsight ...
|
||||
ARG NEXT_PUBLIC_BASE_PATH=""
|
||||
|
||||
# Build Control Plane - run next build first, then custom standalone copy
|
||||
# (The build:standalone script expects a specific path structure that differs in Docker)
|
||||
RUN npm exec -- next build
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.4.11
|
||||
appVersion: "0.4.11"
|
||||
version: 0.4.10
|
||||
appVersion: "0.4.10"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -46,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.4.11"
|
||||
__version__ = "0.4.10"
|
||||
|
||||
@@ -6,7 +6,6 @@ Create Date: 2025-11-27 11:54:19.228030
|
||||
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
@@ -22,96 +21,6 @@ branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _detect_vector_extension() -> str:
|
||||
"""
|
||||
Detect or validate vector extension: 'pgvector', 'vchord', or 'pgvectorscale'.
|
||||
Respects HINDSIGHT_API_VECTOR_EXTENSION env var if set.
|
||||
"""
|
||||
conn = op.get_bind()
|
||||
vector_extension = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
|
||||
|
||||
# Validate configured extension is installed
|
||||
if vector_extension == "pgvectorscale":
|
||||
# pgvectorscale/DiskANN requires pgvector
|
||||
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
|
||||
if not pgvector_check:
|
||||
raise RuntimeError(
|
||||
"DiskANN requires pgvector. Install with: CREATE EXTENSION vector; then vectorscale or pg_diskann CASCADE;"
|
||||
)
|
||||
# Check for either vectorscale (open source) or pg_diskann (Azure)
|
||||
vectorscale_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vectorscale'")).scalar()
|
||||
pg_diskann_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_diskann'")).scalar()
|
||||
|
||||
if vectorscale_check:
|
||||
return "pgvectorscale"
|
||||
elif pg_diskann_check:
|
||||
return "pg_diskann"
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"Configured vector extension 'pgvectorscale' not found. Install either:\n"
|
||||
" - pgvectorscale: CREATE EXTENSION vectorscale CASCADE;\n"
|
||||
" - pg_diskann (Azure): CREATE EXTENSION pg_diskann CASCADE;"
|
||||
)
|
||||
elif vector_extension == "vchord":
|
||||
vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar()
|
||||
if not vchord_check:
|
||||
raise RuntimeError(
|
||||
"Configured vector extension 'vchord' not found. Install it with: CREATE EXTENSION vchord CASCADE;"
|
||||
)
|
||||
return "vchord"
|
||||
elif vector_extension == "pgvector":
|
||||
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
|
||||
if not pgvector_check:
|
||||
raise RuntimeError(
|
||||
"Configured vector extension 'pgvector' not found. Install it with: CREATE EXTENSION vector;"
|
||||
)
|
||||
return "pgvector"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {vector_extension}. Must be 'pgvector', 'vchord', or 'pgvectorscale'"
|
||||
)
|
||||
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
"""
|
||||
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
if text_search_extension == "vchord":
|
||||
# Create vchord_bm25 extension if not exists
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS vchord_bm25 CASCADE")
|
||||
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 = 'vchord_bm25'")).fetchone()
|
||||
if not result:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "vchord"
|
||||
elif text_search_extension == "pg_textsearch":
|
||||
# Create pg_textsearch extension if not exists
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_textsearch CASCADE")
|
||||
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 = 'pg_textsearch'")).fetchone()
|
||||
if not result:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_textsearch"
|
||||
elif text_search_extension == "native":
|
||||
return "native"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade schema - create all tables from scratch."""
|
||||
|
||||
@@ -257,29 +166,11 @@ def upgrade() -> None:
|
||||
)
|
||||
|
||||
# Add search_vector column for full-text search
|
||||
# Type depends on configured text search backend
|
||||
text_search_ext = _detect_text_search_extension()
|
||||
|
||||
if text_search_ext == "vchord":
|
||||
# VectorChord BM25: bm25vector type (no GENERATED - tokenization happens on INSERT)
|
||||
# Note: vchord_bm25 extension creates types in bm25_catalog schema
|
||||
op.execute("""
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector bm25_catalog.bm25vector
|
||||
""")
|
||||
elif text_search_ext == "pg_textsearch":
|
||||
# Timescale pg_textsearch: dummy TEXT column for consistency (indexes operate on base columns directly)
|
||||
op.execute("""
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute("""
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS (to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))) STORED
|
||||
""")
|
||||
op.execute("""
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS (to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))) STORED
|
||||
""")
|
||||
|
||||
op.create_index("idx_memory_units_bank_id", "memory_units", ["bank_id"])
|
||||
op.create_index("idx_memory_units_document_id", "memory_units", ["document_id"])
|
||||
@@ -309,61 +200,19 @@ def upgrade() -> None:
|
||||
["bank_id", sa.text("event_date DESC")],
|
||||
postgresql_where=sa.text("fact_type = 'observation'"),
|
||||
)
|
||||
# Create vector index - conditional based on available extension
|
||||
vector_ext = _detect_vector_extension()
|
||||
op.create_index(
|
||||
"idx_memory_units_embedding",
|
||||
"memory_units",
|
||||
["embedding"],
|
||||
postgresql_using="hnsw",
|
||||
postgresql_ops={"embedding": "vector_cosine_ops"},
|
||||
)
|
||||
|
||||
if vector_ext == "pgvectorscale":
|
||||
# Use DiskANN index for pgvectorscale (disk-based, scalable)
|
||||
op.execute("""
|
||||
CREATE INDEX idx_memory_units_embedding ON memory_units
|
||||
USING diskann (embedding vector_cosine_ops)
|
||||
WITH (num_neighbors = 50)
|
||||
""")
|
||||
elif vector_ext == "pg_diskann":
|
||||
# Use DiskANN index for pg_diskann (Azure)
|
||||
op.execute("""
|
||||
CREATE INDEX idx_memory_units_embedding ON memory_units
|
||||
USING diskann (embedding vector_cosine_ops)
|
||||
WITH (max_neighbors = 50)
|
||||
""")
|
||||
elif vector_ext == "vchord":
|
||||
# Use vchordrq index for vchord (supports high-dimensional embeddings)
|
||||
op.execute("""
|
||||
CREATE INDEX idx_memory_units_embedding ON memory_units
|
||||
USING vchordrq (embedding vector_l2_ops)
|
||||
""")
|
||||
else: # pgvector
|
||||
# Use HNSW index for pgvector
|
||||
op.create_index(
|
||||
"idx_memory_units_embedding",
|
||||
"memory_units",
|
||||
["embedding"],
|
||||
postgresql_using="hnsw",
|
||||
postgresql_ops={"embedding": "vector_cosine_ops"},
|
||||
)
|
||||
|
||||
# Create full-text search index on search_vector
|
||||
# Index type depends on text search backend
|
||||
if text_search_ext == "vchord":
|
||||
# VectorChord BM25 index
|
||||
op.execute("""
|
||||
CREATE INDEX idx_memory_units_text_search ON memory_units
|
||||
USING bm25 (search_vector bm25_catalog.bm25_ops)
|
||||
""")
|
||||
elif text_search_ext == "pg_textsearch":
|
||||
# Timescale pg_textsearch BM25 index on text column
|
||||
# Note: pg_textsearch doesn't support expressions, so we index the main text column
|
||||
op.execute("""
|
||||
CREATE INDEX idx_memory_units_text_search ON memory_units
|
||||
USING bm25(text)
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL GIN index
|
||||
op.execute("""
|
||||
CREATE INDEX idx_memory_units_text_search ON memory_units
|
||||
USING gin(search_vector)
|
||||
""")
|
||||
# Create BM25 full-text search index on search_vector
|
||||
op.execute("""
|
||||
CREATE INDEX idx_memory_units_text_search ON memory_units
|
||||
USING gin(search_vector)
|
||||
""")
|
||||
|
||||
op.execute("""
|
||||
CREATE MATERIALIZED VIEW memory_units_bm25 AS
|
||||
|
||||
@@ -1,70 +0,0 @@
|
||||
"""Add file_storage table for BYTEA-based file storage
|
||||
|
||||
Revision ID: a1b2c3d4e5f6
|
||||
Revises: y0t1u2v3w4x5
|
||||
Create Date: 2026-02-16
|
||||
|
||||
Creates a dedicated table for storing uploaded files using BYTEA.
|
||||
This provides zero-config file storage that "just works" for development
|
||||
and small deployments. For production/scale, use S3-compatible storage.
|
||||
|
||||
Files are stored in a separate table to avoid bloating the documents table.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "a1b2c3d4e5f6"
|
||||
down_revision: str | Sequence[str] | None = "y0t1u2v3w4x5"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Create file_storage table for BYTEA storage."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Create file_storage table (minimal: just key + data)
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE {schema}file_storage (
|
||||
storage_key TEXT PRIMARY KEY,
|
||||
data BYTEA NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Add file tracking columns to documents table
|
||||
op.execute(
|
||||
f"""
|
||||
ALTER TABLE {schema}documents
|
||||
ADD COLUMN IF NOT EXISTS file_storage_key TEXT,
|
||||
ADD COLUMN IF NOT EXISTS file_original_name TEXT,
|
||||
ADD COLUMN IF NOT EXISTS file_content_type TEXT
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Remove file_storage table and related columns."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Drop columns from documents table
|
||||
op.execute(
|
||||
f"""
|
||||
ALTER TABLE {schema}documents
|
||||
DROP COLUMN IF EXISTS file_storage_key,
|
||||
DROP COLUMN IF EXISTS file_original_name,
|
||||
DROP COLUMN IF EXISTS file_content_type
|
||||
"""
|
||||
)
|
||||
|
||||
# Drop file_storage table
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}file_storage")
|
||||
+21
-204
@@ -10,11 +10,9 @@ This migration:
|
||||
3. Adds consolidation tracking columns to the 'banks' table
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "n9i0j1k2l3m4"
|
||||
@@ -29,106 +27,10 @@ def _get_schema_prefix() -> str:
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _detect_vector_extension() -> str:
|
||||
"""
|
||||
Detect or validate vector extension: 'pgvector', 'vchord', or 'pgvectorscale'.
|
||||
Respects HINDSIGHT_API_VECTOR_EXTENSION env var if set.
|
||||
"""
|
||||
conn = op.get_bind()
|
||||
vector_extension = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
|
||||
|
||||
# Validate configured extension is installed
|
||||
if vector_extension == "pgvectorscale":
|
||||
# pgvectorscale/DiskANN requires pgvector
|
||||
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
|
||||
if not pgvector_check:
|
||||
raise RuntimeError(
|
||||
"DiskANN requires pgvector. Install with: CREATE EXTENSION vector; then vectorscale or pg_diskann CASCADE;"
|
||||
)
|
||||
# Check for either vectorscale (open source) or pg_diskann (Azure)
|
||||
vectorscale_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vectorscale'")).scalar()
|
||||
pg_diskann_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_diskann'")).scalar()
|
||||
|
||||
if vectorscale_check:
|
||||
return "pgvectorscale"
|
||||
elif pg_diskann_check:
|
||||
return "pg_diskann"
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"Configured vector extension 'pgvectorscale' not found. Install either:\n"
|
||||
" - pgvectorscale: CREATE EXTENSION vectorscale CASCADE;\n"
|
||||
" - pg_diskann (Azure): CREATE EXTENSION pg_diskann CASCADE;"
|
||||
)
|
||||
elif vector_extension == "vchord":
|
||||
vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar()
|
||||
if not vchord_check:
|
||||
raise RuntimeError(
|
||||
"Configured vector extension 'vchord' not found. Install it with: CREATE EXTENSION vchord CASCADE;"
|
||||
)
|
||||
return "vchord"
|
||||
elif vector_extension == "pgvector":
|
||||
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
|
||||
if not pgvector_check:
|
||||
raise RuntimeError(
|
||||
"Configured vector extension 'pgvector' not found. Install it with: CREATE EXTENSION vector;"
|
||||
)
|
||||
return "pgvector"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {vector_extension}. Must be 'pgvector', 'vchord', or 'pgvectorscale'"
|
||||
)
|
||||
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
"""
|
||||
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
if text_search_extension == "vchord":
|
||||
# Create vchord_bm25 extension if not exists
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS vchord_bm25 CASCADE")
|
||||
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 = 'vchord_bm25'")).fetchone()
|
||||
if not result:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "vchord"
|
||||
elif text_search_extension == "pg_textsearch":
|
||||
# Create pg_textsearch extension if not exists
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_textsearch CASCADE")
|
||||
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 = 'pg_textsearch'")).fetchone()
|
||||
if not result:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_textsearch"
|
||||
elif text_search_extension == "native":
|
||||
return "native"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Create learnings and pinned_reflections tables."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Detect which vector extension is available
|
||||
vector_ext = _detect_vector_extension()
|
||||
|
||||
# Detect which text search extension to use
|
||||
text_search_ext = _detect_text_search_extension()
|
||||
|
||||
# 1. Create learnings table
|
||||
op.execute(f"""
|
||||
CREATE TABLE {schema}learnings (
|
||||
@@ -155,60 +57,18 @@ def upgrade() -> None:
|
||||
|
||||
# Indexes for learnings
|
||||
op.execute(f"CREATE INDEX idx_learnings_bank_id ON {schema}learnings(bank_id)")
|
||||
|
||||
# Create vector index based on detected extension
|
||||
if vector_ext == "pgvectorscale":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_embedding ON {schema}learnings
|
||||
USING diskann (embedding vector_cosine_ops)
|
||||
WITH (num_neighbors = 50)
|
||||
""")
|
||||
elif vector_ext == "pg_diskann":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_embedding ON {schema}learnings
|
||||
USING diskann (embedding vector_cosine_ops)
|
||||
WITH (max_neighbors = 50)
|
||||
""")
|
||||
elif vector_ext == "vchord":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_embedding ON {schema}learnings
|
||||
USING vchordrq (embedding vector_l2_ops)
|
||||
""")
|
||||
else: # pgvector
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_embedding ON {schema}learnings
|
||||
USING hnsw (embedding vector_cosine_ops)
|
||||
""")
|
||||
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_embedding ON {schema}learnings
|
||||
USING hnsw (embedding vector_cosine_ops)
|
||||
""")
|
||||
op.execute(f"CREATE INDEX idx_learnings_tags ON {schema}learnings USING GIN(tags)")
|
||||
|
||||
# Full-text search for learnings
|
||||
if text_search_ext == "vchord":
|
||||
# VectorChord BM25: bm25vector type (no GENERATED - tokenization happens on INSERT)
|
||||
# Note: vchord_bm25 extension creates types in bm25_catalog schema
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}learnings ADD COLUMN search_vector bm25_catalog.bm25vector
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25 (search_vector bm25_catalog.bm25_ops)
|
||||
""")
|
||||
elif text_search_ext == "pg_textsearch":
|
||||
# Timescale pg_textsearch: dummy TEXT column for consistency (indexes operate on base columns directly)
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}learnings ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25(text) WITH (text_config='english')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}learnings ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS (to_tsvector('english', text)) STORED
|
||||
""")
|
||||
op.execute(f"CREATE INDEX idx_learnings_text_search ON {schema}learnings USING gin(search_vector)")
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}learnings ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS (to_tsvector('english', text)) STORED
|
||||
""")
|
||||
op.execute(f"CREATE INDEX idx_learnings_text_search ON {schema}learnings USING gin(search_vector)")
|
||||
|
||||
# 2. Create pinned_reflections table
|
||||
op.execute(f"""
|
||||
@@ -234,64 +94,21 @@ def upgrade() -> None:
|
||||
|
||||
# Indexes for pinned_reflections
|
||||
op.execute(f"CREATE INDEX idx_pinned_reflections_bank_id ON {schema}pinned_reflections(bank_id)")
|
||||
|
||||
# Create vector index based on detected extension
|
||||
if vector_ext == "pgvectorscale":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_embedding ON {schema}pinned_reflections
|
||||
USING diskann (embedding vector_cosine_ops)
|
||||
WITH (num_neighbors = 50)
|
||||
""")
|
||||
elif vector_ext == "pg_diskann":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_embedding ON {schema}pinned_reflections
|
||||
USING diskann (embedding vector_cosine_ops)
|
||||
WITH (max_neighbors = 50)
|
||||
""")
|
||||
elif vector_ext == "vchord":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_embedding ON {schema}pinned_reflections
|
||||
USING vchordrq (embedding vector_l2_ops)
|
||||
""")
|
||||
else: # pgvector
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_embedding ON {schema}pinned_reflections
|
||||
USING hnsw (embedding vector_cosine_ops)
|
||||
""")
|
||||
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_embedding ON {schema}pinned_reflections
|
||||
USING hnsw (embedding vector_cosine_ops)
|
||||
""")
|
||||
op.execute(f"CREATE INDEX idx_pinned_reflections_tags ON {schema}pinned_reflections USING GIN(tags)")
|
||||
|
||||
# Full-text search for pinned_reflections
|
||||
if text_search_ext == "vchord":
|
||||
# VectorChord BM25: bm25vector type (no GENERATED - tokenization happens on INSERT/UPDATE)
|
||||
# Note: vchord_bm25 extension creates types in bm25_catalog schema
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector bm25_catalog.bm25vector
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
|
||||
USING bm25 (search_vector bm25_catalog.bm25_ops)
|
||||
""")
|
||||
elif text_search_ext == "pg_textsearch":
|
||||
# Timescale pg_textsearch: dummy TEXT column for consistency (indexes operate on base columns directly)
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
|
||||
USING bm25(content)
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS (to_tsvector('english', COALESCE(name, '') || ' ' || content)) STORED
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
|
||||
USING gin(search_vector)
|
||||
""")
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS (to_tsvector('english', COALESCE(name, '') || ' ' || content)) STORED
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
|
||||
USING gin(search_vector)
|
||||
""")
|
||||
|
||||
# 3. Add consolidation tracking columns to banks table
|
||||
op.execute(f"""
|
||||
|
||||
@@ -1,64 +0,0 @@
|
||||
"""Add config JSONB column to banks table for hierarchical configuration
|
||||
|
||||
Revision ID: x9s0t1u2v3w4
|
||||
Revises: w8r9s0t1u2v3
|
||||
Create Date: 2026-02-09
|
||||
|
||||
This migration adds a `config` JSONB column to the banks table to support
|
||||
per-bank configuration overrides. This enables hierarchical configuration where:
|
||||
- Global config is loaded from environment variables
|
||||
- Tenant config is provided via TenantExtension
|
||||
- Bank config overrides are stored in banks.config JSONB column
|
||||
|
||||
The config column stores overrides for hierarchical fields (LLM settings,
|
||||
retention parameters, retrieval settings, etc.) in Python field name format.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import context, op
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
|
||||
revision: str = "x9s0t1u2v3w4"
|
||||
down_revision: str | Sequence[str] | None = "w8r9s0t1u2v3"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Add config JSONB column to banks table with GIN index."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Add config column to banks table
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}banks
|
||||
ADD COLUMN config JSONB NOT NULL DEFAULT '{{}}'::jsonb
|
||||
""")
|
||||
|
||||
# Add GIN index for efficient JSONB queries
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_banks_config
|
||||
ON {schema}banks
|
||||
USING gin(config)
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Remove config column and index from banks table."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Drop index first
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_banks_config")
|
||||
|
||||
# Drop column
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}banks
|
||||
DROP COLUMN IF EXISTS config
|
||||
""")
|
||||
-49
@@ -1,49 +0,0 @@
|
||||
"""Add GIN index on async_operations.result_metadata for parent_operation_id queries
|
||||
|
||||
Revision ID: y0t1u2v3w4x5
|
||||
Revises: x9s0t1u2v3w4
|
||||
Create Date: 2026-02-13
|
||||
|
||||
This migration adds a GIN index on the result_metadata JSONB column in the
|
||||
async_operations table to support efficient queries for child operations by
|
||||
parent_operation_id.
|
||||
|
||||
The index enables fast lookups when querying for child operations:
|
||||
SELECT * FROM async_operations
|
||||
WHERE result_metadata::jsonb @> '{"parent_operation_id": "uuid"}'::jsonb
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "y0t1u2v3w4x5"
|
||||
down_revision: str | Sequence[str] | None = "x9s0t1u2v3w4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Add GIN index on result_metadata for efficient parent_operation_id queries."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Add GIN index for JSONB containment queries (@> operator)
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_async_operations_result_metadata
|
||||
ON {schema}async_operations
|
||||
USING gin(result_metadata)
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Remove GIN index on result_metadata."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Drop index
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_async_operations_result_metadata")
|
||||
@@ -13,7 +13,7 @@ from contextlib import asynccontextmanager
|
||||
from datetime import datetime
|
||||
from typing import Any, Literal
|
||||
|
||||
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Query, UploadFile
|
||||
from fastapi import Depends, FastAPI, Header, HTTPException, Query
|
||||
|
||||
from hindsight_api.extensions import AuthenticationError
|
||||
|
||||
@@ -32,45 +32,9 @@ def _parse_metadata(metadata: Any) -> dict[str, Any]:
|
||||
return {}
|
||||
|
||||
|
||||
from typing import Callable
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
|
||||
|
||||
def FieldWithDefault(default_factory: Callable, **kwargs) -> Any:
|
||||
"""
|
||||
Field wrapper that ensures default_factory values appear in OpenAPI schema.
|
||||
|
||||
Pydantic doesn't include default_factory in OpenAPI schemas, causing OpenAPI
|
||||
Generator to make fields Optional with default=None instead of non-optional
|
||||
with the correct default value.
|
||||
|
||||
This wrapper adds json_schema_extra to include the default in the schema.
|
||||
"""
|
||||
# Determine the default value for the schema based on the factory
|
||||
if default_factory is list:
|
||||
schema_default = []
|
||||
elif default_factory is dict:
|
||||
schema_default = {}
|
||||
else:
|
||||
# For custom factories (like IncludeOptions), use empty dict as placeholder
|
||||
schema_default = {}
|
||||
|
||||
# Add or merge json_schema_extra
|
||||
json_extra = kwargs.pop("json_schema_extra", {})
|
||||
if isinstance(json_extra, dict):
|
||||
json_extra["default"] = schema_default
|
||||
else:
|
||||
# If json_schema_extra was a function, we can't merge easily
|
||||
# Fall back to just setting default
|
||||
json_extra = {"default": schema_default}
|
||||
|
||||
return Field(default_factory=default_factory, json_schema_extra=json_extra, **kwargs)
|
||||
|
||||
|
||||
from hindsight_api.config import get_config
|
||||
from hindsight_api.engine.db_utils import acquire_with_retry
|
||||
from hindsight_api.engine.memory_engine import Budget, _get_tiktoken_encoding, fq_table
|
||||
from hindsight_api.engine.reflect.observations import Observation
|
||||
@@ -139,8 +103,8 @@ class RecallRequest(BaseModel):
|
||||
query_timestamp: str | None = Field(
|
||||
default=None, description="ISO format date string (e.g., '2023-05-30T23:40:00')"
|
||||
)
|
||||
include: IncludeOptions = FieldWithDefault(
|
||||
IncludeOptions,
|
||||
include: IncludeOptions = Field(
|
||||
default_factory=IncludeOptions,
|
||||
description="Options for including additional data (entities are included by default)",
|
||||
)
|
||||
tags: list[str] | None = Field(
|
||||
@@ -430,36 +394,6 @@ class RetainRequest(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class FileRetainMetadata(BaseModel):
|
||||
"""Metadata for a single file in file retain request."""
|
||||
|
||||
document_id: str | None = Field(default=None, description="Document ID (auto-generated if not provided)")
|
||||
context: str | None = Field(default=None, description="Context for the file")
|
||||
metadata: dict[str, Any] | None = Field(default=None, description="Additional metadata")
|
||||
tags: list[str] | None = Field(default=None, description="Tags for this file")
|
||||
timestamp: str | None = Field(default=None, description="ISO timestamp")
|
||||
|
||||
|
||||
class FileRetainRequest(BaseModel):
|
||||
"""Request model for file retain endpoint."""
|
||||
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"files_metadata": [
|
||||
{"document_id": "report_2024", "tags": ["quarterly"]},
|
||||
{"context": "meeting notes"},
|
||||
],
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
files_metadata: list[FileRetainMetadata] | None = Field(
|
||||
default=None,
|
||||
description="Metadata for each file (optional, must match number of files if provided)",
|
||||
)
|
||||
|
||||
|
||||
class RetainResponse(BaseModel):
|
||||
"""Response model for retain endpoint."""
|
||||
|
||||
@@ -484,7 +418,7 @@ class RetainResponse(BaseModel):
|
||||
)
|
||||
operation_id: str | None = Field(
|
||||
default=None,
|
||||
description="Operation ID for tracking async operations. Use GET /v1/default/banks/{bank_id}/operations to list operations. Only present when async=true.",
|
||||
description="Operation ID for tracking async operations. Use GET /v1/default/banks/{bank_id}/operations to list operations and find this ID. Only present when async=true.",
|
||||
)
|
||||
usage: TokenUsage | None = Field(
|
||||
default=None,
|
||||
@@ -492,26 +426,6 @@ class RetainResponse(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class FileRetainResponse(BaseModel):
|
||||
"""Response model for file upload endpoint."""
|
||||
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"operation_ids": [
|
||||
"550e8400-e29b-41d4-a716-446655440000",
|
||||
"550e8400-e29b-41d4-a716-446655440001",
|
||||
"550e8400-e29b-41d4-a716-446655440002",
|
||||
],
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
operation_ids: list[str] = Field(
|
||||
description="Operation IDs for tracking file conversion operations. Use GET /v1/default/banks/{bank_id}/operations to list operations."
|
||||
)
|
||||
|
||||
|
||||
class FactsIncludeOptions(BaseModel):
|
||||
"""Options for including facts (based_on) in reflect results."""
|
||||
|
||||
@@ -656,16 +570,18 @@ class ReflectLLMCall(BaseModel):
|
||||
class ReflectBasedOn(BaseModel):
|
||||
"""Evidence the response is based on: memories, mental models, and directives."""
|
||||
|
||||
memories: list[ReflectFact] = FieldWithDefault(list, description="Memory facts used to generate the response")
|
||||
mental_models: list[ReflectMentalModel] = FieldWithDefault(list, description="Mental models used during reflection")
|
||||
directives: list[ReflectDirective] = FieldWithDefault(list, description="Directives applied during reflection")
|
||||
memories: list[ReflectFact] = Field(default_factory=list, description="Memory facts used to generate the response")
|
||||
mental_models: list[ReflectMentalModel] = Field(
|
||||
default_factory=list, description="Mental models used during reflection"
|
||||
)
|
||||
directives: list[ReflectDirective] = Field(default_factory=list, description="Directives applied during reflection")
|
||||
|
||||
|
||||
class ReflectTrace(BaseModel):
|
||||
"""Execution trace of LLM and tool calls during reflection."""
|
||||
|
||||
tool_calls: list[ReflectToolCall] = FieldWithDefault(list, description="Tool calls made during reflection")
|
||||
llm_calls: list[ReflectLLMCall] = FieldWithDefault(list, description="LLM calls made during reflection")
|
||||
tool_calls: list[ReflectToolCall] = Field(default_factory=list, description="Tool calls made during reflection")
|
||||
llm_calls: list[ReflectLLMCall] = Field(default_factory=list, description="LLM calls made during reflection")
|
||||
|
||||
|
||||
class ReflectResponse(BaseModel):
|
||||
@@ -877,55 +793,6 @@ class CreateBankRequest(BaseModel):
|
||||
background: str | None = Field(default=None, description="Deprecated: use mission instead")
|
||||
|
||||
|
||||
class BankConfigUpdate(BaseModel):
|
||||
"""Request model for updating bank configuration."""
|
||||
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"updates": {
|
||||
"llm_model": "claude-sonnet-4-5",
|
||||
"retain_extraction_mode": "verbose",
|
||||
"retain_custom_instructions": "Extract technical details carefully",
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
updates: dict[str, Any] = Field(
|
||||
description="Configuration overrides. Keys can be in Python field format (llm_provider) "
|
||||
"or environment variable format (HINDSIGHT_API_LLM_PROVIDER). "
|
||||
"Only hierarchical fields can be overridden per-bank."
|
||||
)
|
||||
|
||||
|
||||
class BankConfigResponse(BaseModel):
|
||||
"""Response model for bank configuration."""
|
||||
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"bank_id": "my-bank",
|
||||
"config": {
|
||||
"llm_provider": "openai",
|
||||
"llm_model": "gpt-4",
|
||||
"retain_extraction_mode": "verbose",
|
||||
},
|
||||
"overrides": {
|
||||
"llm_model": "gpt-4",
|
||||
"retain_extraction_mode": "verbose",
|
||||
},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
bank_id: str = Field(description="Bank identifier")
|
||||
config: dict[str, Any] = Field(
|
||||
description="Fully resolved configuration with all hierarchical overrides applied (Python field names)"
|
||||
)
|
||||
overrides: dict[str, Any] = Field(description="Bank-specific configuration overrides only (Python field names)")
|
||||
|
||||
|
||||
class GraphDataResponse(BaseModel):
|
||||
"""Response model for graph data endpoint."""
|
||||
|
||||
@@ -1075,7 +942,7 @@ class DocumentResponse(BaseModel):
|
||||
created_at: str
|
||||
updated_at: str
|
||||
memory_unit_count: int
|
||||
tags: list[str] = FieldWithDefault(list, description="Tags associated with this document")
|
||||
tags: list[str] = Field(default_factory=list, description="Tags associated with this document")
|
||||
|
||||
|
||||
class DeleteDocumentResponse(BaseModel):
|
||||
@@ -1199,7 +1066,7 @@ class DirectiveResponse(BaseModel):
|
||||
content: str
|
||||
priority: int = 0
|
||||
is_active: bool = True
|
||||
tags: list[str] = FieldWithDefault(list)
|
||||
tags: list[str] = Field(default_factory=list)
|
||||
created_at: str | None = None
|
||||
updated_at: str | None = None
|
||||
|
||||
@@ -1217,7 +1084,7 @@ class CreateDirectiveRequest(BaseModel):
|
||||
content: str = Field(description="The directive text to inject into prompts")
|
||||
priority: int = Field(default=0, description="Higher priority directives are injected first")
|
||||
is_active: bool = Field(default=True, description="Whether this directive is active")
|
||||
tags: list[str] = FieldWithDefault(list, description="Tags for filtering")
|
||||
tags: list[str] = Field(default_factory=list, description="Tags for filtering")
|
||||
|
||||
|
||||
class UpdateDirectiveRequest(BaseModel):
|
||||
@@ -1254,9 +1121,9 @@ class MentalModelResponse(BaseModel):
|
||||
content: str = Field(
|
||||
description="The mental model content as well-formatted markdown (auto-generated from reflect endpoint)"
|
||||
)
|
||||
tags: list[str] = FieldWithDefault(list)
|
||||
tags: list[str] = Field(default_factory=list)
|
||||
max_tokens: int = Field(default=2048)
|
||||
trigger: MentalModelTrigger = FieldWithDefault(MentalModelTrigger)
|
||||
trigger: MentalModelTrigger = Field(default_factory=MentalModelTrigger)
|
||||
last_refreshed_at: str | None = None
|
||||
created_at: str | None = None
|
||||
reflect_response: dict | None = Field(
|
||||
@@ -1292,9 +1159,9 @@ class CreateMentalModelRequest(BaseModel):
|
||||
)
|
||||
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] = FieldWithDefault(list, description="Tags for scoped visibility")
|
||||
tags: list[str] = Field(default_factory=list, description="Tags for scoped visibility")
|
||||
max_tokens: int = Field(default=2048, ge=256, le=8192, description="Maximum tokens for generated content")
|
||||
trigger: MentalModelTrigger = FieldWithDefault(MentalModelTrigger, description="Trigger settings")
|
||||
trigger: MentalModelTrigger = Field(default_factory=MentalModelTrigger, description="Trigger settings")
|
||||
|
||||
|
||||
class CreateMentalModelResponse(BaseModel):
|
||||
@@ -1407,16 +1274,6 @@ class CancelOperationResponse(BaseModel):
|
||||
operation_id: str
|
||||
|
||||
|
||||
class ChildOperationStatus(BaseModel):
|
||||
"""Status of a child operation (for batch operations)."""
|
||||
|
||||
operation_id: str
|
||||
status: str
|
||||
sub_batch_index: int | None = None
|
||||
items_count: int | None = None
|
||||
error_message: str | None = None
|
||||
|
||||
|
||||
class OperationStatusResponse(BaseModel):
|
||||
"""Response model for getting a single operation status."""
|
||||
|
||||
@@ -1441,13 +1298,6 @@ class OperationStatusResponse(BaseModel):
|
||||
updated_at: str | None = None
|
||||
completed_at: str | None = None
|
||||
error_message: str | None = None
|
||||
result_metadata: dict[str, Any] | None = Field(
|
||||
default=None,
|
||||
description="Internal metadata for debugging. Structure may change without notice. Not for production use.",
|
||||
)
|
||||
child_operations: list[ChildOperationStatus] | None = Field(
|
||||
default=None, description="Child operations for batch operations (if applicable)"
|
||||
)
|
||||
|
||||
|
||||
class AsyncOperationSubmitResponse(BaseModel):
|
||||
@@ -1472,8 +1322,6 @@ class FeaturesInfo(BaseModel):
|
||||
observations: bool = Field(description="Whether observations (auto-consolidation) are enabled")
|
||||
mcp: bool = Field(description="Whether MCP (Model Context Protocol) server is enabled")
|
||||
worker: bool = Field(description="Whether the background worker is enabled")
|
||||
bank_config_api: bool = Field(description="Whether per-bank configuration API is enabled")
|
||||
file_upload_api: bool = Field(description="Whether file upload/conversion API is enabled")
|
||||
|
||||
|
||||
class VersionResponse(BaseModel):
|
||||
@@ -1487,8 +1335,6 @@ class VersionResponse(BaseModel):
|
||||
"observations": False,
|
||||
"mcp": True,
|
||||
"worker": True,
|
||||
"bank_config_api": False,
|
||||
"file_upload_api": True,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -1645,9 +1491,6 @@ def create_app(
|
||||
logging.info("Memory system closed")
|
||||
|
||||
from hindsight_api import __version__
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
config = get_config()
|
||||
|
||||
app = FastAPI(
|
||||
title="Hindsight HTTP API",
|
||||
@@ -1661,7 +1504,6 @@ def create_app(
|
||||
"url": "https://www.apache.org/licenses/LICENSE-2.0.html",
|
||||
},
|
||||
lifespan=lifespan,
|
||||
root_path=config.base_path,
|
||||
)
|
||||
|
||||
# IMPORTANT: Set memory on app.state immediately, don't wait for lifespan
|
||||
@@ -1768,22 +1610,17 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
Returns version info and feature flags that can be used by clients
|
||||
to determine which capabilities are available.
|
||||
|
||||
Note: observations flag shows the global default. Individual banks
|
||||
may override this setting via bank-specific configuration.
|
||||
"""
|
||||
from hindsight_api import __version__
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
config = _get_raw_config()
|
||||
config = get_config()
|
||||
return VersionResponse(
|
||||
api_version=__version__,
|
||||
features=FeaturesInfo(
|
||||
observations=config.enable_observations,
|
||||
mcp=config.mcp_enabled,
|
||||
worker=config.worker_enabled,
|
||||
bank_config_api=config.enable_bank_config_api,
|
||||
file_upload_api=config.enable_file_upload_api,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -3437,112 +3274,6 @@ def _register_routes(app: FastAPI):
|
||||
logger.error(f"Error in DELETE /v1/default/banks/{bank_id}/observations: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/config",
|
||||
response_model=BankConfigResponse,
|
||||
summary="Get bank configuration",
|
||||
description="Get fully resolved configuration for a bank including all hierarchical overrides (global → tenant → bank). "
|
||||
"The 'config' field contains all resolved config values. The 'overrides' field shows only bank-specific overrides.",
|
||||
operation_id="get_bank_config",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_get_bank_config(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
"""Get configuration for a bank with all hierarchical overrides applied."""
|
||||
if not get_config().enable_bank_config_api:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="Bank configuration API is disabled. Set HINDSIGHT_API_ENABLE_BANK_CONFIG_API=true to enable.",
|
||||
)
|
||||
try:
|
||||
# Get resolved config from config resolver
|
||||
config_dict = await app.state.memory._config_resolver.get_bank_config(bank_id, request_context)
|
||||
|
||||
# Get bank-specific overrides only
|
||||
bank_overrides = await app.state.memory._config_resolver._load_bank_config(bank_id)
|
||||
|
||||
return BankConfigResponse(bank_id=bank_id, config=config_dict, overrides=bank_overrides)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in GET /v1/default/banks/{bank_id}/config: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.patch(
|
||||
"/v1/default/banks/{bank_id}/config",
|
||||
response_model=BankConfigResponse,
|
||||
summary="Update bank configuration",
|
||||
description="Update configuration overrides for a bank. Only hierarchical fields can be overridden (LLM settings, retention parameters, etc.). "
|
||||
"Keys can be provided in Python field format (llm_provider) or environment variable format (HINDSIGHT_API_LLM_PROVIDER).",
|
||||
operation_id="update_bank_config",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_update_bank_config(
|
||||
bank_id: str, request: BankConfigUpdate, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Update configuration overrides for a bank."""
|
||||
if not get_config().enable_bank_config_api:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="Bank configuration API is disabled. Set HINDSIGHT_API_ENABLE_BANK_CONFIG_API=true to enable.",
|
||||
)
|
||||
try:
|
||||
# Update config via config resolver (validates configurable fields and permissions)
|
||||
await app.state.memory._config_resolver.update_bank_config(bank_id, request.updates, request_context)
|
||||
|
||||
# Return updated config
|
||||
config_dict = await app.state.memory._config_resolver.get_bank_config(bank_id, request_context)
|
||||
bank_overrides = await app.state.memory._config_resolver._load_bank_config(bank_id)
|
||||
|
||||
return BankConfigResponse(bank_id=bank_id, config=config_dict, overrides=bank_overrides)
|
||||
except ValueError as e:
|
||||
# Validation error (e.g., trying to override static field)
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in PATCH /v1/default/banks/{bank_id}/config: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.delete(
|
||||
"/v1/default/banks/{bank_id}/config",
|
||||
response_model=BankConfigResponse,
|
||||
summary="Reset bank configuration",
|
||||
description="Reset bank configuration to defaults by removing all bank-specific overrides. "
|
||||
"The bank will then use global and tenant-level configuration only.",
|
||||
operation_id="reset_bank_config",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_reset_bank_config(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
"""Reset bank configuration to defaults (remove all overrides)."""
|
||||
if not get_config().enable_bank_config_api:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="Bank configuration API is disabled. Set HINDSIGHT_API_ENABLE_BANK_CONFIG_API=true to enable.",
|
||||
)
|
||||
try:
|
||||
# Reset config via config resolver
|
||||
await app.state.memory._config_resolver.reset_bank_config(bank_id)
|
||||
|
||||
# Return updated config (should match defaults now)
|
||||
config_dict = await app.state.memory._config_resolver.get_bank_config(bank_id, request_context)
|
||||
bank_overrides = await app.state.memory._config_resolver._load_bank_config(bank_id)
|
||||
|
||||
return BankConfigResponse(bank_id=bank_id, config=config_dict, overrides=bank_overrides)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in DELETE /v1/default/banks/{bank_id}/config: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post(
|
||||
"/v1/default/banks/{bank_id}/consolidate",
|
||||
response_model=ConsolidationResponse,
|
||||
@@ -3633,21 +3364,6 @@ def _register_routes(app: FastAPI):
|
||||
}
|
||||
)
|
||||
else:
|
||||
# Check if batch API is enabled - if so, require async mode
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
config = get_config()
|
||||
if config.retain_batch_enabled:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(
|
||||
"Batch API is enabled (HINDSIGHT_API_RETAIN_BATCH_ENABLED=true) but async=false. "
|
||||
"Batch operations can take several minutes to hours and will timeout in synchronous mode. "
|
||||
"Please set async=true in your request to use background processing, or disable batch API "
|
||||
"by setting HINDSIGHT_API_RETAIN_BATCH_ENABLED=false in your environment."
|
||||
),
|
||||
)
|
||||
|
||||
# Synchronous processing: wait for completion (record metrics)
|
||||
with metrics.record_operation("retain", bank_id=bank_id, source="api"):
|
||||
result, usage = await app.state.memory.retain_batch_async(
|
||||
@@ -3685,147 +3401,6 @@ def _register_routes(app: FastAPI):
|
||||
logger.error(f"Error in /v1/default/banks/{bank_id}/memories (retain): {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post(
|
||||
"/v1/default/banks/{bank_id}/files/retain",
|
||||
response_model=FileRetainResponse,
|
||||
summary="Convert files to memories",
|
||||
description="Upload files (PDF, DOCX, etc.), convert them to markdown, and retain as memories.\n\n"
|
||||
"This endpoint handles file upload, conversion, and memory creation in a single operation.\n\n"
|
||||
"**Features:**\n"
|
||||
"- Supports PDF, DOCX, PPTX, XLSX, images (with OCR), audio (with transcription)\n"
|
||||
"- Automatic file-to-markdown conversion using pluggable parsers\n"
|
||||
"- Files stored in object storage (PostgreSQL by default, S3 for production)\n"
|
||||
"- Each file becomes a separate document with optional metadata/tags\n"
|
||||
"- Always processes asynchronously — returns operation IDs immediately\n\n"
|
||||
"**The system automatically:**\n"
|
||||
"1. Stores uploaded files in object storage\n"
|
||||
"2. Converts files to markdown\n"
|
||||
"3. Creates document records with file metadata\n"
|
||||
"4. Extracts facts and creates memory units (same as regular retain)\n\n"
|
||||
"Use the operations endpoint to monitor progress.\n\n"
|
||||
"**Request format:** multipart/form-data with:\n"
|
||||
"- `files`: One or more files to upload\n"
|
||||
"- `request`: JSON string with FileRetainRequest model (files_metadata)\n\n"
|
||||
"**Note:** File parser is configured server-side via `HINDSIGHT_API_FILE_PARSER` (default: markitdown).",
|
||||
operation_id="file_retain",
|
||||
tags=["Files"],
|
||||
)
|
||||
async def api_file_retain(
|
||||
bank_id: str,
|
||||
files: list[UploadFile] = File(..., description="Files to upload and convert"),
|
||||
request: str = Form(..., description="JSON string with FileRetainRequest model"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Upload and convert files to memories."""
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
config = get_config()
|
||||
|
||||
# Check if file upload API is enabled
|
||||
if not config.enable_file_upload_api:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="File upload API is disabled. Set HINDSIGHT_API_ENABLE_FILE_UPLOAD_API=true to enable.",
|
||||
)
|
||||
|
||||
try:
|
||||
# Parse request JSON
|
||||
try:
|
||||
request_data = FileRetainRequest.model_validate_json(request)
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Invalid request JSON: {str(e)}",
|
||||
)
|
||||
|
||||
# Validate file count
|
||||
if len(files) > config.file_conversion_max_batch_size:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Too many files. Maximum {config.file_conversion_max_batch_size} files per request.",
|
||||
)
|
||||
|
||||
# Validate files_metadata count matches files count if provided
|
||||
if request_data.files_metadata and len(request_data.files_metadata) != len(files):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"files_metadata count ({len(request_data.files_metadata)}) must match files count ({len(files)})",
|
||||
)
|
||||
|
||||
# Prepare file items and calculate total batch size
|
||||
file_items = []
|
||||
total_batch_size = 0
|
||||
|
||||
for i, file in enumerate(files):
|
||||
# Read file content to check size
|
||||
file_content = await file.read()
|
||||
size = len(file_content)
|
||||
total_batch_size += size
|
||||
|
||||
# Create a temporary file-like object from the bytes
|
||||
import io
|
||||
|
||||
file_obj = io.BytesIO(file_content)
|
||||
|
||||
# Create a mock UploadFile with the necessary attributes
|
||||
class FileWrapper:
|
||||
def __init__(self, content, filename, content_type):
|
||||
self._content = content
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
self._buffer = io.BytesIO(content)
|
||||
|
||||
async def read(self):
|
||||
return self._content
|
||||
|
||||
wrapped_file = FileWrapper(file_content, file.filename, file.content_type)
|
||||
|
||||
# Get per-file metadata
|
||||
file_meta = request_data.files_metadata[i] if request_data.files_metadata else FileRetainMetadata()
|
||||
doc_id = file_meta.document_id or f"file_{uuid.uuid4()}"
|
||||
|
||||
item = {
|
||||
"file": wrapped_file,
|
||||
"document_id": doc_id,
|
||||
"context": file_meta.context,
|
||||
"metadata": file_meta.metadata or {},
|
||||
"tags": file_meta.tags or [],
|
||||
"timestamp": file_meta.timestamp,
|
||||
}
|
||||
file_items.append(item)
|
||||
|
||||
# Check total batch size after processing all files
|
||||
if total_batch_size > config.file_conversion_max_batch_size_bytes:
|
||||
total_mb = total_batch_size / (1024 * 1024)
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Total batch size ({total_mb:.1f}MB) exceeds maximum of {config.file_conversion_max_batch_size_mb}MB",
|
||||
)
|
||||
|
||||
result = await app.state.memory.submit_async_file_retain(
|
||||
bank_id=bank_id,
|
||||
file_items=file_items,
|
||||
parser=config.file_parser,
|
||||
document_tags=None,
|
||||
request_context=request_context,
|
||||
)
|
||||
return FileRetainResponse.model_validate(
|
||||
{
|
||||
"operation_ids": result["operation_ids"],
|
||||
}
|
||||
)
|
||||
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in /v1/default/banks/{bank_id}/files/retain: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.delete(
|
||||
"/v1/default/banks/{bank_id}/memories",
|
||||
response_model=DeleteResponse,
|
||||
|
||||
@@ -114,39 +114,9 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
logger.info(f"Loading MCP extension: {mcp_extension.__class__.__name__}")
|
||||
mcp_extension.register_tools(mcp, memory)
|
||||
|
||||
# Make all tools tolerant of extra arguments from LLMs (e.g., "explanation")
|
||||
_make_tools_tolerant(mcp)
|
||||
|
||||
return mcp
|
||||
|
||||
|
||||
def _make_tools_tolerant(mcp: FastMCP) -> None:
|
||||
"""Wrap all tool run methods to strip unknown arguments before validation.
|
||||
|
||||
LLMs frequently add extra fields like "explanation" or "reasoning" to tool calls.
|
||||
FastMCP's Pydantic TypeAdapter rejects these with "Unexpected keyword argument".
|
||||
This wraps each tool's run() to filter arguments to only known parameters.
|
||||
"""
|
||||
try:
|
||||
for name, tool in mcp._tool_manager._tools.items():
|
||||
if hasattr(tool, "parameters") and tool.parameters:
|
||||
allowed = set(tool.parameters.get("properties", {}).keys())
|
||||
original_run = tool.run
|
||||
|
||||
async def _tolerant_run(arguments, _allowed=allowed, _orig=original_run):
|
||||
extra_keys = set(arguments.keys()) - _allowed
|
||||
if extra_keys:
|
||||
logger.debug(f"Stripping unknown arguments from tool call: {extra_keys}")
|
||||
arguments = {k: v for k, v in arguments.items() if k in _allowed}
|
||||
return await _orig(arguments)
|
||||
|
||||
# FunctionTool is a Pydantic model with extra='forbid', so use
|
||||
# object.__setattr__ to bypass Pydantic's setter validation.
|
||||
object.__setattr__(tool, "run", _tolerant_run)
|
||||
except (AttributeError, KeyError) as e:
|
||||
logger.warning(f"Could not make tools tolerant of extra arguments: {e}")
|
||||
|
||||
|
||||
class MCPMiddleware:
|
||||
"""ASGI middleware that intercepts MCP requests and routes to appropriate MCP server.
|
||||
|
||||
@@ -172,11 +142,6 @@ class MCPMiddleware:
|
||||
- No bank management tools (list_banks, create_bank)
|
||||
- Recommended for agent isolation
|
||||
|
||||
Bank ID resolution priority:
|
||||
1. URL path (e.g., /mcp/{bank_id}/) → single-bank mode
|
||||
2. X-Bank-Id header → multi-bank mode
|
||||
3. HINDSIGHT_MCP_BANK_ID env var → multi-bank mode (default: "default")
|
||||
|
||||
Examples:
|
||||
# Single-bank mode (recommended for agent isolation)
|
||||
claude mcp add --transport http my-agent http://localhost:8888/mcp/my-agent-bank/ \\
|
||||
@@ -277,25 +242,20 @@ class MCPMiddleware:
|
||||
_current_schema.set(tenant_context.schema_name) if tenant_context and tenant_context.schema_name else None
|
||||
)
|
||||
|
||||
# Resolve bank_id: path takes priority over header.
|
||||
# Path = user's explicit connection endpoint (e.g., /mcp/my-bank/).
|
||||
# X-Bank-Id header = per-request override for multi-bank mode only.
|
||||
bank_id = None
|
||||
# Try to get bank_id from header first (for Claude Code compatibility)
|
||||
bank_id = self._get_header(scope, "X-Bank-Id")
|
||||
bank_id_from_path = False
|
||||
new_path = path
|
||||
|
||||
# First, try to extract from path: /{bank_id}/...
|
||||
if path.startswith("/") and len(path) > 1:
|
||||
# If no header, try to extract from path: /{bank_id}/...
|
||||
new_path = path
|
||||
if not bank_id and path.startswith("/") and len(path) > 1:
|
||||
parts = path[1:].split("/", 1)
|
||||
if parts[0]:
|
||||
# First segment looks like a bank_id
|
||||
bank_id = parts[0]
|
||||
bank_id_from_path = True
|
||||
new_path = "/" + parts[1] if len(parts) > 1 else "/"
|
||||
|
||||
# If no path-based bank_id, try X-Bank-Id header (multi-bank mode)
|
||||
if not bank_id:
|
||||
bank_id = self._get_header(scope, "X-Bank-Id")
|
||||
|
||||
# Fall back to default bank_id
|
||||
if not bank_id:
|
||||
bank_id = DEFAULT_BANK_ID
|
||||
|
||||
@@ -86,8 +86,6 @@ def print_startup_info(
|
||||
reranker_provider: str,
|
||||
mcp_enabled: bool = False,
|
||||
version: str | None = None,
|
||||
vector_extension: str | None = None,
|
||||
text_search_extension: str | None = None,
|
||||
):
|
||||
"""Print styled startup information."""
|
||||
print(color_start("Starting Hindsight API..."))
|
||||
@@ -98,8 +96,6 @@ def print_startup_info(
|
||||
print(f" {dim('LLM:')} {color(f'{llm_provider} / {llm_model}', 0.6)}")
|
||||
print(f" {dim('Embeddings:')} {color(embeddings_provider, 0.8)}")
|
||||
print(f" {dim('Reranker:')} {color(reranker_provider, 1.0)}")
|
||||
extensions = f"{vector_extension or 'default'} (vector) / {text_search_extension or 'default'} (text)"
|
||||
print(f" {dim('Extensions:')} {color(extensions, 0.4)}")
|
||||
if mcp_enabled:
|
||||
print(f" {dim('MCP:')} {color_end('enabled at /mcp')}")
|
||||
print()
|
||||
|
||||
@@ -8,9 +8,8 @@ import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from dataclasses import dataclass, field, fields
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from dotenv import find_dotenv, load_dotenv
|
||||
|
||||
@@ -19,103 +18,6 @@ load_dotenv(find_dotenv(usecwd=True), override=True)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConfigFieldAccessError(AttributeError):
|
||||
"""Raised when trying to access a bank-configurable field from global config."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class StaticConfigProxy:
|
||||
"""
|
||||
Proxy that wraps HindsightConfig and only allows access to static (non-configurable) fields.
|
||||
|
||||
Raises ConfigFieldAccessError when trying to access configurable fields that vary per-bank.
|
||||
Forces developers to use get_resolved_config(bank_id, context) for bank-specific settings.
|
||||
"""
|
||||
|
||||
def __init__(self, config: "HindsightConfig"):
|
||||
object.__setattr__(self, "_config", config)
|
||||
object.__setattr__(self, "_configurable_fields", HindsightConfig.get_configurable_fields())
|
||||
|
||||
def __getattribute__(self, name: str):
|
||||
if name.startswith("_"):
|
||||
return object.__getattribute__(self, name)
|
||||
|
||||
configurable_fields = object.__getattribute__(self, "_configurable_fields")
|
||||
if name in configurable_fields:
|
||||
raise ConfigFieldAccessError(
|
||||
f"Field '{name}' is bank-configurable and cannot be accessed from global config. "
|
||||
f"Use ConfigResolver.resolve_full_config(bank_id, context) to get bank-specific config. "
|
||||
f"This prevents accidentally using global defaults when bank-specific overrides exist."
|
||||
)
|
||||
|
||||
config = object.__getattribute__(self, "_config")
|
||||
return getattr(config, name)
|
||||
|
||||
def __setattr__(self, name: str, value):
|
||||
raise AttributeError("Config is read-only. Modifications must go through ConfigResolver.")
|
||||
|
||||
|
||||
# Configuration field markers for hierarchical configuration
|
||||
def hierarchical(default_value):
|
||||
"""
|
||||
Mark a config field as hierarchical (can be overridden per-tenant/bank).
|
||||
|
||||
Hierarchical fields can be customized at the tenant or bank level via database
|
||||
configuration. Examples: LLM settings, retention parameters, retrieval settings.
|
||||
"""
|
||||
return field(default=default_value, metadata={"hierarchical": True})
|
||||
|
||||
|
||||
def static(default_value):
|
||||
"""
|
||||
Mark a config field as static (server-level only, cannot be overridden).
|
||||
|
||||
Static fields are infrastructure-level settings that affect the entire server
|
||||
and cannot vary per tenant or bank. Examples: database URL, API port, worker settings.
|
||||
"""
|
||||
return field(default=default_value, metadata={"hierarchical": False})
|
||||
|
||||
|
||||
# Configuration key normalization utilities
|
||||
def normalize_config_key(key: str) -> str:
|
||||
"""
|
||||
Convert environment variable format to Python field name format.
|
||||
|
||||
Examples:
|
||||
HINDSIGHT_API_LLM_PROVIDER -> llm_provider
|
||||
LLM_MODEL -> llm_model
|
||||
llm_model -> llm_model (already normalized)
|
||||
|
||||
Args:
|
||||
key: Environment variable name or Python field name
|
||||
|
||||
Returns:
|
||||
Normalized Python field name (lowercase snake_case)
|
||||
"""
|
||||
if key.startswith("HINDSIGHT_API_"):
|
||||
key = key[len("HINDSIGHT_API_") :]
|
||||
return key.lower()
|
||||
|
||||
|
||||
def normalize_config_dict(config: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Normalize all keys in a config dict to Python field names.
|
||||
|
||||
Allows users to provide config overrides in either format:
|
||||
- Python field format: {"llm_provider": "openai"}
|
||||
- Env var format: {"HINDSIGHT_API_LLM_PROVIDER": "openai"}
|
||||
|
||||
Args:
|
||||
config: Dict with env var or Python field names as keys
|
||||
|
||||
Returns:
|
||||
Dict with all keys normalized to Python field names
|
||||
"""
|
||||
return {normalize_config_key(k): v for k, v in config.items()}
|
||||
|
||||
|
||||
# Environment variable names
|
||||
ENV_DATABASE_URL = "HINDSIGHT_API_DATABASE_URL"
|
||||
ENV_DATABASE_SCHEMA = "HINDSIGHT_API_DATABASE_SCHEMA"
|
||||
@@ -129,11 +31,6 @@ 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"
|
||||
ENV_LLM_OPENAI_SERVICE_TIER = "HINDSIGHT_API_LLM_OPENAI_SERVICE_TIER"
|
||||
|
||||
# Defaults for service tiers
|
||||
DEFAULT_LLM_GROQ_SERVICE_TIER = "auto" # "on_demand", "flex", or "auto"
|
||||
DEFAULT_LLM_OPENAI_SERVICE_TIER = None # None (default) or "flex" (50% cheaper)
|
||||
|
||||
# Per-operation LLM configuration (optional, falls back to global LLM config)
|
||||
ENV_RETAIN_LLM_PROVIDER = "HINDSIGHT_API_RETAIN_LLM_PROVIDER"
|
||||
@@ -194,14 +91,6 @@ ENV_RERANKER_LITELLM_API_BASE = "HINDSIGHT_API_RERANKER_LITELLM_API_BASE"
|
||||
ENV_RERANKER_LITELLM_API_KEY = "HINDSIGHT_API_RERANKER_LITELLM_API_KEY"
|
||||
ENV_RERANKER_LITELLM_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_MODEL"
|
||||
|
||||
# LiteLLM SDK configuration (direct API access, no proxy needed)
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_KEY"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_MODEL = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_BASE = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_BASE"
|
||||
ENV_RERANKER_LITELLM_SDK_API_KEY = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY"
|
||||
ENV_RERANKER_LITELLM_SDK_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL"
|
||||
ENV_RERANKER_LITELLM_SDK_API_BASE = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_BASE"
|
||||
|
||||
# Deprecated: Legacy shared LiteLLM config (for backward compatibility)
|
||||
ENV_LITELLM_API_BASE = "HINDSIGHT_API_LITELLM_API_BASE"
|
||||
ENV_LITELLM_API_KEY = "HINDSIGHT_API_LITELLM_API_KEY"
|
||||
@@ -218,17 +107,12 @@ ENV_RERANKER_MAX_CANDIDATES = "HINDSIGHT_API_RERANKER_MAX_CANDIDATES"
|
||||
ENV_RERANKER_FLASHRANK_MODEL = "HINDSIGHT_API_RERANKER_FLASHRANK_MODEL"
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR"
|
||||
|
||||
ENV_VECTOR_EXTENSION = "HINDSIGHT_API_VECTOR_EXTENSION"
|
||||
ENV_TEXT_SEARCH_EXTENSION = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION"
|
||||
|
||||
ENV_HOST = "HINDSIGHT_API_HOST"
|
||||
ENV_PORT = "HINDSIGHT_API_PORT"
|
||||
ENV_BASE_PATH = "HINDSIGHT_API_BASE_PATH"
|
||||
ENV_LOG_LEVEL = "HINDSIGHT_API_LOG_LEVEL"
|
||||
ENV_LOG_FORMAT = "HINDSIGHT_API_LOG_FORMAT"
|
||||
ENV_WORKERS = "HINDSIGHT_API_WORKERS"
|
||||
ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
|
||||
ENV_ENABLE_BANK_CONFIG_API = "HINDSIGHT_API_ENABLE_BANK_CONFIG_API"
|
||||
ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER"
|
||||
ENV_MPFP_TOP_K_NEIGHBORS = "HINDSIGHT_API_MPFP_TOP_K_NEIGHBORS"
|
||||
ENV_RECALL_MAX_CONCURRENT = "HINDSIGHT_API_RECALL_MAX_CONCURRENT"
|
||||
@@ -255,29 +139,6 @@ 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_BATCH_TOKENS = "HINDSIGHT_API_RETAIN_BATCH_TOKENS"
|
||||
ENV_RETAIN_BATCH_ENABLED = "HINDSIGHT_API_RETAIN_BATCH_ENABLED"
|
||||
ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS = "HINDSIGHT_API_RETAIN_BATCH_POLL_INTERVAL_SECONDS"
|
||||
|
||||
# File storage configuration
|
||||
ENV_FILE_STORAGE_TYPE = "HINDSIGHT_API_FILE_STORAGE_TYPE"
|
||||
ENV_FILE_STORAGE_S3_BUCKET = "HINDSIGHT_API_FILE_STORAGE_S3_BUCKET"
|
||||
ENV_FILE_STORAGE_S3_REGION = "HINDSIGHT_API_FILE_STORAGE_S3_REGION"
|
||||
ENV_FILE_STORAGE_S3_ENDPOINT = "HINDSIGHT_API_FILE_STORAGE_S3_ENDPOINT"
|
||||
ENV_FILE_STORAGE_S3_ACCESS_KEY_ID = "HINDSIGHT_API_FILE_STORAGE_S3_ACCESS_KEY_ID"
|
||||
ENV_FILE_STORAGE_S3_SECRET_ACCESS_KEY = "HINDSIGHT_API_FILE_STORAGE_S3_SECRET_ACCESS_KEY"
|
||||
ENV_FILE_STORAGE_GCS_BUCKET = "HINDSIGHT_API_FILE_STORAGE_GCS_BUCKET"
|
||||
ENV_FILE_STORAGE_GCS_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_FILE_STORAGE_GCS_SERVICE_ACCOUNT_KEY"
|
||||
ENV_FILE_STORAGE_AZURE_CONTAINER = "HINDSIGHT_API_FILE_STORAGE_AZURE_CONTAINER"
|
||||
ENV_FILE_STORAGE_AZURE_ACCOUNT_NAME = "HINDSIGHT_API_FILE_STORAGE_AZURE_ACCOUNT_NAME"
|
||||
ENV_FILE_STORAGE_AZURE_ACCOUNT_KEY = "HINDSIGHT_API_FILE_STORAGE_AZURE_ACCOUNT_KEY"
|
||||
ENV_FILE_PARSER = "HINDSIGHT_API_FILE_PARSER"
|
||||
ENV_FILE_PARSER_IRIS_TOKEN = "HINDSIGHT_API_FILE_PARSER_IRIS_TOKEN"
|
||||
ENV_FILE_PARSER_IRIS_ORG_ID = "HINDSIGHT_API_FILE_PARSER_IRIS_ORG_ID"
|
||||
ENV_FILE_CONVERSION_MAX_BATCH_SIZE_MB = "HINDSIGHT_API_FILE_CONVERSION_MAX_BATCH_SIZE_MB"
|
||||
ENV_FILE_CONVERSION_MAX_BATCH_SIZE = "HINDSIGHT_API_FILE_CONVERSION_MAX_BATCH_SIZE"
|
||||
ENV_ENABLE_FILE_UPLOAD_API = "HINDSIGHT_API_ENABLE_FILE_UPLOAD_API"
|
||||
ENV_FILE_DELETE_AFTER_RETAIN = "HINDSIGHT_API_FILE_DELETE_AFTER_RETAIN"
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
ENV_ENABLE_OBSERVATIONS = "HINDSIGHT_API_ENABLE_OBSERVATIONS"
|
||||
@@ -362,29 +223,17 @@ DEFAULT_RERANKER_FLASHRANK_CACHE_DIR = None # Use default cache directory
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL = "embed-english-v3.0"
|
||||
DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
|
||||
|
||||
# Vector extension (pgvector, vchord, or pgvectorscale)
|
||||
DEFAULT_VECTOR_EXTENSION = "pgvector" # Options: "pgvector", "vchord", "pgvectorscale"
|
||||
|
||||
# Text search extension (native PostgreSQL, vchord BM25, or Timescale pg_textsearch)
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord", "pg_textsearch"
|
||||
|
||||
# LiteLLM defaults
|
||||
DEFAULT_LITELLM_API_BASE = "http://localhost:4000"
|
||||
DEFAULT_EMBEDDINGS_LITELLM_MODEL = "text-embedding-3-small"
|
||||
DEFAULT_RERANKER_LITELLM_MODEL = "cohere/rerank-english-v3.0"
|
||||
|
||||
# LiteLLM SDK defaults
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL = "cohere/embed-english-v3.0"
|
||||
DEFAULT_RERANKER_LITELLM_SDK_MODEL = "cohere/rerank-english-v3.0"
|
||||
|
||||
DEFAULT_HOST = "0.0.0.0"
|
||||
DEFAULT_PORT = 8888
|
||||
DEFAULT_BASE_PATH = "" # Empty string = root path
|
||||
DEFAULT_LOG_LEVEL = "info"
|
||||
DEFAULT_LOG_FORMAT = "text" # Options: "text", "json"
|
||||
DEFAULT_WORKERS = 1
|
||||
DEFAULT_MCP_ENABLED = True
|
||||
DEFAULT_ENABLE_BANK_CONFIG_API = False # Disabled by default for security
|
||||
DEFAULT_GRAPH_RETRIEVER = "link_expansion" # Options: "link_expansion", "mpfp", "bfs"
|
||||
DEFAULT_MPFP_TOP_K_NEIGHBORS = 20 # Fan-out limit per node in MPFP graph traversal
|
||||
DEFAULT_RECALL_MAX_CONCURRENT = 32 # Max concurrent recall operations per worker
|
||||
@@ -399,17 +248,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_BATCH_TOKENS = 10_000 # ~40KB of text # Max chars per sub-batch for async retain auto-splitting
|
||||
DEFAULT_RETAIN_BATCH_ENABLED = False # Use LLM Batch API for fact extraction (only when async=True)
|
||||
DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60 # Batch API polling interval in seconds
|
||||
|
||||
# File storage defaults
|
||||
DEFAULT_FILE_STORAGE_TYPE = "native" # PostgreSQL BYTEA storage
|
||||
DEFAULT_FILE_PARSER = "markitdown" # File parser to use (markitdown is the only supported parser)
|
||||
DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE_MB = 100 # Max total batch size in MB (all files combined)
|
||||
DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE = 10 # Max files per batch upload
|
||||
DEFAULT_ENABLE_FILE_UPLOAD_API = True # Enable file upload endpoint
|
||||
DEFAULT_FILE_DELETE_AFTER_RETAIN = True # Delete file bytes after retain (saves storage)
|
||||
|
||||
# Observations defaults (consolidated knowledge from facts)
|
||||
DEFAULT_ENABLE_OBSERVATIONS = True # Observations enabled by default
|
||||
@@ -520,8 +358,6 @@ class HindsightConfig:
|
||||
# Database
|
||||
database_url: str
|
||||
database_schema: str
|
||||
vector_extension: str # "pgvector" or "vchord"
|
||||
text_search_extension: str # "native" or "vchord"
|
||||
|
||||
# LLM (default, used as fallback for per-operation config)
|
||||
llm_provider: str
|
||||
@@ -533,8 +369,6 @@ class HindsightConfig:
|
||||
llm_initial_backoff: float
|
||||
llm_max_backoff: float
|
||||
llm_timeout: float
|
||||
llm_groq_service_tier: str # Groq: "on_demand", "flex", or "auto"
|
||||
llm_openai_service_tier: str | None # OpenAI: None (default) or "flex" (50% cheaper)
|
||||
|
||||
# Vertex AI configuration
|
||||
llm_vertexai_project_id: str | None
|
||||
@@ -585,9 +419,6 @@ class HindsightConfig:
|
||||
embeddings_litellm_api_base: str
|
||||
embeddings_litellm_api_key: str | None
|
||||
embeddings_litellm_model: str
|
||||
embeddings_litellm_sdk_api_key: str | None
|
||||
embeddings_litellm_sdk_model: str
|
||||
embeddings_litellm_sdk_api_base: str | None
|
||||
|
||||
# Reranker
|
||||
reranker_provider: str
|
||||
@@ -605,18 +436,13 @@ class HindsightConfig:
|
||||
reranker_litellm_api_base: str
|
||||
reranker_litellm_api_key: str | None
|
||||
reranker_litellm_model: str
|
||||
reranker_litellm_sdk_api_key: str | None
|
||||
reranker_litellm_sdk_model: str
|
||||
reranker_litellm_sdk_api_base: str | None
|
||||
|
||||
# Server
|
||||
host: str
|
||||
port: int
|
||||
base_path: str
|
||||
log_level: str
|
||||
log_format: str
|
||||
mcp_enabled: bool
|
||||
enable_bank_config_api: bool
|
||||
|
||||
# Recall
|
||||
graph_retriever: str
|
||||
@@ -631,29 +457,6 @@ class HindsightConfig:
|
||||
retain_extract_causal_links: bool
|
||||
retain_extraction_mode: str
|
||||
retain_custom_instructions: str | None
|
||||
retain_batch_tokens: int
|
||||
retain_batch_enabled: bool
|
||||
retain_batch_poll_interval_seconds: int
|
||||
|
||||
# File storage (static - server-level only)
|
||||
file_storage_type: str # "native" (PostgreSQL) or "s3" (S3-compatible)
|
||||
file_storage_s3_bucket: str | None # S3 bucket name (required for s3 storage)
|
||||
file_storage_s3_region: str | None # S3 region (optional, uses SDK default)
|
||||
file_storage_s3_endpoint: str | None # S3 endpoint URL (for MinIO, R2, etc.)
|
||||
file_storage_s3_access_key_id: str | None # S3 access key (optional, uses env/IAM)
|
||||
file_storage_s3_secret_access_key: str | None # S3 secret key (optional, uses env/IAM)
|
||||
file_storage_gcs_bucket: str | None # GCS bucket name (required for gcs storage)
|
||||
file_storage_gcs_service_account_key: str | None # GCS service account key JSON (optional, uses ADC)
|
||||
file_storage_azure_container: str | None # Azure container name (required for azure storage)
|
||||
file_storage_azure_account_name: str | None # Azure storage account name
|
||||
file_storage_azure_account_key: str | None # Azure storage account key
|
||||
file_parser: str # File parser to use (e.g., "markitdown", "iris")
|
||||
file_parser_iris_token: str | None # Vectorize API token for iris parser (VECTORIZE_TOKEN)
|
||||
file_parser_iris_org_id: str | None # Vectorize org ID for iris parser (VECTORIZE_ORG_ID)
|
||||
file_conversion_max_batch_size_mb: int # Max total batch size in MB (all files combined)
|
||||
file_conversion_max_batch_size: int # Max files per request
|
||||
enable_file_upload_api: bool
|
||||
file_delete_after_retain: bool
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations: bool
|
||||
@@ -692,120 +495,8 @@ class HindsightConfig:
|
||||
otel_service_name: str
|
||||
otel_deployment_environment: str
|
||||
|
||||
# Class-level sets for configuration categorization
|
||||
|
||||
# CREDENTIAL_FIELDS: Never exposed via API, never configurable per-tenant/bank
|
||||
_CREDENTIAL_FIELDS = {
|
||||
# API Keys
|
||||
"llm_api_key",
|
||||
"retain_llm_api_key",
|
||||
"reflect_llm_api_key",
|
||||
"consolidation_llm_api_key",
|
||||
# Base URLs (could expose infrastructure)
|
||||
"llm_base_url",
|
||||
"retain_llm_base_url",
|
||||
"reflect_llm_base_url",
|
||||
"consolidation_llm_base_url",
|
||||
"embeddings_tei_base_url",
|
||||
"reranker_tei_base_url",
|
||||
"reranker_cohere_base_url",
|
||||
# Service Account Keys
|
||||
"llm_vertexai_service_account_key",
|
||||
# File storage credentials
|
||||
"file_storage_s3_access_key_id",
|
||||
"file_storage_s3_secret_access_key",
|
||||
"file_storage_gcs_service_account_key",
|
||||
"file_storage_azure_account_key",
|
||||
# File parser credentials
|
||||
"file_parser_iris_token",
|
||||
}
|
||||
|
||||
# CONFIGURABLE_FIELDS: Safe behavioral settings that can be customized per-tenant/bank
|
||||
# These fields are manually tagged as safe to expose and modify.
|
||||
# Excludes credentials, infrastructure config, provider/model selection, and performance tuning.
|
||||
_CONFIGURABLE_FIELDS = {
|
||||
# Retention settings (behavioral)
|
||||
"retain_chunk_size",
|
||||
"retain_extraction_mode",
|
||||
"retain_custom_instructions",
|
||||
# Consolidation settings
|
||||
"enable_observations",
|
||||
}
|
||||
|
||||
@property
|
||||
def file_conversion_max_batch_size_bytes(self) -> int:
|
||||
"""Get maximum total batch size in bytes."""
|
||||
return self.file_conversion_max_batch_size_mb * 1024 * 1024
|
||||
|
||||
@classmethod
|
||||
def get_configurable_fields(cls) -> set[str]:
|
||||
"""
|
||||
Get set of field names that are configurable per-tenant/bank via API.
|
||||
|
||||
Configurable fields are manually tagged behavioral settings that are safe
|
||||
to expose and modify (e.g., retain_chunk_size, custom_instructions).
|
||||
Excludes credentials, infrastructure config, and provider/model selection.
|
||||
|
||||
Returns:
|
||||
Set of configurable field names
|
||||
"""
|
||||
return cls._CONFIGURABLE_FIELDS.copy()
|
||||
|
||||
@classmethod
|
||||
def get_credential_fields(cls) -> set[str]:
|
||||
"""
|
||||
Get set of field names that are credentials (NEVER exposed via API).
|
||||
|
||||
Credential fields include API keys, base URLs, and service account keys.
|
||||
These must never be returned in API responses or accepted in updates.
|
||||
|
||||
Returns:
|
||||
Set of credential field names
|
||||
"""
|
||||
return cls._CREDENTIAL_FIELDS.copy()
|
||||
|
||||
@classmethod
|
||||
def get_hierarchical_fields(cls) -> set[str]:
|
||||
"""
|
||||
DEPRECATED: Use get_configurable_fields() instead.
|
||||
|
||||
Kept for backward compatibility during migration.
|
||||
"""
|
||||
return cls.get_configurable_fields()
|
||||
|
||||
@classmethod
|
||||
def get_static_fields(cls) -> set[str]:
|
||||
"""
|
||||
Get set of field names that are static (server-level only).
|
||||
|
||||
Static fields are infrastructure-level settings that cannot vary
|
||||
per tenant or bank. These include database config, API port, worker settings, etc.
|
||||
Also includes credential fields which are never configurable.
|
||||
|
||||
Returns:
|
||||
Set of static field names
|
||||
"""
|
||||
# Get all field names from dataclass
|
||||
all_fields = {f.name for f in fields(cls)}
|
||||
# Static fields = all fields - configurable fields
|
||||
return all_fields - cls._CONFIGURABLE_FIELDS
|
||||
|
||||
def validate(self) -> None:
|
||||
"""Validate configuration values and raise errors for invalid combinations."""
|
||||
# Validate vector_extension
|
||||
valid_extensions = ("pgvector", "vchord", "pgvectorscale")
|
||||
if self.vector_extension not in valid_extensions:
|
||||
raise ValueError(
|
||||
f"Invalid vector_extension: {self.vector_extension}. Must be one of: {', '.join(valid_extensions)}"
|
||||
)
|
||||
|
||||
# Validate text_search_extension
|
||||
valid_text_search = ("native", "vchord", "pg_textsearch")
|
||||
if self.text_search_extension not in valid_text_search:
|
||||
raise ValueError(
|
||||
f"Invalid text_search_extension: {self.text_search_extension}. Must be one of: {', '.join(valid_text_search)}"
|
||||
)
|
||||
|
||||
# RETAIN_MAX_COMPLETION_TOKENS must be greater than RETAIN_CHUNK_SIZE
|
||||
# to ensure the LLM has enough output capacity to extract facts from chunks
|
||||
if self.retain_max_completion_tokens <= self.retain_chunk_size:
|
||||
@@ -831,8 +522,6 @@ class HindsightConfig:
|
||||
# Database
|
||||
database_url=os.getenv(ENV_DATABASE_URL, DEFAULT_DATABASE_URL),
|
||||
database_schema=os.getenv(ENV_DATABASE_SCHEMA, DEFAULT_DATABASE_SCHEMA),
|
||||
vector_extension=os.getenv(ENV_VECTOR_EXTENSION, DEFAULT_VECTOR_EXTENSION).lower(),
|
||||
text_search_extension=os.getenv(ENV_TEXT_SEARCH_EXTENSION, DEFAULT_TEXT_SEARCH_EXTENSION).lower(),
|
||||
# LLM
|
||||
llm_provider=llm_provider,
|
||||
llm_api_key=os.getenv(ENV_LLM_API_KEY),
|
||||
@@ -843,8 +532,6 @@ class HindsightConfig:
|
||||
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))),
|
||||
llm_groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
|
||||
llm_openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
|
||||
# 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),
|
||||
@@ -943,12 +630,6 @@ class HindsightConfig:
|
||||
or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE),
|
||||
embeddings_litellm_api_key=os.getenv(ENV_EMBEDDINGS_LITELLM_API_KEY) or os.getenv(ENV_LITELLM_API_KEY),
|
||||
embeddings_litellm_model=os.getenv(ENV_EMBEDDINGS_LITELLM_MODEL, DEFAULT_EMBEDDINGS_LITELLM_MODEL),
|
||||
# LiteLLM SDK embeddings (direct API access)
|
||||
embeddings_litellm_sdk_api_key=os.getenv(ENV_EMBEDDINGS_LITELLM_SDK_API_KEY),
|
||||
embeddings_litellm_sdk_model=os.getenv(
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_MODEL, DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL
|
||||
),
|
||||
embeddings_litellm_sdk_api_base=os.getenv(ENV_EMBEDDINGS_LITELLM_SDK_API_BASE) or None,
|
||||
# Reranker
|
||||
reranker_provider=os.getenv(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER),
|
||||
reranker_local_model=os.getenv(ENV_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_LOCAL_MODEL),
|
||||
@@ -978,19 +659,12 @@ class HindsightConfig:
|
||||
or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE),
|
||||
reranker_litellm_api_key=os.getenv(ENV_RERANKER_LITELLM_API_KEY) or os.getenv(ENV_LITELLM_API_KEY),
|
||||
reranker_litellm_model=os.getenv(ENV_RERANKER_LITELLM_MODEL, DEFAULT_RERANKER_LITELLM_MODEL),
|
||||
# LiteLLM SDK reranker (direct API access)
|
||||
reranker_litellm_sdk_api_key=os.getenv(ENV_RERANKER_LITELLM_SDK_API_KEY),
|
||||
reranker_litellm_sdk_model=os.getenv(ENV_RERANKER_LITELLM_SDK_MODEL, DEFAULT_RERANKER_LITELLM_SDK_MODEL),
|
||||
reranker_litellm_sdk_api_base=os.getenv(ENV_RERANKER_LITELLM_SDK_API_BASE) or None,
|
||||
# Server
|
||||
host=os.getenv(ENV_HOST, DEFAULT_HOST),
|
||||
port=int(os.getenv(ENV_PORT, DEFAULT_PORT)),
|
||||
base_path=os.getenv(ENV_BASE_PATH, DEFAULT_BASE_PATH),
|
||||
log_level=os.getenv(ENV_LOG_LEVEL, DEFAULT_LOG_LEVEL),
|
||||
log_format=os.getenv(ENV_LOG_FORMAT, DEFAULT_LOG_FORMAT).lower(),
|
||||
mcp_enabled=os.getenv(ENV_MCP_ENABLED, str(DEFAULT_MCP_ENABLED)).lower() == "true",
|
||||
enable_bank_config_api=os.getenv(ENV_ENABLE_BANK_CONFIG_API, str(DEFAULT_ENABLE_BANK_CONFIG_API)).lower()
|
||||
== "true",
|
||||
# Recall
|
||||
graph_retriever=os.getenv(ENV_GRAPH_RETRIEVER, DEFAULT_GRAPH_RETRIEVER),
|
||||
mpfp_top_k_neighbors=int(os.getenv(ENV_MPFP_TOP_K_NEIGHBORS, str(DEFAULT_MPFP_TOP_K_NEIGHBORS))),
|
||||
@@ -1017,39 +691,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_batch_tokens=int(os.getenv(ENV_RETAIN_BATCH_TOKENS, str(DEFAULT_RETAIN_BATCH_TOKENS))),
|
||||
retain_batch_enabled=os.getenv(ENV_RETAIN_BATCH_ENABLED, str(DEFAULT_RETAIN_BATCH_ENABLED)).lower()
|
||||
== "true",
|
||||
retain_batch_poll_interval_seconds=int(
|
||||
os.getenv(ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS, str(DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS))
|
||||
),
|
||||
# File storage
|
||||
file_storage_type=os.getenv(ENV_FILE_STORAGE_TYPE, DEFAULT_FILE_STORAGE_TYPE),
|
||||
file_storage_s3_bucket=os.getenv(ENV_FILE_STORAGE_S3_BUCKET) or None,
|
||||
file_storage_s3_region=os.getenv(ENV_FILE_STORAGE_S3_REGION) or None,
|
||||
file_storage_s3_endpoint=os.getenv(ENV_FILE_STORAGE_S3_ENDPOINT) or None,
|
||||
file_storage_s3_access_key_id=os.getenv(ENV_FILE_STORAGE_S3_ACCESS_KEY_ID) or None,
|
||||
file_storage_s3_secret_access_key=os.getenv(ENV_FILE_STORAGE_S3_SECRET_ACCESS_KEY) or None,
|
||||
file_storage_gcs_bucket=os.getenv(ENV_FILE_STORAGE_GCS_BUCKET) or None,
|
||||
file_storage_gcs_service_account_key=os.getenv(ENV_FILE_STORAGE_GCS_SERVICE_ACCOUNT_KEY) or None,
|
||||
file_storage_azure_container=os.getenv(ENV_FILE_STORAGE_AZURE_CONTAINER) or None,
|
||||
file_storage_azure_account_name=os.getenv(ENV_FILE_STORAGE_AZURE_ACCOUNT_NAME) or None,
|
||||
file_storage_azure_account_key=os.getenv(ENV_FILE_STORAGE_AZURE_ACCOUNT_KEY) or None,
|
||||
file_parser=os.getenv(ENV_FILE_PARSER, DEFAULT_FILE_PARSER),
|
||||
file_parser_iris_token=os.getenv(ENV_FILE_PARSER_IRIS_TOKEN) or None,
|
||||
file_parser_iris_org_id=os.getenv(ENV_FILE_PARSER_IRIS_ORG_ID) or None,
|
||||
file_conversion_max_batch_size_mb=int(
|
||||
os.getenv(ENV_FILE_CONVERSION_MAX_BATCH_SIZE_MB, str(DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE_MB))
|
||||
),
|
||||
file_conversion_max_batch_size=int(
|
||||
os.getenv(ENV_FILE_CONVERSION_MAX_BATCH_SIZE, str(DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE))
|
||||
),
|
||||
enable_file_upload_api=os.getenv(ENV_ENABLE_FILE_UPLOAD_API, str(DEFAULT_ENABLE_FILE_UPLOAD_API)).lower()
|
||||
== "true",
|
||||
file_delete_after_retain=os.getenv(
|
||||
ENV_FILE_DELETE_AFTER_RETAIN, str(DEFAULT_FILE_DELETE_AFTER_RETAIN)
|
||||
).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(
|
||||
@@ -1164,35 +805,8 @@ class HindsightConfig:
|
||||
_config_cache: HindsightConfig | None = None
|
||||
|
||||
|
||||
def get_config() -> StaticConfigProxy:
|
||||
"""
|
||||
Get global configuration with ONLY static (non-configurable) fields accessible.
|
||||
|
||||
This returns a proxy that prevents access to bank-configurable fields
|
||||
(like enable_observations, retain_chunk_size, etc.).
|
||||
|
||||
For bank-specific configuration, use:
|
||||
config_resolver.resolve_full_config(bank_id, context)
|
||||
|
||||
This design prevents accidentally using global defaults when bank-specific
|
||||
overrides exist.
|
||||
|
||||
Returns:
|
||||
StaticConfigProxy that only exposes static infrastructure fields
|
||||
|
||||
Raises:
|
||||
ConfigFieldAccessError: If you try to access a bank-configurable field
|
||||
"""
|
||||
return StaticConfigProxy(_get_raw_config())
|
||||
|
||||
|
||||
def _get_raw_config() -> HindsightConfig:
|
||||
"""
|
||||
Get raw config (internal use only).
|
||||
|
||||
INTERNAL USE ONLY. Do not use this directly in application code.
|
||||
Use get_config() for static fields or ConfigResolver.resolve_full_config() for bank-specific config.
|
||||
"""
|
||||
def get_config() -> HindsightConfig:
|
||||
"""Get the cached configuration, loading from environment on first call."""
|
||||
global _config_cache
|
||||
if _config_cache is None:
|
||||
_config_cache = HindsightConfig.from_env()
|
||||
|
||||
@@ -1,274 +0,0 @@
|
||||
"""
|
||||
Configuration resolution with hierarchical overrides.
|
||||
|
||||
Resolves config values through the hierarchy:
|
||||
Global (env vars) → Tenant config (via extension) → Bank config (database)
|
||||
|
||||
Config values are resolved on every request to ensure consistency across
|
||||
multiple API servers.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import asdict
|
||||
from typing import Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
from hindsight_api.config import HindsightConfig, _get_raw_config, normalize_config_dict
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConfigResolver:
|
||||
"""Resolves hierarchical configuration with tenant/bank overrides."""
|
||||
|
||||
def __init__(self, pool: asyncpg.Pool, tenant_extension: TenantExtension | None = None):
|
||||
"""
|
||||
Initialize config resolver.
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
tenant_extension: Optional tenant extension for tenant-level config and permissions
|
||||
"""
|
||||
self.pool = pool
|
||||
self.tenant_extension = tenant_extension
|
||||
self._global_config = _get_raw_config()
|
||||
self._configurable_fields = HindsightConfig.get_configurable_fields()
|
||||
self._credential_fields = HindsightConfig.get_credential_fields()
|
||||
|
||||
async def resolve_full_config(self, bank_id: str, context: RequestContext | None = None) -> HindsightConfig:
|
||||
"""
|
||||
Resolve full HindsightConfig for a bank with hierarchical overrides applied.
|
||||
|
||||
This is for INTERNAL USE ONLY. Returns the complete config object with all fields
|
||||
including credentials and static fields. Use get_bank_config() for API responses.
|
||||
|
||||
Resolution order:
|
||||
1. Global config (from environment variables)
|
||||
2. Tenant config overrides (from TenantExtension.get_tenant_config())
|
||||
3. Bank config overrides (from banks.config JSONB)
|
||||
|
||||
Args:
|
||||
bank_id: Bank identifier
|
||||
context: Request context for tenant config resolution
|
||||
|
||||
Returns:
|
||||
Complete HindsightConfig with hierarchical overrides applied
|
||||
"""
|
||||
# Start with global config (all fields)
|
||||
config_dict = asdict(self._global_config)
|
||||
|
||||
# Load tenant config overrides (if tenant extension available)
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
tenant_overrides = await self.tenant_extension.get_tenant_config(context)
|
||||
if tenant_overrides:
|
||||
# Normalize keys and filter to configurable fields only
|
||||
normalized_tenant = normalize_config_dict(tenant_overrides)
|
||||
configurable_tenant = {k: v for k, v in normalized_tenant.items() if k in self._configurable_fields}
|
||||
config_dict.update(configurable_tenant)
|
||||
logger.debug(
|
||||
f"Applied tenant config overrides for bank {bank_id}: {list(configurable_tenant.keys())}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load tenant config for bank {bank_id}: {e}")
|
||||
|
||||
# Load bank config overrides
|
||||
bank_overrides = await self._load_bank_config(bank_id)
|
||||
if bank_overrides:
|
||||
config_dict.update(bank_overrides)
|
||||
logger.debug(f"Applied bank config overrides for bank {bank_id}: {list(bank_overrides.keys())}")
|
||||
|
||||
# Return full config object (dataclass doesn't have __init__ that accepts kwargs, so we update the object)
|
||||
# Create a new config instance by copying the global config and updating fields
|
||||
resolved_config = HindsightConfig(**config_dict)
|
||||
return resolved_config
|
||||
|
||||
async def get_bank_config(self, bank_id: str, context: RequestContext | None = None) -> dict[str, Any]:
|
||||
"""
|
||||
Get fully resolved config for a bank (filtered by permissions).
|
||||
|
||||
Resolution order:
|
||||
1. Global config (from environment variables)
|
||||
2. Tenant config overrides (from TenantExtension.get_tenant_config())
|
||||
3. Bank config overrides (from banks.config JSONB)
|
||||
|
||||
Note: Config is resolved on every call (not cached) to ensure consistency
|
||||
across multiple API servers.
|
||||
|
||||
SECURITY:
|
||||
- Only returns configurable fields (excludes static/infrastructure fields)
|
||||
- Filters out ALL credential fields (API keys, base URLs, etc.)
|
||||
- Further filtered by tenant/bank permissions if extension provides them
|
||||
|
||||
Args:
|
||||
bank_id: Bank identifier
|
||||
context: Request context for tenant config resolution and permissions
|
||||
|
||||
Returns:
|
||||
Dict of allowed configurable fields only (never includes credentials or static fields)
|
||||
"""
|
||||
# Resolve full config with all hierarchical overrides
|
||||
resolved_config = await self.resolve_full_config(bank_id, context)
|
||||
config_dict = asdict(resolved_config)
|
||||
|
||||
# SECURITY: Filter to only configurable fields (exclude static/infrastructure)
|
||||
filtered = {k: v for k, v in config_dict.items() if k in self._configurable_fields}
|
||||
|
||||
# SECURITY: Remove ALL credential fields (API keys, base URLs, etc.)
|
||||
filtered = {k: v for k, v in filtered.items() if k not in self._credential_fields}
|
||||
|
||||
# PERMISSIONS: Further filter based on tenant/bank permissions
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
filtered = {k: v for k, v in filtered.items() if k in allowed_fields}
|
||||
logger.debug(
|
||||
f"Applied permission filter for bank {bank_id}: allowed={len(allowed_fields)} fields, "
|
||||
f"returned={len(filtered)} fields"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load permissions for bank {bank_id}: {e}")
|
||||
|
||||
return filtered
|
||||
|
||||
async def _load_bank_config(self, bank_id: str) -> dict[str, Any]:
|
||||
"""
|
||||
Load bank config overrides from banks.config JSONB column.
|
||||
|
||||
Args:
|
||||
bank_id: Bank identifier
|
||||
|
||||
Returns:
|
||||
Dict of config overrides (only configurable fields, normalized keys)
|
||||
"""
|
||||
try:
|
||||
async with self.pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT config FROM banks WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if row and row["config"]:
|
||||
config_data = row["config"]
|
||||
|
||||
# Handle case where JSONB is returned as JSON string
|
||||
if isinstance(config_data, str):
|
||||
config_data = json.loads(config_data)
|
||||
|
||||
# Normalize keys (handle both env var format and Python field format)
|
||||
normalized = normalize_config_dict(config_data)
|
||||
|
||||
# Only return overrides for configurable fields
|
||||
return {k: v for k, v in normalized.items() if k in self._configurable_fields}
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load bank config for {bank_id}: {e}")
|
||||
|
||||
return {}
|
||||
|
||||
async def update_bank_config(
|
||||
self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Update bank configuration overrides (with permission checking).
|
||||
|
||||
Args:
|
||||
bank_id: Bank identifier
|
||||
updates: Dict of config field names to new values.
|
||||
Keys can be in env var format (HINDSIGHT_API_LLM_PROVIDER)
|
||||
or Python field format (llm_provider).
|
||||
Only configurable fields are allowed.
|
||||
context: Request context for permission checking
|
||||
|
||||
Raises:
|
||||
ValueError: If attempting to override invalid/disallowed fields
|
||||
"""
|
||||
# Normalize keys
|
||||
normalized_updates = normalize_config_dict(updates)
|
||||
|
||||
# SECURITY: Reject credential fields explicitly
|
||||
credential_attempts = set(normalized_updates.keys()) & self._credential_fields
|
||||
if credential_attempts:
|
||||
raise ValueError(
|
||||
f"Cannot set credential fields via API: {sorted(credential_attempts)}. "
|
||||
f"Credentials (API keys, base URLs) must be set at server level only."
|
||||
)
|
||||
|
||||
# Validate all fields are configurable
|
||||
invalid_fields = set(normalized_updates.keys()) - self._configurable_fields
|
||||
if invalid_fields:
|
||||
static_fields = HindsightConfig.get_static_fields()
|
||||
invalid_static = invalid_fields & static_fields
|
||||
if invalid_static:
|
||||
raise ValueError(
|
||||
f"Cannot override static (server-level) fields: {sorted(invalid_static)}. "
|
||||
f"Only configurable fields can be overridden per-bank. "
|
||||
f"Configurable fields include: {sorted(list(self._configurable_fields)[:10])}... "
|
||||
f"(total: {len(self._configurable_fields)} fields)"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown configuration fields: {sorted(invalid_fields)}. "
|
||||
f"Valid configurable fields: {sorted(list(self._configurable_fields)[:10])}..."
|
||||
)
|
||||
|
||||
# PERMISSIONS: Check tenant/bank permissions
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
disallowed = set(normalized_updates.keys()) - allowed_fields
|
||||
if disallowed:
|
||||
raise ValueError(
|
||||
f"Not allowed to modify fields: {sorted(disallowed)}. "
|
||||
f"Your permissions allow: {sorted(list(allowed_fields)[:10])}..."
|
||||
if allowed_fields
|
||||
else "Not allowed to modify fields: {sorted(disallowed)}. "
|
||||
"Your permissions do not allow any config modifications."
|
||||
)
|
||||
except ValueError:
|
||||
raise # Re-raise permission errors
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to check permissions for bank {bank_id}: {e}")
|
||||
# Continue without permission check (fail open for backward compatibility)
|
||||
|
||||
# Merge with existing config (JSONB || operator)
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE banks
|
||||
SET config = config || $1::jsonb,
|
||||
updated_at = now()
|
||||
WHERE bank_id = $2
|
||||
""",
|
||||
json.dumps(normalized_updates),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
logger.info(f"Updated bank config for {bank_id}: {list(normalized_updates.keys())}")
|
||||
|
||||
async def reset_bank_config(self, bank_id: str) -> None:
|
||||
"""
|
||||
Reset bank configuration to defaults (remove all overrides).
|
||||
|
||||
Args:
|
||||
bank_id: Bank identifier
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE banks
|
||||
SET config = '{}'::jsonb,
|
||||
updated_at = now()
|
||||
WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
logger.info(f"Reset bank config for {bank_id} to defaults")
|
||||
@@ -18,7 +18,6 @@ import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ...config import get_config
|
||||
from ..memory_engine import fq_table
|
||||
from ..retain import embedding_utils
|
||||
from .prompts import (
|
||||
@@ -83,8 +82,9 @@ async def run_consolidation_job(
|
||||
Returns:
|
||||
Dict with consolidation results
|
||||
"""
|
||||
# Resolve bank-specific config with hierarchical overrides
|
||||
config = await memory_engine._config_resolver.resolve_full_config(bank_id, request_context)
|
||||
from ...config import get_config
|
||||
|
||||
config = get_config()
|
||||
perf = ConsolidationPerfLog(bank_id)
|
||||
max_memories_per_batch = config.consolidation_batch_size
|
||||
|
||||
@@ -1016,34 +1016,15 @@ async def _create_observation_directly(
|
||||
|
||||
t0 = time.time()
|
||||
observation_id = uuid.uuid4()
|
||||
|
||||
# Query varies based on text search backend
|
||||
config = get_config()
|
||||
if config.text_search_extension == "vchord":
|
||||
# VectorChord: manually tokenize and insert search_vector
|
||||
query = f"""
|
||||
INSERT INTO {fq_table("memory_units")} (
|
||||
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history,
|
||||
tags, event_date, occurred_start, occurred_end, mentioned_at, search_vector
|
||||
)
|
||||
VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8, $9, $10,
|
||||
tokenize($3, 'llmlingua2')::bm25_catalog.bm25vector)
|
||||
RETURNING id
|
||||
"""
|
||||
else: # native or pg_textsearch
|
||||
# Native PostgreSQL: search_vector is GENERATED ALWAYS, don't include it
|
||||
# pg_textsearch: indexes operate on base columns directly, don't populate search_vector
|
||||
query = f"""
|
||||
INSERT INTO {fq_table("memory_units")} (
|
||||
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history,
|
||||
tags, event_date, occurred_start, occurred_end, mentioned_at
|
||||
)
|
||||
VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8, $9, $10)
|
||||
RETURNING id
|
||||
"""
|
||||
|
||||
row = await conn.fetchrow(
|
||||
query,
|
||||
f"""
|
||||
INSERT INTO {fq_table("memory_units")} (
|
||||
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history,
|
||||
tags, event_date, occurred_start, occurred_end, mentioned_at
|
||||
)
|
||||
VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8, $9, $10)
|
||||
RETURNING id
|
||||
""",
|
||||
observation_id,
|
||||
bank_id,
|
||||
observation_text,
|
||||
|
||||
@@ -21,7 +21,6 @@ from ..config import (
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL,
|
||||
DEFAULT_RERANKER_LITELLM_MODEL,
|
||||
DEFAULT_RERANKER_LITELLM_SDK_MODEL,
|
||||
DEFAULT_RERANKER_LOCAL_FORCE_CPU,
|
||||
DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_LOCAL_MODEL,
|
||||
@@ -33,7 +32,6 @@ from ..config import (
|
||||
ENV_RERANKER_COHERE_MODEL,
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
ENV_RERANKER_FLASHRANK_MODEL,
|
||||
ENV_RERANKER_LITELLM_SDK_API_KEY,
|
||||
ENV_RERANKER_LOCAL_FORCE_CPU,
|
||||
ENV_RERANKER_LOCAL_MAX_CONCURRENT,
|
||||
ENV_RERANKER_LOCAL_MODEL,
|
||||
@@ -830,126 +828,6 @@ class LiteLLMCrossEncoder(CrossEncoderModel):
|
||||
return all_scores
|
||||
|
||||
|
||||
class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
LiteLLM SDK cross-encoder for direct API integration.
|
||||
|
||||
Supports reranking via LiteLLM SDK without requiring a proxy server.
|
||||
Supported providers: Cohere, DeepInfra, Together AI, HuggingFace, Jina AI, Voyage AI, AWS Bedrock.
|
||||
|
||||
Example model names:
|
||||
- cohere/rerank-english-v3.0
|
||||
- deepinfra/Qwen3-reranker-8B
|
||||
- together_ai/Salesforce/Llama-Rank-V1
|
||||
- huggingface/BAAI/bge-reranker-v2-m3
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_RERANKER_LITELLM_SDK_MODEL,
|
||||
api_base: str | None = None,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
"""
|
||||
Initialize LiteLLM SDK cross-encoder client.
|
||||
|
||||
Args:
|
||||
api_key: API key for the reranking provider
|
||||
model: Model name with provider prefix (e.g., "deepinfra/Qwen3-reranker-8B")
|
||||
api_base: Custom base URL for API (optional)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.api_base = api_base
|
||||
self.timeout = timeout
|
||||
self._initialized = False
|
||||
self._litellm = None # Will be set during initialization
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "litellm-sdk"
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the LiteLLM SDK client."""
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
try:
|
||||
import litellm
|
||||
|
||||
self._litellm = litellm # Store reference
|
||||
except ImportError:
|
||||
raise ImportError("litellm is required for LiteLLMSDKCrossEncoder. Install it with: pip install litellm")
|
||||
|
||||
api_base_msg = f" at {self.api_base}" if self.api_base else ""
|
||||
logger.info(f"Reranker: initializing LiteLLM SDK provider with model {self.model}{api_base_msg}")
|
||||
|
||||
self._initialized = True
|
||||
logger.info("Reranker: LiteLLM SDK provider initialized")
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
Score query-document pairs using the LiteLLM SDK.
|
||||
|
||||
Args:
|
||||
pairs: List of (query, document) tuples to score
|
||||
|
||||
Returns:
|
||||
List of relevance scores
|
||||
"""
|
||||
if not self._initialized:
|
||||
raise RuntimeError("Reranker not initialized. Call initialize() first.")
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
# Group pairs by query for efficient batching
|
||||
# LiteLLM rerank expects one query with multiple documents
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
texts = [text for _, text in indexed_texts]
|
||||
indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
# Build kwargs for rerank call
|
||||
rerank_kwargs = {
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": texts,
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_base:
|
||||
rerank_kwargs["api_base"] = self.api_base
|
||||
|
||||
response = await self._litellm.arerank(**rerank_kwargs)
|
||||
|
||||
# Map scores back to original positions
|
||||
# Response format: RerankResponse with results list
|
||||
# Each result is a TypedDict with "index" and "relevance_score"
|
||||
if hasattr(response, "results") and response.results:
|
||||
for result in response.results:
|
||||
# Results are TypedDicts, use dict-style access
|
||||
original_idx = result["index"]
|
||||
score = result.get("relevance_score", result.get("score", 0.0))
|
||||
all_scores[indices[original_idx]] = score
|
||||
elif isinstance(response, list):
|
||||
# Direct list of scores (unlikely but defensive)
|
||||
for i, score in enumerate(response):
|
||||
all_scores[indices[i]] = score
|
||||
else:
|
||||
logger.warning(f"Unexpected response format from LiteLLM rerank: {type(response)}")
|
||||
|
||||
return all_scores
|
||||
|
||||
|
||||
def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
"""
|
||||
Create a CrossEncoderModel instance based on configuration.
|
||||
@@ -999,20 +877,9 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=config.reranker_litellm_api_key,
|
||||
model=config.reranker_litellm_model,
|
||||
)
|
||||
elif provider == "litellm-sdk":
|
||||
api_key = config.reranker_litellm_sdk_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_RERANKER_LITELLM_SDK_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'litellm-sdk'"
|
||||
)
|
||||
return LiteLLMSDKCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_litellm_sdk_model,
|
||||
api_base=config.reranker_litellm_sdk_api_base,
|
||||
)
|
||||
elif provider == "rrf":
|
||||
return RRFPassthroughCrossEncoder()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'flashrank', 'litellm', 'litellm-sdk', 'rrf'"
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'flashrank', 'litellm', 'rrf'"
|
||||
)
|
||||
|
||||
@@ -19,7 +19,6 @@ import httpx
|
||||
from ..config import (
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LITELLM_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
@@ -27,7 +26,6 @@ from ..config import (
|
||||
DEFAULT_EMBEDDINGS_PROVIDER,
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY,
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY,
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL,
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
@@ -722,148 +720,6 @@ class LiteLLMEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
class LiteLLMSDKEmbeddings(Embeddings):
|
||||
"""
|
||||
LiteLLM SDK embeddings for direct API integration.
|
||||
|
||||
Supports embeddings via LiteLLM SDK without requiring a proxy server.
|
||||
Supported providers: Cohere, OpenAI, Azure OpenAI, HuggingFace, Voyage AI, Together AI, etc.
|
||||
|
||||
Example model names:
|
||||
- cohere/embed-english-v3.0
|
||||
- openai/text-embedding-3-small
|
||||
- together_ai/togethercomputer/m2-bert-80M-8k-retrieval
|
||||
- voyage/voyage-2
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
api_base: str | None = None,
|
||||
batch_size: int = 100,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
"""
|
||||
Initialize LiteLLM SDK embeddings client.
|
||||
|
||||
Args:
|
||||
api_key: API key for the embedding provider
|
||||
model: Model name with provider prefix (e.g., "cohere/embed-english-v3.0")
|
||||
api_base: Custom base URL for API (optional)
|
||||
batch_size: Maximum batch size for embedding requests (default: 100)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.api_base = api_base
|
||||
self.batch_size = batch_size
|
||||
self.timeout = timeout
|
||||
self._litellm = None # Will be set during initialization
|
||||
self._dimension: int | None = None
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "litellm-sdk"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
if self._dimension is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
return self._dimension
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the LiteLLM SDK client and detect dimension."""
|
||||
if self._litellm is not None:
|
||||
return
|
||||
|
||||
try:
|
||||
import litellm
|
||||
|
||||
self._litellm = litellm # Store reference
|
||||
except ImportError:
|
||||
raise ImportError("litellm is required for LiteLLMSDKEmbeddings. Install it with: pip install litellm")
|
||||
|
||||
api_base_msg = f" at {self.api_base}" if self.api_base else ""
|
||||
logger.info(f"Embeddings: initializing LiteLLM SDK provider with model {self.model}{api_base_msg}")
|
||||
|
||||
# Do a test embedding to detect dimension
|
||||
try:
|
||||
# Build kwargs for embedding call
|
||||
embed_kwargs = {
|
||||
"model": self.model,
|
||||
"input": ["test"],
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_base:
|
||||
embed_kwargs["api_base"] = self.api_base
|
||||
|
||||
# Use async embedding method (standard in litellm)
|
||||
response = await self._litellm.aembedding(**embed_kwargs)
|
||||
|
||||
# Extract dimension from response
|
||||
if response.data and len(response.data) > 0:
|
||||
self._dimension = len(response.data[0]["embedding"])
|
||||
else:
|
||||
raise RuntimeError(f"Unable to detect embedding dimension for model {self.model}")
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to initialize LiteLLM SDK embeddings: {e}")
|
||||
|
||||
logger.info(f"Embeddings: LiteLLM SDK provider initialized (model: {self.model}, dim: {self._dimension})")
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""
|
||||
Generate embeddings using the LiteLLM SDK.
|
||||
|
||||
Args:
|
||||
texts: List of text strings to encode
|
||||
|
||||
Returns:
|
||||
List of embedding vectors (one per input text)
|
||||
"""
|
||||
if self._litellm is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
all_embeddings = []
|
||||
|
||||
# Process in batches
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
|
||||
try:
|
||||
# Build kwargs for embedding call
|
||||
embed_kwargs = {
|
||||
"model": self.model,
|
||||
"input": batch,
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_base:
|
||||
embed_kwargs["api_base"] = self.api_base
|
||||
|
||||
# Use sync embedding (litellm doesn't have async in thread-safe way)
|
||||
response = self._litellm.embedding(**embed_kwargs)
|
||||
|
||||
# Extract embeddings from response
|
||||
# Sort by index to ensure correct order
|
||||
batch_embeddings = sorted(response.data, key=lambda x: x.get("index", 0))
|
||||
all_embeddings.extend([e["embedding"] for e in batch_embeddings])
|
||||
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
logger.error(
|
||||
f"Error in LiteLLM embedding for batch starting at index {i}: {e}\n"
|
||||
f"Traceback: {traceback.format_exc()}"
|
||||
)
|
||||
raise
|
||||
|
||||
return all_embeddings
|
||||
|
||||
|
||||
def create_embeddings_from_env() -> Embeddings:
|
||||
"""
|
||||
Create an Embeddings instance based on configuration.
|
||||
@@ -915,19 +771,7 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
api_key=config.embeddings_litellm_api_key,
|
||||
model=config.embeddings_litellm_model,
|
||||
)
|
||||
elif provider == "litellm-sdk":
|
||||
api_key = config.embeddings_litellm_sdk_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_LITELLM_SDK_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'litellm-sdk'"
|
||||
)
|
||||
return LiteLLMSDKEmbeddings(
|
||||
api_key=api_key,
|
||||
model=config.embeddings_litellm_sdk_model,
|
||||
api_base=config.embeddings_litellm_sdk_api_base,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown embeddings provider: {provider}. "
|
||||
f"Supported: 'local', 'tei', 'openai', 'cohere', 'litellm', 'litellm-sdk'"
|
||||
f"Unknown embeddings provider: {provider}. Supported: 'local', 'tei', 'openai', 'cohere', 'litellm'"
|
||||
)
|
||||
|
||||
@@ -48,7 +48,6 @@ class MemoryEngineInterface(ABC):
|
||||
contents: list[dict[str, Any]],
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
document_tags: list[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Retain a batch of memory items.
|
||||
@@ -56,9 +55,8 @@ class MemoryEngineInterface(ABC):
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
contents: List of content dicts with 'content', optional 'event_date',
|
||||
'context', 'metadata', 'document_id', and per-item 'tags'.
|
||||
'context', 'metadata', 'document_id'.
|
||||
request_context: Request context for authentication.
|
||||
document_tags: Optional tags applied to all items in the batch.
|
||||
|
||||
Returns:
|
||||
Dict with processing results.
|
||||
@@ -563,7 +561,6 @@ class MemoryEngineInterface(ABC):
|
||||
contents: list[dict[str, Any]],
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
document_tags: list[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Submit a batch retain operation to run asynchronously.
|
||||
@@ -572,7 +569,6 @@ class MemoryEngineInterface(ABC):
|
||||
bank_id: The memory bank ID.
|
||||
contents: List of content dicts to retain.
|
||||
request_context: Request context for authentication.
|
||||
document_tags: Optional tags applied to all items in the async batch.
|
||||
|
||||
Returns:
|
||||
Dict with operation_id and items_count.
|
||||
|
||||
@@ -128,67 +128,6 @@ class LLMInterface(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
async def supports_batch_api(self) -> bool:
|
||||
"""
|
||||
Check if this provider supports batch API operations.
|
||||
|
||||
Returns:
|
||||
True if provider supports submit_batch/get_batch_status/retrieve_batch_results
|
||||
"""
|
||||
return False
|
||||
|
||||
async def submit_batch(
|
||||
self,
|
||||
requests: list[dict[str, Any]],
|
||||
endpoint: str = "/v1/chat/completions",
|
||||
completion_window: str = "24h",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Submit a batch of requests to the provider's batch API.
|
||||
|
||||
Args:
|
||||
requests: List of request dicts in JSONL format (custom_id, method, url, body)
|
||||
endpoint: API endpoint for the batch (e.g., "/v1/chat/completions")
|
||||
completion_window: Completion window (e.g., "24h")
|
||||
|
||||
Returns:
|
||||
Dict with batch metadata: {"batch_id": str, "status": str, ...}
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If provider doesn't support batch API
|
||||
"""
|
||||
raise NotImplementedError(f"Batch API not supported for provider: {self.provider}")
|
||||
|
||||
async def get_batch_status(self, batch_id: str) -> dict[str, Any]:
|
||||
"""
|
||||
Get the status of a batch job.
|
||||
|
||||
Args:
|
||||
batch_id: Batch identifier returned from submit_batch
|
||||
|
||||
Returns:
|
||||
Dict with status info: {"batch_id": str, "status": str, "completed_at": str, ...}
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If provider doesn't support batch API
|
||||
"""
|
||||
raise NotImplementedError(f"Batch API not supported for provider: {self.provider}")
|
||||
|
||||
async def retrieve_batch_results(self, batch_id: str) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Retrieve completed batch results.
|
||||
|
||||
Args:
|
||||
batch_id: Batch identifier returned from submit_batch
|
||||
|
||||
Returns:
|
||||
List of result dicts (one per request, matched by custom_id)
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If provider doesn't support batch API
|
||||
"""
|
||||
raise NotImplementedError(f"Batch API not supported for provider: {self.provider}")
|
||||
|
||||
@abstractmethod
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (close connections, etc.)."""
|
||||
|
||||
@@ -67,7 +67,6 @@ def create_llm_provider(
|
||||
model: str,
|
||||
reasoning_effort: str,
|
||||
groq_service_tier: str | None = None,
|
||||
openai_service_tier: str | None = None,
|
||||
vertexai_project_id: str | None = None,
|
||||
vertexai_region: str | None = None,
|
||||
vertexai_credentials: Any = None,
|
||||
@@ -81,8 +80,7 @@ def create_llm_provider(
|
||||
base_url: Base URL for the API.
|
||||
model: Model name.
|
||||
reasoning_effort: Reasoning effort level for supported providers.
|
||||
groq_service_tier: Groq service tier (for Groq provider) - "on_demand", "flex", or "auto".
|
||||
openai_service_tier: OpenAI service tier (for OpenAI provider) - None (default) or "flex" (50% cheaper).
|
||||
groq_service_tier: Groq service tier (for Groq provider).
|
||||
vertexai_project_id: Vertex AI project ID (for VertexAI provider).
|
||||
vertexai_region: Vertex AI region (for VertexAI provider).
|
||||
vertexai_credentials: Vertex AI credentials object (for VertexAI provider).
|
||||
@@ -158,7 +156,6 @@ def create_llm_provider(
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
groq_service_tier=groq_service_tier,
|
||||
openai_service_tier=openai_service_tier,
|
||||
)
|
||||
|
||||
else:
|
||||
@@ -180,7 +177,6 @@ class LLMProvider:
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
groq_service_tier: str | None = None,
|
||||
openai_service_tier: str | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize LLM provider.
|
||||
@@ -191,17 +187,15 @@ class LLMProvider:
|
||||
base_url: Base URL for the API.
|
||||
model: Model name.
|
||||
reasoning_effort: Reasoning effort level for supported providers.
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto") - from config.
|
||||
openai_service_tier: OpenAI service tier (None or "flex") - from config.
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto"). Default: None (uses Groq's default).
|
||||
"""
|
||||
self.provider = provider.lower()
|
||||
self.api_key = api_key
|
||||
self.base_url = base_url
|
||||
self.model = model
|
||||
self.reasoning_effort = reasoning_effort
|
||||
# Service tiers from hierarchical config (not env vars)
|
||||
self.groq_service_tier = groq_service_tier
|
||||
self.openai_service_tier = openai_service_tier
|
||||
# Default to 'auto' for best performance, users can override to 'on_demand' for free tier
|
||||
self.groq_service_tier = groq_service_tier or os.getenv(ENV_LLM_GROQ_SERVICE_TIER, "auto")
|
||||
|
||||
# Validate provider
|
||||
valid_providers = [
|
||||
@@ -278,7 +272,6 @@ class LLMProvider:
|
||||
model=self.model,
|
||||
reasoning_effort=self.reasoning_effort,
|
||||
groq_service_tier=self.groq_service_tier,
|
||||
openai_service_tier=self.openai_service_tier,
|
||||
vertexai_project_id=vertexai_project_id,
|
||||
vertexai_region=vertexai_region,
|
||||
vertexai_credentials=vertexai_credentials,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,69 +0,0 @@
|
||||
"""
|
||||
Typed metadata models for async operations.
|
||||
|
||||
These dataclasses define the structure of result_metadata for different operation types.
|
||||
The metadata is exposed in the API for debugging purposes and may change without notice.
|
||||
"""
|
||||
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchRetainParentMetadata:
|
||||
"""Metadata for parent batch_retain operations (when split into sub-batches)."""
|
||||
|
||||
items_count: int
|
||||
total_tokens: int
|
||||
num_sub_batches: int
|
||||
is_parent: bool = True
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchRetainChildMetadata:
|
||||
"""Metadata for child batch_retain operations (individual sub-batches)."""
|
||||
|
||||
items_count: int
|
||||
parent_operation_id: str
|
||||
sub_batch_index: int
|
||||
total_sub_batches: int
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainMetadata:
|
||||
"""Metadata for regular retain operations (non-batched, deprecated async path)."""
|
||||
|
||||
items_count: int
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConsolidationMetadata:
|
||||
"""Metadata for consolidation operations."""
|
||||
|
||||
# Currently empty, but structure for future fields
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RefreshMentalModelMetadata:
|
||||
"""Metadata for mental model refresh operations."""
|
||||
|
||||
mental_model_id: str
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
@@ -1,62 +0,0 @@
|
||||
"""File parser implementations."""
|
||||
|
||||
from .base import FileParser, UnsupportedFileTypeError
|
||||
from .iris import IrisParser
|
||||
from .markitdown import MarkitdownParser
|
||||
|
||||
__all__ = ["FileParser", "UnsupportedFileTypeError", "IrisParser", "MarkitdownParser", "FileParserRegistry"]
|
||||
|
||||
|
||||
class FileParserRegistry:
|
||||
"""Registry for file parsers with auto-detection."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize empty parser registry."""
|
||||
self._parsers: dict[str, FileParser] = {}
|
||||
|
||||
def register(self, parser: FileParser):
|
||||
"""
|
||||
Register a parser.
|
||||
|
||||
Args:
|
||||
parser: FileParser instance
|
||||
"""
|
||||
self._parsers[parser.name()] = parser
|
||||
|
||||
def get_parser(
|
||||
self,
|
||||
name: str | None,
|
||||
filename: str,
|
||||
content_type: str | None = None,
|
||||
) -> FileParser:
|
||||
"""
|
||||
Get parser by name or auto-detect.
|
||||
|
||||
Args:
|
||||
name: Parser name (e.g., "markitdown") or None for auto-detect
|
||||
filename: File name for auto-detection
|
||||
content_type: MIME type (optional)
|
||||
|
||||
Returns:
|
||||
FileParser instance
|
||||
|
||||
Raises:
|
||||
ValueError: If no suitable parser found
|
||||
"""
|
||||
if name:
|
||||
# Explicit parser requested — return it directly, let the parser
|
||||
# raise UnsupportedFileTypeError from convert() if needed
|
||||
if name not in self._parsers:
|
||||
raise ValueError(f"Parser '{name}' not found. Available: {list(self._parsers.keys())}")
|
||||
return self._parsers[name]
|
||||
|
||||
# Auto-detect parser
|
||||
for parser in self._parsers.values():
|
||||
if parser.supports(filename, content_type):
|
||||
return parser
|
||||
|
||||
raise ValueError(f"No parser found for {filename}. Available parsers: {list(self._parsers.keys())}")
|
||||
|
||||
def list_parsers(self) -> list[str]:
|
||||
"""Get list of registered parser names."""
|
||||
return list(self._parsers.keys())
|
||||
@@ -1,58 +0,0 @@
|
||||
"""Abstract base class for file parsers."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class UnsupportedFileTypeError(Exception):
|
||||
"""Raised by a parser when it does not support the given file type."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class FileParser(ABC):
|
||||
"""Abstract base for file to markdown parsers."""
|
||||
|
||||
@abstractmethod
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
"""
|
||||
Parse file to markdown.
|
||||
|
||||
Args:
|
||||
file_data: Raw file bytes
|
||||
filename: Original filename (used for format detection)
|
||||
|
||||
Returns:
|
||||
Markdown content as string
|
||||
|
||||
Raises:
|
||||
UnsupportedFileTypeError: If the file type is not supported by this parser
|
||||
RuntimeError: If parsing fails for another reason
|
||||
"""
|
||||
pass
|
||||
|
||||
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||
"""
|
||||
Check if parser supports this file type.
|
||||
|
||||
Override this for local/static extension-based filtering.
|
||||
Parsers that delegate to a remote service should leave this as True
|
||||
and raise UnsupportedFileTypeError from convert() instead.
|
||||
|
||||
Args:
|
||||
filename: File name (used for extension check)
|
||||
content_type: MIME type (optional)
|
||||
|
||||
Returns:
|
||||
True if this parser can handle the file (default: True)
|
||||
"""
|
||||
return True
|
||||
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
"""
|
||||
Get parser name.
|
||||
|
||||
Returns:
|
||||
Parser name (e.g., "markitdown")
|
||||
"""
|
||||
pass
|
||||
@@ -1,137 +0,0 @@
|
||||
"""Iris parser implementation using the Vectorize Iris HTTP API."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import mimetypes
|
||||
import time
|
||||
|
||||
import httpx
|
||||
|
||||
from .base import FileParser, UnsupportedFileTypeError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_IRIS_BASE_URL = "https://api.vectorize.io/v1"
|
||||
_DEFAULT_POLL_INTERVAL = 2.0 # seconds
|
||||
_DEFAULT_TIMEOUT = 300.0 # seconds
|
||||
|
||||
|
||||
class IrisParser(FileParser):
|
||||
"""
|
||||
Iris file parser using the Vectorize Iris cloud extraction service.
|
||||
|
||||
Uploads files to the Vectorize Iris API, starts an extraction job,
|
||||
and polls until the text is ready. The API determines which file types
|
||||
are supported — UnsupportedFileTypeError is raised if the file is rejected.
|
||||
|
||||
Authentication:
|
||||
Requires HINDSIGHT_API_FILE_PARSER_IRIS_TOKEN and
|
||||
HINDSIGHT_API_FILE_PARSER_IRIS_ORG_ID environment variables,
|
||||
or pass them explicitly via the constructor.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
token: str,
|
||||
org_id: str,
|
||||
poll_interval: float = _DEFAULT_POLL_INTERVAL,
|
||||
timeout: float = _DEFAULT_TIMEOUT,
|
||||
):
|
||||
"""
|
||||
Initialize iris parser.
|
||||
|
||||
Args:
|
||||
token: Vectorize API token
|
||||
org_id: Vectorize organization ID
|
||||
poll_interval: Seconds between status poll requests (default: 2)
|
||||
timeout: Maximum seconds to wait for extraction (default: 300)
|
||||
"""
|
||||
self._token = token
|
||||
self._org_id = org_id
|
||||
self._poll_interval = poll_interval
|
||||
self._timeout = timeout
|
||||
self._auth_headers = {"Authorization": f"Bearer {token}"}
|
||||
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
"""
|
||||
Parse file to text using the Vectorize Iris API.
|
||||
|
||||
Raises:
|
||||
UnsupportedFileTypeError: If the Iris API rejects the file type (4xx)
|
||||
RuntimeError: If extraction fails for another reason
|
||||
"""
|
||||
content_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
# Step 1: Request a presigned upload URL
|
||||
init_resp = await client.post(
|
||||
f"{_IRIS_BASE_URL}/org/{self._org_id}/files",
|
||||
headers=self._auth_headers,
|
||||
json={"name": filename, "contentType": content_type},
|
||||
)
|
||||
_raise_for_status(init_resp, filename, "file upload init")
|
||||
init_data = init_resp.json()
|
||||
file_id: str = init_data["fileId"]
|
||||
upload_url: str = init_data["uploadUrl"]
|
||||
|
||||
# Step 2: Upload the file bytes to the presigned URL (no auth header)
|
||||
upload_resp = await client.put(
|
||||
upload_url,
|
||||
content=file_data,
|
||||
headers={"Content-Type": content_type},
|
||||
)
|
||||
_raise_for_status(upload_resp, filename, "file upload")
|
||||
|
||||
# Step 3: Start extraction
|
||||
extract_resp = await client.post(
|
||||
f"{_IRIS_BASE_URL}/org/{self._org_id}/extraction",
|
||||
headers=self._auth_headers,
|
||||
json={"fileId": file_id},
|
||||
)
|
||||
_raise_for_status(extract_resp, filename, "start extraction")
|
||||
extraction_id: str = extract_resp.json()["extractionId"]
|
||||
|
||||
# Step 4: Poll until ready or timeout
|
||||
deadline = time.monotonic() + self._timeout
|
||||
while True:
|
||||
status_resp = await client.get(
|
||||
f"{_IRIS_BASE_URL}/org/{self._org_id}/extraction/{extraction_id}",
|
||||
headers=self._auth_headers,
|
||||
)
|
||||
_raise_for_status(status_resp, filename, "poll extraction status")
|
||||
status_data = status_resp.json()
|
||||
|
||||
if status_data.get("ready"):
|
||||
data = status_data.get("data", {})
|
||||
if not data.get("success"):
|
||||
error = data.get("error", "unknown error")
|
||||
raise RuntimeError(f"Iris extraction failed for '{filename}': {error}")
|
||||
text = data.get("text")
|
||||
if not text:
|
||||
raise RuntimeError(f"No content extracted from '{filename}'")
|
||||
return text
|
||||
|
||||
if time.monotonic() >= deadline:
|
||||
raise RuntimeError(f"Iris extraction timed out after {self._timeout}s for '{filename}'")
|
||||
|
||||
await asyncio.sleep(self._poll_interval)
|
||||
|
||||
def name(self) -> str:
|
||||
"""Get parser name."""
|
||||
return "iris"
|
||||
|
||||
|
||||
def _raise_for_status(response: httpx.Response, filename: str, step: str) -> None:
|
||||
"""
|
||||
Raise an appropriate error including the response body on HTTP errors.
|
||||
|
||||
Raises UnsupportedFileTypeError for 4xx responses (file rejected by the API),
|
||||
RuntimeError for other HTTP errors.
|
||||
"""
|
||||
if not response.is_error:
|
||||
return
|
||||
body = response.text or "<empty>"
|
||||
msg = f"Iris API error during {step} for '{filename}': {response.status_code} {response.reason_phrase} — {body}"
|
||||
if response.is_client_error:
|
||||
raise UnsupportedFileTypeError(msg)
|
||||
raise RuntimeError(msg)
|
||||
@@ -1,109 +0,0 @@
|
||||
"""Markitdown parser implementation."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from .base import FileParser
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MarkitdownParser(FileParser):
|
||||
"""
|
||||
Markitdown file parser.
|
||||
|
||||
Uses Microsoft's markitdown library to convert various file formats
|
||||
to markdown including PDF, Office docs, images (via OCR), audio, HTML.
|
||||
|
||||
Supported formats:
|
||||
- PDF (.pdf)
|
||||
- Word (.docx, .doc)
|
||||
- PowerPoint (.pptx, .ppt)
|
||||
- Excel (.xlsx, .xls)
|
||||
- Images (.jpg, .jpeg, .png) - with OCR
|
||||
- HTML (.html, .htm)
|
||||
- Text (.txt, .md)
|
||||
- Audio (.mp3, .wav) - with transcription
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize markitdown parser."""
|
||||
# Lazy import to avoid requiring markitdown for all users
|
||||
try:
|
||||
from markitdown import MarkItDown
|
||||
|
||||
self._markitdown = MarkItDown()
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"markitdown package is required for file parsing. Install with: pip install markitdown"
|
||||
) from e
|
||||
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
"""Parse file to markdown using markitdown."""
|
||||
# markitdown is synchronous, so we run it in executor to avoid blocking
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, self._convert_sync, file_data, filename)
|
||||
|
||||
def _convert_sync(self, file_data: bytes, filename: str) -> str:
|
||||
"""Synchronous parsing (runs in thread pool)."""
|
||||
# Write to temp file (markitdown requires file path)
|
||||
with tempfile.NamedTemporaryFile(suffix=Path(filename).suffix, delete=False) as tmp:
|
||||
tmp.write(file_data)
|
||||
tmp_path = tmp.name
|
||||
|
||||
try:
|
||||
# Parse using markitdown
|
||||
result = self._markitdown.convert(tmp_path)
|
||||
|
||||
if not result or not result.text_content:
|
||||
raise RuntimeError(f"No content extracted from '{filename}'")
|
||||
|
||||
return result.text_content
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Markitdown parsing failed for {filename}: {e}")
|
||||
raise RuntimeError(f"Failed to parse '{filename}': {e}") from e
|
||||
|
||||
finally:
|
||||
# Clean up temp file
|
||||
try:
|
||||
Path(tmp_path).unlink()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||
"""Check if markitdown supports this file type."""
|
||||
# Supported extensions (from markitdown docs)
|
||||
supported_extensions = {
|
||||
# Documents
|
||||
".pdf",
|
||||
".docx",
|
||||
".doc",
|
||||
".pptx",
|
||||
".ppt",
|
||||
".xlsx",
|
||||
".xls",
|
||||
# Images (with OCR)
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".png",
|
||||
# Web
|
||||
".html",
|
||||
".htm",
|
||||
# Text
|
||||
".txt",
|
||||
".md",
|
||||
".csv",
|
||||
# Audio (with transcription)
|
||||
".mp3",
|
||||
".wav",
|
||||
}
|
||||
|
||||
ext = Path(filename).suffix.lower()
|
||||
return ext in supported_extensions
|
||||
|
||||
def name(self) -> str:
|
||||
"""Get parser name."""
|
||||
return "markitdown"
|
||||
@@ -16,7 +16,6 @@ Features:
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
@@ -97,9 +96,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
if self.provider in ("openai", "groq") and not self.api_key:
|
||||
raise ValueError(f"API key is required for {self.provider}")
|
||||
|
||||
# Service tier configuration (from config, not env vars)
|
||||
self.groq_service_tier = groq_service_tier
|
||||
self.openai_service_tier = kwargs.get("openai_service_tier")
|
||||
# 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)))
|
||||
@@ -784,140 +782,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
raise last_exception
|
||||
raise RuntimeError("Ollama call failed after all retries")
|
||||
|
||||
async def supports_batch_api(self) -> bool:
|
||||
"""Check if this provider supports batch API operations."""
|
||||
# Only OpenAI and Groq support batch API
|
||||
return self.provider in ("openai", "groq")
|
||||
|
||||
async def submit_batch(
|
||||
self,
|
||||
requests: list[dict[str, Any]],
|
||||
endpoint: str = "/v1/chat/completions",
|
||||
completion_window: str = "24h",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Submit a batch of requests to OpenAI/Groq Batch API.
|
||||
|
||||
Args:
|
||||
requests: List of request dicts with custom_id, method, url, body
|
||||
endpoint: API endpoint (e.g., "/v1/chat/completions")
|
||||
completion_window: Completion window (e.g., "24h")
|
||||
|
||||
Returns:
|
||||
Dict with batch metadata including batch_id
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If provider doesn't support batch API
|
||||
"""
|
||||
if not await self.supports_batch_api():
|
||||
raise NotImplementedError(f"Batch API not supported for provider: {self.provider}")
|
||||
|
||||
logger.info(f"Submitting batch with {len(requests)} requests to {self.provider}")
|
||||
|
||||
# Format requests as JSONL
|
||||
jsonl_content = "\n".join(json.dumps(req) for req in requests)
|
||||
|
||||
# Upload file to provider (wrap in BytesIO with filename)
|
||||
file_bytes = io.BytesIO(jsonl_content.encode("utf-8"))
|
||||
file_bytes.name = "batch_input.jsonl" # OpenAI SDK needs a filename
|
||||
|
||||
file_response = await self._client.files.create(
|
||||
file=file_bytes,
|
||||
purpose="batch",
|
||||
)
|
||||
|
||||
logger.debug(f"Uploaded batch file: {file_response.id}")
|
||||
|
||||
# Create batch
|
||||
batch_response = await self._client.batches.create(
|
||||
input_file_id=file_response.id,
|
||||
endpoint=endpoint,
|
||||
completion_window=completion_window,
|
||||
)
|
||||
|
||||
logger.info(f"Batch submitted: {batch_response.id}, status={batch_response.status}")
|
||||
|
||||
return {
|
||||
"batch_id": batch_response.id,
|
||||
"status": batch_response.status,
|
||||
"input_file_id": file_response.id,
|
||||
"created_at": batch_response.created_at,
|
||||
"request_count": len(requests),
|
||||
}
|
||||
|
||||
async def get_batch_status(self, batch_id: str) -> dict[str, Any]:
|
||||
"""
|
||||
Get the status of a batch job.
|
||||
|
||||
Args:
|
||||
batch_id: Batch identifier
|
||||
|
||||
Returns:
|
||||
Dict with status info (batch_id, status, completed_at, etc.)
|
||||
"""
|
||||
if not await self.supports_batch_api():
|
||||
raise NotImplementedError(f"Batch API not supported for provider: {self.provider}")
|
||||
|
||||
batch = await self._client.batches.retrieve(batch_id)
|
||||
|
||||
result = {
|
||||
"batch_id": batch.id,
|
||||
"status": batch.status,
|
||||
"created_at": batch.created_at,
|
||||
"request_counts": {
|
||||
"total": batch.request_counts.total if batch.request_counts else 0,
|
||||
"completed": batch.request_counts.completed if batch.request_counts else 0,
|
||||
"failed": batch.request_counts.failed if batch.request_counts else 0,
|
||||
},
|
||||
}
|
||||
|
||||
if batch.completed_at:
|
||||
result["completed_at"] = batch.completed_at
|
||||
if batch.output_file_id:
|
||||
result["output_file_id"] = batch.output_file_id
|
||||
if batch.error_file_id:
|
||||
result["error_file_id"] = batch.error_file_id
|
||||
if batch.errors:
|
||||
result["errors"] = batch.errors
|
||||
|
||||
return result
|
||||
|
||||
async def retrieve_batch_results(self, batch_id: str) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Retrieve completed batch results.
|
||||
|
||||
Args:
|
||||
batch_id: Batch identifier
|
||||
|
||||
Returns:
|
||||
List of result dicts (one per request, matched by custom_id)
|
||||
"""
|
||||
if not await self.supports_batch_api():
|
||||
raise NotImplementedError(f"Batch API not supported for provider: {self.provider}")
|
||||
|
||||
# Get batch status
|
||||
batch = await self._client.batches.retrieve(batch_id)
|
||||
|
||||
if batch.status != "completed":
|
||||
raise ValueError(f"Batch {batch_id} is not completed yet (status: {batch.status})")
|
||||
|
||||
if not batch.output_file_id:
|
||||
raise ValueError(f"Batch {batch_id} has no output file")
|
||||
|
||||
# Download results file
|
||||
logger.debug(f"Downloading results for batch {batch_id} from file {batch.output_file_id}")
|
||||
file_content = await self._client.files.content(batch.output_file_id)
|
||||
|
||||
# Parse JSONL results
|
||||
results = []
|
||||
for line in file_content.text.strip().split("\n"):
|
||||
if line:
|
||||
results.append(json.loads(line))
|
||||
|
||||
logger.info(f"Retrieved {len(results)} results for batch {batch_id}")
|
||||
|
||||
return results
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (close OpenAI client connections)."""
|
||||
if hasattr(self, "_client") and self._client:
|
||||
|
||||
@@ -695,20 +695,44 @@ Example: "Lost job → couldn't pay rent → moved apartment"
|
||||
- Fact 2: Moved apartment, causal_relations: [{target_index: 1, relation_type: "caused_by"}]"""
|
||||
|
||||
|
||||
def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
async def _extract_facts_from_chunk(
|
||||
chunk: str,
|
||||
chunk_index: int,
|
||||
total_chunks: int,
|
||||
event_date: datetime,
|
||||
context: str,
|
||||
llm_config: "LLMConfig",
|
||||
agent_name: str = None,
|
||||
) -> tuple[list[dict[str, str]], TokenUsage]:
|
||||
"""
|
||||
Build extraction prompt and response schema based on config.
|
||||
Extract facts from a single chunk (internal helper for parallel processing).
|
||||
|
||||
Returns:
|
||||
Tuple of (prompt, response_schema)
|
||||
Note: event_date parameter is kept for backward compatibility but not used in prompt.
|
||||
The LLM extracts temporal information from the context string instead.
|
||||
"""
|
||||
import logging
|
||||
|
||||
from openai import BadRequestError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Determine which fact types to extract
|
||||
# Note: We use "assistant" in the prompt but convert to "bank" for storage
|
||||
fact_types_instruction = "Extract ONLY 'world' and 'assistant' type facts."
|
||||
|
||||
# Check config for extraction mode and causal link extraction
|
||||
config = get_config()
|
||||
extraction_mode = config.retain_extraction_mode
|
||||
extract_causal_links = config.retain_extract_causal_links
|
||||
|
||||
# Select base prompt based on extraction mode
|
||||
if extraction_mode == "custom":
|
||||
# Custom mode: inject user-provided guidelines
|
||||
if not config.retain_custom_instructions:
|
||||
logger.warning(
|
||||
"extraction_mode='custom' but HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS not set. "
|
||||
"Falling back to 'concise' mode."
|
||||
)
|
||||
base_prompt = CONCISE_FACT_EXTRACTION_PROMPT
|
||||
prompt = base_prompt.format(fact_types_instruction=fact_types_instruction)
|
||||
else:
|
||||
@@ -724,26 +748,33 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
base_prompt = CONCISE_FACT_EXTRACTION_PROMPT
|
||||
prompt = base_prompt.format(fact_types_instruction=fact_types_instruction)
|
||||
|
||||
# Add causal relationships section if enabled
|
||||
# Build the full prompt with or without causal relationships section
|
||||
# Select appropriate response schema based on extraction mode and causal links
|
||||
if extract_causal_links:
|
||||
prompt = prompt + CAUSAL_RELATIONSHIPS_SECTION
|
||||
response_schema = FactExtractionResponseVerbose if extraction_mode == "verbose" else FactExtractionResponse
|
||||
if extraction_mode == "verbose":
|
||||
response_schema = FactExtractionResponseVerbose
|
||||
else:
|
||||
response_schema = FactExtractionResponse
|
||||
else:
|
||||
response_schema = FactExtractionResponseNoCausal
|
||||
|
||||
return prompt, response_schema
|
||||
|
||||
|
||||
def _build_user_message(chunk: str, chunk_index: int, total_chunks: int, event_date: datetime, context: str) -> str:
|
||||
"""Build user message for fact extraction."""
|
||||
from .orchestrator import parse_datetime_flexible
|
||||
# Retry logic for JSON validation errors
|
||||
max_retries = 2
|
||||
last_error = None
|
||||
|
||||
# Sanitize input text to prevent Unicode encoding errors (e.g., unpaired surrogates)
|
||||
sanitized_chunk = _sanitize_text(chunk)
|
||||
sanitized_context = _sanitize_text(context) if context else "none"
|
||||
event_date = parse_datetime_flexible(event_date)
|
||||
event_date_formatted = event_date.strftime("%A, %B %d, %Y")
|
||||
|
||||
return f"""Extract facts from the following text 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.
|
||||
|
||||
Chunk: {chunk_index + 1}/{total_chunks}
|
||||
Event Date: {event_date_formatted} ({event_date.isoformat()})
|
||||
@@ -752,70 +783,6 @@ Context: {sanitized_context}
|
||||
Text:
|
||||
{sanitized_chunk}"""
|
||||
|
||||
|
||||
def _build_request_body(llm_config, config, prompt: str, user_message: str, response_schema: type) -> dict:
|
||||
"""Build request body for LLM API call."""
|
||||
request_body = {
|
||||
"model": llm_config.model,
|
||||
"messages": [{"role": "system", "content": prompt}, {"role": "user", "content": user_message}],
|
||||
"temperature": 0.1,
|
||||
}
|
||||
|
||||
# Add max_completion_tokens if configured
|
||||
if config.retain_max_completion_tokens:
|
||||
request_body["max_completion_tokens"] = config.retain_max_completion_tokens
|
||||
|
||||
# Add service_tier for OpenAI Flex Processing
|
||||
if llm_config.provider == "openai" and llm_config._provider_impl.openai_service_tier:
|
||||
request_body["service_tier"] = llm_config._provider_impl.openai_service_tier
|
||||
|
||||
# Add response_format (JSON schema)
|
||||
if hasattr(response_schema, "model_json_schema"):
|
||||
schema = response_schema.model_json_schema()
|
||||
request_body["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {"name": "facts", "schema": schema},
|
||||
}
|
||||
|
||||
return request_body
|
||||
|
||||
|
||||
async def _extract_facts_from_chunk(
|
||||
chunk: str,
|
||||
chunk_index: int,
|
||||
total_chunks: int,
|
||||
event_date: datetime,
|
||||
context: str,
|
||||
llm_config: "LLMConfig",
|
||||
config,
|
||||
agent_name: str = None,
|
||||
) -> tuple[list[dict[str, str]], TokenUsage]:
|
||||
"""
|
||||
Extract facts from a single chunk (internal helper for parallel processing).
|
||||
|
||||
Note: event_date parameter is kept for backward compatibility but not used in prompt.
|
||||
The LLM extracts temporal information from the context string instead.
|
||||
"""
|
||||
import logging
|
||||
|
||||
from openai import BadRequestError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Build prompt and schema using helper function
|
||||
prompt, response_schema = _build_extraction_prompt_and_schema(config)
|
||||
|
||||
# Check config for extraction mode and causal link extraction
|
||||
extraction_mode = config.retain_extraction_mode
|
||||
extract_causal_links = config.retain_extract_causal_links
|
||||
|
||||
# Build user message using helper function
|
||||
user_message = _build_user_message(chunk, chunk_index, total_chunks, event_date, context)
|
||||
|
||||
# Retry logic for JSON validation errors
|
||||
max_retries = 2
|
||||
last_error = None
|
||||
|
||||
usage = TokenUsage() # Track cumulative usage across retries
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
@@ -1088,7 +1055,6 @@ async def _extract_facts_with_auto_split(
|
||||
event_date: datetime,
|
||||
context: str,
|
||||
llm_config: LLMConfig,
|
||||
config,
|
||||
agent_name: str = None,
|
||||
) -> tuple[list[dict[str, str]], TokenUsage]:
|
||||
"""
|
||||
@@ -1104,7 +1070,6 @@ async def _extract_facts_with_auto_split(
|
||||
event_date: Reference date for temporal information
|
||||
context: Context about the conversation/document
|
||||
llm_config: LLM configuration to use
|
||||
config: Resolved HindsightConfig for this bank
|
||||
agent_name: Optional agent name (memory owner)
|
||||
|
||||
Returns:
|
||||
@@ -1123,7 +1088,6 @@ async def _extract_facts_with_auto_split(
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
config=config,
|
||||
agent_name=agent_name,
|
||||
)
|
||||
except OutputTooLongError:
|
||||
@@ -1168,7 +1132,6 @@ async def _extract_facts_with_auto_split(
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
config=config,
|
||||
agent_name=agent_name,
|
||||
),
|
||||
_extract_facts_with_auto_split(
|
||||
@@ -1178,7 +1141,6 @@ async def _extract_facts_with_auto_split(
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
config=config,
|
||||
agent_name=agent_name,
|
||||
),
|
||||
]
|
||||
@@ -1202,7 +1164,6 @@ async def extract_facts_from_text(
|
||||
event_date: datetime,
|
||||
llm_config: LLMConfig,
|
||||
agent_name: str,
|
||||
config,
|
||||
context: str = "",
|
||||
) -> tuple[list[Fact], list[tuple[str, int]], TokenUsage]:
|
||||
"""
|
||||
@@ -1217,10 +1178,9 @@ async def extract_facts_from_text(
|
||||
Args:
|
||||
text: Input text (conversation, article, etc.)
|
||||
event_date: Reference date for resolving relative times
|
||||
context: Context about the conversation/document
|
||||
llm_config: LLM configuration to use
|
||||
agent_name: Agent name (memory owner)
|
||||
config: Resolved HindsightConfig for this bank
|
||||
context: Context about the conversation/document
|
||||
|
||||
Returns:
|
||||
Tuple of (facts, chunks, usage) where:
|
||||
@@ -1228,6 +1188,7 @@ async def extract_facts_from_text(
|
||||
- chunks: List of tuples (chunk_text, fact_count) for each chunk
|
||||
- usage: Aggregated token usage across all LLM calls
|
||||
"""
|
||||
config = get_config()
|
||||
chunks = chunk_text(text, max_chars=config.retain_chunk_size)
|
||||
|
||||
# Log chunk count before starting LLM requests
|
||||
@@ -1246,7 +1207,6 @@ async def extract_facts_from_text(
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
config=config,
|
||||
agent_name=agent_name,
|
||||
)
|
||||
for i, chunk in enumerate(chunks)
|
||||
@@ -1278,420 +1238,8 @@ logger = logging.getLogger(__name__)
|
||||
SECONDS_PER_FACT = 10
|
||||
|
||||
|
||||
async def extract_facts_from_contents_batch_api(
|
||||
contents: list[RetainContent],
|
||||
llm_config,
|
||||
agent_name: str,
|
||||
config,
|
||||
pool=None,
|
||||
operation_id: str | None = None,
|
||||
schema: str | None = None,
|
||||
) -> tuple[list[ExtractedFactType], list[ChunkMetadata], TokenUsage]:
|
||||
"""
|
||||
Extract facts using LLM Batch API (OpenAI/Groq).
|
||||
|
||||
Submits all chunks as a single batch, polls until complete, then processes results.
|
||||
Only called when config.retain_batch_enabled=True.
|
||||
|
||||
Args:
|
||||
contents: List of RetainContent objects to process
|
||||
llm_config: LLM configuration with batch API support
|
||||
agent_name: Name of the agent
|
||||
config: Resolved HindsightConfig for this bank
|
||||
pool: Database connection pool (for storing batch state)
|
||||
operation_id: Async operation ID (for crash recovery)
|
||||
schema: Database schema (for multi-tenant support)
|
||||
|
||||
Returns:
|
||||
Tuple of (extracted_facts, chunks_metadata, usage)
|
||||
"""
|
||||
if not contents:
|
||||
return [], [], TokenUsage()
|
||||
|
||||
logger.info(f"Using Batch API for fact extraction ({len(contents)} contents)")
|
||||
|
||||
# Check config for extraction mode and causal link extraction (used throughout)
|
||||
extraction_mode = config.retain_extraction_mode
|
||||
extract_causal_links = config.retain_extract_causal_links
|
||||
|
||||
# Check if provider supports batch API
|
||||
if not await llm_config._provider_impl.supports_batch_api():
|
||||
logger.warning(f"Batch API not supported for provider {llm_config.provider}, falling back to sync mode")
|
||||
return await extract_facts_from_contents(contents, llm_config, agent_name, config, pool, operation_id, schema)
|
||||
|
||||
# Check if we're resuming an existing batch (crash recovery)
|
||||
batch_id = None
|
||||
if operation_id and pool:
|
||||
from ..task_backend import fq_table
|
||||
|
||||
table = fq_table("async_operations", schema)
|
||||
row = await pool.fetchrow(
|
||||
f"SELECT result_metadata FROM {table} WHERE operation_id = $1",
|
||||
operation_id,
|
||||
)
|
||||
|
||||
if row and row["result_metadata"]:
|
||||
metadata = row["result_metadata"]
|
||||
if isinstance(metadata, str):
|
||||
metadata = json.loads(metadata)
|
||||
batch_id = metadata.get("batch_id")
|
||||
|
||||
if batch_id:
|
||||
logger.info(f"Resuming existing batch: batch_id={batch_id} (crash recovery)")
|
||||
|
||||
# Step 1: Chunk all contents and build batch requests (skip if resuming)
|
||||
all_chunks_info = [] # List of (chunk_text, content_index, chunk_index_in_content, event_date, context)
|
||||
batch_requests = []
|
||||
|
||||
# Build prompt and schema once (same for all chunks)
|
||||
prompt, response_schema = _build_extraction_prompt_and_schema(config)
|
||||
|
||||
for content_index, item in enumerate(contents):
|
||||
chunks = chunk_text(item.content, max_chars=config.retain_chunk_size)
|
||||
|
||||
for chunk_index_in_content, chunk in enumerate(chunks):
|
||||
all_chunks_info.append((chunk, content_index, chunk_index_in_content, item.event_date, item.context))
|
||||
|
||||
# Build batch request for this chunk
|
||||
custom_id = f"chunk_{len(all_chunks_info) - 1}" # Global chunk index
|
||||
|
||||
# Build user message using helper function
|
||||
user_message = _build_user_message(
|
||||
chunk, chunk_index_in_content, len(chunks), item.event_date, item.context
|
||||
)
|
||||
|
||||
# Build request body using helper function
|
||||
request_body = _build_request_body(llm_config, config, prompt, user_message, response_schema)
|
||||
|
||||
batch_requests.append(
|
||||
{"custom_id": custom_id, "method": "POST", "url": "/v1/chat/completions", "body": request_body}
|
||||
)
|
||||
|
||||
if not batch_requests and not batch_id: # No requests and not resuming
|
||||
return [], [], TokenUsage()
|
||||
|
||||
# Step 2: Submit batch (skip if resuming)
|
||||
if not batch_id:
|
||||
logger.info(f"Submitting batch with {len(batch_requests)} chunk requests")
|
||||
|
||||
batch_metadata = await llm_config._provider_impl.submit_batch(batch_requests)
|
||||
batch_id = batch_metadata["batch_id"]
|
||||
|
||||
logger.info(f"Batch submitted: {batch_id}, polling every {config.retain_batch_poll_interval_seconds}s")
|
||||
|
||||
# CRITICAL: Store minimal batch state in operation metadata for crash recovery
|
||||
# This allows resuming polling if worker restarts
|
||||
if operation_id and pool:
|
||||
batch_state = {
|
||||
"batch_id": batch_id,
|
||||
"batch_provider": llm_config.provider,
|
||||
"chunk_count": len(batch_requests),
|
||||
}
|
||||
|
||||
# Update operation result_metadata
|
||||
from ..task_backend import fq_table
|
||||
|
||||
table = fq_table("async_operations", schema)
|
||||
await pool.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET result_metadata = result_metadata || $1::jsonb, updated_at = now()
|
||||
WHERE operation_id = $2
|
||||
""",
|
||||
json.dumps(batch_state),
|
||||
operation_id,
|
||||
)
|
||||
logger.info(f"Stored batch state for operation {operation_id} (crash recovery enabled)")
|
||||
else:
|
||||
logger.info(f"Resuming polling for existing batch: {batch_id}")
|
||||
|
||||
# Step 3: Poll until complete
|
||||
import time
|
||||
|
||||
start_time = time.time()
|
||||
while True:
|
||||
status_info = await llm_config._provider_impl.get_batch_status(batch_id)
|
||||
status = status_info["status"]
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
logger.info(
|
||||
f"Batch {batch_id}: status={status}, "
|
||||
f"completed={status_info['request_counts']['completed']}/{status_info['request_counts']['total']}, "
|
||||
f"elapsed={elapsed:.0f}s"
|
||||
)
|
||||
|
||||
if status == "completed":
|
||||
break
|
||||
elif status in ("failed", "expired", "cancelled"):
|
||||
error_msg = status_info.get("errors", "Unknown error")
|
||||
raise RuntimeError(f"Batch {batch_id} failed with status {status}: {error_msg}")
|
||||
|
||||
# Wait before polling again
|
||||
await asyncio.sleep(config.retain_batch_poll_interval_seconds)
|
||||
|
||||
logger.info(f"Batch {batch_id} completed in {elapsed:.0f}s, retrieving results")
|
||||
|
||||
# Step 4: Retrieve results
|
||||
batch_results = await llm_config._provider_impl.retrieve_batch_results(batch_id)
|
||||
|
||||
# Map results by custom_id
|
||||
results_by_id = {result["custom_id"]: result for result in batch_results}
|
||||
|
||||
# Step 5: Parse results into facts (same as sync mode)
|
||||
all_facts_from_llm = []
|
||||
chunks_metadata = []
|
||||
total_usage = TokenUsage()
|
||||
|
||||
for chunk_idx, (chunk_content, content_index, chunk_index_in_content, event_date, context) in enumerate(
|
||||
all_chunks_info
|
||||
):
|
||||
custom_id = f"chunk_{chunk_idx}"
|
||||
result = results_by_id.get(custom_id)
|
||||
|
||||
if not result:
|
||||
logger.warning(f"Missing result for {custom_id}, skipping")
|
||||
chunks_metadata.append(
|
||||
ChunkMetadata(
|
||||
chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
# Check for errors
|
||||
if result.get("error"):
|
||||
logger.error(f"Error in {custom_id}: {result['error']}")
|
||||
chunks_metadata.append(
|
||||
ChunkMetadata(
|
||||
chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
# Extract response
|
||||
response_body = result.get("response", {}).get("body", {})
|
||||
choices = response_body.get("choices", [])
|
||||
|
||||
if not choices:
|
||||
logger.warning(f"No choices in response for {custom_id}")
|
||||
chunks_metadata.append(
|
||||
ChunkMetadata(
|
||||
chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
# Parse JSON content
|
||||
message = choices[0].get("message", {})
|
||||
content_str = message.get("content", "{}")
|
||||
|
||||
try:
|
||||
extraction_response_json = json.loads(content_str)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"Failed to parse JSON for {custom_id}: {e}")
|
||||
chunks_metadata.append(
|
||||
ChunkMetadata(
|
||||
chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
# Parse facts (reuse existing logic from _extract_facts_from_chunk)
|
||||
raw_facts = extraction_response_json.get("facts", [])
|
||||
chunk_facts = []
|
||||
|
||||
for i, llm_fact in enumerate(raw_facts):
|
||||
if not isinstance(llm_fact, dict):
|
||||
continue
|
||||
|
||||
def get_value(field_name):
|
||||
value = llm_fact.get(field_name)
|
||||
if value and value != "" and value != [] and value != {} and str(value).upper() != "N/A":
|
||||
return value
|
||||
return None
|
||||
|
||||
what = get_value("what")
|
||||
if not what:
|
||||
what = get_value("factual_core")
|
||||
if not what:
|
||||
continue
|
||||
|
||||
when = get_value("when")
|
||||
who = get_value("who")
|
||||
why = get_value("why")
|
||||
|
||||
# Critical field: fact_type
|
||||
original_fact_type = llm_fact.get("fact_type")
|
||||
fact_type = original_fact_type
|
||||
|
||||
# Convert "assistant" → "experience"
|
||||
if fact_type == "assistant":
|
||||
fact_type = "experience"
|
||||
|
||||
# Validate fact_type
|
||||
if fact_type not in ["world", "experience", "opinion"]:
|
||||
fact_kind = llm_fact.get("fact_kind")
|
||||
if fact_kind == "assistant":
|
||||
fact_type = "experience"
|
||||
elif fact_kind in ["world", "experience", "opinion"]:
|
||||
fact_type = fact_kind
|
||||
else:
|
||||
fact_type = "world"
|
||||
|
||||
# Build combined fact text
|
||||
combined_parts = [what]
|
||||
if when:
|
||||
combined_parts.append(f"When: {when}")
|
||||
if who:
|
||||
combined_parts.append(f"Involving: {who}")
|
||||
if why:
|
||||
combined_parts.append(why)
|
||||
combined_text = " | ".join(combined_parts)
|
||||
|
||||
# Temporal fields
|
||||
fact_data = {}
|
||||
fact_kind = llm_fact.get("fact_kind", "conversation")
|
||||
if fact_kind not in ["conversation", "event", "other"]:
|
||||
fact_kind = "conversation"
|
||||
|
||||
if fact_kind == "event":
|
||||
occurred_start = get_value("occurred_start")
|
||||
occurred_end = get_value("occurred_end")
|
||||
|
||||
if not occurred_start:
|
||||
fact_data["occurred_start"] = _infer_temporal_date(combined_text, event_date)
|
||||
else:
|
||||
fact_data["occurred_start"] = occurred_start
|
||||
|
||||
if occurred_end:
|
||||
fact_data["occurred_end"] = occurred_end
|
||||
elif fact_data.get("occurred_start"):
|
||||
fact_data["occurred_end"] = fact_data["occurred_start"]
|
||||
|
||||
# Entities
|
||||
entities = get_value("entities")
|
||||
if entities:
|
||||
validated_entities = []
|
||||
for ent in entities:
|
||||
if isinstance(ent, str):
|
||||
validated_entities.append(Entity(text=ent))
|
||||
elif isinstance(ent, dict) and "text" in ent:
|
||||
try:
|
||||
validated_entities.append(Entity.model_validate(ent))
|
||||
except Exception:
|
||||
pass
|
||||
if validated_entities:
|
||||
fact_data["entities"] = validated_entities
|
||||
|
||||
# Causal relations
|
||||
if extract_causal_links:
|
||||
validated_relations = []
|
||||
causal_relations_raw = get_value("causal_relations")
|
||||
if causal_relations_raw:
|
||||
for rel in causal_relations_raw:
|
||||
if not isinstance(rel, dict):
|
||||
continue
|
||||
target_idx = rel.get("target_index")
|
||||
relation_type = rel.get("relation_type")
|
||||
strength = rel.get("strength", 1.0)
|
||||
|
||||
if target_idx is None or relation_type is None:
|
||||
continue
|
||||
if target_idx < 0 or target_idx >= i:
|
||||
continue
|
||||
|
||||
try:
|
||||
validated_relations.append(
|
||||
CausalRelation(
|
||||
target_fact_index=target_idx, relation_type=relation_type, strength=strength
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if validated_relations:
|
||||
fact_data["causal_relations"] = validated_relations
|
||||
|
||||
# Always set mentioned_at
|
||||
fact_data["mentioned_at"] = event_date.isoformat()
|
||||
|
||||
try:
|
||||
fact = Fact(fact=combined_text, fact_type=fact_type, **fact_data)
|
||||
chunk_facts.append(fact)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create Fact model for fact {i}: {e}")
|
||||
continue
|
||||
|
||||
all_facts_from_llm.extend(chunk_facts)
|
||||
chunks_metadata.append(
|
||||
ChunkMetadata(
|
||||
chunk_text=chunk_content,
|
||||
fact_count=len(chunk_facts),
|
||||
content_index=content_index,
|
||||
chunk_index=chunk_idx,
|
||||
)
|
||||
)
|
||||
|
||||
# Track token usage
|
||||
usage_data = response_body.get("usage", {})
|
||||
if usage_data:
|
||||
total_usage = total_usage + TokenUsage(
|
||||
input_tokens=usage_data.get("prompt_tokens", 0),
|
||||
output_tokens=usage_data.get("completion_tokens", 0),
|
||||
total_tokens=usage_data.get("total_tokens", 0),
|
||||
)
|
||||
|
||||
# Step 6: Convert to ExtractedFact objects with proper chunk mapping
|
||||
# Group facts by chunk
|
||||
facts_by_chunk = [] # List of (chunk_metadata, [facts])
|
||||
fact_start_idx = 0
|
||||
|
||||
for chunk_meta in chunks_metadata:
|
||||
chunk_facts = all_facts_from_llm[fact_start_idx : fact_start_idx + chunk_meta.fact_count]
|
||||
facts_by_chunk.append((chunk_meta, chunk_facts))
|
||||
fact_start_idx += chunk_meta.fact_count
|
||||
|
||||
# Now convert to ExtractedFactType
|
||||
extracted_facts = []
|
||||
global_fact_idx = 0
|
||||
|
||||
for chunk_meta, chunk_facts in facts_by_chunk:
|
||||
content = contents[chunk_meta.content_index]
|
||||
|
||||
for fact_from_llm in chunk_facts:
|
||||
extracted_fact = ExtractedFactType(
|
||||
fact_text=fact_from_llm.fact,
|
||||
fact_type=fact_from_llm.fact_type,
|
||||
entities=[e.text for e in (fact_from_llm.entities or [])],
|
||||
occurred_start=_parse_datetime(fact_from_llm.occurred_start) if fact_from_llm.occurred_start else None,
|
||||
occurred_end=_parse_datetime(fact_from_llm.occurred_end) if fact_from_llm.occurred_end else None,
|
||||
causal_relations=_convert_causal_relations(fact_from_llm.causal_relations or [], global_fact_idx),
|
||||
content_index=chunk_meta.content_index,
|
||||
chunk_index=chunk_meta.chunk_index,
|
||||
context=content.context,
|
||||
mentioned_at=content.event_date,
|
||||
metadata=content.metadata,
|
||||
tags=content.tags,
|
||||
)
|
||||
|
||||
extracted_facts.append(extracted_fact)
|
||||
global_fact_idx += 1
|
||||
|
||||
# Step 7: Add temporal offsets
|
||||
_add_temporal_offsets(extracted_facts, contents)
|
||||
|
||||
logger.info(f"Batch API extracted {len(extracted_facts)} facts from {len(all_chunks_info)} chunks")
|
||||
|
||||
return extracted_facts, chunks_metadata, total_usage
|
||||
|
||||
|
||||
async def extract_facts_from_contents(
|
||||
contents: list[RetainContent],
|
||||
llm_config,
|
||||
agent_name: str,
|
||||
config,
|
||||
pool=None,
|
||||
operation_id: str | None = None,
|
||||
schema: str | None = None,
|
||||
contents: list[RetainContent], llm_config, agent_name: str
|
||||
) -> tuple[list[ExtractedFactType], list[ChunkMetadata], TokenUsage]:
|
||||
"""
|
||||
Extract facts from multiple content items in parallel.
|
||||
@@ -1702,16 +1250,10 @@ async def extract_facts_from_contents(
|
||||
3. Adds time offsets to preserve fact ordering within each content
|
||||
4. Returns typed ExtractedFact and ChunkMetadata objects
|
||||
|
||||
Routes to batch API mode if config.retain_batch_enabled=True.
|
||||
|
||||
Args:
|
||||
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)
|
||||
config: Resolved HindsightConfig for this bank
|
||||
pool: Database connection pool (passed to batch API for state storage)
|
||||
operation_id: Async operation ID (passed to batch API for crash recovery)
|
||||
schema: Database schema (passed to batch API for multi-tenant support)
|
||||
|
||||
Returns:
|
||||
Tuple of (extracted_facts, chunks_metadata, usage)
|
||||
@@ -1719,12 +1261,6 @@ async def extract_facts_from_contents(
|
||||
if not contents:
|
||||
return [], [], TokenUsage()
|
||||
|
||||
# Route to batch API if enabled
|
||||
if config.retain_batch_enabled:
|
||||
return await extract_facts_from_contents_batch_api(
|
||||
contents, llm_config, agent_name, config, pool, operation_id, schema
|
||||
)
|
||||
|
||||
# Step 1: Create parallel fact extraction tasks
|
||||
fact_extraction_tasks = []
|
||||
for item in contents:
|
||||
@@ -1736,7 +1272,6 @@ async def extract_facts_from_contents(
|
||||
context=item.context,
|
||||
llm_config=llm_config,
|
||||
agent_name=agent_name,
|
||||
config=config,
|
||||
)
|
||||
fact_extraction_tasks.append(task)
|
||||
|
||||
|
||||
@@ -7,7 +7,6 @@ Handles insertion of facts into the database.
|
||||
import json
|
||||
import logging
|
||||
|
||||
from ...config import get_config
|
||||
from ..memory_engine import fq_table
|
||||
from .fact_extraction import _sanitize_text
|
||||
from .types import ProcessedFact
|
||||
@@ -71,59 +70,28 @@ async def insert_facts_batch(
|
||||
|
||||
# Batch insert all facts
|
||||
# Note: tags are passed as JSON strings and converted back to varchar[] via jsonb_array_elements_text + array_agg
|
||||
# Query varies based on text search backend
|
||||
config = get_config()
|
||||
if config.text_search_extension == "vchord":
|
||||
# VectorChord: manually tokenize and insert search_vector
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags_json)
|
||||
)
|
||||
INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags, search_vector)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
tokenize(COALESCE(text, '') || ' ' || COALESCE(context, ''), 'llmlingua2')::bm25_catalog.bm25vector
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
else: # native or pg_textsearch
|
||||
# Native PostgreSQL: search_vector is GENERATED ALWAYS, don't include it
|
||||
# pg_textsearch: indexes operate on base columns directly, don't populate search_vector
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags_json)
|
||||
)
|
||||
INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
)
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
|
||||
results = await conn.fetch(
|
||||
query,
|
||||
f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags_json)
|
||||
)
|
||||
INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
)
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
""",
|
||||
bank_id,
|
||||
fact_texts,
|
||||
embeddings,
|
||||
|
||||
@@ -76,14 +76,11 @@ async def retain_batch(
|
||||
duplicate_checker_fn,
|
||||
bank_id: str,
|
||||
contents_dicts: list[RetainContentDict],
|
||||
config,
|
||||
document_id: str | None = None,
|
||||
is_first_batch: bool = True,
|
||||
fact_type_override: str | None = None,
|
||||
confidence_score: float | None = None,
|
||||
document_tags: list[str] | None = None,
|
||||
operation_id: str | None = None,
|
||||
schema: str | None = None,
|
||||
) -> tuple[list[list[str]], TokenUsage]:
|
||||
"""
|
||||
Process a batch of content through the retain pipeline.
|
||||
@@ -97,7 +94,6 @@ async def retain_batch(
|
||||
duplicate_checker_fn: Function to check for duplicate facts
|
||||
bank_id: Bank identifier
|
||||
contents_dicts: List of content dictionaries
|
||||
config: Resolved HindsightConfig for this bank
|
||||
document_id: Optional document ID
|
||||
is_first_batch: Whether this is the first batch
|
||||
fact_type_override: Override fact type for all facts
|
||||
@@ -148,9 +144,7 @@ async def retain_batch(
|
||||
# Step 1: Extract facts from all contents
|
||||
step_start = time.time()
|
||||
|
||||
extracted_facts, chunks, usage = await fact_extraction.extract_facts_from_contents(
|
||||
contents, llm_config, agent_name, config, pool, operation_id, schema
|
||||
)
|
||||
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"
|
||||
)
|
||||
|
||||
@@ -13,10 +13,12 @@ from .reranking import CrossEncoderReranker
|
||||
from .retrieval import (
|
||||
ParallelRetrievalResult,
|
||||
get_default_graph_retriever,
|
||||
retrieve_parallel,
|
||||
set_default_graph_retriever,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"retrieve_parallel",
|
||||
"get_default_graph_retriever",
|
||||
"set_default_graph_retriever",
|
||||
"ParallelRetrievalResult",
|
||||
|
||||
@@ -85,6 +85,116 @@ def set_default_graph_retriever(retriever: GraphRetriever) -> None:
|
||||
_default_graph_retriever = retriever
|
||||
|
||||
|
||||
async def retrieve_semantic(
|
||||
conn,
|
||||
query_emb_str: str,
|
||||
bank_id: str,
|
||||
fact_type: str,
|
||||
limit: int,
|
||||
tags: list[str] | None = None,
|
||||
) -> list[RetrievalResult]:
|
||||
"""
|
||||
Semantic retrieval via vector similarity.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
query_emb_str: Query embedding as string
|
||||
agent_id: bank ID
|
||||
fact_type: Fact type to filter
|
||||
limit: Maximum results to return
|
||||
tags: Optional list of tags for visibility filtering (OR matching)
|
||||
|
||||
Returns:
|
||||
List of RetrievalResult objects
|
||||
"""
|
||||
from .tags import TagsMatch, build_tags_where_clause_simple
|
||||
|
||||
tags_clause = build_tags_where_clause_simple(tags, 5)
|
||||
params = [query_emb_str, bank_id, fact_type, limit]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
|
||||
results = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
|
||||
1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND fact_type = $3
|
||||
AND (1 - (embedding <=> $1::vector)) >= 0.3
|
||||
{tags_clause}
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $4
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
return [RetrievalResult.from_db_row(dict(r)) for r in results]
|
||||
|
||||
|
||||
async def retrieve_bm25(
|
||||
conn,
|
||||
query_text: str,
|
||||
bank_id: str,
|
||||
fact_type: str,
|
||||
limit: int,
|
||||
tags: list[str] | None = None,
|
||||
) -> list[RetrievalResult]:
|
||||
"""
|
||||
BM25 keyword retrieval via full-text search.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
query_text: Query text
|
||||
agent_id: bank ID
|
||||
fact_type: Fact type to filter
|
||||
limit: Maximum results to return
|
||||
tags: Optional list of tags for visibility filtering (OR matching)
|
||||
|
||||
Returns:
|
||||
List of RetrievalResult objects
|
||||
"""
|
||||
import re
|
||||
|
||||
from .tags import TagsMatch, build_tags_where_clause_simple
|
||||
|
||||
# Sanitize query text: remove special characters that have meaning in tsquery
|
||||
# Keep only alphanumeric characters and spaces
|
||||
sanitized_text = re.sub(r"[^\w\s]", " ", query_text.lower())
|
||||
|
||||
# Split and filter empty strings
|
||||
tokens = [token for token in sanitized_text.split() if token]
|
||||
|
||||
if not tokens:
|
||||
# If no valid tokens, return empty results
|
||||
return []
|
||||
|
||||
# Convert query to tsquery using OR for more flexible matching
|
||||
# This prevents empty results when some terms are missing
|
||||
query_tsquery = " | ".join(tokens)
|
||||
|
||||
tags_clause = build_tags_where_clause_simple(tags, 5)
|
||||
params = [query_tsquery, bank_id, fact_type, limit]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
|
||||
results = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
|
||||
ts_rank_cd(search_vector, to_tsquery('english', $1)) AS bm25_score
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = $3
|
||||
AND search_vector @@ to_tsquery('english', $1)
|
||||
{tags_clause}
|
||||
ORDER BY bm25_score DESC
|
||||
LIMIT $4
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
return [RetrievalResult.from_db_row(dict(r)) for r in results]
|
||||
|
||||
|
||||
async def retrieve_semantic_bm25_combined(
|
||||
conn,
|
||||
query_emb_str: str,
|
||||
@@ -158,41 +268,18 @@ async def retrieve_semantic_bm25_combined(
|
||||
result_dict[ft][0].append(RetrievalResult.from_db_row(row))
|
||||
return result_dict
|
||||
|
||||
# Build BM25 query based on text search backend
|
||||
config = get_config()
|
||||
query_tsquery = " | ".join(tokens)
|
||||
|
||||
# Build tags clause - param 6 if tags provided
|
||||
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
|
||||
|
||||
# Build backend-specific BM25 parts
|
||||
if config.text_search_extension == "vchord":
|
||||
# VectorChord BM25: use <&> operator with to_bm25query and tokenize
|
||||
# Note: VectorChord scores are negative (higher = better, so -1 > -10)
|
||||
bm25_score_expr = "search_vector <&> to_bm25query('idx_memory_units_text_search', tokenize($5, 'llmlingua2'))"
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = "" # No additional WHERE filter for vchord
|
||||
params = [query_emb_str, bank_id, fact_types, limit, query_text] # Pass raw query_text for tokenization
|
||||
elif config.text_search_extension == "pg_textsearch":
|
||||
# Timescale pg_textsearch: use <@> operator with to_bm25query
|
||||
# Note: pg_textsearch scores are negative (lower/more negative = better, so -10 > -1)
|
||||
# We negate the score to maintain API consistency (higher = better)
|
||||
bm25_score_expr = "-(text <@> to_bm25query($5, 'idx_memory_units_text_search'))"
|
||||
bm25_order_by = "text <@> to_bm25query($5, 'idx_memory_units_text_search') ASC"
|
||||
bm25_where_filter = "" # No additional WHERE filter for pg_textsearch
|
||||
params = [query_emb_str, bank_id, fact_types, limit, query_text]
|
||||
else: # native
|
||||
# Native PostgreSQL: use ts_rank_cd with to_tsquery
|
||||
query_tsquery = " | ".join(tokens)
|
||||
bm25_score_expr = "ts_rank_cd(search_vector, to_tsquery('english', $5))"
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = "AND search_vector @@ to_tsquery('english', $5)"
|
||||
params = [query_emb_str, bank_id, fact_types, limit, query_tsquery]
|
||||
|
||||
params = [query_emb_str, bank_id, fact_types, limit, query_tsquery]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
|
||||
# Single query template with backend-specific parts injected
|
||||
query = f"""
|
||||
# Combined CTE query for both semantic and BM25 across all fact types
|
||||
# Uses window functions to limit per fact_type per method
|
||||
results = await conn.fetch(
|
||||
f"""
|
||||
WITH semantic_ranked AS (
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
|
||||
1 - (embedding <=> $1::vector) AS similarity,
|
||||
@@ -209,13 +296,13 @@ async def retrieve_semantic_bm25_combined(
|
||||
bm25_ranked AS (
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
|
||||
NULL::float AS similarity,
|
||||
{bm25_score_expr} AS bm25_score,
|
||||
ts_rank_cd(search_vector, to_tsquery('english', $5)) AS bm25_score,
|
||||
'bm25' AS source,
|
||||
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY {bm25_order_by}) AS rn
|
||||
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY ts_rank_cd(search_vector, to_tsquery('english', $5)) DESC) AS rn
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = ANY($3)
|
||||
{bm25_where_filter}
|
||||
AND search_vector @@ to_tsquery('english', $5)
|
||||
{tags_clause}
|
||||
),
|
||||
semantic AS (
|
||||
@@ -231,11 +318,9 @@ async def retrieve_semantic_bm25_combined(
|
||||
SELECT * FROM semantic
|
||||
UNION ALL
|
||||
SELECT * FROM bm25
|
||||
"""
|
||||
|
||||
# Combined CTE query for both semantic and BM25 across all fact types
|
||||
# Uses window functions to limit per fact_type per method
|
||||
results = await conn.fetch(query, *params)
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
# Group results by fact_type and source
|
||||
result_dict: dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]] = {ft: ([], []) for ft in fact_types}
|
||||
@@ -476,6 +561,623 @@ async def retrieve_temporal_combined(
|
||||
return results_by_ft
|
||||
|
||||
|
||||
async def retrieve_temporal(
|
||||
conn,
|
||||
query_emb_str: str,
|
||||
bank_id: str,
|
||||
fact_type: str,
|
||||
start_date: datetime,
|
||||
end_date: datetime,
|
||||
budget: int,
|
||||
semantic_threshold: float = 0.1,
|
||||
tags: list[str] | None = None,
|
||||
) -> list[RetrievalResult]:
|
||||
"""
|
||||
Temporal retrieval with spreading activation.
|
||||
|
||||
Strategy:
|
||||
1. Find entry points (facts in date range with semantic relevance)
|
||||
2. Spread through temporal links to related facts
|
||||
3. Score by temporal proximity + semantic similarity + link weight
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
query_emb_str: Query embedding as string
|
||||
agent_id: bank ID
|
||||
fact_type: Fact type to filter
|
||||
start_date: Start of time range
|
||||
end_date: End of time range
|
||||
budget: Node budget for spreading
|
||||
semantic_threshold: Minimum semantic similarity to include
|
||||
tags: Optional list of tags for visibility filtering (OR matching)
|
||||
|
||||
Returns:
|
||||
List of RetrievalResult objects with temporal scores
|
||||
"""
|
||||
|
||||
# Ensure start_date and end_date are timezone-aware (UTC) to match database datetimes
|
||||
if start_date.tzinfo is None:
|
||||
start_date = start_date.replace(tzinfo=UTC)
|
||||
if end_date.tzinfo is None:
|
||||
end_date = end_date.replace(tzinfo=UTC)
|
||||
|
||||
from .tags import TagsMatch, build_tags_where_clause_simple
|
||||
|
||||
tags_clause = build_tags_where_clause_simple(tags, 7)
|
||||
params = [query_emb_str, bank_id, fact_type, start_date, end_date, semantic_threshold]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
|
||||
entry_points = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
|
||||
1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = $3
|
||||
AND embedding IS NOT NULL
|
||||
AND (
|
||||
-- Match if occurred range overlaps with query range
|
||||
(occurred_start IS NOT NULL AND occurred_end IS NOT NULL
|
||||
AND occurred_start <= $5 AND occurred_end >= $4)
|
||||
OR
|
||||
-- Match if mentioned_at falls within query range
|
||||
(mentioned_at IS NOT NULL AND mentioned_at BETWEEN $4 AND $5)
|
||||
OR
|
||||
-- Match if any occurred date is set and overlaps (even if only start or end is set)
|
||||
(occurred_start IS NOT NULL AND occurred_start BETWEEN $4 AND $5)
|
||||
OR
|
||||
(occurred_end IS NOT NULL AND occurred_end BETWEEN $4 AND $5)
|
||||
)
|
||||
AND (1 - (embedding <=> $1::vector)) >= $6
|
||||
{tags_clause}
|
||||
ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC, (embedding <=> $1::vector) ASC
|
||||
LIMIT 10
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
if not entry_points:
|
||||
return []
|
||||
|
||||
# Calculate temporal scores for entry points
|
||||
total_days = (end_date - start_date).total_seconds() / 86400
|
||||
mid_date = start_date + (end_date - start_date) / 2 # Calculate once for all comparisons
|
||||
results = []
|
||||
visited = set()
|
||||
|
||||
for ep in entry_points:
|
||||
unit_id = str(ep["id"])
|
||||
visited.add(unit_id)
|
||||
|
||||
# Calculate temporal proximity using the most relevant date
|
||||
# Priority: occurred_start/end (event time) > mentioned_at (mention time)
|
||||
best_date = None
|
||||
if ep["occurred_start"] is not None and ep["occurred_end"] is not None:
|
||||
# Use midpoint of occurred range
|
||||
best_date = ep["occurred_start"] + (ep["occurred_end"] - ep["occurred_start"]) / 2
|
||||
elif ep["occurred_start"] is not None:
|
||||
best_date = ep["occurred_start"]
|
||||
elif ep["occurred_end"] is not None:
|
||||
best_date = ep["occurred_end"]
|
||||
elif ep["mentioned_at"] is not None:
|
||||
best_date = ep["mentioned_at"]
|
||||
|
||||
# Temporal proximity score (closer to range center = higher score)
|
||||
if best_date:
|
||||
days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
|
||||
temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
||||
else:
|
||||
temporal_proximity = 0.5 # Fallback if no dates (shouldn't happen due to WHERE clause)
|
||||
|
||||
# Create RetrievalResult with temporal scores
|
||||
ep_result = RetrievalResult.from_db_row(dict(ep))
|
||||
ep_result.temporal_score = temporal_proximity
|
||||
ep_result.temporal_proximity = temporal_proximity
|
||||
results.append(ep_result)
|
||||
|
||||
# Spread through temporal links using BATCHED neighbor fetching
|
||||
# Map node_id -> (semantic_sim, temporal_score) for propagation
|
||||
node_scores = {str(ep["id"]): (ep["similarity"], 1.0) for ep in entry_points}
|
||||
frontier = list(node_scores.keys()) # Current batch of nodes to expand
|
||||
budget_remaining = budget - len(entry_points)
|
||||
batch_size = 20 # Process this many nodes per DB query
|
||||
|
||||
while frontier and budget_remaining > 0:
|
||||
# Take a batch from frontier
|
||||
batch_ids = frontier[:batch_size]
|
||||
frontier = frontier[batch_size:]
|
||||
|
||||
# Batch fetch all neighbors for this batch of nodes
|
||||
neighbors = await conn.fetch(
|
||||
f"""
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.embedding, mu.fact_type, mu.document_id, mu.chunk_id,
|
||||
ml.weight, ml.link_type, ml.from_unit_id,
|
||||
1 - (mu.embedding <=> $1::vector) AS similarity
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = ANY($2::uuid[])
|
||||
AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents')
|
||||
AND ml.weight >= 0.1
|
||||
AND mu.fact_type = $3
|
||||
AND mu.embedding IS NOT NULL
|
||||
AND (1 - (mu.embedding <=> $1::vector)) >= $4
|
||||
ORDER BY ml.weight DESC
|
||||
LIMIT $5
|
||||
""",
|
||||
query_emb_str,
|
||||
batch_ids,
|
||||
fact_type,
|
||||
semantic_threshold,
|
||||
batch_size * 10, # Allow up to 10 neighbors per node in batch
|
||||
)
|
||||
|
||||
for n in neighbors:
|
||||
neighbor_id = str(n["id"])
|
||||
if neighbor_id in visited:
|
||||
continue
|
||||
|
||||
visited.add(neighbor_id)
|
||||
budget_remaining -= 1
|
||||
|
||||
# Get parent's scores for propagation
|
||||
parent_id = str(n["from_unit_id"])
|
||||
_, parent_temporal_score = node_scores.get(parent_id, (0.5, 0.5))
|
||||
|
||||
# Calculate temporal score for neighbor using best available date
|
||||
neighbor_best_date = None
|
||||
if n["occurred_start"] is not None and n["occurred_end"] is not None:
|
||||
neighbor_best_date = n["occurred_start"] + (n["occurred_end"] - n["occurred_start"]) / 2
|
||||
elif n["occurred_start"] is not None:
|
||||
neighbor_best_date = n["occurred_start"]
|
||||
elif n["occurred_end"] is not None:
|
||||
neighbor_best_date = n["occurred_end"]
|
||||
elif n["mentioned_at"] is not None:
|
||||
neighbor_best_date = n["mentioned_at"]
|
||||
|
||||
if neighbor_best_date:
|
||||
days_from_mid = abs((neighbor_best_date - mid_date).total_seconds() / 86400)
|
||||
neighbor_temporal_proximity = (
|
||||
1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
||||
)
|
||||
else:
|
||||
neighbor_temporal_proximity = 0.3 # Lower score if no temporal data
|
||||
|
||||
# Boost causal links (same as graph retrieval)
|
||||
link_type = n["link_type"]
|
||||
if link_type in ("causes", "caused_by"):
|
||||
causal_boost = 2.0
|
||||
elif link_type in ("enables", "prevents"):
|
||||
causal_boost = 1.5
|
||||
else:
|
||||
causal_boost = 1.0
|
||||
|
||||
# Propagate temporal score through links (decay, with causal boost)
|
||||
propagated_temporal = parent_temporal_score * n["weight"] * causal_boost * 0.7
|
||||
|
||||
# Combined temporal score
|
||||
combined_temporal = max(neighbor_temporal_proximity, propagated_temporal)
|
||||
|
||||
# Create RetrievalResult with temporal scores
|
||||
neighbor_result = RetrievalResult.from_db_row(dict(n))
|
||||
neighbor_result.temporal_score = combined_temporal
|
||||
neighbor_result.temporal_proximity = neighbor_temporal_proximity
|
||||
results.append(neighbor_result)
|
||||
|
||||
# Track scores for propagation and add to frontier
|
||||
if budget_remaining > 0 and combined_temporal > 0.2:
|
||||
node_scores[neighbor_id] = (n["similarity"], combined_temporal)
|
||||
frontier.append(neighbor_id)
|
||||
|
||||
if budget_remaining <= 0:
|
||||
break
|
||||
|
||||
return results
|
||||
|
||||
|
||||
async def retrieve_parallel(
|
||||
pool,
|
||||
query_text: str,
|
||||
query_embedding_str: str,
|
||||
bank_id: str,
|
||||
fact_type: str,
|
||||
thinking_budget: int,
|
||||
question_date: datetime | None = None,
|
||||
query_analyzer: Optional["QueryAnalyzer"] = None,
|
||||
graph_retriever: GraphRetriever | None = None,
|
||||
temporal_constraint: tuple | None = None, # Pre-extracted temporal constraint
|
||||
tags: list[str] | None = None, # Visibility scope tags for filtering
|
||||
) -> ParallelRetrievalResult:
|
||||
"""
|
||||
Run 3-way or 4-way parallel retrieval (adds temporal if detected).
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
query_text: Query text
|
||||
query_embedding_str: Query embedding as string
|
||||
bank_id: Bank ID
|
||||
fact_type: Fact type to filter
|
||||
thinking_budget: Budget for graph traversal and retrieval limits
|
||||
question_date: Optional date when question was asked (for temporal filtering)
|
||||
query_analyzer: Query analyzer to use (defaults to TransformerQueryAnalyzer)
|
||||
graph_retriever: Graph retrieval strategy (defaults to configured retriever)
|
||||
temporal_constraint: Pre-extracted temporal constraint (optional)
|
||||
tags: Optional list of tags for visibility filtering (OR matching)
|
||||
|
||||
Returns:
|
||||
ParallelRetrievalResult with semantic, bm25, graph, temporal results and timings
|
||||
"""
|
||||
retriever = graph_retriever or get_default_graph_retriever()
|
||||
|
||||
# Use optimized parallel path for MPFP and LinkExpansion (runs all methods truly in parallel)
|
||||
# BFS uses legacy path that extracts temporal constraint upfront
|
||||
if retriever.name in ("mpfp", "link_expansion"):
|
||||
return await _retrieve_parallel_mpfp(
|
||||
pool,
|
||||
query_text,
|
||||
query_embedding_str,
|
||||
bank_id,
|
||||
fact_type,
|
||||
thinking_budget,
|
||||
temporal_constraint,
|
||||
retriever,
|
||||
question_date,
|
||||
query_analyzer,
|
||||
tags=tags,
|
||||
)
|
||||
else:
|
||||
# For BFS, extract temporal constraint upfront (legacy path)
|
||||
if temporal_constraint is None:
|
||||
from .temporal_extraction import extract_temporal_constraint
|
||||
|
||||
temporal_constraint = extract_temporal_constraint(
|
||||
query_text, reference_date=question_date, analyzer=query_analyzer
|
||||
)
|
||||
return await _retrieve_parallel_bfs(
|
||||
pool,
|
||||
query_text,
|
||||
query_embedding_str,
|
||||
bank_id,
|
||||
fact_type,
|
||||
thinking_budget,
|
||||
temporal_constraint,
|
||||
retriever,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _TimedResult:
|
||||
"""Internal result with timing."""
|
||||
|
||||
results: list[RetrievalResult]
|
||||
time: float
|
||||
conn_wait: float = 0.0 # Connection acquisition wait time
|
||||
|
||||
|
||||
async def _retrieve_parallel_mpfp(
|
||||
pool,
|
||||
query_text: str,
|
||||
query_embedding_str: str,
|
||||
bank_id: str,
|
||||
fact_type: str,
|
||||
thinking_budget: int,
|
||||
temporal_constraint: tuple | None,
|
||||
retriever: GraphRetriever,
|
||||
question_date: datetime | None = None,
|
||||
query_analyzer=None,
|
||||
tags: list[str] | None = None,
|
||||
) -> ParallelRetrievalResult:
|
||||
"""
|
||||
MPFP retrieval with true parallelization.
|
||||
|
||||
All methods run independently in parallel:
|
||||
- Semantic: vector similarity search
|
||||
- BM25: keyword search
|
||||
- Graph: MPFP traversal (does its own semantic seeds internally)
|
||||
- Temporal: date extraction (if needed) + date-range search
|
||||
|
||||
Temporal extraction runs IN PARALLEL with other retrievals, so even if
|
||||
dateparser is slow, it doesn't block semantic/BM25/graph.
|
||||
"""
|
||||
import time
|
||||
|
||||
async def run_semantic() -> _TimedResult:
|
||||
"""Independent semantic retrieval."""
|
||||
start = time.time()
|
||||
acquire_start = time.time()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
conn_wait = time.time() - acquire_start
|
||||
results = await retrieve_semantic(
|
||||
conn, query_embedding_str, bank_id, fact_type, limit=thinking_budget, tags=tags
|
||||
)
|
||||
return _TimedResult(results, time.time() - start, conn_wait)
|
||||
|
||||
async def run_bm25() -> _TimedResult:
|
||||
"""Independent BM25 retrieval."""
|
||||
start = time.time()
|
||||
acquire_start = time.time()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
conn_wait = time.time() - acquire_start
|
||||
results = await retrieve_bm25(conn, query_text, bank_id, fact_type, limit=thinking_budget, tags=tags)
|
||||
return _TimedResult(results, time.time() - start, conn_wait)
|
||||
|
||||
async def run_graph() -> tuple[list[RetrievalResult], float, MPFPTimings | None]:
|
||||
"""Independent graph retrieval - does its own semantic seeds."""
|
||||
start = time.time()
|
||||
|
||||
# MPFP does its own semantic seeds via _find_semantic_seeds
|
||||
# Note: temporal_seeds not used here to avoid dependency on temporal extraction
|
||||
results, mpfp_timing = await retriever.retrieve(
|
||||
pool=pool,
|
||||
query_embedding_str=query_embedding_str,
|
||||
bank_id=bank_id,
|
||||
fact_type=fact_type,
|
||||
budget=thinking_budget,
|
||||
query_text=query_text,
|
||||
semantic_seeds=None, # Let MPFP find its own seeds
|
||||
temporal_seeds=None, # Don't wait for temporal extraction
|
||||
tags=tags,
|
||||
)
|
||||
return results, time.time() - start, mpfp_timing
|
||||
|
||||
@dataclass
|
||||
class _TemporalWithConstraint:
|
||||
"""Temporal results with the extracted constraint."""
|
||||
|
||||
results: list[RetrievalResult]
|
||||
time: float
|
||||
constraint: tuple | None
|
||||
extraction_time: float # Time spent in query analyzer (dateparser)
|
||||
conn_wait: float = 0.0 # Connection acquisition wait time
|
||||
|
||||
async def run_temporal_with_extraction() -> _TemporalWithConstraint:
|
||||
"""
|
||||
Extract temporal constraint AND run temporal retrieval.
|
||||
|
||||
This runs in parallel with semantic/BM25/graph, so dateparser
|
||||
latency doesn't block other retrievals.
|
||||
"""
|
||||
start = time.time()
|
||||
|
||||
# Use pre-provided constraint if available
|
||||
tc = temporal_constraint
|
||||
extraction_time = 0.0
|
||||
|
||||
# Otherwise extract from query (this is the potentially slow dateparser call)
|
||||
if tc is None:
|
||||
from .temporal_extraction import extract_temporal_constraint
|
||||
|
||||
extraction_start = time.time()
|
||||
tc = extract_temporal_constraint(query_text, reference_date=question_date, analyzer=query_analyzer)
|
||||
extraction_time = time.time() - extraction_start
|
||||
|
||||
# If no temporal constraint found, return empty (but still report extraction time)
|
||||
if tc is None:
|
||||
return _TemporalWithConstraint([], time.time() - start, None, extraction_time, 0.0)
|
||||
|
||||
# Run temporal retrieval with the extracted constraint
|
||||
tc_start, tc_end = tc
|
||||
acquire_start = time.time()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
conn_wait = time.time() - acquire_start
|
||||
results = await retrieve_temporal(
|
||||
conn,
|
||||
query_embedding_str,
|
||||
bank_id,
|
||||
fact_type,
|
||||
tc_start,
|
||||
tc_end,
|
||||
budget=thinking_budget,
|
||||
semantic_threshold=0.1,
|
||||
)
|
||||
return _TemporalWithConstraint(results, time.time() - start, tc, extraction_time, conn_wait)
|
||||
|
||||
# Run ALL methods in parallel (including temporal extraction!)
|
||||
semantic_result, bm25_result, graph_result, temporal_result = await asyncio.gather(
|
||||
run_semantic(),
|
||||
run_bm25(),
|
||||
run_graph(),
|
||||
run_temporal_with_extraction(),
|
||||
)
|
||||
graph_results, graph_time, mpfp_timing = graph_result
|
||||
|
||||
# Compute max connection wait across all methods (graph handles its own connections)
|
||||
max_conn_wait = max(semantic_result.conn_wait, bm25_result.conn_wait, temporal_result.conn_wait)
|
||||
|
||||
return ParallelRetrievalResult(
|
||||
semantic=semantic_result.results,
|
||||
bm25=bm25_result.results,
|
||||
graph=graph_results,
|
||||
temporal=temporal_result.results if temporal_result.results else None,
|
||||
timings={
|
||||
"semantic": semantic_result.time,
|
||||
"bm25": bm25_result.time,
|
||||
"graph": graph_time,
|
||||
"temporal": temporal_result.time,
|
||||
"temporal_extraction": temporal_result.extraction_time,
|
||||
},
|
||||
temporal_constraint=temporal_result.constraint,
|
||||
mpfp_timings=[mpfp_timing] if mpfp_timing else [],
|
||||
max_conn_wait=max_conn_wait,
|
||||
)
|
||||
|
||||
|
||||
async def _get_temporal_entry_points(
|
||||
conn,
|
||||
query_embedding_str: str,
|
||||
bank_id: str,
|
||||
fact_type: str,
|
||||
start_date: datetime,
|
||||
end_date: datetime,
|
||||
limit: int = 20,
|
||||
semantic_threshold: float = 0.1,
|
||||
) -> list[RetrievalResult]:
|
||||
"""Get temporal entry points (facts in date range with semantic relevance)."""
|
||||
|
||||
if start_date.tzinfo is None:
|
||||
start_date = start_date.replace(tzinfo=UTC)
|
||||
if end_date.tzinfo is None:
|
||||
end_date = end_date.replace(tzinfo=UTC)
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
embedding, fact_type, document_id, chunk_id,
|
||||
1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = $3
|
||||
AND embedding IS NOT NULL
|
||||
AND (
|
||||
(occurred_start IS NOT NULL AND occurred_end IS NOT NULL
|
||||
AND occurred_start <= $5 AND occurred_end >= $4)
|
||||
OR (mentioned_at IS NOT NULL AND mentioned_at BETWEEN $4 AND $5)
|
||||
OR (occurred_start IS NOT NULL AND occurred_start BETWEEN $4 AND $5)
|
||||
OR (occurred_end IS NOT NULL AND occurred_end BETWEEN $4 AND $5)
|
||||
)
|
||||
AND (1 - (embedding <=> $1::vector)) >= $6
|
||||
ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC,
|
||||
(embedding <=> $1::vector) ASC
|
||||
LIMIT $7
|
||||
""",
|
||||
query_embedding_str,
|
||||
bank_id,
|
||||
fact_type,
|
||||
start_date,
|
||||
end_date,
|
||||
semantic_threshold,
|
||||
limit,
|
||||
)
|
||||
|
||||
results = []
|
||||
total_days = max((end_date - start_date).total_seconds() / 86400, 1)
|
||||
mid_date = start_date + (end_date - start_date) / 2
|
||||
|
||||
for row in rows:
|
||||
result = RetrievalResult.from_db_row(dict(row))
|
||||
|
||||
# Calculate temporal proximity score
|
||||
best_date = None
|
||||
if row["occurred_start"] and row["occurred_end"]:
|
||||
best_date = row["occurred_start"] + (row["occurred_end"] - row["occurred_start"]) / 2
|
||||
elif row["occurred_start"]:
|
||||
best_date = row["occurred_start"]
|
||||
elif row["occurred_end"]:
|
||||
best_date = row["occurred_end"]
|
||||
elif row["mentioned_at"]:
|
||||
best_date = row["mentioned_at"]
|
||||
|
||||
if best_date:
|
||||
days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
|
||||
result.temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0)
|
||||
else:
|
||||
result.temporal_proximity = 0.5
|
||||
|
||||
result.temporal_score = result.temporal_proximity
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
async def _retrieve_parallel_bfs(
|
||||
pool,
|
||||
query_text: str,
|
||||
query_embedding_str: str,
|
||||
bank_id: str,
|
||||
fact_type: str,
|
||||
thinking_budget: int,
|
||||
temporal_constraint: tuple | None,
|
||||
retriever: GraphRetriever,
|
||||
tags: list[str] | None = None,
|
||||
) -> ParallelRetrievalResult:
|
||||
"""BFS retrieval: all methods run in parallel (original behavior)."""
|
||||
import time
|
||||
|
||||
async def run_semantic() -> _TimedResult:
|
||||
start = time.time()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
results = await retrieve_semantic(
|
||||
conn, query_embedding_str, bank_id, fact_type, limit=thinking_budget, tags=tags
|
||||
)
|
||||
return _TimedResult(results, time.time() - start)
|
||||
|
||||
async def run_bm25() -> _TimedResult:
|
||||
start = time.time()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
results = await retrieve_bm25(conn, query_text, bank_id, fact_type, limit=thinking_budget, tags=tags)
|
||||
return _TimedResult(results, time.time() - start)
|
||||
|
||||
async def run_graph() -> _TimedResult:
|
||||
start = time.time()
|
||||
results, _ = await retriever.retrieve(
|
||||
pool=pool,
|
||||
query_embedding_str=query_embedding_str,
|
||||
bank_id=bank_id,
|
||||
fact_type=fact_type,
|
||||
budget=thinking_budget,
|
||||
query_text=query_text,
|
||||
tags=tags,
|
||||
)
|
||||
return _TimedResult(results, time.time() - start)
|
||||
|
||||
async def run_temporal(tc_start, tc_end) -> _TimedResult:
|
||||
start = time.time()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
results = await retrieve_temporal(
|
||||
conn,
|
||||
query_embedding_str,
|
||||
bank_id,
|
||||
fact_type,
|
||||
tc_start,
|
||||
tc_end,
|
||||
budget=thinking_budget,
|
||||
semantic_threshold=0.1,
|
||||
tags=tags,
|
||||
)
|
||||
return _TimedResult(results, time.time() - start)
|
||||
|
||||
if temporal_constraint:
|
||||
tc_start, tc_end = temporal_constraint
|
||||
semantic_r, bm25_r, graph_r, temporal_r = await asyncio.gather(
|
||||
run_semantic(),
|
||||
run_bm25(),
|
||||
run_graph(),
|
||||
run_temporal(tc_start, tc_end),
|
||||
)
|
||||
return ParallelRetrievalResult(
|
||||
semantic=semantic_r.results,
|
||||
bm25=bm25_r.results,
|
||||
graph=graph_r.results,
|
||||
temporal=temporal_r.results,
|
||||
timings={
|
||||
"semantic": semantic_r.time,
|
||||
"bm25": bm25_r.time,
|
||||
"graph": graph_r.time,
|
||||
"temporal": temporal_r.time,
|
||||
},
|
||||
temporal_constraint=temporal_constraint,
|
||||
)
|
||||
else:
|
||||
semantic_r, bm25_r, graph_r = await asyncio.gather(
|
||||
run_semantic(),
|
||||
run_bm25(),
|
||||
run_graph(),
|
||||
)
|
||||
return ParallelRetrievalResult(
|
||||
semantic=semantic_r.results,
|
||||
bm25=bm25_r.results,
|
||||
graph=graph_r.results,
|
||||
temporal=None,
|
||||
timings={
|
||||
"semantic": semantic_r.time,
|
||||
"bm25": bm25_r.time,
|
||||
"graph": graph_r.time,
|
||||
},
|
||||
temporal_constraint=None,
|
||||
)
|
||||
|
||||
|
||||
async def retrieve_all_fact_types_parallel(
|
||||
pool,
|
||||
query_text: str,
|
||||
|
||||
@@ -1,77 +0,0 @@
|
||||
"""File storage backends for uploaded files."""
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from .base import FileStorage
|
||||
from .postgresql import PostgreSQLFileStorage
|
||||
|
||||
__all__ = ["FileStorage", "PostgreSQLFileStorage", "create_file_storage"]
|
||||
|
||||
|
||||
def create_file_storage(
|
||||
storage_type: str,
|
||||
pool_getter: Callable | None = None,
|
||||
schema: str | None = None,
|
||||
**kwargs,
|
||||
) -> FileStorage:
|
||||
"""
|
||||
Create file storage backend based on configuration.
|
||||
|
||||
Args:
|
||||
storage_type: "native" (PostgreSQL BYTEA) or "s3" (S3-compatible object storage)
|
||||
pool_getter: Database pool getter (required for native)
|
||||
schema: Database schema (for native multi-tenant)
|
||||
**kwargs: Additional args passed to storage backend
|
||||
|
||||
Returns:
|
||||
FileStorage instance
|
||||
|
||||
Raises:
|
||||
ValueError: If storage_type is unknown or required args are missing
|
||||
"""
|
||||
if storage_type == "native":
|
||||
if not pool_getter:
|
||||
raise ValueError("pool_getter required for native (PostgreSQL) storage")
|
||||
return PostgreSQLFileStorage(pool_getter=pool_getter, schema=schema)
|
||||
elif storage_type == "s3":
|
||||
from ...config import get_config
|
||||
from .s3 import S3FileStorage
|
||||
|
||||
config = get_config()
|
||||
bucket = config.file_storage_s3_bucket
|
||||
if not bucket:
|
||||
raise ValueError("HINDSIGHT_API_FILE_STORAGE_S3_BUCKET is required for S3 storage")
|
||||
return S3FileStorage(
|
||||
bucket=bucket,
|
||||
region=config.file_storage_s3_region,
|
||||
endpoint=config.file_storage_s3_endpoint,
|
||||
access_key_id=config.file_storage_s3_access_key_id,
|
||||
secret_access_key=config.file_storage_s3_secret_access_key,
|
||||
)
|
||||
elif storage_type == "gcs":
|
||||
from ...config import get_config
|
||||
from .gcs import GCSFileStorage
|
||||
|
||||
config = get_config()
|
||||
bucket = config.file_storage_gcs_bucket
|
||||
if not bucket:
|
||||
raise ValueError("HINDSIGHT_API_FILE_STORAGE_GCS_BUCKET is required for GCS storage")
|
||||
return GCSFileStorage(
|
||||
bucket=bucket,
|
||||
service_account_key=config.file_storage_gcs_service_account_key,
|
||||
)
|
||||
elif storage_type == "azure":
|
||||
from ...config import get_config
|
||||
from .azure import AzureFileStorage
|
||||
|
||||
config = get_config()
|
||||
container = config.file_storage_azure_container
|
||||
if not container:
|
||||
raise ValueError("HINDSIGHT_API_FILE_STORAGE_AZURE_CONTAINER is required for Azure storage")
|
||||
return AzureFileStorage(
|
||||
container_name=container,
|
||||
account_name=config.file_storage_azure_account_name,
|
||||
account_key=config.file_storage_azure_account_key,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown storage type: {storage_type}. Supported: 'native', 's3', 'gcs', 'azure'.")
|
||||
@@ -1,62 +0,0 @@
|
||||
"""Azure Blob Storage backend using obstore."""
|
||||
|
||||
import logging
|
||||
from datetime import timedelta
|
||||
|
||||
import obstore as obs
|
||||
from obstore.store import AzureStore
|
||||
|
||||
from .base import FileStorage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AzureFileStorage(FileStorage):
|
||||
"""
|
||||
Azure Blob Storage backend.
|
||||
|
||||
Uses obstore (Rust-backed) for high-throughput async access to Azure Blob Storage.
|
||||
Supports account key, SAS token, and default Azure credentials.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
container_name: str,
|
||||
account_name: str | None = None,
|
||||
account_key: str | None = None,
|
||||
):
|
||||
kwargs: dict = {}
|
||||
if account_name:
|
||||
kwargs["account_name"] = account_name
|
||||
if account_key:
|
||||
kwargs["account_key"] = account_key
|
||||
|
||||
self._store = AzureStore(container_name, **kwargs)
|
||||
logger.info(f"Initialized Azure file storage: container={container_name}, account={account_name}")
|
||||
|
||||
async def store(self, file_data: bytes, key: str, metadata: dict[str, str] | None = None) -> str:
|
||||
await obs.put_async(self._store, key, file_data)
|
||||
logger.debug(f"Stored file {key} ({len(file_data)} bytes) in Azure")
|
||||
return key
|
||||
|
||||
async def retrieve(self, key: str) -> bytes:
|
||||
try:
|
||||
response = await obs.get_async(self._store, key)
|
||||
return await response.bytes_async()
|
||||
except Exception as e:
|
||||
if "not found" in str(e).lower() or "BlobNotFound" in str(e):
|
||||
raise FileNotFoundError(f"File not found: {key}") from e
|
||||
raise
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
await obs.delete_async(self._store, key)
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
try:
|
||||
await obs.head_async(self._store, key)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def get_download_url(self, key: str, expires_in: int = 3600) -> str:
|
||||
return await obs.sign_async(self._store, "GET", key, timedelta(seconds=expires_in))
|
||||
@@ -1,83 +0,0 @@
|
||||
"""Abstract base class for file storage backends."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class FileStorage(ABC):
|
||||
"""Abstract base for file storage backends."""
|
||||
|
||||
@abstractmethod
|
||||
async def store(
|
||||
self,
|
||||
file_data: bytes,
|
||||
key: str,
|
||||
metadata: dict[str, str] | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Store file and return storage key.
|
||||
|
||||
Args:
|
||||
file_data: Raw file bytes
|
||||
key: Storage key (e.g., "banks/{bank_id}/files/{file_id}.pdf")
|
||||
metadata: Optional metadata to store with file
|
||||
|
||||
Returns:
|
||||
Storage key that can be used to retrieve the file
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def retrieve(self, key: str) -> bytes:
|
||||
"""
|
||||
Retrieve file by storage key.
|
||||
|
||||
Args:
|
||||
key: Storage key
|
||||
|
||||
Returns:
|
||||
File data as bytes
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If file does not exist
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def delete(self, key: str) -> None:
|
||||
"""
|
||||
Delete file by storage key.
|
||||
|
||||
Args:
|
||||
key: Storage key
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def exists(self, key: str) -> bool:
|
||||
"""
|
||||
Check if file exists.
|
||||
|
||||
Args:
|
||||
key: Storage key
|
||||
|
||||
Returns:
|
||||
True if file exists, False otherwise
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_download_url(self, key: str, expires_in: int = 3600) -> str:
|
||||
"""
|
||||
Get a URL for downloading the file.
|
||||
|
||||
For PostgreSQL storage, this might be a relative API path.
|
||||
For S3, this would be a pre-signed URL.
|
||||
|
||||
Args:
|
||||
key: Storage key
|
||||
expires_in: Expiration time in seconds (may be ignored for some backends)
|
||||
|
||||
Returns:
|
||||
Download URL or path
|
||||
"""
|
||||
pass
|
||||
@@ -1,59 +0,0 @@
|
||||
"""Google Cloud Storage backend using obstore."""
|
||||
|
||||
import logging
|
||||
from datetime import timedelta
|
||||
|
||||
import obstore as obs
|
||||
from obstore.store import GCSStore
|
||||
|
||||
from .base import FileStorage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GCSFileStorage(FileStorage):
|
||||
"""
|
||||
Google Cloud Storage backend.
|
||||
|
||||
Uses obstore (Rust-backed) for high-throughput async access to GCS.
|
||||
Supports Application Default Credentials, service account keys, and explicit credentials.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bucket: str,
|
||||
service_account_key: str | None = None,
|
||||
):
|
||||
kwargs: dict = {}
|
||||
if service_account_key:
|
||||
kwargs["service_account_key"] = service_account_key
|
||||
|
||||
self._store = GCSStore(bucket, **kwargs)
|
||||
logger.info(f"Initialized GCS file storage: bucket={bucket}")
|
||||
|
||||
async def store(self, file_data: bytes, key: str, metadata: dict[str, str] | None = None) -> str:
|
||||
await obs.put_async(self._store, key, file_data)
|
||||
logger.debug(f"Stored file {key} ({len(file_data)} bytes) in GCS")
|
||||
return key
|
||||
|
||||
async def retrieve(self, key: str) -> bytes:
|
||||
try:
|
||||
response = await obs.get_async(self._store, key)
|
||||
return await response.bytes_async()
|
||||
except Exception as e:
|
||||
if "not found" in str(e).lower():
|
||||
raise FileNotFoundError(f"File not found: {key}") from e
|
||||
raise
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
await obs.delete_async(self._store, key)
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
try:
|
||||
await obs.head_async(self._store, key)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def get_download_url(self, key: str, expires_in: int = 3600) -> str:
|
||||
return await obs.sign_async(self._store, "GET", key, timedelta(seconds=expires_in))
|
||||
@@ -1,139 +0,0 @@
|
||||
"""PostgreSQL BYTEA-based file storage (default, zero-config)."""
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import asyncpg
|
||||
|
||||
from .base import FileStorage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def fq_table(table: str, schema: str | None = None) -> str:
|
||||
"""Get fully-qualified table name with optional schema prefix."""
|
||||
if schema:
|
||||
return f'"{schema}".{table}'
|
||||
return table
|
||||
|
||||
|
||||
class PostgreSQLFileStorage(FileStorage):
|
||||
"""
|
||||
PostgreSQL BYTEA-based file storage.
|
||||
|
||||
Stores files directly in PostgreSQL using BYTEA columns.
|
||||
This is the default storage backend - zero configuration required!
|
||||
|
||||
Pros:
|
||||
- Works out of the box (no external dependencies)
|
||||
- Transactional consistency with database
|
||||
- Simple backups (included in pg_dump)
|
||||
- Good performance for <10MB files
|
||||
|
||||
Cons:
|
||||
- Database bloat for large/many files
|
||||
- Not ideal for distributed deployments
|
||||
- Higher cost than object storage at scale
|
||||
|
||||
For production/scale, consider S3FileStorage instead.
|
||||
"""
|
||||
|
||||
def __init__(self, pool_getter: Callable[[], "asyncpg.Pool"], schema: str | None = None):
|
||||
"""
|
||||
Initialize PostgreSQL file storage.
|
||||
|
||||
Args:
|
||||
pool_getter: Function that returns asyncpg connection pool
|
||||
schema: Database schema (for multi-tenant support)
|
||||
"""
|
||||
self._pool_getter = pool_getter
|
||||
self._schema = schema
|
||||
|
||||
async def store(
|
||||
self,
|
||||
file_data: bytes,
|
||||
key: str,
|
||||
metadata: dict[str, str] | None = None,
|
||||
) -> str:
|
||||
"""Store file in PostgreSQL."""
|
||||
pool = self._pool_getter()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("file_storage", self._schema)}
|
||||
(storage_key, data)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT (storage_key) DO UPDATE SET
|
||||
data = EXCLUDED.data
|
||||
""",
|
||||
key,
|
||||
file_data,
|
||||
)
|
||||
|
||||
logger.debug(f"Stored file {key} ({len(file_data)} bytes) in PostgreSQL")
|
||||
return key
|
||||
|
||||
async def retrieve(self, key: str) -> bytes:
|
||||
"""Retrieve file from PostgreSQL."""
|
||||
pool = self._pool_getter()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT data FROM {fq_table("file_storage", self._schema)}
|
||||
WHERE storage_key = $1
|
||||
""",
|
||||
key,
|
||||
)
|
||||
|
||||
if not row:
|
||||
raise FileNotFoundError(f"File not found: {key}")
|
||||
|
||||
return bytes(row["data"])
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
"""Delete file from PostgreSQL."""
|
||||
pool = self._pool_getter()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {fq_table("file_storage", self._schema)}
|
||||
WHERE storage_key = $1
|
||||
""",
|
||||
key,
|
||||
)
|
||||
|
||||
# Check if anything was deleted
|
||||
if result == "DELETE 0":
|
||||
logger.warning(f"Attempted to delete non-existent file: {key}")
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
"""Check if file exists in PostgreSQL."""
|
||||
pool = self._pool_getter()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT 1 FROM {fq_table("file_storage", self._schema)}
|
||||
WHERE storage_key = $1
|
||||
""",
|
||||
key,
|
||||
)
|
||||
|
||||
return row is not None
|
||||
|
||||
async def get_download_url(self, key: str, expires_in: int = 3600) -> str:
|
||||
"""
|
||||
Get download URL for PostgreSQL-stored file.
|
||||
|
||||
Returns an API endpoint path (not a pre-signed URL since the file
|
||||
is stored in the database). The expires_in parameter is ignored
|
||||
for PostgreSQL storage.
|
||||
"""
|
||||
# Return API path for download endpoint
|
||||
# (expires_in ignored for database storage - auth handled at API level)
|
||||
return f"/v1/default/files/download/{key}"
|
||||
@@ -1,71 +0,0 @@
|
||||
"""S3 object storage backend using obstore."""
|
||||
|
||||
import logging
|
||||
from datetime import timedelta
|
||||
|
||||
import obstore as obs
|
||||
from obstore.store import S3Store
|
||||
|
||||
from .base import FileStorage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class S3FileStorage(FileStorage):
|
||||
"""
|
||||
S3-compatible object storage backend.
|
||||
|
||||
Uses obstore (Rust-backed) for high-throughput async access to
|
||||
Amazon S3, MinIO, Cloudflare R2, and other S3-compliant APIs.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bucket: str,
|
||||
region: str | None = None,
|
||||
endpoint: str | None = None,
|
||||
access_key_id: str | None = None,
|
||||
secret_access_key: str | None = None,
|
||||
):
|
||||
kwargs: dict = {}
|
||||
if region:
|
||||
kwargs["region"] = region
|
||||
if endpoint:
|
||||
kwargs["endpoint"] = endpoint
|
||||
# Allow plain HTTP for local S3-compatible services (MinIO, LocalStack, etc.)
|
||||
if endpoint.startswith("http://"):
|
||||
kwargs["allow_http"] = True
|
||||
if access_key_id:
|
||||
kwargs["access_key_id"] = access_key_id
|
||||
if secret_access_key:
|
||||
kwargs["secret_access_key"] = secret_access_key
|
||||
|
||||
self._store = S3Store(bucket, **kwargs)
|
||||
logger.info(f"Initialized S3 file storage: bucket={bucket}, region={region}, endpoint={endpoint}")
|
||||
|
||||
async def store(self, file_data: bytes, key: str, metadata: dict[str, str] | None = None) -> str:
|
||||
await obs.put_async(self._store, key, file_data)
|
||||
logger.debug(f"Stored file {key} ({len(file_data)} bytes) in S3")
|
||||
return key
|
||||
|
||||
async def retrieve(self, key: str) -> bytes:
|
||||
try:
|
||||
response = await obs.get_async(self._store, key)
|
||||
return await response.bytes_async()
|
||||
except Exception as e:
|
||||
if "not found" in str(e).lower() or "NoSuchKey" in str(e):
|
||||
raise FileNotFoundError(f"File not found: {key}") from e
|
||||
raise
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
await obs.delete_async(self._store, key)
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
try:
|
||||
await obs.head_async(self._store, key)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def get_download_url(self, key: str, expires_in: int = 3600) -> str:
|
||||
return await obs.sign_async(self._store, "GET", key, timedelta(seconds=expires_in))
|
||||
@@ -19,7 +19,6 @@ async def extract_facts(
|
||||
context: str = "",
|
||||
llm_config: "LLMConfig" = None,
|
||||
agent_name: str = None,
|
||||
config=None,
|
||||
) -> 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
|
||||
config: HindsightConfig to use (defaults to global config if not provided)
|
||||
|
||||
Returns:
|
||||
Tuple of (facts, chunks) where:
|
||||
@@ -49,19 +47,12 @@ async def extract_facts(
|
||||
if not text or not text.strip():
|
||||
return [], []
|
||||
|
||||
# Use provided config or fall back to global config
|
||||
if config is None:
|
||||
from ..config import _get_raw_config
|
||||
|
||||
config = _get_raw_config()
|
||||
|
||||
facts, chunks, _ = await extract_facts_from_text(
|
||||
text,
|
||||
event_date,
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name=agent_name,
|
||||
config=config,
|
||||
context=context,
|
||||
)
|
||||
|
||||
if not facts:
|
||||
|
||||
@@ -96,13 +96,7 @@ class DefaultExtensionContext(ExtensionContext):
|
||||
|
||||
async def run_migration(self, schema: str) -> None:
|
||||
"""Run migrations for a specific schema."""
|
||||
from hindsight_api.config import get_config
|
||||
from hindsight_api.migrations import (
|
||||
ensure_embedding_dimension,
|
||||
ensure_text_search_extension,
|
||||
ensure_vector_extension,
|
||||
run_migrations,
|
||||
)
|
||||
from hindsight_api.migrations import ensure_embedding_dimension, run_migrations
|
||||
|
||||
# Prefer getting URL from memory engine (handles pg0 case where URL is set after init)
|
||||
db_url = self._database_url
|
||||
@@ -113,9 +107,6 @@ class DefaultExtensionContext(ExtensionContext):
|
||||
|
||||
run_migrations(db_url, schema=schema)
|
||||
|
||||
# Get config for vector extension setting
|
||||
config = get_config()
|
||||
|
||||
# Ensure embedding column dimension matches the model's dimension
|
||||
# This is needed because migrations create columns with default dimension
|
||||
if self._memory_engine is not None:
|
||||
@@ -123,15 +114,7 @@ class DefaultExtensionContext(ExtensionContext):
|
||||
if embeddings is not None:
|
||||
dimension = getattr(embeddings, "dimension", None)
|
||||
if dimension is not None:
|
||||
ensure_embedding_dimension(
|
||||
db_url, dimension, schema=schema, vector_extension=config.vector_extension
|
||||
)
|
||||
|
||||
# Ensure vector indexes match the configured extension
|
||||
ensure_vector_extension(db_url, vector_extension=config.vector_extension, schema=schema)
|
||||
|
||||
# Ensure text search columns/indexes match the configured extension
|
||||
ensure_text_search_extension(db_url, text_search_extension=config.text_search_extension, schema=schema)
|
||||
ensure_embedding_dimension(db_url, dimension, schema=schema)
|
||||
|
||||
def get_memory_engine(self) -> "MemoryEngineInterface":
|
||||
"""Get the memory engine interface."""
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
from hindsight_api.models import RequestContext
|
||||
@@ -89,54 +88,6 @@ class TenantExtension(Extension, ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
async def get_tenant_config(self, context: RequestContext) -> dict[str, Any]:
|
||||
"""
|
||||
Get tenant-specific configuration overrides.
|
||||
|
||||
This method is called during hierarchical configuration resolution to get
|
||||
tenant-level config overrides. The returned dict should contain Python field
|
||||
names (lowercase snake_case) as keys, not environment variable names.
|
||||
|
||||
Example:
|
||||
{"llm_model": "gpt-4", "retain_extraction_mode": "verbose"}
|
||||
|
||||
The default implementation returns an empty dict (no tenant-specific config).
|
||||
Override this method in custom extensions to provide tenant-specific configuration.
|
||||
|
||||
Args:
|
||||
context: The request context containing tenant information.
|
||||
|
||||
Returns:
|
||||
Dict of config field names to values (only configurable fields).
|
||||
Empty dict if no tenant-specific config.
|
||||
"""
|
||||
return {}
|
||||
|
||||
async def get_allowed_config_fields(self, context: RequestContext, bank_id: str) -> set[str] | None:
|
||||
"""
|
||||
Get set of config fields that this tenant/bank is allowed to modify.
|
||||
|
||||
This method controls which configurable fields can be modified via the bank config API.
|
||||
It enables fine-grained permission control per tenant or per bank.
|
||||
|
||||
Examples:
|
||||
- Return None: Allow all configurable fields (default)
|
||||
- Return {"retain_chunk_size", "retain_custom_instructions"}: Allow only these fields
|
||||
- Return set(): Allow no modifications (read-only)
|
||||
|
||||
The default implementation returns None (all configurable fields allowed).
|
||||
Override this method in custom extensions to implement custom permission logic.
|
||||
|
||||
Args:
|
||||
context: The request context containing tenant information.
|
||||
bank_id: The bank identifier for per-bank permissions.
|
||||
|
||||
Returns:
|
||||
Set of allowed field names, or None to allow all configurable fields.
|
||||
Returned fields must be a subset of HindsightConfig.get_configurable_fields().
|
||||
"""
|
||||
return None
|
||||
|
||||
async def authenticate_mcp(self, context: RequestContext) -> TenantContext:
|
||||
"""
|
||||
Authenticate MCP requests.
|
||||
|
||||
@@ -23,7 +23,7 @@ import uvicorn
|
||||
from . import MemoryEngine, __version__
|
||||
from .api import create_app
|
||||
from .banner import print_banner
|
||||
from .config import DEFAULT_WORKERS, ENV_WORKERS, HindsightConfig, _get_raw_config
|
||||
from .config import DEFAULT_WORKERS, ENV_WORKERS, HindsightConfig, get_config
|
||||
from .daemon import (
|
||||
DEFAULT_DAEMON_PORT,
|
||||
DEFAULT_IDLE_TIMEOUT,
|
||||
@@ -68,7 +68,7 @@ def main():
|
||||
global _memory
|
||||
|
||||
# Load configuration from environment (for CLI args defaults)
|
||||
config = _get_raw_config()
|
||||
config = get_config()
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="hindsight-api",
|
||||
@@ -155,8 +155,6 @@ def main():
|
||||
config = HindsightConfig(
|
||||
database_url=config.database_url,
|
||||
database_schema=config.database_schema,
|
||||
vector_extension=config.vector_extension,
|
||||
text_search_extension=config.text_search_extension,
|
||||
llm_provider=config.llm_provider,
|
||||
llm_api_key=config.llm_api_key,
|
||||
llm_model=config.llm_model,
|
||||
@@ -166,8 +164,6 @@ def main():
|
||||
llm_initial_backoff=config.llm_initial_backoff,
|
||||
llm_max_backoff=config.llm_max_backoff,
|
||||
llm_timeout=config.llm_timeout,
|
||||
llm_groq_service_tier=config.llm_groq_service_tier,
|
||||
llm_openai_service_tier=config.llm_openai_service_tier,
|
||||
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,
|
||||
@@ -210,9 +206,6 @@ def main():
|
||||
embeddings_litellm_api_base=config.embeddings_litellm_api_base,
|
||||
embeddings_litellm_api_key=config.embeddings_litellm_api_key,
|
||||
embeddings_litellm_model=config.embeddings_litellm_model,
|
||||
embeddings_litellm_sdk_api_key=config.embeddings_litellm_sdk_api_key,
|
||||
embeddings_litellm_sdk_model=config.embeddings_litellm_sdk_model,
|
||||
embeddings_litellm_sdk_api_base=config.embeddings_litellm_sdk_api_base,
|
||||
reranker_provider=config.reranker_provider,
|
||||
reranker_local_model=config.reranker_local_model,
|
||||
reranker_local_force_cpu=config.reranker_local_force_cpu,
|
||||
@@ -228,16 +221,11 @@ def main():
|
||||
reranker_litellm_api_base=config.reranker_litellm_api_base,
|
||||
reranker_litellm_api_key=config.reranker_litellm_api_key,
|
||||
reranker_litellm_model=config.reranker_litellm_model,
|
||||
reranker_litellm_sdk_api_key=config.reranker_litellm_sdk_api_key,
|
||||
reranker_litellm_sdk_model=config.reranker_litellm_sdk_model,
|
||||
reranker_litellm_sdk_api_base=config.reranker_litellm_sdk_api_base,
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
base_path=config.base_path,
|
||||
log_level=args.log_level,
|
||||
log_format=config.log_format,
|
||||
mcp_enabled=config.mcp_enabled,
|
||||
enable_bank_config_api=config.enable_bank_config_api,
|
||||
graph_retriever=config.graph_retriever,
|
||||
mpfp_top_k_neighbors=config.mpfp_top_k_neighbors,
|
||||
recall_max_concurrent=config.recall_max_concurrent,
|
||||
@@ -247,27 +235,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_batch_tokens=config.retain_batch_tokens,
|
||||
retain_batch_enabled=config.retain_batch_enabled,
|
||||
retain_batch_poll_interval_seconds=config.retain_batch_poll_interval_seconds,
|
||||
file_storage_type=config.file_storage_type,
|
||||
file_storage_s3_bucket=config.file_storage_s3_bucket,
|
||||
file_storage_s3_region=config.file_storage_s3_region,
|
||||
file_storage_s3_endpoint=config.file_storage_s3_endpoint,
|
||||
file_storage_s3_access_key_id=config.file_storage_s3_access_key_id,
|
||||
file_storage_s3_secret_access_key=config.file_storage_s3_secret_access_key,
|
||||
file_storage_gcs_bucket=config.file_storage_gcs_bucket,
|
||||
file_storage_gcs_service_account_key=config.file_storage_gcs_service_account_key,
|
||||
file_storage_azure_container=config.file_storage_azure_container,
|
||||
file_storage_azure_account_name=config.file_storage_azure_account_name,
|
||||
file_storage_azure_account_key=config.file_storage_azure_account_key,
|
||||
file_parser=config.file_parser,
|
||||
file_parser_iris_token=config.file_parser_iris_token,
|
||||
file_parser_iris_org_id=config.file_parser_iris_org_id,
|
||||
file_conversion_max_batch_size_mb=config.file_conversion_max_batch_size_mb,
|
||||
file_conversion_max_batch_size=config.file_conversion_max_batch_size,
|
||||
enable_file_upload_api=config.enable_file_upload_api,
|
||||
file_delete_after_retain=config.file_delete_after_retain,
|
||||
enable_observations=config.enable_observations,
|
||||
consolidation_batch_size=config.consolidation_batch_size,
|
||||
consolidation_max_tokens=config.consolidation_max_tokens,
|
||||
@@ -369,7 +336,6 @@ def main():
|
||||
"proxy_headers": args.proxy_headers,
|
||||
"ws": "wsproto", # Use wsproto instead of websockets to avoid deprecation warnings
|
||||
"loop": loop_impl, # Explicitly set event loop implementation
|
||||
"timeout_keep_alive": 30, # Exceed aiohttp's 15s client timeout so the client always closes first
|
||||
}
|
||||
|
||||
# Add optional parameters if provided
|
||||
@@ -398,8 +364,6 @@ def main():
|
||||
reranker_provider=config.reranker_provider,
|
||||
mcp_enabled=config.mcp_enabled,
|
||||
version=__version__,
|
||||
vector_extension=config.vector_extension,
|
||||
text_search_extension=config.text_search_extension,
|
||||
)
|
||||
|
||||
# Start idle checker in daemon mode
|
||||
|
||||
@@ -33,69 +33,6 @@ logger = logging.getLogger(__name__)
|
||||
MIGRATION_LOCK_ID = 123456789
|
||||
|
||||
|
||||
def _detect_vector_extension(conn, vector_extension: str = "pgvector") -> str:
|
||||
"""
|
||||
Validate vector extension: 'pgvector', 'vchord', or 'pgvectorscale'.
|
||||
|
||||
Args:
|
||||
conn: SQLAlchemy connection object
|
||||
vector_extension: Configured extension ("pgvector", "vchord", or "pgvectorscale")
|
||||
|
||||
Returns:
|
||||
"pgvector", "vchord", "pgvectorscale", or "pg_diskann"
|
||||
|
||||
Raises:
|
||||
RuntimeError: If configured extension is not installed
|
||||
"""
|
||||
# Verify the configured extension is installed
|
||||
if vector_extension == "pgvectorscale":
|
||||
# pgvectorscale/DiskANN requires pgvector to be installed first
|
||||
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
|
||||
if not pgvector_check:
|
||||
raise RuntimeError(
|
||||
"DiskANN (pgvectorscale/pg_diskann) requires pgvector to be installed. "
|
||||
"Install it with: CREATE EXTENSION vector; then CREATE EXTENSION vectorscale CASCADE; (or pg_diskann on Azure)"
|
||||
)
|
||||
|
||||
# Check for either vectorscale (open source) or pg_diskann (Azure)
|
||||
vectorscale_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vectorscale'")).scalar()
|
||||
pg_diskann_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_diskann'")).scalar()
|
||||
|
||||
if vectorscale_check:
|
||||
logger.debug("Using vector extension: pgvectorscale (DiskANN)")
|
||||
return "pgvectorscale"
|
||||
elif pg_diskann_check:
|
||||
logger.debug("Using vector extension: pg_diskann (Azure DiskANN)")
|
||||
return "pg_diskann" # Return distinct name for parameter handling
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"Configured vector extension 'pgvectorscale' not found. "
|
||||
"Install either:\n"
|
||||
" - pgvectorscale (open source): CREATE EXTENSION vectorscale CASCADE;\n"
|
||||
" - pg_diskann (Azure): CREATE EXTENSION pg_diskann CASCADE;"
|
||||
)
|
||||
elif vector_extension == "vchord":
|
||||
vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar()
|
||||
if not vchord_check:
|
||||
raise RuntimeError(
|
||||
"Configured vector extension 'vchord' not found. Install it with: CREATE EXTENSION vchord CASCADE;"
|
||||
)
|
||||
logger.debug("Using configured vector extension: vchord")
|
||||
return "vchord"
|
||||
elif vector_extension == "pgvector":
|
||||
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
|
||||
if not pgvector_check:
|
||||
raise RuntimeError(
|
||||
"Configured vector extension 'pgvector' not found. Install it with: CREATE EXTENSION vector;"
|
||||
)
|
||||
logger.debug("Using configured vector extension: pgvector")
|
||||
return "pgvector"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid vector_extension: {vector_extension}. Must be 'pgvector', 'vchord', or 'pgvectorscale'"
|
||||
)
|
||||
|
||||
|
||||
def _get_schema_lock_id(schema: str) -> int:
|
||||
"""
|
||||
Generate a unique advisory lock ID for a schema.
|
||||
@@ -305,48 +242,6 @@ def run_migrations(
|
||||
"Please install it with: CREATE EXTENSION vector;"
|
||||
) from e
|
||||
|
||||
# If using pgvectorscale, ensure vectorscale extension is also installed
|
||||
vector_extension = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
|
||||
if vector_extension == "pgvectorscale":
|
||||
logger.debug("Checking pgvectorscale (vectorscale) extension availability...")
|
||||
|
||||
vectorscale_check = conn.execute(
|
||||
text("SELECT 1 FROM pg_extension WHERE extname = 'vectorscale'")
|
||||
).scalar()
|
||||
|
||||
if vectorscale_check:
|
||||
logger.info("pgvectorscale extension already installed")
|
||||
else:
|
||||
# Extension doesn't exist - try to install
|
||||
logger.info("pgvectorscale extension not found, attempting to install...")
|
||||
try:
|
||||
conn.execute(text("CREATE EXTENSION vectorscale CASCADE"))
|
||||
conn.commit()
|
||||
logger.info("pgvectorscale extension installed successfully")
|
||||
except Exception as e:
|
||||
# Installation failed - check one more time in case another process installed it
|
||||
conn.rollback()
|
||||
vectorscale_recheck = conn.execute(
|
||||
text("SELECT 1 FROM pg_extension WHERE extname = 'vectorscale'")
|
||||
).fetchone()
|
||||
|
||||
if vectorscale_recheck:
|
||||
logger.warning(
|
||||
"Could not install pgvectorscale extension (permission denied?), "
|
||||
"but extension exists. Continuing..."
|
||||
)
|
||||
else:
|
||||
# Extension truly doesn't exist and we can't install it
|
||||
logger.error(
|
||||
f"pgvectorscale extension is not installed and cannot be installed: {e}. "
|
||||
f"Please ensure pgvectorscale is installed by a database administrator. "
|
||||
f"See: https://github.com/timescale/pgvectorscale#installation"
|
||||
)
|
||||
raise RuntimeError(
|
||||
"pgvectorscale extension is required but not installed. "
|
||||
"Please install it with: CREATE EXTENSION vectorscale CASCADE;"
|
||||
) from e
|
||||
|
||||
# Run migrations while holding the lock
|
||||
_run_migrations_internal(database_url, script_location, schema=schema)
|
||||
finally:
|
||||
@@ -429,7 +324,6 @@ def ensure_embedding_dimension(
|
||||
database_url: str,
|
||||
required_dimension: int,
|
||||
schema: str | None = None,
|
||||
vector_extension: str = "pgvector",
|
||||
) -> None:
|
||||
"""
|
||||
Ensure the embedding column dimension matches the model's dimension.
|
||||
@@ -444,7 +338,6 @@ def ensure_embedding_dimension(
|
||||
database_url: SQLAlchemy database URL
|
||||
required_dimension: The embedding dimension required by the model
|
||||
schema: Target PostgreSQL schema name (None for public)
|
||||
vector_extension: Configured vector extension ("pgvector" or "vchord")
|
||||
|
||||
Raises:
|
||||
RuntimeError: If dimension mismatch with existing data
|
||||
@@ -468,10 +361,6 @@ def ensure_embedding_dimension(
|
||||
logger.debug(f"memory_units table does not exist in schema '{schema_name}', skipping dimension check")
|
||||
return
|
||||
|
||||
# Detect which vector extension is available
|
||||
vector_ext = _detect_vector_extension(conn, vector_extension)
|
||||
logger.info(f"Using vector extension: {vector_ext}")
|
||||
|
||||
# Get current column dimension from pg_attribute
|
||||
# pgvector stores dimension in atttypmod
|
||||
current_dim = conn.execute(
|
||||
@@ -519,7 +408,8 @@ def ensure_embedding_dimension(
|
||||
# Table is empty, safe to alter column
|
||||
logger.info(f"Altering embedding column dimension from {current_dimension} to {required_dimension}")
|
||||
|
||||
# Drop existing vector index (works for both HNSW and vchordrq)
|
||||
# Drop the HNSW index on embedding column if it exists
|
||||
# Only drop indexes that use 'hnsw' and reference the 'embedding' column
|
||||
conn.execute(
|
||||
text(f"""
|
||||
DO $$
|
||||
@@ -529,7 +419,7 @@ def ensure_embedding_dimension(
|
||||
SELECT indexname FROM pg_indexes
|
||||
WHERE schemaname = '{schema_name}'
|
||||
AND tablename = 'memory_units'
|
||||
AND (indexdef LIKE '%hnsw%' OR indexdef LIKE '%vchordrq%')
|
||||
AND indexdef LIKE '%hnsw%'
|
||||
AND indexdef LIKE '%embedding%'
|
||||
LOOP
|
||||
EXECUTE 'DROP INDEX IF EXISTS {schema_name}.' || idx_name;
|
||||
@@ -544,452 +434,15 @@ def ensure_embedding_dimension(
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
# Recreate index with appropriate type based on detected extension
|
||||
if vector_ext == "pgvectorscale":
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_embedding_diskann
|
||||
ON {schema_name}.memory_units
|
||||
USING diskann (embedding vector_cosine_ops)
|
||||
WITH (num_neighbors = 50)
|
||||
""")
|
||||
)
|
||||
logger.info(f"Created DiskANN index for {required_dimension}-dimensional embeddings")
|
||||
elif vector_ext == "vchord":
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_embedding_vchordrq
|
||||
ON {schema_name}.memory_units
|
||||
USING vchordrq (embedding vector_l2_ops)
|
||||
""")
|
||||
)
|
||||
logger.info(f"Created vchordrq index for {required_dimension}-dimensional embeddings")
|
||||
else: # pgvector
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_embedding_hnsw
|
||||
ON {schema_name}.memory_units
|
||||
USING hnsw (embedding vector_cosine_ops)
|
||||
WITH (m = 16, ef_construction = 64)
|
||||
""")
|
||||
)
|
||||
logger.info(f"Created HNSW index for {required_dimension}-dimensional embeddings")
|
||||
# Recreate the HNSW index
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_embedding_hnsw
|
||||
ON {schema_name}.memory_units
|
||||
USING hnsw (embedding vector_cosine_ops)
|
||||
WITH (m = 16, ef_construction = 64)
|
||||
""")
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
logger.info(f"Successfully changed embedding dimension to {required_dimension}")
|
||||
|
||||
|
||||
def ensure_vector_extension(
|
||||
database_url: str,
|
||||
vector_extension: str = "pgvector",
|
||||
schema: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Ensure the vector indexes match the configured vector extension.
|
||||
|
||||
This function checks the current vector index type in the database
|
||||
and adjusts it if necessary:
|
||||
- If index type matches configured extension: no action needed
|
||||
- If they differ and tables are empty: drop old indexes, recreate with new type
|
||||
- If they differ and tables have data: raise error with migration guidance
|
||||
|
||||
Args:
|
||||
database_url: SQLAlchemy database URL
|
||||
vector_extension: Configured vector extension ("pgvector" or "vchord")
|
||||
schema: Target PostgreSQL schema name (None for public)
|
||||
|
||||
Raises:
|
||||
RuntimeError: If extension mismatch with existing data
|
||||
"""
|
||||
schema_name = schema or "public"
|
||||
|
||||
engine = create_engine(database_url)
|
||||
with engine.connect() as conn:
|
||||
# Detect which vector extension should be used
|
||||
target_ext = _detect_vector_extension(conn, vector_extension)
|
||||
logger.info(f"Target vector extension: {target_ext}")
|
||||
|
||||
# Tables with vector indexes to check
|
||||
tables_to_check = [
|
||||
("memory_units", "idx_memory_units_embedding"),
|
||||
("learnings", "idx_learnings_embedding"),
|
||||
("pinned_reflections", "idx_pinned_reflections_embedding"),
|
||||
]
|
||||
|
||||
# Determine target index type
|
||||
if target_ext in ("pgvectorscale", "pg_diskann"):
|
||||
target_index_type = "diskann"
|
||||
elif target_ext == "vchord":
|
||||
target_index_type = "vchordrq"
|
||||
else:
|
||||
target_index_type = "hnsw"
|
||||
|
||||
mismatched_tables = []
|
||||
tables_with_data = []
|
||||
|
||||
for table_name, index_name in tables_to_check:
|
||||
# Check if table exists
|
||||
table_exists = conn.execute(
|
||||
text("""
|
||||
SELECT EXISTS (
|
||||
SELECT 1 FROM information_schema.tables
|
||||
WHERE table_schema = :schema AND table_name = :table_name
|
||||
)
|
||||
"""),
|
||||
{"schema": schema_name, "table_name": table_name},
|
||||
).scalar()
|
||||
|
||||
if not table_exists:
|
||||
logger.debug(f"Table {table_name} does not exist in schema '{schema_name}', skipping")
|
||||
continue
|
||||
|
||||
# Check current index type by querying pg_indexes
|
||||
current_index_info = conn.execute(
|
||||
text("""
|
||||
SELECT indexdef
|
||||
FROM pg_indexes
|
||||
WHERE schemaname = :schema
|
||||
AND tablename = :table_name
|
||||
AND indexname LIKE :index_pattern
|
||||
"""),
|
||||
{"schema": schema_name, "table_name": table_name, "index_pattern": "%embedding%"},
|
||||
).fetchone()
|
||||
|
||||
if not current_index_info:
|
||||
logger.warning(f"No embedding index found for {table_name}, will create it")
|
||||
mismatched_tables.append((table_name, index_name, None))
|
||||
continue
|
||||
|
||||
indexdef = current_index_info[0].lower()
|
||||
if "diskann" in indexdef:
|
||||
current_index_type = "diskann"
|
||||
elif "vchordrq" in indexdef:
|
||||
current_index_type = "vchordrq"
|
||||
elif "hnsw" in indexdef:
|
||||
current_index_type = "hnsw"
|
||||
else:
|
||||
logger.warning(f"Unknown index type for {table_name}: {indexdef}")
|
||||
continue
|
||||
|
||||
# Check if index type matches target
|
||||
if current_index_type != target_index_type:
|
||||
logger.info(
|
||||
f"Index type mismatch on {table_name}: current={current_index_type}, target={target_index_type}"
|
||||
)
|
||||
mismatched_tables.append((table_name, index_name, current_index_type))
|
||||
|
||||
# Check if table has data
|
||||
row_count = conn.execute(
|
||||
text(f"SELECT COUNT(*) FROM {schema_name}.{table_name} WHERE embedding IS NOT NULL")
|
||||
).scalar()
|
||||
|
||||
if row_count > 0:
|
||||
tables_with_data.append((table_name, row_count))
|
||||
else:
|
||||
logger.debug(f"Index type OK for {table_name}: {current_index_type}")
|
||||
|
||||
# If no mismatches, we're done
|
||||
if not mismatched_tables:
|
||||
logger.debug(f"All vector indexes match configured extension: {target_ext}")
|
||||
return
|
||||
|
||||
# If there's data in any mismatched table, raise error
|
||||
if tables_with_data:
|
||||
table_list = ", ".join([f"{table}({count} rows)" for table, count in tables_with_data])
|
||||
# Map index type back to extension name for error message
|
||||
current_ext_name = {"diskann": "pgvectorscale", "vchordrq": "vchord", "hnsw": "pgvector"}.get(
|
||||
current_index_type, current_index_type
|
||||
)
|
||||
|
||||
raise RuntimeError(
|
||||
f"Cannot change vector extension from {current_index_type} to {target_index_type}: "
|
||||
f"the following tables contain data: {table_list}. "
|
||||
f"To change vector extension, you must either:\n"
|
||||
f" 1. Re-embed all data: DELETE FROM {schema_name}.memory_units; "
|
||||
f"DELETE FROM {schema_name}.learnings; DELETE FROM {schema_name}.pinned_reflections; then restart\n"
|
||||
f" 2. Use the current vector extension (set HINDSIGHT_API_VECTOR_EXTENSION='{current_ext_name}')"
|
||||
)
|
||||
|
||||
# Tables are empty, safe to recreate indexes
|
||||
logger.info(f"Recreating vector indexes for {target_ext}")
|
||||
|
||||
for table_name, index_name, current_type in mismatched_tables:
|
||||
# Drop existing index if it exists
|
||||
if current_type:
|
||||
logger.info(f"Dropping {current_type} index on {table_name}")
|
||||
conn.execute(text(f"DROP INDEX IF EXISTS {schema_name}.{index_name}"))
|
||||
|
||||
# Create new index with appropriate type
|
||||
if target_ext == "pgvectorscale":
|
||||
logger.info(f"Creating DiskANN index on {table_name} (pgvectorscale)")
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX IF NOT EXISTS {index_name}
|
||||
ON {schema_name}.{table_name}
|
||||
USING diskann (embedding vector_cosine_ops)
|
||||
WITH (num_neighbors = 50)
|
||||
""")
|
||||
)
|
||||
elif target_ext == "pg_diskann":
|
||||
logger.info(f"Creating DiskANN index on {table_name} (pg_diskann/Azure)")
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX IF NOT EXISTS {index_name}
|
||||
ON {schema_name}.{table_name}
|
||||
USING diskann (embedding vector_cosine_ops)
|
||||
WITH (max_neighbors = 50)
|
||||
""")
|
||||
)
|
||||
elif target_ext == "vchord":
|
||||
logger.info(f"Creating vchordrq index on {table_name}")
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX IF NOT EXISTS {index_name}
|
||||
ON {schema_name}.{table_name}
|
||||
USING vchordrq (embedding vector_l2_ops)
|
||||
""")
|
||||
)
|
||||
else: # pgvector
|
||||
logger.info(f"Creating HNSW index on {table_name}")
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX IF NOT EXISTS {index_name}
|
||||
ON {schema_name}.{table_name}
|
||||
USING hnsw (embedding vector_cosine_ops)
|
||||
WITH (m = 16, ef_construction = 64)
|
||||
""")
|
||||
)
|
||||
|
||||
conn.commit()
|
||||
logger.info(f"Successfully migrated vector indexes to {target_ext}")
|
||||
|
||||
|
||||
def ensure_text_search_extension(
|
||||
database_url: str,
|
||||
text_search_extension: str = "native",
|
||||
schema: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Ensure the text search columns and indexes match the configured extension.
|
||||
|
||||
This function checks the current search_vector column type and index type
|
||||
in the database and adjusts them if necessary:
|
||||
- If they match configured extension: no action needed
|
||||
- If they differ and tables are empty: drop old column/index, recreate with new type
|
||||
- If they differ and tables have data: raise error with migration guidance
|
||||
|
||||
Args:
|
||||
database_url: SQLAlchemy database URL
|
||||
text_search_extension: Configured text search extension ("native" or "vchord")
|
||||
schema: Target PostgreSQL schema name (None for public)
|
||||
|
||||
Raises:
|
||||
RuntimeError: If extension mismatch with existing data
|
||||
"""
|
||||
schema_name = schema or "public"
|
||||
|
||||
engine = create_engine(database_url)
|
||||
with engine.connect() as conn:
|
||||
# Tables with search_vector columns to check
|
||||
tables_to_check = [
|
||||
"memory_units",
|
||||
"reflections", # Renamed from pinned_reflections in p1k2l3m4n5o6 migration
|
||||
]
|
||||
|
||||
# Determine target column type and index type
|
||||
if text_search_extension == "vchord":
|
||||
target_column_type = "bm25vector"
|
||||
target_index_type = "bm25"
|
||||
elif text_search_extension == "pg_textsearch":
|
||||
target_column_type = "text"
|
||||
target_index_type = "bm25"
|
||||
else: # native
|
||||
target_column_type = "tsvector"
|
||||
target_index_type = "gin"
|
||||
|
||||
mismatched_tables = []
|
||||
tables_with_data = []
|
||||
|
||||
for table_name in tables_to_check:
|
||||
# Check if table exists
|
||||
table_exists = conn.execute(
|
||||
text("""
|
||||
SELECT EXISTS (
|
||||
SELECT 1 FROM information_schema.tables
|
||||
WHERE table_schema = :schema AND table_name = :table_name
|
||||
)
|
||||
"""),
|
||||
{"schema": schema_name, "table_name": table_name},
|
||||
).scalar()
|
||||
|
||||
if not table_exists:
|
||||
logger.debug(f"Table {table_name} does not exist in schema '{schema_name}', skipping")
|
||||
continue
|
||||
|
||||
# Get current column type from information_schema
|
||||
current_column_info = conn.execute(
|
||||
text("""
|
||||
SELECT data_type, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table_name
|
||||
AND column_name = 'search_vector'
|
||||
"""),
|
||||
{"schema": schema_name, "table_name": table_name},
|
||||
).fetchone()
|
||||
|
||||
if not current_column_info:
|
||||
logger.warning(f"No search_vector column found for {table_name}, will create it")
|
||||
mismatched_tables.append((table_name, None, None))
|
||||
continue
|
||||
|
||||
# Check column type (udt_name contains the actual type: tsvector, bm25vector, etc.)
|
||||
current_column_type = current_column_info[1] # udt_name
|
||||
|
||||
# Get current index type
|
||||
current_index_info = conn.execute(
|
||||
text("""
|
||||
SELECT am.amname
|
||||
FROM pg_indexes pi
|
||||
JOIN pg_class c ON c.relname = pi.indexname
|
||||
JOIN pg_am am ON am.oid = c.relam
|
||||
WHERE pi.schemaname = :schema
|
||||
AND pi.tablename = :table_name
|
||||
AND pi.indexname LIKE '%text_search%'
|
||||
"""),
|
||||
{"schema": schema_name, "table_name": table_name},
|
||||
).fetchone()
|
||||
|
||||
current_index_type = current_index_info[0] if current_index_info else None
|
||||
|
||||
# Check if column and index types match target
|
||||
column_matches = current_column_type == target_column_type
|
||||
index_matches = current_index_type == target_index_type if current_index_type else False
|
||||
|
||||
if not (column_matches and index_matches):
|
||||
logger.info(
|
||||
f"Text search mismatch on {table_name}: "
|
||||
f"column={current_column_type} (want {target_column_type}), "
|
||||
f"index={current_index_type} (want {target_index_type})"
|
||||
)
|
||||
mismatched_tables.append((table_name, current_column_type, current_index_type))
|
||||
|
||||
# Check if table has data
|
||||
row_count = conn.execute(text(f"SELECT COUNT(*) FROM {schema_name}.{table_name}")).scalar()
|
||||
|
||||
if row_count > 0:
|
||||
tables_with_data.append((table_name, row_count))
|
||||
else:
|
||||
logger.debug(f"Text search OK for {table_name}: {current_column_type}/{current_index_type}")
|
||||
|
||||
# If no mismatches, we're done
|
||||
if not mismatched_tables:
|
||||
logger.debug(f"All text search columns/indexes match configured extension: {text_search_extension}")
|
||||
return
|
||||
|
||||
# If there's data in any mismatched table, raise error
|
||||
if tables_with_data:
|
||||
table_list = ", ".join([f"{table}({count} rows)" for table, count in tables_with_data])
|
||||
# Detect current extension from column type
|
||||
current_col_type = mismatched_tables[0][1]
|
||||
if current_col_type == "tsvector":
|
||||
current_ext = "native"
|
||||
elif current_col_type == "bm25vector":
|
||||
current_ext = "vchord"
|
||||
elif current_col_type == "text":
|
||||
current_ext = "pg_textsearch"
|
||||
else:
|
||||
current_ext = "unknown"
|
||||
raise RuntimeError(
|
||||
f"Cannot change text search extension from {current_ext} to {text_search_extension}: "
|
||||
f"the following tables contain data: {table_list}. "
|
||||
f"To change text search extension, you must either:\n"
|
||||
f" 1. Clear all data: DELETE FROM {schema_name}.memory_units; "
|
||||
f"DELETE FROM {schema_name}.reflections; then restart\n"
|
||||
f" 2. Use the current text search extension (set HINDSIGHT_API_TEXT_SEARCH_EXTENSION='{current_ext}')"
|
||||
)
|
||||
|
||||
# Tables are empty, safe to recreate columns/indexes
|
||||
logger.info(f"Recreating text search columns/indexes for {text_search_extension}")
|
||||
|
||||
for table_name, current_col_type, current_idx_type in mismatched_tables:
|
||||
# Drop existing index if it exists
|
||||
if current_idx_type:
|
||||
logger.info(f"Dropping {current_idx_type} index on {table_name}")
|
||||
conn.execute(
|
||||
text(f"""
|
||||
DROP INDEX IF EXISTS {schema_name}.idx_{table_name.replace(".", "_")}_text_search
|
||||
""")
|
||||
)
|
||||
|
||||
# Drop existing column if it exists
|
||||
if current_col_type:
|
||||
logger.info(f"Dropping {current_col_type} column on {table_name}")
|
||||
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} DROP COLUMN IF EXISTS search_vector"))
|
||||
|
||||
# Create new column with appropriate type
|
||||
if text_search_extension == "vchord":
|
||||
logger.info(f"Creating bm25vector column on {table_name}")
|
||||
# Note: vchord_bm25 extension creates types in bm25_catalog schema
|
||||
conn.execute(
|
||||
text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector bm25_catalog.bm25vector")
|
||||
)
|
||||
|
||||
# Create BM25 index
|
||||
logger.info(f"Creating BM25 index on {table_name}")
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
|
||||
ON {schema_name}.{table_name}
|
||||
USING bm25 (search_vector bm25_catalog.bm25_ops)
|
||||
""")
|
||||
)
|
||||
elif text_search_extension == "pg_textsearch":
|
||||
logger.info(f"Creating TEXT column on {table_name}")
|
||||
# Dummy TEXT column for consistency (indexes operate on base columns)
|
||||
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector TEXT"))
|
||||
|
||||
# Create BM25 index on expression
|
||||
logger.info(f"Creating BM25 index on {table_name}")
|
||||
# Different expression for each table
|
||||
if table_name == "memory_units":
|
||||
index_expr = "(COALESCE(text, '') || ' ' || COALESCE(context, ''))"
|
||||
else: # reflections
|
||||
index_expr = "(COALESCE(name, '') || ' ' || content)"
|
||||
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
|
||||
ON {schema_name}.{table_name}
|
||||
USING bm25({index_expr})
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
)
|
||||
else: # native
|
||||
logger.info(f"Creating tsvector column on {table_name}")
|
||||
# Different GENERATED expression for each table
|
||||
if table_name == "memory_units":
|
||||
generated_expr = "to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))"
|
||||
else: # reflections
|
||||
generated_expr = "to_tsvector('english', COALESCE(name, '') || ' ' || content)"
|
||||
|
||||
conn.execute(
|
||||
text(f"""
|
||||
ALTER TABLE {schema_name}.{table_name}
|
||||
ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS ({generated_expr}) STORED
|
||||
""")
|
||||
)
|
||||
|
||||
# Create GIN index
|
||||
logger.info(f"Creating GIN index on {table_name}")
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
|
||||
ON {schema_name}.{table_name}
|
||||
USING gin(search_vector)
|
||||
""")
|
||||
)
|
||||
|
||||
conn.commit()
|
||||
logger.info(f"Successfully migrated text search to {text_search_extension}")
|
||||
|
||||
@@ -376,12 +376,7 @@ class WorkerPoller:
|
||||
del self._in_flight_by_type[operation_type]
|
||||
|
||||
async def _execute_task_inner(self, task: ClaimedTask):
|
||||
"""Inner task execution with error handling.
|
||||
|
||||
Note: The executor (MemoryEngine.execute_task) handles status marking internally
|
||||
(marking operations as completed/failed and handling retries). This method should
|
||||
NOT override those status updates.
|
||||
"""
|
||||
"""Inner task execution with error handling."""
|
||||
task_type = task.task_dict.get("type", "unknown")
|
||||
bank_id = task.task_dict.get("bank_id", "unknown")
|
||||
|
||||
@@ -391,12 +386,12 @@ class WorkerPoller:
|
||||
if task.schema:
|
||||
task.task_dict["_schema"] = task.schema
|
||||
await self._executor(task.task_dict)
|
||||
logger.debug(f"Task {task.operation_id} execution finished")
|
||||
await self._mark_completed(task.operation_id, task.schema)
|
||||
logger.debug(f"Task {task.operation_id} completed successfully")
|
||||
except Exception as e:
|
||||
# The executor should handle its own errors, but if an unexpected exception
|
||||
# propagates (e.g., from schema setup), log it as a warning
|
||||
logger.error(f"Task {task.operation_id} raised unexpected exception: {e}")
|
||||
traceback.print_exc()
|
||||
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)
|
||||
|
||||
async def recover_own_tasks(self) -> int:
|
||||
"""
|
||||
@@ -406,8 +401,6 @@ class WorkerPoller:
|
||||
On startup, we reset any tasks stuck in 'processing' for this worker_id
|
||||
back to 'pending' so they can be picked up again.
|
||||
|
||||
Also recovers batch API operations that were in-flight.
|
||||
|
||||
If tenant_extension is configured, recovers across all tenant schemas.
|
||||
|
||||
Returns:
|
||||
@@ -420,16 +413,11 @@ class WorkerPoller:
|
||||
try:
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
# First, recover batch API operations (before resetting worker tasks)
|
||||
batch_count = await self._recover_batch_operations(schema)
|
||||
total_count += batch_count
|
||||
|
||||
# Then reset normal worker tasks
|
||||
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 AND result_metadata->>'batch_id' IS NULL
|
||||
WHERE status = 'processing' AND worker_id = $1
|
||||
""",
|
||||
self._worker_id,
|
||||
)
|
||||
@@ -446,80 +434,6 @@ class WorkerPoller:
|
||||
logger.info(f"Worker {self._worker_id} recovered {total_count} stale tasks from previous run")
|
||||
return total_count
|
||||
|
||||
async def _recover_batch_operations(self, schema: str | None) -> int:
|
||||
"""
|
||||
Recover batch API operations that were in-flight when worker crashed.
|
||||
|
||||
Finds operations with batch_id in metadata and re-submits them as tasks
|
||||
so polling can resume.
|
||||
|
||||
Args:
|
||||
schema: Database schema to recover from
|
||||
|
||||
Returns:
|
||||
Number of batch operations recovered
|
||||
"""
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
try:
|
||||
# Find operations with batch_id in metadata (batch API operations)
|
||||
rows = await self._pool.fetch(
|
||||
f"""
|
||||
SELECT operation_id, task_payload, result_metadata
|
||||
FROM {table}
|
||||
WHERE status = 'processing'
|
||||
AND result_metadata ? 'batch_id'
|
||||
AND task_payload IS NOT NULL
|
||||
"""
|
||||
)
|
||||
|
||||
if not rows:
|
||||
return 0
|
||||
|
||||
recovered = 0
|
||||
for row in rows:
|
||||
operation_id = str(row["operation_id"])
|
||||
task_payload = row["task_payload"]
|
||||
result_metadata = row["result_metadata"]
|
||||
|
||||
# Parse metadata
|
||||
if isinstance(result_metadata, str):
|
||||
result_metadata = json.loads(result_metadata)
|
||||
|
||||
batch_id = result_metadata.get("batch_id")
|
||||
batch_provider = result_metadata.get("batch_provider", "openai")
|
||||
|
||||
logger.info(
|
||||
f"Recovering batch operation: operation_id={operation_id}, batch_id={batch_id}, provider={batch_provider}"
|
||||
)
|
||||
|
||||
# Parse task_payload
|
||||
if isinstance(task_payload, str):
|
||||
task_dict = json.loads(task_payload)
|
||||
else:
|
||||
task_dict = task_payload
|
||||
|
||||
# Mark operation as ready for re-processing
|
||||
# Reset to pending with task_payload intact so worker picks it up again
|
||||
await self._pool.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
operation_id,
|
||||
)
|
||||
|
||||
recovered += 1
|
||||
logger.info(f"Batch operation {operation_id} reset to pending for re-processing")
|
||||
|
||||
return recovered
|
||||
|
||||
except Exception as e:
|
||||
schema_display = f'"{schema}"' if schema else str(schema)
|
||||
logger.error(f"Failed to recover batch operations for schema {schema_display}: {e}")
|
||||
return 0
|
||||
|
||||
async def run(self):
|
||||
"""
|
||||
Main polling loop with fire-and-forget task execution.
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api"
|
||||
version = "0.4.11"
|
||||
version = "0.4.10"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -42,9 +42,6 @@ dependencies = [
|
||||
"typer>=0.9.0",
|
||||
"cohere>=5.0.0",
|
||||
"flashrank>=0.2.0",
|
||||
"litellm>=1.0.0",
|
||||
"markitdown[pdf,docx,pptx,xlsx,xls]>=0.1.4", # File to markdown conversion
|
||||
"obstore>=0.4.0", # S3/GCS/Azure object storage client (Rust-backed)
|
||||
# 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
|
||||
@@ -67,7 +64,6 @@ test = [
|
||||
"pytest-timeout>=2.4.0",
|
||||
"pytest-xdist>=3.0.0",
|
||||
"filelock>=3.20.1", # TOCTOU race condition fix
|
||||
"testcontainers>=4.0.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
@@ -117,7 +113,6 @@ dev = [
|
||||
"filelock>=3.20.1", # TOCTOU race condition fix
|
||||
"ruff>=0.8.0",
|
||||
"ty>=0.0.1",
|
||||
"testcontainers>=4.0.0",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
|
||||
@@ -1,423 +0,0 @@
|
||||
"""Test async batch retain with smart batching and parent-child operations."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.extensions import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplicate_document_ids_rejected_async(memory, request_context):
|
||||
"""Test that async retain rejects batches with duplicate document_ids."""
|
||||
bank_id = "test_duplicate_async"
|
||||
contents = [
|
||||
{"content": "First item", "document_id": "doc1"},
|
||||
{"content": "Second item", "document_id": "doc2"},
|
||||
{"content": "Third item", "document_id": "doc1"}, # Duplicate!
|
||||
]
|
||||
|
||||
# Should raise ValueError due to duplicate document_ids
|
||||
with pytest.raises(ValueError, match="duplicate document_ids.*doc1"):
|
||||
await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplicate_document_ids_rejected_sync(memory, request_context):
|
||||
"""Test that sync retain also rejects batches with duplicate document_ids."""
|
||||
bank_id = "test_duplicate_sync"
|
||||
contents = [
|
||||
{"content": "First item", "document_id": "doc1"},
|
||||
{"content": "Second item", "document_id": "doc1"}, # Duplicate!
|
||||
]
|
||||
|
||||
# Should raise ValueError due to duplicate document_ids
|
||||
with pytest.raises(ValueError, match="duplicate document_ids.*doc1"):
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_small_async_batch_no_splitting(memory, request_context):
|
||||
"""Test that small async batches create parent with single child (simplified code path)."""
|
||||
bank_id = "test_small_async"
|
||||
contents = [{"content": "Alice works at Google", "document_id": f"doc{i}"} for i in range(5)]
|
||||
|
||||
# Calculate total chars (should be well under threshold)
|
||||
total_chars = sum(len(item["content"]) for item in contents)
|
||||
assert total_chars < 10_000, "Test batch should be small"
|
||||
|
||||
# Submit async retain
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify we got an operation_id back
|
||||
assert "operation_id" in result
|
||||
assert "items_count" in result
|
||||
assert result["items_count"] == 5
|
||||
|
||||
operation_id = result["operation_id"]
|
||||
|
||||
# Wait for task to complete (SyncTaskBackend executes immediately)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Check operation status
|
||||
status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=operation_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should be a parent operation with single child (simplified code path)
|
||||
assert status["status"] == "completed"
|
||||
assert status["operation_type"] == "batch_retain"
|
||||
assert "child_operations" in status
|
||||
assert status["result_metadata"]["num_sub_batches"] == 1 # Single sub-batch
|
||||
assert len(status["child_operations"]) == 1
|
||||
assert status["child_operations"][0]["status"] == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_large_async_batch_auto_splits(memory, request_context):
|
||||
"""Test that large async batches automatically split into sub-batches with parent operation."""
|
||||
from hindsight_api.engine.memory_engine import count_tokens
|
||||
|
||||
bank_id = "test_large_async"
|
||||
|
||||
# Create a large batch that exceeds the threshold (10k tokens default)
|
||||
# Repeating "A"s gets heavily compressed by tokenizer, use varied content
|
||||
# Use ~22k chars per item = ~5.5k tokens per item, 2 items = ~11k tokens total (exceeds 10k)
|
||||
large_content = "The quick brown fox jumps over the lazy dog. " * 500 # ~22k chars = ~5.5k tokens
|
||||
contents = [{"content": large_content + f" item {i}", "document_id": f"doc{i}"} for i in range(2)]
|
||||
|
||||
# Calculate total tokens (should exceed threshold)
|
||||
total_tokens = sum(count_tokens(item["content"]) for item in contents)
|
||||
assert total_tokens > 10_000, "Test batch should exceed threshold"
|
||||
|
||||
# Submit async retain
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify we got an operation_id back
|
||||
assert "operation_id" in result
|
||||
assert "items_count" in result
|
||||
assert result["items_count"] == 2
|
||||
|
||||
parent_operation_id = result["operation_id"]
|
||||
|
||||
# Wait for tasks to complete
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Check parent operation status
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=parent_operation_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should be a parent operation with children
|
||||
assert parent_status["operation_type"] == "batch_retain"
|
||||
assert "child_operations" in parent_status
|
||||
assert "num_sub_batches" in parent_status["result_metadata"]
|
||||
assert parent_status["result_metadata"]["num_sub_batches"] >= 2 # Should split into at least 2 batches
|
||||
assert parent_status["result_metadata"]["items_count"] == 2
|
||||
|
||||
# Verify child operations
|
||||
child_ops = parent_status["child_operations"]
|
||||
assert len(child_ops) >= 2, "Should have at least 2 child operations"
|
||||
|
||||
# All children should be completed (SyncTaskBackend executes immediately)
|
||||
for child in child_ops:
|
||||
assert child["status"] == "completed"
|
||||
assert child["sub_batch_index"] is not None
|
||||
assert child["items_count"] > 0
|
||||
|
||||
# Parent status should be aggregated as "completed"
|
||||
assert parent_status["status"] == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parent_operation_status_aggregation_pending(memory, request_context):
|
||||
"""Test that parent operation shows 'pending' when children are pending."""
|
||||
bank_id = "test_parent_pending"
|
||||
pool = await memory._get_pool()
|
||||
|
||||
# Manually create a parent operation
|
||||
parent_id = uuid.uuid4()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
parent_id,
|
||||
bank_id,
|
||||
"batch_retain",
|
||||
json.dumps({"items_count": 20, "num_sub_batches": 2, "is_parent": True}),
|
||||
"pending",
|
||||
)
|
||||
|
||||
# Create 2 child operations - one completed, one pending
|
||||
child1_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child1_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 1,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"completed",
|
||||
)
|
||||
|
||||
child2_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child2_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 2,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"pending",
|
||||
)
|
||||
|
||||
# Check parent status
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=str(parent_id),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Parent should aggregate as "pending" since one child is still pending
|
||||
assert parent_status["status"] == "pending"
|
||||
assert len(parent_status["child_operations"]) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parent_operation_status_aggregation_failed(memory, request_context):
|
||||
"""Test that parent operation shows 'failed' when any child fails."""
|
||||
bank_id = "test_parent_failed"
|
||||
pool = await memory._get_pool()
|
||||
|
||||
# Manually create a parent operation
|
||||
parent_id = uuid.uuid4()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
parent_id,
|
||||
bank_id,
|
||||
"batch_retain",
|
||||
json.dumps({"items_count": 20, "num_sub_batches": 2, "is_parent": True}),
|
||||
"pending",
|
||||
)
|
||||
|
||||
# Create 2 child operations - one completed, one failed
|
||||
child1_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child1_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 1,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"completed",
|
||||
)
|
||||
|
||||
child2_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status, error_message)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
""",
|
||||
child2_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 2,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"failed",
|
||||
"Test error",
|
||||
)
|
||||
|
||||
# Check parent status
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=str(parent_id),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Parent should aggregate as "failed" since one child failed
|
||||
assert parent_status["status"] == "failed"
|
||||
assert len(parent_status["child_operations"]) == 2
|
||||
|
||||
# Verify child with error is included
|
||||
failed_child = [c for c in parent_status["child_operations"] if c["status"] == "failed"][0]
|
||||
assert failed_child["error_message"] == "Test error"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parent_operation_status_aggregation_completed(memory, request_context):
|
||||
"""Test that parent operation shows 'completed' when all children are completed."""
|
||||
bank_id = "test_parent_completed"
|
||||
pool = await memory._get_pool()
|
||||
|
||||
# Manually create a parent operation
|
||||
parent_id = uuid.uuid4()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
parent_id,
|
||||
bank_id,
|
||||
"batch_retain",
|
||||
json.dumps({"items_count": 20, "num_sub_batches": 2, "is_parent": True}),
|
||||
"pending",
|
||||
)
|
||||
|
||||
# Create 2 child operations - both completed
|
||||
child1_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child1_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 1,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"completed",
|
||||
)
|
||||
|
||||
child2_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child2_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 2,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"completed",
|
||||
)
|
||||
|
||||
# Check parent status
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=str(parent_id),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Parent should aggregate as "completed" since all children are completed
|
||||
assert parent_status["status"] == "completed"
|
||||
assert len(parent_status["child_operations"]) == 2
|
||||
assert all(c["status"] == "completed" for c in parent_status["child_operations"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_retain_batch_tokens_respected(memory, request_context):
|
||||
"""Test that the retain_batch_tokens config setting is respected."""
|
||||
from hindsight_api.config import get_config
|
||||
from hindsight_api.engine.memory_engine import count_tokens
|
||||
|
||||
bank_id = "test_config_batch_tokens"
|
||||
config = get_config()
|
||||
|
||||
# Check that config has the retain_batch_tokens setting
|
||||
assert hasattr(config, "retain_batch_tokens")
|
||||
assert config.retain_batch_tokens > 0
|
||||
|
||||
# Create a batch that's just under the threshold
|
||||
# Use content that produces roughly half the token limit per item
|
||||
content_size = config.retain_batch_tokens * 2 # chars (rough estimate: 1 token ~= 4 chars)
|
||||
contents = [{"content": "A" * content_size, "document_id": f"doc{i}"} for i in range(2)]
|
||||
|
||||
total_tokens = sum(count_tokens(item["content"]) for item in contents)
|
||||
# Should be equal to threshold (boundary case, no splitting since we use > not >=)
|
||||
assert total_tokens <= config.retain_batch_tokens
|
||||
|
||||
# Submit - should NOT split
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Wait for completion
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Check status - should be a parent with single child (even for small batches)
|
||||
status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=result["operation_id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Even small batches use parent-child pattern now (simpler code path)
|
||||
assert "child_operations" in status
|
||||
assert status["result_metadata"]["num_sub_batches"] == 1
|
||||
@@ -1,93 +0,0 @@
|
||||
"""Unit tests for async retain tag propagation."""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_async_retain_includes_document_tags_in_task_payload():
|
||||
"""submit_async_retain should include document_tags in queued task payload."""
|
||||
engine = MemoryEngine.__new__(MemoryEngine)
|
||||
engine._initialized = True
|
||||
engine._authenticate_tenant = AsyncMock()
|
||||
engine._submit_async_operation = AsyncMock(return_value={"operation_id": "op-1"})
|
||||
|
||||
# Mock the pool and connection for parent operation creation
|
||||
mock_conn = AsyncMock()
|
||||
mock_conn.execute = AsyncMock()
|
||||
mock_conn.transaction = MagicMock()
|
||||
mock_conn.transaction.return_value.__aenter__ = AsyncMock()
|
||||
mock_conn.transaction.return_value.__aexit__ = AsyncMock()
|
||||
|
||||
mock_pool = AsyncMock()
|
||||
mock_pool.acquire = AsyncMock(return_value=mock_conn)
|
||||
mock_pool.release = AsyncMock()
|
||||
|
||||
engine._get_pool = AsyncMock(return_value=mock_pool)
|
||||
|
||||
request_context = RequestContext(tenant_id="tenant-a", api_key_id="key-a")
|
||||
contents = [{"content": "Async retain payload test."}]
|
||||
document_tags = ["scope:tools", "user:alice"]
|
||||
|
||||
result = await MemoryEngine.submit_async_retain(
|
||||
engine,
|
||||
bank_id="bank-1",
|
||||
contents=contents,
|
||||
document_tags=document_tags,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Check result structure
|
||||
assert "operation_id" in result
|
||||
assert "items_count" in result
|
||||
assert result["items_count"] == 1
|
||||
|
||||
# Verify authentication was called
|
||||
engine._authenticate_tenant.assert_awaited_once_with(request_context)
|
||||
|
||||
# Verify child operation was submitted
|
||||
engine._submit_async_operation.assert_awaited_once()
|
||||
|
||||
# Verify child operation payload contains document_tags
|
||||
kwargs = engine._submit_async_operation.await_args.kwargs
|
||||
assert kwargs["bank_id"] == "bank-1"
|
||||
assert kwargs["operation_type"] == "retain"
|
||||
assert kwargs["task_type"] == "batch_retain"
|
||||
assert kwargs["task_payload"]["contents"] == contents
|
||||
assert kwargs["task_payload"]["document_tags"] == document_tags
|
||||
assert kwargs["task_payload"]["_tenant_id"] == "tenant-a"
|
||||
assert kwargs["task_payload"]["_api_key_id"] == "key-a"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_batch_retain_forwards_document_tags_to_retain_batch_async():
|
||||
"""Worker handler should forward document_tags from task payload."""
|
||||
engine = MemoryEngine.__new__(MemoryEngine)
|
||||
engine._initialized = True
|
||||
engine.retain_batch_async = AsyncMock(return_value={"items_count": 1})
|
||||
|
||||
task_dict = {
|
||||
"bank_id": "bank-1",
|
||||
"contents": [{"content": "Forward tags test."}],
|
||||
"document_tags": ["scope:client"],
|
||||
"_tenant_id": "tenant-a",
|
||||
"_api_key_id": "key-a",
|
||||
}
|
||||
|
||||
await MemoryEngine._handle_batch_retain(engine, task_dict)
|
||||
|
||||
engine.retain_batch_async.assert_awaited_once()
|
||||
kwargs = engine.retain_batch_async.await_args.kwargs
|
||||
assert kwargs["bank_id"] == "bank-1"
|
||||
assert kwargs["contents"] == task_dict["contents"]
|
||||
assert kwargs["document_tags"] == ["scope:client"]
|
||||
|
||||
request_context = kwargs["request_context"]
|
||||
assert request_context.internal is True
|
||||
assert request_context.user_initiated is True
|
||||
assert request_context.tenant_id == "tenant-a"
|
||||
assert request_context.api_key_id == "key-a"
|
||||
@@ -1,189 +0,0 @@
|
||||
"""
|
||||
Integration test for API base path support.
|
||||
|
||||
Tests that the API works correctly when deployed with a base path (e.g., /hindsight)
|
||||
for reverse proxy deployments.
|
||||
"""
|
||||
import os
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import httpx
|
||||
from hindsight_api.api import create_app
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client_with_base_path(memory):
|
||||
"""Create an async test client for the FastAPI app with a base path."""
|
||||
# Set base path in environment
|
||||
base_path = "/hindsight"
|
||||
os.environ["HINDSIGHT_API_BASE_PATH"] = base_path
|
||||
|
||||
# Clear config cache to force reload with new base_path
|
||||
clear_config_cache()
|
||||
|
||||
# Memory is already initialized by the conftest fixture (with migrations)
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
|
||||
# Use base_url with base path
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport,
|
||||
base_url=f"http://test{base_path}"
|
||||
) as client:
|
||||
yield client
|
||||
|
||||
# Cleanup: unset base path
|
||||
os.environ.pop("HINDSIGHT_API_BASE_PATH", None)
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client_without_base_path(memory):
|
||||
"""Create an async test client for the FastAPI app without a base path (root)."""
|
||||
# Ensure no base path is set
|
||||
os.environ.pop("HINDSIGHT_API_BASE_PATH", None)
|
||||
clear_config_cache()
|
||||
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_health_endpoint(api_client_with_base_path):
|
||||
"""Test that health endpoint works with base path."""
|
||||
# With base path set to /hindsight, health should be at /hindsight/health
|
||||
# But since our client base_url is already http://test/hindsight, we request /health
|
||||
response = await api_client_with_base_path.get("/health")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "status" in data
|
||||
assert data["status"] in ["ok", "healthy"] # Accept both formats
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_banks_endpoint(api_client_with_base_path):
|
||||
"""Test that banks endpoint works with base path."""
|
||||
response = await api_client_with_base_path.get("/v1/default/banks")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "banks" in data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_openapi_schema(api_client_with_base_path):
|
||||
"""Test that OpenAPI schema includes correct base path in servers."""
|
||||
response = await api_client_with_base_path.get("/openapi.json")
|
||||
assert response.status_code == 200
|
||||
openapi_schema = response.json()
|
||||
|
||||
# Check that servers array includes base path
|
||||
assert "servers" in openapi_schema
|
||||
servers = openapi_schema["servers"]
|
||||
assert len(servers) > 0
|
||||
# FastAPI should set server URL to the root_path
|
||||
assert servers[0]["url"] == "/hindsight"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_docs_redirect(api_client_with_base_path):
|
||||
"""Test that /docs redirects correctly with base path."""
|
||||
# FastAPI docs endpoint should work
|
||||
response = await api_client_with_base_path.get("/docs", follow_redirects=False)
|
||||
# Should either return 200 (direct) or 307 (redirect to trailing slash)
|
||||
assert response.status_code in [200, 307]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_metrics(api_client_with_base_path):
|
||||
"""Test that metrics endpoint works with base path."""
|
||||
response = await api_client_with_base_path.get("/metrics")
|
||||
assert response.status_code == 200
|
||||
# Metrics should be in Prometheus format
|
||||
assert "# HELP" in response.text or "# TYPE" in response.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_full_workflow(api_client_with_base_path):
|
||||
"""
|
||||
Test a full retain/recall workflow with base path.
|
||||
|
||||
This ensures that all memory operations work correctly when the API
|
||||
is deployed with a base path.
|
||||
"""
|
||||
bank_id = "test_base_path_bank"
|
||||
|
||||
# 1. Create/get bank
|
||||
response = await api_client_with_base_path.get(f"/v1/default/banks/{bank_id}/profile")
|
||||
assert response.status_code == 200
|
||||
|
||||
# 2. Store a memory
|
||||
response = await api_client_with_base_path.post(
|
||||
f"/v1/default/banks/{bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "The API supports base path deployment for reverse proxy use cases.",
|
||||
"context": "testing base path feature"
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["success"] is True
|
||||
|
||||
# 3. Recall the memory
|
||||
response = await api_client_with_base_path.post(
|
||||
f"/v1/default/banks/{bank_id}/memories/recall",
|
||||
json={
|
||||
"query": "base path support"
|
||||
}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
recall_result = response.json()
|
||||
# API returns "results" not "memories"
|
||||
assert "results" in recall_result
|
||||
assert len(recall_result["results"]) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_without_base_path_still_works(api_client_without_base_path):
|
||||
"""
|
||||
Regression test: ensure default behavior (no base path) still works.
|
||||
|
||||
This test verifies that when HINDSIGHT_API_BASE_PATH is not set,
|
||||
the API works at the root path as before.
|
||||
"""
|
||||
# Health check at root
|
||||
response = await api_client_without_base_path.get("/health")
|
||||
assert response.status_code == 200
|
||||
|
||||
# Banks endpoint at root
|
||||
response = await api_client_without_base_path.get("/v1/default/banks")
|
||||
assert response.status_code == 200
|
||||
|
||||
# OpenAPI schema should have empty or "/" server path
|
||||
response = await api_client_without_base_path.get("/openapi.json")
|
||||
assert response.status_code == 200
|
||||
openapi_schema = response.json()
|
||||
servers = openapi_schema.get("servers", [])
|
||||
if servers:
|
||||
# Server URL should be empty string (root) or "/"
|
||||
assert servers[0]["url"] in ["", "/"]
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="MCP endpoint routing with base path needs investigation")
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_mcp_endpoint(api_client_with_base_path):
|
||||
"""Test that MCP endpoint is accessible with base path."""
|
||||
bank_id = "test_mcp_bank"
|
||||
|
||||
# MCP endpoint should be mounted at /mcp/{bank_id}/
|
||||
# The MCP server uses a different protocol, so just check the root exists
|
||||
response = await api_client_with_base_path.get(f"/mcp/{bank_id}/")
|
||||
# MCP may return various status codes, but should not be 404 (not found)
|
||||
# Accept 405 (method not allowed), 400 (bad request), etc.
|
||||
assert response.status_code != 404, "MCP endpoint should exist"
|
||||
@@ -1,508 +0,0 @@
|
||||
"""
|
||||
Test OpenAI Batch API integration for retain fact extraction.
|
||||
|
||||
Tests cover:
|
||||
- Normal batch API flow (submit, poll, complete)
|
||||
- Crash recovery (resume from existing batch_id)
|
||||
- Provider fallback (when batch API not supported)
|
||||
- Worker recovery on restart
|
||||
"""
|
||||
import pytest
|
||||
import asyncio
|
||||
import logging
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.retain.fact_extraction import (
|
||||
extract_facts_from_contents_batch_api,
|
||||
extract_facts_from_contents,
|
||||
RetainContent,
|
||||
)
|
||||
from hindsight_api.config import HindsightConfig
|
||||
from hindsight_api.engine.llm_wrapper import create_llm_provider
|
||||
from hindsight_api.worker.poller import WorkerPoller
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_llm_config():
|
||||
"""Create a mock LLM config with batch API support."""
|
||||
mock = MagicMock()
|
||||
mock.provider = "openai"
|
||||
mock.model = "gpt-4o-mini"
|
||||
mock._provider_impl = AsyncMock()
|
||||
return mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_contents():
|
||||
"""Create test content for fact extraction."""
|
||||
return [
|
||||
RetainContent(
|
||||
content="Alice is a senior software engineer at TechCorp. She specializes in distributed systems.",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
context="team overview",
|
||||
),
|
||||
RetainContent(
|
||||
content="Bob joined the team last month as a junior developer. He is learning React.",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
context="team overview",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def hindsight_config():
|
||||
"""Create test config with batch API enabled."""
|
||||
config = HindsightConfig.from_env()
|
||||
config.retain_batch_enabled = True
|
||||
config.retain_batch_poll_interval_seconds = 1 # Fast polling for tests
|
||||
config.retain_chunk_size = 4000
|
||||
config.retain_extraction_mode = "concise"
|
||||
config.retain_extract_causal_links = False
|
||||
return config
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_normal_flow(mock_llm_config, test_contents, hindsight_config, memory, request_context):
|
||||
"""Test normal batch API flow: submit, poll, complete."""
|
||||
bank_id = f"test_batch_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Mock batch API responses
|
||||
batch_id = "batch_test123"
|
||||
|
||||
# Mock supports_batch_api
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)
|
||||
|
||||
# Mock submit_batch - returns batch metadata
|
||||
mock_llm_config._provider_impl.submit_batch = AsyncMock(
|
||||
return_value={
|
||||
"batch_id": batch_id,
|
||||
"status": "validating",
|
||||
"request_counts": {"total": 2, "completed": 0, "failed": 0},
|
||||
}
|
||||
)
|
||||
|
||||
# Mock get_batch_status - simulate polling sequence
|
||||
status_sequence = [
|
||||
{"status": "in_progress", "request_counts": {"total": 2, "completed": 1, "failed": 0}},
|
||||
{"status": "completed", "request_counts": {"total": 2, "completed": 2, "failed": 0}},
|
||||
]
|
||||
mock_llm_config._provider_impl.get_batch_status = AsyncMock(side_effect=status_sequence)
|
||||
|
||||
# Mock retrieve_batch_results - returns fact extraction results
|
||||
mock_results = [
|
||||
{
|
||||
"custom_id": "chunk_0",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": json.dumps({
|
||||
"facts": [
|
||||
{
|
||||
"what": "Alice is a senior software engineer at TechCorp",
|
||||
"when": "present",
|
||||
"where": "TechCorp",
|
||||
"who": "Alice",
|
||||
"why": "Professional background information",
|
||||
"fact_type": "world",
|
||||
"fact_kind": "conversation",
|
||||
}
|
||||
]
|
||||
})
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "chunk_1",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": json.dumps({
|
||||
"facts": [
|
||||
{
|
||||
"what": "Bob joined the team last month as a junior developer",
|
||||
"when": "last month",
|
||||
"where": "team",
|
||||
"who": "Bob",
|
||||
"why": "New team member information",
|
||||
"fact_type": "world",
|
||||
"fact_kind": "conversation",
|
||||
}
|
||||
]
|
||||
})
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
}
|
||||
},
|
||||
},
|
||||
]
|
||||
mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(return_value=mock_results)
|
||||
|
||||
# Call batch API extraction
|
||||
facts, chunks, usage = await extract_facts_from_contents_batch_api(
|
||||
contents=test_contents,
|
||||
llm_config=mock_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=hindsight_config,
|
||||
pool=None, # No DB pool for this test
|
||||
operation_id=None,
|
||||
schema=None,
|
||||
)
|
||||
|
||||
# Verify results
|
||||
assert len(facts) == 2, "Should extract 2 facts (one per chunk)"
|
||||
# Facts are ExtractedFact objects with .fact_text field
|
||||
assert "Alice" in facts[0].fact_text and "senior software engineer" in facts[0].fact_text
|
||||
assert "Bob" in facts[1].fact_text and "junior developer" in facts[1].fact_text
|
||||
|
||||
# Verify chunks metadata
|
||||
assert len(chunks) == 2, "Should have 2 chunks metadata"
|
||||
assert chunks[0].fact_count == 1
|
||||
assert chunks[1].fact_count == 1
|
||||
|
||||
# Verify token usage
|
||||
assert usage.input_tokens == 200 # 100 per chunk
|
||||
assert usage.output_tokens == 100 # 50 per chunk
|
||||
assert usage.total_tokens == 300
|
||||
|
||||
# Verify API calls
|
||||
mock_llm_config._provider_impl.submit_batch.assert_called_once()
|
||||
assert mock_llm_config._provider_impl.get_batch_status.call_count == 2
|
||||
mock_llm_config._provider_impl.retrieve_batch_results.assert_called_once_with(batch_id)
|
||||
|
||||
logger.info("✅ Normal batch API flow test passed")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
try:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_crash_recovery(mock_llm_config, test_contents, hindsight_config, memory, request_context):
|
||||
"""Test crash recovery: resume polling from existing batch_id."""
|
||||
bank_id = f"test_crash_{datetime.now(timezone.utc).timestamp()}"
|
||||
operation_id = str(uuid.uuid4()) # Must be UUID for async_operations table
|
||||
|
||||
try:
|
||||
# Ensure bank exists
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Setup: Store batch_id in async_operations table (simulates partial execution)
|
||||
batch_id = "batch_recovered_456"
|
||||
pool = memory._pool
|
||||
schema = request_context.tenant_id
|
||||
|
||||
from hindsight_api.engine.task_backend import fq_table
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
# Create operation with batch_id already stored
|
||||
await pool.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (operation_id, operation_type, bank_id, status, result_metadata)
|
||||
VALUES ($1, 'retain', $2, 'processing', $3::jsonb)
|
||||
""",
|
||||
operation_id,
|
||||
bank_id,
|
||||
json.dumps({
|
||||
"batch_id": batch_id,
|
||||
"batch_provider": "openai",
|
||||
"chunk_count": 2,
|
||||
}),
|
||||
)
|
||||
|
||||
# Mock batch API responses for resume scenario
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)
|
||||
|
||||
# Mock get_batch_status - batch already in progress
|
||||
mock_llm_config._provider_impl.get_batch_status = AsyncMock(
|
||||
return_value={
|
||||
"status": "completed",
|
||||
"request_counts": {"total": 2, "completed": 2, "failed": 0},
|
||||
}
|
||||
)
|
||||
|
||||
# Mock retrieve_batch_results
|
||||
mock_results = [
|
||||
{
|
||||
"custom_id": "chunk_0",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": json.dumps({
|
||||
"facts": [
|
||||
{
|
||||
"what": "Alice is a senior software engineer",
|
||||
"when": "present",
|
||||
"where": "TechCorp",
|
||||
"who": "Alice",
|
||||
"why": "Background",
|
||||
"fact_type": "world",
|
||||
"fact_kind": "conversation",
|
||||
}
|
||||
]
|
||||
})
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "chunk_1",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": json.dumps({
|
||||
"facts": [
|
||||
{
|
||||
"what": "Bob is a junior developer",
|
||||
"when": "last month",
|
||||
"where": "team",
|
||||
"who": "Bob",
|
||||
"why": "New member",
|
||||
"fact_type": "world",
|
||||
"fact_kind": "conversation",
|
||||
}
|
||||
]
|
||||
})
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
}
|
||||
},
|
||||
},
|
||||
]
|
||||
mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(return_value=mock_results)
|
||||
|
||||
# Call batch API extraction with operation_id (crash recovery scenario)
|
||||
facts, chunks, usage = await extract_facts_from_contents_batch_api(
|
||||
contents=test_contents,
|
||||
llm_config=mock_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=hindsight_config,
|
||||
pool=pool,
|
||||
operation_id=operation_id, # Provides crash recovery context
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Verify results
|
||||
assert len(facts) == 2, "Should extract 2 facts after recovery"
|
||||
|
||||
# CRITICAL: Verify submit_batch was NOT called (because batch_id already exists)
|
||||
mock_llm_config._provider_impl.submit_batch.assert_not_called()
|
||||
|
||||
# Verify get_batch_status WAS called (polling resumed)
|
||||
mock_llm_config._provider_impl.get_batch_status.assert_called()
|
||||
|
||||
# Verify retrieve_batch_results was called with the recovered batch_id
|
||||
mock_llm_config._provider_impl.retrieve_batch_results.assert_called_once_with(batch_id)
|
||||
|
||||
logger.info("✅ Crash recovery test passed - resumed polling without re-submission")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
try:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_fallback_unsupported_provider(mock_llm_config, test_contents, hindsight_config):
|
||||
"""Test fallback to sync mode when provider doesn't support batch API."""
|
||||
|
||||
# Mock provider that doesn't support batch API
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=False)
|
||||
mock_llm_config.provider = "groq" # Example of provider
|
||||
|
||||
# Patch the sync mode function to verify it's called
|
||||
with patch(
|
||||
"hindsight_api.engine.retain.fact_extraction.extract_facts_from_contents"
|
||||
) as mock_sync_extract:
|
||||
mock_sync_extract.return_value = ([], [], MagicMock())
|
||||
|
||||
# Call batch API extraction (should fallback to sync)
|
||||
await extract_facts_from_contents_batch_api(
|
||||
contents=test_contents,
|
||||
llm_config=mock_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=hindsight_config,
|
||||
pool=None,
|
||||
operation_id=None,
|
||||
schema=None,
|
||||
)
|
||||
|
||||
# Verify fallback occurred
|
||||
mock_sync_extract.assert_called_once()
|
||||
|
||||
# Verify batch API methods were NOT called
|
||||
mock_llm_config._provider_impl.submit_batch.assert_not_called()
|
||||
|
||||
logger.info("✅ Fallback to sync mode test passed")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_batch_recovery(memory, request_context):
|
||||
"""Test that WorkerPoller._recover_batch_operations finds and resets orphaned batches."""
|
||||
bank_id = f"test_worker_recovery_{datetime.now(timezone.utc).timestamp()}"
|
||||
operation_id = str(uuid.uuid4()) # Must be UUID for async_operations table
|
||||
|
||||
try:
|
||||
# Ensure bank exists
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
pool = memory._pool
|
||||
schema = request_context.tenant_id
|
||||
|
||||
from hindsight_api.engine.task_backend import fq_table
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
# Create orphaned batch operation (simulates worker crash during polling)
|
||||
batch_id = "batch_orphaned_999"
|
||||
task_payload = {
|
||||
"operation_type": "retain",
|
||||
"bank_id": bank_id,
|
||||
"contents": [{"content": "test", "event_date": "2024-01-15T00:00:00Z"}],
|
||||
}
|
||||
|
||||
await pool.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (operation_id, operation_type, bank_id, status, worker_id, result_metadata, task_payload)
|
||||
VALUES ($1, 'retain', $2, 'processing', 'worker_crashed', $3::jsonb, $4::jsonb)
|
||||
""",
|
||||
operation_id,
|
||||
bank_id,
|
||||
json.dumps({
|
||||
"batch_id": batch_id,
|
||||
"batch_provider": "openai",
|
||||
"chunk_count": 1,
|
||||
}),
|
||||
json.dumps(task_payload),
|
||||
)
|
||||
|
||||
# Create WorkerPoller
|
||||
from hindsight_api.extensions.builtin.tenant import DefaultTenantExtension
|
||||
tenant_extension = DefaultTenantExtension(config={"schema": schema} if schema else {})
|
||||
|
||||
poller = WorkerPoller(
|
||||
pool=pool,
|
||||
worker_id="test_worker_recovery",
|
||||
executor=memory,
|
||||
poll_interval_ms=100,
|
||||
max_retries=3,
|
||||
schema=schema,
|
||||
tenant_extension=tenant_extension,
|
||||
max_slots=5,
|
||||
consolidation_max_slots=2,
|
||||
)
|
||||
|
||||
# Run recovery
|
||||
recovered_count = await poller._recover_batch_operations(schema)
|
||||
|
||||
# Verify recovery
|
||||
assert recovered_count == 1, "Should recover 1 batch operation"
|
||||
|
||||
# Verify operation was reset to pending
|
||||
row = await pool.fetchrow(
|
||||
f"SELECT status, worker_id FROM {table} WHERE operation_id = $1",
|
||||
operation_id,
|
||||
)
|
||||
|
||||
assert row["status"] == "pending", "Operation should be reset to pending"
|
||||
assert row["worker_id"] is None, "Worker ID should be cleared"
|
||||
|
||||
logger.info("✅ Worker batch recovery test passed")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
try:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_via_extract_facts_from_contents(
|
||||
mock_llm_config, test_contents, hindsight_config, memory, request_context
|
||||
):
|
||||
"""Test that extract_facts_from_contents routes to batch API when enabled."""
|
||||
bank_id = f"test_routing_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Enable batch API in config
|
||||
hindsight_config.retain_batch_enabled = True
|
||||
|
||||
# Mock batch API support
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)
|
||||
mock_llm_config._provider_impl.submit_batch = AsyncMock(
|
||||
return_value={"batch_id": "batch_123", "status": "validating", "request_counts": {}}
|
||||
)
|
||||
mock_llm_config._provider_impl.get_batch_status = AsyncMock(
|
||||
return_value={"status": "completed", "request_counts": {"total": 1, "completed": 1, "failed": 0}}
|
||||
)
|
||||
mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"custom_id": "chunk_0",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": json.dumps({"facts": []})
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
# Call main extract_facts_from_contents (should route to batch API)
|
||||
facts, chunks, usage = await extract_facts_from_contents(
|
||||
contents=test_contents,
|
||||
llm_config=mock_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=hindsight_config,
|
||||
pool=None,
|
||||
operation_id=None,
|
||||
schema=None,
|
||||
)
|
||||
|
||||
# Verify batch API was called
|
||||
mock_llm_config._provider_impl.submit_batch.assert_called_once()
|
||||
|
||||
logger.info("✅ Routing to batch API test passed")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
try:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -1,263 +0,0 @@
|
||||
"""
|
||||
Real integration test for OpenAI Batch API.
|
||||
|
||||
This test makes REAL API calls to OpenAI and measures actual timing.
|
||||
It will be slow (minutes to hours) depending on OpenAI's queue.
|
||||
|
||||
To run:
|
||||
pytest tests/test_batch_api_integration.py -v -s
|
||||
|
||||
To skip in CI:
|
||||
Add @pytest.mark.skip at the test level
|
||||
"""
|
||||
import pytest
|
||||
import os
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from dotenv import load_dotenv
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.retain.fact_extraction import (
|
||||
extract_facts_from_contents_batch_api,
|
||||
RetainContent,
|
||||
)
|
||||
from hindsight_api.config import HindsightConfig
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Load .env file for API keys
|
||||
load_dotenv()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def openai_api_key():
|
||||
"""Get OpenAI API key from environment."""
|
||||
# Try both current and commented keys from .env
|
||||
api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY")
|
||||
|
||||
# Check if it's an OpenAI key (starts with sk-proj- or sk-)
|
||||
if not api_key or not api_key.startswith("sk-"):
|
||||
# Try the OpenAI-specific env var (if set separately)
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
|
||||
if not api_key or not api_key.startswith("sk-"):
|
||||
pytest.skip("OpenAI API key not found in environment. Set OPENAI_API_KEY or uncomment OpenAI config in .env")
|
||||
|
||||
return api_key
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def real_llm_config(openai_api_key):
|
||||
"""Create real LLM config for OpenAI."""
|
||||
# Create config with OpenAI settings
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
# Use LLMProvider wrapper (which creates _provider_impl internally)
|
||||
llm_config = LLMProvider(
|
||||
provider="openai",
|
||||
api_key=openai_api_key,
|
||||
base_url="https://api.openai.com/v1",
|
||||
model="gpt-4o-mini", # Fast, cheap model for testing
|
||||
reasoning_effort="medium", # Required parameter
|
||||
)
|
||||
|
||||
return llm_config
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_contents_real():
|
||||
"""Create realistic test content for fact extraction."""
|
||||
return [
|
||||
RetainContent(
|
||||
content="""
|
||||
Alice is a senior software engineer at TechCorp, where she has been working for 5 years.
|
||||
She specializes in distributed systems and microservices architecture. Alice graduated
|
||||
from MIT with a degree in Computer Science in 2015. She is known for writing clean,
|
||||
well-documented code and mentoring junior developers.
|
||||
""",
|
||||
event_date=datetime(2024, 1, 15, 10, 30, tzinfo=timezone.utc),
|
||||
context="team member profile",
|
||||
),
|
||||
RetainContent(
|
||||
content="""
|
||||
Bob joined TechCorp last month as a junior developer. He is learning React and Node.js
|
||||
and recently completed his first feature, which was a user authentication flow. Bob
|
||||
graduated from Berkeley with a degree in Computer Science in 2023. He is enthusiastic
|
||||
and asks great questions during code reviews.
|
||||
""",
|
||||
event_date=datetime(2024, 1, 15, 10, 30, tzinfo=timezone.utc),
|
||||
context="team member profile",
|
||||
),
|
||||
RetainContent(
|
||||
content="""
|
||||
The team uses Kubernetes for container orchestration and deploys to AWS. They follow
|
||||
agile methodologies with two-week sprints. Code reviews are mandatory before merging
|
||||
any pull request. The team meets every morning for a 15-minute standup to discuss
|
||||
progress and blockers.
|
||||
""",
|
||||
event_date=datetime(2024, 1, 15, 10, 30, tzinfo=timezone.utc),
|
||||
context="team processes",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def integration_config():
|
||||
"""Create config for integration test."""
|
||||
config = HindsightConfig.from_env()
|
||||
config.retain_batch_enabled = True
|
||||
config.retain_batch_poll_interval_seconds = 30 # Poll every 30 seconds (reasonable for real API)
|
||||
config.retain_chunk_size = 4000
|
||||
config.retain_extraction_mode = "concise"
|
||||
config.retain_extract_causal_links = False
|
||||
return config
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Real API test - takes minutes and costs money. Run manually with: pytest tests/test_batch_api_integration.py::test_real_openai_batch_api -v -s")
|
||||
@pytest.mark.integration # Mark as integration test
|
||||
@pytest.mark.slow # Mark as slow test
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_openai_batch_api(real_llm_config, test_contents_real, integration_config, memory, request_context):
|
||||
"""
|
||||
REAL integration test: Submit actual batch to OpenAI and measure timing.
|
||||
|
||||
WARNING: This test:
|
||||
- Makes real API calls to OpenAI
|
||||
- Will take minutes to hours to complete
|
||||
- Costs money (though very little with gpt-4o-mini)
|
||||
- Requires valid OpenAI API key
|
||||
|
||||
To skip this test:
|
||||
pytest tests/test_batch_api_integration.py --skip-integration
|
||||
"""
|
||||
bank_id = f"test_real_batch_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
logger.info("=" * 80)
|
||||
logger.info("STARTING REAL OPENAI BATCH API INTEGRATION TEST")
|
||||
logger.info("=" * 80)
|
||||
logger.info(f"Test contents: {len(test_contents_real)} items")
|
||||
logger.info(f"Poll interval: {integration_config.retain_batch_poll_interval_seconds}s")
|
||||
logger.info(f"Model: {real_llm_config.model}")
|
||||
logger.info("This may take several minutes to hours depending on OpenAI's queue...")
|
||||
logger.info("=" * 80)
|
||||
|
||||
try:
|
||||
# Ensure bank exists
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Get database pool and schema for crash recovery testing
|
||||
pool = memory._pool
|
||||
schema = request_context.tenant_id
|
||||
|
||||
# Track overall timing
|
||||
test_start_time = time.time()
|
||||
|
||||
# Call REAL batch API extraction
|
||||
logger.info("\n📤 Submitting batch to OpenAI...")
|
||||
|
||||
facts, chunks, usage = await extract_facts_from_contents_batch_api(
|
||||
contents=test_contents_real,
|
||||
llm_config=real_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=integration_config,
|
||||
pool=pool,
|
||||
operation_id=None, # No crash recovery for this test
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
test_end_time = time.time()
|
||||
total_duration = test_end_time - test_start_time
|
||||
|
||||
# Log results
|
||||
logger.info("\n" + "=" * 80)
|
||||
logger.info("✅ BATCH COMPLETED SUCCESSFULLY")
|
||||
logger.info("=" * 80)
|
||||
logger.info(f"Total duration: {total_duration:.1f} seconds ({total_duration/60:.1f} minutes)")
|
||||
logger.info(f"Facts extracted: {len(facts)}")
|
||||
logger.info(f"Chunks processed: {len(chunks)}")
|
||||
logger.info(f"Token usage: {usage.input_tokens} input + {usage.output_tokens} output = {usage.total_tokens} total")
|
||||
logger.info(f"Estimated cost: ${(usage.input_tokens * 0.00015 / 1000 + usage.output_tokens * 0.0006 / 1000):.4f}")
|
||||
logger.info("=" * 80)
|
||||
|
||||
# Log sample facts
|
||||
logger.info("\n📋 Sample extracted facts:")
|
||||
for i, fact in enumerate(facts[:5]): # Show first 5 facts
|
||||
logger.info(f"\nFact {i+1}:")
|
||||
logger.info(f" Type: {fact.fact_type}")
|
||||
logger.info(f" Text: {fact.fact_text[:100]}...")
|
||||
logger.info(f" Entities: {fact.entities}")
|
||||
|
||||
# Verify results
|
||||
assert len(facts) > 0, "Should extract at least some facts"
|
||||
assert len(chunks) == len(test_contents_real), f"Should have {len(test_contents_real)} chunks"
|
||||
assert usage.total_tokens > 0, "Should have token usage"
|
||||
|
||||
# Verify fact structure
|
||||
for fact in facts:
|
||||
assert hasattr(fact, "fact_text"), "Fact should have fact_text"
|
||||
assert hasattr(fact, "fact_type"), "Fact should have fact_type"
|
||||
assert fact.fact_type in ["world", "experience", "opinion"], f"Invalid fact_type: {fact.fact_type}"
|
||||
|
||||
logger.info("\n✅ All assertions passed!")
|
||||
|
||||
# Write timing report to file for later analysis
|
||||
report_path = "/tmp/openai_batch_api_timing_report.txt"
|
||||
with open(report_path, "w") as f:
|
||||
f.write(f"OpenAI Batch API Integration Test Report\n")
|
||||
f.write(f"={'=' * 60}\n\n")
|
||||
f.write(f"Test Date: {datetime.now(timezone.utc).isoformat()}\n")
|
||||
f.write(f"Model: {real_llm_config.model}\n")
|
||||
f.write(f"Contents: {len(test_contents_real)} items\n")
|
||||
f.write(f"Poll Interval: {integration_config.retain_batch_poll_interval_seconds}s\n\n")
|
||||
f.write(f"Results:\n")
|
||||
f.write(f" Total Duration: {total_duration:.1f}s ({total_duration/60:.1f} min)\n")
|
||||
f.write(f" Facts Extracted: {len(facts)}\n")
|
||||
f.write(f" Chunks Processed: {len(chunks)}\n")
|
||||
f.write(f" Token Usage: {usage.total_tokens} ({usage.input_tokens} in + {usage.output_tokens} out)\n")
|
||||
f.write(f" Estimated Cost: ${(usage.input_tokens * 0.00015 / 1000 + usage.output_tokens * 0.0006 / 1000):.4f}\n")
|
||||
|
||||
logger.info(f"\n📄 Timing report written to: {report_path}")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
try:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
logger.info(f"\n🧹 Cleaned up test bank: {bank_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to cleanup bank: {e}")
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Real API test - requires Groq API key. Run manually if needed.")
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_batch_supports_groq(integration_config):
|
||||
"""
|
||||
Test that Groq also supports batch API (if configured).
|
||||
|
||||
Groq has the same batch API interface as OpenAI.
|
||||
"""
|
||||
groq_api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY")
|
||||
|
||||
if not groq_api_key or not groq_api_key.startswith("gsk_"):
|
||||
pytest.skip("Groq API key not found in environment")
|
||||
|
||||
llm_config = LLMProvider(
|
||||
provider="groq",
|
||||
api_key=groq_api_key,
|
||||
base_url="https://api.groq.com/openai/v1",
|
||||
model="llama-3.1-8b-instant",
|
||||
reasoning_effort="medium",
|
||||
)
|
||||
|
||||
# Check if Groq supports batch API
|
||||
supports_batch = await llm_config._provider_impl.supports_batch_api()
|
||||
|
||||
logger.info(f"Groq batch API support: {supports_batch}")
|
||||
|
||||
# Groq should support batch API (same interface as OpenAI)
|
||||
assert supports_batch, "Groq should support batch API"
|
||||
|
||||
logger.info("✅ Groq batch API support confirmed")
|
||||
@@ -1,38 +0,0 @@
|
||||
"""
|
||||
Test validation for batch API + synchronous retain.
|
||||
|
||||
When HINDSIGHT_API_RETAIN_BATCH_ENABLED=true, synchronous retain operations
|
||||
should be rejected with a 400 error since they will timeout.
|
||||
"""
|
||||
|
||||
import os
|
||||
import pytest
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
from hindsight_api.config import HindsightConfig
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_validation(memory, request_context):
|
||||
"""
|
||||
Test that attempting synchronous retain with batch API enabled
|
||||
raises an error at the HTTP layer.
|
||||
|
||||
This test verifies the validation logic exists - actual HTTP testing
|
||||
would require full FastAPI app setup.
|
||||
"""
|
||||
# Create config with batch API enabled
|
||||
config = HindsightConfig.from_env()
|
||||
config.retain_batch_enabled = True
|
||||
config.retain_batch_poll_interval_seconds = 1
|
||||
|
||||
# Verify the validation exists in memory engine
|
||||
# The actual HTTP validation happens in http.py api_retain()
|
||||
# This test documents the expected behavior
|
||||
|
||||
assert config.retain_batch_enabled is True
|
||||
assert config.retain_batch_poll_interval_seconds == 1
|
||||
|
||||
# When batch API is enabled and async=false, the HTTP endpoint
|
||||
# should return 400 with message:
|
||||
# "Batch API is enabled (HINDSIGHT_API_RETAIN_BATCH_ENABLED=true) but async=false"
|
||||
@@ -12,7 +12,6 @@ from datetime import datetime
|
||||
import pytest
|
||||
|
||||
from hindsight_api import LLMConfig
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
|
||||
|
||||
|
||||
@@ -45,7 +44,6 @@ class TestCausalRelationsValidation:
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -90,7 +88,6 @@ class TestCausalRelationsValidation:
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -127,7 +124,6 @@ class TestCausalRelationsValidation:
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract facts about the causal chain"
|
||||
@@ -177,7 +173,6 @@ class TestCausalRelationsValidation:
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract facts"
|
||||
@@ -214,7 +209,6 @@ class TestCausalRelationsValidation:
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
)
|
||||
|
||||
# Verify relation types are all backward-looking
|
||||
|
||||
@@ -10,7 +10,6 @@ from datetime import datetime
|
||||
import pytest
|
||||
|
||||
from hindsight_api import LLMConfig
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
|
||||
|
||||
|
||||
@@ -38,8 +37,7 @@ After searching for weeks, I finally found a cheaper apartment in Brooklyn.
|
||||
llm_config = LLMConfig.for_memory()
|
||||
|
||||
facts, _, _ = await extract_facts_from_text(
|
||||
text=text, event_date=datetime(2024, 3, 15), context=context, llm_config=llm_config, agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
text=text, event_date=datetime(2024, 3, 15), context=context, llm_config=llm_config, agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) >= 3, f"Should extract at least 3 facts from the causal chain. Got {len(facts)}"
|
||||
@@ -108,8 +106,7 @@ The renovation took three months and cost $15,000.
|
||||
llm_config = LLMConfig.for_memory()
|
||||
|
||||
facts, _, _ = await extract_facts_from_text(
|
||||
text=text, event_date=datetime(2024, 6, 1), context=context, llm_config=llm_config, agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
text=text, event_date=datetime(2024, 6, 1), context=context, llm_config=llm_config, agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) >= 4, f"Should extract at least 4 facts. Got {len(facts)}"
|
||||
@@ -139,8 +136,7 @@ Machine learning fascinated me so much that I changed my career to data science.
|
||||
llm_config = LLMConfig.for_memory()
|
||||
|
||||
facts, _, _ = await extract_facts_from_text(
|
||||
text=text, event_date=datetime(2024, 1, 1), context=context, llm_config=llm_config, agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
text=text, event_date=datetime(2024, 1, 1), context=context, llm_config=llm_config, agent_name="TestUser"
|
||||
)
|
||||
|
||||
# Check no fact references itself
|
||||
@@ -167,8 +163,7 @@ The new role enabled me to lead a team of engineers.
|
||||
llm_config = LLMConfig.for_memory()
|
||||
|
||||
facts, _, _ = await extract_facts_from_text(
|
||||
text=text, event_date=datetime(2024, 2, 15), context=context, llm_config=llm_config, agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
text=text, event_date=datetime(2024, 2, 15), context=context, llm_config=llm_config, agent_name="TestUser"
|
||||
)
|
||||
|
||||
# Validate all indices (must reference PREVIOUS facts only)
|
||||
@@ -195,8 +190,7 @@ Reduced spending somewhat affected local businesses.
|
||||
llm_config = LLMConfig.for_memory()
|
||||
|
||||
facts, _, _ = await extract_facts_from_text(
|
||||
text=text, event_date=datetime(2024, 4, 1), context=context, llm_config=llm_config, agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
text=text, event_date=datetime(2024, 4, 1), context=context, llm_config=llm_config, agent_name="TestUser"
|
||||
)
|
||||
|
||||
for i, fact in enumerate(facts):
|
||||
|
||||
@@ -21,9 +21,9 @@ from hindsight_api.engine.reflect.tools import (
|
||||
@pytest.fixture(autouse=True)
|
||||
def enable_observations():
|
||||
"""Enable observations for all tests in this module."""
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
config = _get_raw_config()
|
||||
config = get_config()
|
||||
original_value = config.enable_observations
|
||||
config.enable_observations = True
|
||||
yield
|
||||
@@ -563,26 +563,25 @@ class TestConsolidationDisabled:
|
||||
self, memory: MemoryEngine, request_context
|
||||
):
|
||||
"""Test that consolidation returns disabled status when enable_observations is False."""
|
||||
from unittest.mock import patch
|
||||
|
||||
bank_id = f"test-consolidation-disabled-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Create the bank
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
# Disable observations for this bank via bank config
|
||||
await memory._config_resolver.update_bank_config(
|
||||
bank_id=bank_id,
|
||||
updates={"enable_observations": False},
|
||||
context=request_context,
|
||||
)
|
||||
# Disable observations via config
|
||||
with patch("hindsight_api.config.get_config") as mock_config:
|
||||
mock_config.return_value.enable_observations = False
|
||||
|
||||
result = await run_consolidation_job(
|
||||
memory_engine=memory,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
result = await run_consolidation_job(
|
||||
memory_engine=memory,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result["status"] == "disabled"
|
||||
assert result["bank_id"] == bank_id
|
||||
assert result["status"] == "disabled"
|
||||
assert result["bank_id"] == bank_id
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -135,179 +135,3 @@ async def test_memory_without_document(memory, request_context):
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_persisted_with_zero_facts(memory, request_context):
|
||||
"""
|
||||
Test that documents are persisted even when zero facts are extracted.
|
||||
|
||||
This is a regression test for issue #324 where documents with no extractable
|
||||
facts were reported as disappearing from the system.
|
||||
"""
|
||||
bank_id = f"test_zero_facts_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
document_id = "doc-zero-facts"
|
||||
|
||||
# Retain content that produces zero facts (gibberish/random characters)
|
||||
units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="xyzabc123 !!!### @@@ $$$", # Random characters unlikely to produce facts
|
||||
context="Test zero facts",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should return empty unit list (no facts extracted)
|
||||
assert len(units) == 0, "Should extract zero facts from gibberish content"
|
||||
|
||||
# But document should still be persisted and retrievable
|
||||
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc is not None, "Document should be persisted even with zero facts"
|
||||
assert doc["id"] == document_id
|
||||
assert doc["bank_id"] == bank_id
|
||||
assert doc["memory_unit_count"] == 0, "Should have zero memory units"
|
||||
assert len(doc["original_text"]) > 0, "Should have non-zero text length"
|
||||
assert "xyzabc123" in doc["original_text"], "Should contain original content"
|
||||
|
||||
# Document should also appear in list
|
||||
docs_list = await memory.list_documents(
|
||||
bank_id=bank_id,
|
||||
search_query=None,
|
||||
limit=100,
|
||||
offset=0,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert docs_list["total"] == 1, "Document should appear in list"
|
||||
assert any(d["id"] == document_id for d in docs_list["items"]), "Document should be in items"
|
||||
|
||||
listed_doc = next(d for d in docs_list["items"] if d["id"] == document_id)
|
||||
assert listed_doc["memory_unit_count"] == 0, "Listed document should show zero memory units"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_persisted_with_zero_facts_batch(memory, request_context):
|
||||
"""
|
||||
Test that documents are persisted with zero facts in batch retain operations.
|
||||
|
||||
This tests the async batch code path to ensure it also handles zero facts correctly.
|
||||
"""
|
||||
bank_id = f"test_zero_facts_batch_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Mix of content: some produces facts, some produces zero facts
|
||||
contents = [
|
||||
{
|
||||
"content": "Alice works at Google",
|
||||
"document_id": "doc-with-facts",
|
||||
},
|
||||
{
|
||||
"content": "!@# $$$ %%% ^^^ &&& ***", # Gibberish - zero facts expected
|
||||
"document_id": "doc-zero-facts",
|
||||
},
|
||||
]
|
||||
|
||||
unit_ids = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# First content should produce facts, second should not
|
||||
assert len(unit_ids[0]) > 0, "First content should produce facts"
|
||||
assert len(unit_ids[1]) == 0, "Second content should produce zero facts"
|
||||
|
||||
# Both documents should be persisted
|
||||
doc_with_facts = await memory.get_document("doc-with-facts", bank_id, request_context=request_context)
|
||||
assert doc_with_facts is not None
|
||||
assert doc_with_facts["memory_unit_count"] > 0
|
||||
|
||||
doc_zero_facts = await memory.get_document("doc-zero-facts", bank_id, request_context=request_context)
|
||||
assert doc_zero_facts is not None, "Document with zero facts should be persisted"
|
||||
assert doc_zero_facts["memory_unit_count"] == 0, "Should have zero memory units"
|
||||
assert "!@#" in doc_zero_facts["original_text"]
|
||||
|
||||
# Both should appear in list
|
||||
docs_list = await memory.list_documents(
|
||||
bank_id=bank_id,
|
||||
search_query=None,
|
||||
limit=100,
|
||||
offset=0,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert docs_list["total"] == 2, "Both documents should appear in list"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_persisted_with_zero_facts_async_submit(memory, request_context):
|
||||
"""
|
||||
Test that documents are persisted with zero facts in fire-and-forget async retain.
|
||||
|
||||
This tests the submit_async_retain (background task) code path to ensure it also
|
||||
handles zero facts correctly.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
bank_id = f"test_zero_facts_async_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Submit async retain with gibberish content
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{
|
||||
"content": "!@# $$$ %%% ^^^ &&& ***", # Gibberish - zero facts expected
|
||||
"document_id": "doc-async-zero-facts",
|
||||
}
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
operation_id = result["operation_id"]
|
||||
assert operation_id is not None, "Should return operation_id"
|
||||
|
||||
# Wait for background task to complete
|
||||
max_wait = 60 # 60 seconds max
|
||||
wait_interval = 0.5
|
||||
elapsed = 0
|
||||
|
||||
while elapsed < max_wait:
|
||||
await asyncio.sleep(wait_interval)
|
||||
elapsed += wait_interval
|
||||
|
||||
# Check if document exists
|
||||
doc = await memory.get_document(
|
||||
"doc-async-zero-facts", bank_id, request_context=request_context
|
||||
)
|
||||
if doc is not None:
|
||||
break
|
||||
|
||||
# Document should be persisted even with zero facts
|
||||
assert doc is not None, "Document should be persisted after async task completes"
|
||||
assert doc["id"] == "doc-async-zero-facts"
|
||||
assert doc["memory_unit_count"] == 0, "Should have zero memory units"
|
||||
assert "!@#" in doc["original_text"]
|
||||
|
||||
# Document should appear in list
|
||||
docs_list = await memory.list_documents(
|
||||
bank_id=bank_id,
|
||||
search_query=None,
|
||||
limit=100,
|
||||
offset=0,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert docs_list["total"] == 1, "Document should appear in list"
|
||||
assert any(d["id"] == "doc-async-zero-facts" for d in docs_list["items"])
|
||||
|
||||
listed_doc = next(d for d in docs_list["items"] if d["id"] == "doc-async-zero-facts")
|
||||
assert listed_doc["memory_unit_count"] == 0, "Listed document should show zero memory units"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -8,7 +8,7 @@ from datetime import datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config import get_config, clear_config_cache, _get_raw_config
|
||||
from hindsight_api.config import get_config, clear_config_cache
|
||||
from hindsight_api.engine.llm_wrapper import LLMConfig
|
||||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
|
||||
|
||||
@@ -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",
|
||||
config=_get_raw_config(),
|
||||
)
|
||||
|
||||
duration = time.time() - start_time
|
||||
|
||||
@@ -11,7 +11,6 @@ from datetime import datetime
|
||||
import pytest
|
||||
|
||||
from hindsight_api import LLMConfig
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
|
||||
|
||||
|
||||
@@ -45,8 +44,7 @@ I ran into my neighbor Sarah who mentioned she's planning a trip to Italy next m
|
||||
event_date=datetime(2024, 6, 15),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
input_length = len(text)
|
||||
@@ -90,8 +88,7 @@ User: Perfect, I'll make a reservation for Saturday at 7pm.
|
||||
event_date=datetime(2024, 6, 15),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
input_length = len(text)
|
||||
@@ -147,8 +144,7 @@ I edited about 20 photos from my recent trip to the mountains.
|
||||
event_date=datetime(2024, 4, 15),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
input_length = len(text)
|
||||
@@ -212,8 +208,7 @@ I edited about 20 photos from my recent trip to the mountains.
|
||||
event_date=datetime(2023, 5, 8), # Date from locomo dataset
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name=data["conversation"]["speaker_a"],
|
||||
config=_get_raw_config(),
|
||||
agent_name=data["conversation"]["speaker_a"]
|
||||
)
|
||||
|
||||
# Calculate ratios
|
||||
@@ -274,8 +269,7 @@ I'm planning to visit Japan next year.
|
||||
event_date=datetime(2024, 6, 15),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
# Count approximate number of statements (sentences)
|
||||
|
||||
@@ -17,7 +17,6 @@ from datetime import UTC, datetime
|
||||
import pytest
|
||||
|
||||
from hindsight_api import LLMConfig
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
|
||||
|
||||
# =============================================================================
|
||||
@@ -49,8 +48,7 @@ Marcus felt anxious about the upcoming interview.
|
||||
event_date=datetime(2024, 11, 13),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -82,8 +80,7 @@ The music was so loud I could barely hear myself think.
|
||||
event_date=datetime(2024, 11, 13),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -116,8 +113,7 @@ Maybe we should reconsider the timeline.
|
||||
event_date=datetime(2024, 11, 13),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -150,8 +146,7 @@ I'm unable to attend the conference due to scheduling conflicts.
|
||||
event_date=datetime(2024, 11, 13),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -183,8 +178,7 @@ Unlike last year, we're ahead of schedule.
|
||||
event_date=datetime(2024, 11, 13),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -217,8 +211,7 @@ She's enthusiastic about the opportunity.
|
||||
event_date=datetime(2024, 11, 13),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -251,8 +244,7 @@ I'm planning to switch careers because I'm not fulfilled in my current role.
|
||||
event_date=datetime(2024, 11, 13),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -289,8 +281,7 @@ Family is the most important thing to her.
|
||||
event_date=datetime(2024, 11, 13),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -324,8 +315,7 @@ I prefer presenting in person rather than virtually because I can read the room
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -382,8 +372,7 @@ I'm planning to visit Tokyo next month.
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -434,8 +423,7 @@ with a concert surrounded by music, joy and the warm summer breeze.
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="Melanie",
|
||||
config=_get_raw_config(),
|
||||
agent_name="Melanie"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -505,8 +493,7 @@ It was a beautiful day and I plan to make this a regular habit.
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -560,8 +547,7 @@ It was a beautiful day and I plan to make this a regular habit.
|
||||
event_date=reference_date,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
context="Personal diary",
|
||||
config=_get_raw_config(),
|
||||
context="Personal diary"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -591,8 +577,7 @@ It was a beautiful day and I plan to make this a regular habit.
|
||||
event_date=reference_date,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
context="General info",
|
||||
config=_get_raw_config(),
|
||||
context="General info"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -619,8 +604,7 @@ It was a beautiful day and I plan to make this a regular habit.
|
||||
event_date=reference_date,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
context="Calendar events",
|
||||
config=_get_raw_config(),
|
||||
context="Calendar events"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -671,8 +655,7 @@ great time! Every time I see it, I can't help but smile.
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="Deborah",
|
||||
config=_get_raw_config(),
|
||||
agent_name="Deborah"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -722,8 +705,7 @@ I've learned so much from it.
|
||||
event_date=datetime(2024, 11, 13),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -792,8 +774,7 @@ Jamie: Congratulations! I'd love to read it.
|
||||
event_date=datetime(2024, 11, 13),
|
||||
llm_config=llm_config,
|
||||
agent_name="Marcus",
|
||||
context=context,
|
||||
config=_get_raw_config(),
|
||||
context=context
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact from the transcript"
|
||||
@@ -838,8 +819,7 @@ We presented our findings to the team yesterday.
|
||||
event_date=datetime(2024, 11, 13),
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
context=context,
|
||||
config=_get_raw_config(),
|
||||
context=context
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract facts"
|
||||
@@ -874,8 +854,7 @@ Jamie: [teasing] We'll see who's right, my Niners pick is solid.
|
||||
event_date=datetime(2024, 11, 14),
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
agent_name=agent_name,
|
||||
config=_get_raw_config(),
|
||||
agent_name=agent_name
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
@@ -941,8 +920,7 @@ so the algorithm learns to box out. See you next week!
|
||||
event_date=datetime(2024, 11, 13),
|
||||
llm_config=llm_config,
|
||||
agent_name="Marcus",
|
||||
context=context,
|
||||
config=_get_raw_config(),
|
||||
context=context
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
|
||||
@@ -1,553 +0,0 @@
|
||||
"""
|
||||
End-to-end tests for file retain (upload, convert, retain) functionality.
|
||||
"""
|
||||
|
||||
import io
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_pdf_content():
|
||||
"""Create a simple PDF-like content for testing."""
|
||||
# This is a minimal PDF that markitdown can parse
|
||||
return b"""%PDF-1.4
|
||||
1 0 obj
|
||||
<<
|
||||
/Type /Catalog
|
||||
/Pages 2 0 R
|
||||
>>
|
||||
endobj
|
||||
2 0 obj
|
||||
<<
|
||||
/Type /Pages
|
||||
/Kids [3 0 R]
|
||||
/Count 1
|
||||
>>
|
||||
endobj
|
||||
3 0 obj
|
||||
<<
|
||||
/Type /Page
|
||||
/Parent 2 0 R
|
||||
/MediaBox [0 0 612 792]
|
||||
/Contents 4 0 R
|
||||
/Resources <<
|
||||
/Font <<
|
||||
/F1 <<
|
||||
/Type /Font
|
||||
/Subtype /Type1
|
||||
/BaseFont /Helvetica
|
||||
>>
|
||||
>>
|
||||
>>
|
||||
>>
|
||||
endobj
|
||||
4 0 obj
|
||||
<<
|
||||
/Length 44
|
||||
>>
|
||||
stream
|
||||
BT
|
||||
/F1 12 Tf
|
||||
100 700 Td
|
||||
(Test Document) Tj
|
||||
ET
|
||||
endstream
|
||||
endobj
|
||||
xref
|
||||
0 5
|
||||
0000000000 65535 f
|
||||
0000000009 00000 n
|
||||
0000000058 00000 n
|
||||
0000000115 00000 n
|
||||
0000000317 00000 n
|
||||
trailer
|
||||
<<
|
||||
/Size 5
|
||||
/Root 1 0 R
|
||||
>>
|
||||
startxref
|
||||
410
|
||||
%%EOF
|
||||
"""
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_txt_content():
|
||||
"""Create simple text content."""
|
||||
return b"This is a test document.\nIt contains some important information.\nAlice works at Google."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_basic(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test basic file upload and conversion."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
# Create a bank first
|
||||
bank_response = await client.put("/v1/default/banks/test-file-bank", json={"name": "Test File Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
# Upload file
|
||||
request_data = {
|
||||
"document_tags": ["test"],
|
||||
"async": True,
|
||||
}
|
||||
|
||||
files = {"files": ("test.txt", sample_txt_content, "text/plain")}
|
||||
data = {"request": json.dumps(request_data)}
|
||||
|
||||
response = await client.post(
|
||||
"/v1/default/banks/test-file-bank/files/retain",
|
||||
files=files,
|
||||
data=data,
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "operation_ids" in result
|
||||
assert len(result["operation_ids"]) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_with_metadata(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test file upload with per-file metadata."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
# Create bank
|
||||
bank_response = await client.put("/v1/default/banks/test-file-meta-bank", json={"name": "Test Meta Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
# Upload file with metadata
|
||||
request_data = {
|
||||
"document_tags": ["work", "reports"],
|
||||
"async": True,
|
||||
"files_metadata": [
|
||||
{
|
||||
"document_id": "test_doc_123",
|
||||
"context": "quarterly report",
|
||||
"metadata": {"author": "Alice", "year": "2024"},
|
||||
"tags": ["Q1"],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
files = {"files": ("report.txt", sample_txt_content, "text/plain")}
|
||||
data = {"request": json.dumps(request_data)}
|
||||
|
||||
response = await client.post(
|
||||
"/v1/default/banks/test-file-meta-bank/files/retain",
|
||||
files=files,
|
||||
data=data,
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "operation_ids" in result
|
||||
assert len(result["operation_ids"]) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_multiple_files(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test uploading multiple files at once."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
# Create bank
|
||||
bank_response = await client.put("/v1/default/banks/test-multi-file-bank", json={"name": "Test Multi Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
# Upload multiple files
|
||||
request_data = {
|
||||
"async": True,
|
||||
"files_metadata": [
|
||||
{"document_id": "doc1", "tags": ["file1"]},
|
||||
{"document_id": "doc2", "tags": ["file2"]},
|
||||
],
|
||||
}
|
||||
|
||||
content1 = b"First document content"
|
||||
content2 = b"Second document content"
|
||||
|
||||
files = [
|
||||
("files", ("file1.txt", content1, "text/plain")),
|
||||
("files", ("file2.txt", content2, "text/plain")),
|
||||
]
|
||||
data = {"request": json.dumps(request_data)}
|
||||
|
||||
response = await client.post(
|
||||
"/v1/default/banks/test-multi-file-bank/files/retain",
|
||||
files=files,
|
||||
data=data,
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "operation_ids" in result
|
||||
assert len(result["operation_ids"]) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_validation_errors(memory_no_llm_verify):
|
||||
"""Test validation errors."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
# Create bank
|
||||
bank_response = await client.put("/v1/default/banks/test-validation-bank", json={"name": "Test Validation Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
# Test: metadata count mismatch
|
||||
request_data = {
|
||||
"async": True,
|
||||
"files_metadata": [
|
||||
{"document_id": "doc1"},
|
||||
{"document_id": "doc2"}, # 2 metadata entries
|
||||
],
|
||||
}
|
||||
|
||||
files = {"files": ("file1.txt", b"content", "text/plain")} # But only 1 file
|
||||
data = {"request": json.dumps(request_data)}
|
||||
|
||||
response = await client.post(
|
||||
"/v1/default/banks/test-validation-bank/files/retain",
|
||||
files=files,
|
||||
data=data,
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "files_metadata count" in response.json()["detail"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_no_files(memory_no_llm_verify):
|
||||
"""Test error when no files provided."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
# Create bank
|
||||
bank_response = await client.put("/v1/default/banks/test-no-files-bank", json={"name": "Test No Files Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
request_data = {
|
||||
"async": True,
|
||||
}
|
||||
|
||||
# No files provided
|
||||
data = {"request": json.dumps(request_data)}
|
||||
|
||||
response = await client.post(
|
||||
"/v1/default/banks/test-no-files-bank/files/retain",
|
||||
data=data,
|
||||
)
|
||||
|
||||
# FastAPI will return 422 for missing required field
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_sync_not_supported(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test that file retain is always async (sync is not supported)."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
# Create bank
|
||||
bank_response = await client.put("/v1/default/banks/test-sync-bank", json={"name": "Test Sync Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
# File retain is always async - just verify it succeeds and returns operation_ids
|
||||
files = {"files": ("test.txt", sample_txt_content, "text/plain")}
|
||||
data = {"request": json.dumps({})}
|
||||
|
||||
response = await client.post(
|
||||
"/v1/default/banks/test-sync-bank/files/retain",
|
||||
files=files,
|
||||
data=data,
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "operation_ids" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_storage_postgresql(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test file storage in PostgreSQL."""
|
||||
# Test that files are stored and retrieved correctly
|
||||
storage = memory_no_llm_verify._file_storage
|
||||
|
||||
# Store a file
|
||||
key = "test/file1.txt"
|
||||
stored_key = await storage.store(
|
||||
file_data=sample_txt_content,
|
||||
key=key,
|
||||
metadata={"content_type": "text/plain"},
|
||||
)
|
||||
|
||||
assert stored_key == key
|
||||
|
||||
# Retrieve the file
|
||||
retrieved = await storage.retrieve(key)
|
||||
assert retrieved == sample_txt_content
|
||||
|
||||
# Check if file exists
|
||||
exists = await storage.exists(key)
|
||||
assert exists is True
|
||||
|
||||
# Delete the file
|
||||
await storage.delete(key)
|
||||
|
||||
# Check file no longer exists
|
||||
exists_after = await storage.exists(key)
|
||||
assert exists_after is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_markitdown_converter():
|
||||
"""Test markitdown parser."""
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
parser = MarkitdownParser()
|
||||
|
||||
# Test simple text file
|
||||
text_content = b"This is a test document.\nWith multiple lines."
|
||||
result = await parser.convert(text_content, "test.txt")
|
||||
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
assert "test document" in result.lower() or "multiple lines" in result.lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_converter_registry():
|
||||
"""Test file parser registry."""
|
||||
from hindsight_api.engine.parsers import FileParserRegistry, MarkitdownParser
|
||||
|
||||
registry = FileParserRegistry()
|
||||
parser = MarkitdownParser()
|
||||
registry.register(parser)
|
||||
|
||||
# Test get by name
|
||||
retrieved = registry.get_parser("markitdown", "test.txt")
|
||||
assert retrieved is parser
|
||||
|
||||
# Test auto-detection
|
||||
auto = registry.get_parser(None, "test.pdf")
|
||||
assert auto is parser
|
||||
|
||||
# Test unsupported format
|
||||
with pytest.raises(ValueError, match="No parser found"):
|
||||
registry.get_parser(None, "test.xyz")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_conversion_creates_separate_retain_operation(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test that file conversion and retain are two separate async operations.
|
||||
|
||||
The file_convert_retain task should:
|
||||
1. Convert the file to markdown
|
||||
2. In a single transaction: create a separate 'retain' operation AND mark itself as 'completed'
|
||||
3. Free the worker slot immediately after conversion
|
||||
|
||||
The retain then runs as its own task. This prevents deadlocks where file conversion
|
||||
tasks hold worker slots while waiting for inline retain to finish.
|
||||
"""
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
bank_id = "test_file_two_phase_bank"
|
||||
|
||||
context = RequestContext(internal=True)
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id, request_context=context)
|
||||
|
||||
class MockFile:
|
||||
def __init__(self, content, filename, content_type):
|
||||
self.content = content
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
|
||||
async def read(self):
|
||||
return self.content
|
||||
|
||||
mock_file = MockFile(sample_txt_content, "test.txt", "text/plain")
|
||||
|
||||
file_items = [
|
||||
{
|
||||
"file": mock_file,
|
||||
"document_id": "test_doc_two_phase",
|
||||
"context": "test context",
|
||||
"metadata": {"source": "test"},
|
||||
"tags": ["test_tag"],
|
||||
"timestamp": None,
|
||||
}
|
||||
]
|
||||
|
||||
result = await memory_no_llm_verify.submit_async_file_retain(
|
||||
bank_id=bank_id,
|
||||
file_items=file_items,
|
||||
parser="markitdown",
|
||||
document_tags=["two_phase_test"],
|
||||
request_context=context,
|
||||
)
|
||||
|
||||
assert "operation_ids" in result
|
||||
assert len(result["operation_ids"]) == 1
|
||||
convert_operation_id = result["operation_ids"][0]
|
||||
|
||||
import asyncio
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
pool = await memory_no_llm_verify._get_pool()
|
||||
from hindsight_api.engine.memory_engine import get_current_schema
|
||||
|
||||
schema = get_current_schema()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
# 1. The file_convert_retain operation must be completed
|
||||
convert_op = await conn.fetchrow(
|
||||
f"SELECT status, operation_type FROM {schema}.async_operations WHERE operation_id = $1",
|
||||
convert_operation_id,
|
||||
)
|
||||
assert convert_op is not None
|
||||
assert convert_op["operation_type"] == "file_convert_retain"
|
||||
assert convert_op["status"] == "completed", (
|
||||
f"file_convert_retain should be 'completed' after conversion, got '{convert_op['status']}'"
|
||||
)
|
||||
|
||||
# 2. A separate retain operation must have been created
|
||||
retain_op = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT status, operation_type
|
||||
FROM {schema}.async_operations
|
||||
WHERE bank_id = $1 AND operation_type = 'retain' AND operation_id != $2
|
||||
""",
|
||||
bank_id,
|
||||
convert_operation_id,
|
||||
)
|
||||
assert retain_op is not None, "A separate 'retain' operation should have been created by file conversion"
|
||||
# With SyncTaskBackend the retain runs immediately, so it should be completed
|
||||
assert retain_op["status"] == "completed"
|
||||
|
||||
# 3. The document should exist with file metadata and retained content
|
||||
doc = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT id, original_text, file_original_name, file_content_type
|
||||
FROM {schema}.documents
|
||||
WHERE id = $1 AND bank_id = $2
|
||||
""",
|
||||
"test_doc_two_phase",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
assert doc is not None
|
||||
assert doc["file_original_name"] == "test.txt"
|
||||
assert doc["file_content_type"] == "text/plain"
|
||||
assert doc["original_text"] is not None
|
||||
assert len(doc["original_text"]) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_conversion_failure_sets_status_to_failed(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test that when file conversion fails, the operation status is set to 'failed' not 'completed'."""
|
||||
from hindsight_api.engine.parsers.base import FileParser
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
bank_id = "test_file_failure_bank"
|
||||
|
||||
# Create a mock parser that always fails
|
||||
class FailingParser(FileParser):
|
||||
"""Mock parser that raises an error."""
|
||||
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
# Simulate conversion failure
|
||||
raise RuntimeError(f"Failed to convert '{filename}': Mock conversion error")
|
||||
|
||||
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||
return filename.endswith(".fail")
|
||||
|
||||
def name(self) -> str:
|
||||
return "failing_converter"
|
||||
|
||||
# Register the failing parser
|
||||
failing_converter = FailingParser()
|
||||
memory_no_llm_verify._parser_registry.register(failing_converter)
|
||||
|
||||
# Create bank
|
||||
context = RequestContext(internal=True)
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id, request_context=context)
|
||||
|
||||
# Create mock file
|
||||
class MockFile:
|
||||
def __init__(self, content, filename, content_type):
|
||||
self.content = content
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
|
||||
async def read(self):
|
||||
return self.content
|
||||
|
||||
mock_file = MockFile(sample_txt_content, "test.fail", "application/octet-stream")
|
||||
|
||||
file_items = [
|
||||
{
|
||||
"file": mock_file,
|
||||
"document_id": "test_doc_fail",
|
||||
"context": None,
|
||||
"metadata": {},
|
||||
"tags": [],
|
||||
"timestamp": None,
|
||||
}
|
||||
]
|
||||
|
||||
# Submit async file retain with failing parser
|
||||
result = await memory_no_llm_verify.submit_async_file_retain(
|
||||
bank_id=bank_id,
|
||||
file_items=file_items,
|
||||
parser="failing_converter",
|
||||
document_tags=None,
|
||||
request_context=context,
|
||||
)
|
||||
|
||||
assert "operation_ids" in result
|
||||
assert len(result["operation_ids"]) == 1
|
||||
operation_id = result["operation_ids"][0]
|
||||
|
||||
# Wait for async processing (with SyncTaskBackend, this is immediate)
|
||||
import asyncio
|
||||
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
# Check operation status - should be 'failed' not 'completed'
|
||||
pool = await memory_no_llm_verify._get_pool()
|
||||
from hindsight_api.engine.memory_engine import get_current_schema
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
operation = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT status, error_message
|
||||
FROM {get_current_schema()}.async_operations
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
operation_id,
|
||||
)
|
||||
|
||||
assert operation is not None, f"Operation {operation_id} not found"
|
||||
assert operation["status"] == "failed", f"Expected status 'failed' but got '{operation['status']}'"
|
||||
assert operation["error_message"] is not None
|
||||
assert "Mock conversion error" in operation["error_message"]
|
||||
assert "test.fail" in operation["error_message"]
|
||||
@@ -1,252 +0,0 @@
|
||||
"""
|
||||
Integration tests for S3FileStorage against a SeaweedFS Docker container.
|
||||
|
||||
SeaweedFS (Apache 2.0) provides an S3-compatible API via `weed server -s3`.
|
||||
Requires Docker to be running. Tests are skipped automatically if Docker is unavailable.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
from testcontainers.core.container import DockerContainer
|
||||
|
||||
_has_testcontainers = True
|
||||
except ImportError:
|
||||
_has_testcontainers = False
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.skipif(not _has_testcontainers, reason="testcontainers not installed"),
|
||||
]
|
||||
|
||||
SEAWEEDFS_S3_PORT = 8333
|
||||
TEST_BUCKET = "hindsight-test"
|
||||
ACCESS_KEY = "test_access_key"
|
||||
SECRET_KEY = "test_secret_key"
|
||||
|
||||
# SeaweedFS S3 IAM config granting full access to our test credentials
|
||||
_S3_CONFIG = {
|
||||
"identities": [
|
||||
{
|
||||
"name": "test-user",
|
||||
"credentials": [{"accessKey": ACCESS_KEY, "secretKey": SECRET_KEY}],
|
||||
"actions": ["Admin", "Read", "Write", "List"],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def _docker_available() -> bool:
|
||||
"""Check if Docker daemon is running."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["docker", "info"],
|
||||
capture_output=True,
|
||||
timeout=5,
|
||||
)
|
||||
return result.returncode == 0
|
||||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||||
return False
|
||||
|
||||
|
||||
def _wait_for_seaweedfs(endpoint: str, timeout: int = 30) -> None:
|
||||
"""Poll SeaweedFS S3 endpoint until ready."""
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
resp = httpx.get(endpoint, timeout=2)
|
||||
# 200 = no auth, 403 = auth enabled but gateway is up — either means ready
|
||||
if resp.status_code in (200, 403):
|
||||
logger.info("SeaweedFS S3 is ready at %s", endpoint)
|
||||
return
|
||||
except httpx.HTTPError:
|
||||
pass
|
||||
time.sleep(0.5)
|
||||
raise TimeoutError(f"SeaweedFS did not become ready at {endpoint} within {timeout}s")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def seaweedfs_container():
|
||||
"""Start a SeaweedFS container for the test module, shared across all tests.
|
||||
|
||||
Mounts an s3.json config file to set up S3 credentials for the test user.
|
||||
"""
|
||||
if not _docker_available():
|
||||
pytest.skip("Docker is not available")
|
||||
|
||||
# Write S3 IAM config to a temp file that persists for the module scope
|
||||
s3_config_file = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False)
|
||||
json.dump(_S3_CONFIG, s3_config_file)
|
||||
s3_config_file.flush()
|
||||
|
||||
container = (
|
||||
DockerContainer(image="chrislusf/seaweedfs:latest")
|
||||
.with_exposed_ports(SEAWEEDFS_S3_PORT)
|
||||
.with_volume_mapping(s3_config_file.name, "/etc/seaweedfs/s3.json", "ro")
|
||||
.with_command(
|
||||
f"server -s3 -s3.port={SEAWEEDFS_S3_PORT} -s3.config=/etc/seaweedfs/s3.json -ip.bind=0.0.0.0"
|
||||
)
|
||||
)
|
||||
|
||||
container.start()
|
||||
|
||||
try:
|
||||
host = container.get_container_host_ip()
|
||||
port = container.get_exposed_port(SEAWEEDFS_S3_PORT)
|
||||
endpoint = f"http://{host}:{port}"
|
||||
|
||||
_wait_for_seaweedfs(endpoint)
|
||||
|
||||
# Create test bucket using obstore (proper SigV4 signing)
|
||||
import obstore as obs
|
||||
from obstore.store import S3Store
|
||||
|
||||
admin_store = S3Store(
|
||||
TEST_BUCKET,
|
||||
endpoint=endpoint,
|
||||
region="us-east-1",
|
||||
access_key_id=ACCESS_KEY,
|
||||
secret_access_key=SECRET_KEY,
|
||||
allow_http=True,
|
||||
)
|
||||
# SeaweedFS auto-creates buckets on first write
|
||||
obs.put(admin_store, ".bucket-init", b"")
|
||||
obs.delete(admin_store, ".bucket-init")
|
||||
logger.info("Test bucket '%s' is ready", TEST_BUCKET)
|
||||
|
||||
yield {
|
||||
"endpoint": endpoint,
|
||||
"access_key": ACCESS_KEY,
|
||||
"secret_key": SECRET_KEY,
|
||||
"bucket": TEST_BUCKET,
|
||||
}
|
||||
finally:
|
||||
container.stop()
|
||||
import os
|
||||
|
||||
os.unlink(s3_config_file.name)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def s3_storage(seaweedfs_container):
|
||||
"""Create an S3FileStorage instance pointing at the SeaweedFS container."""
|
||||
from hindsight_api.engine.storage.s3 import S3FileStorage
|
||||
|
||||
return S3FileStorage(
|
||||
bucket=seaweedfs_container["bucket"],
|
||||
region="us-east-1",
|
||||
endpoint=seaweedfs_container["endpoint"],
|
||||
access_key_id=seaweedfs_container["access_key"],
|
||||
secret_access_key=seaweedfs_container["secret_key"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_storage_store_and_retrieve(s3_storage):
|
||||
"""Store a file, retrieve it, verify bytes match."""
|
||||
content = b"Hello, SeaweedFS! This is a test file."
|
||||
key = f"test/{uuid.uuid4()}.txt"
|
||||
|
||||
stored_key = await s3_storage.store(
|
||||
file_data=content,
|
||||
key=key,
|
||||
metadata={"content_type": "text/plain"},
|
||||
)
|
||||
assert stored_key == key
|
||||
|
||||
retrieved = await s3_storage.retrieve(key)
|
||||
assert retrieved == content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_storage_exists_and_delete(s3_storage):
|
||||
"""Store, check exists=True, delete, check exists=False."""
|
||||
content = b"File to be deleted."
|
||||
key = f"test/{uuid.uuid4()}.txt"
|
||||
|
||||
await s3_storage.store(file_data=content, key=key)
|
||||
|
||||
assert await s3_storage.exists(key) is True
|
||||
|
||||
await s3_storage.delete(key)
|
||||
|
||||
assert await s3_storage.exists(key) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_storage_file_not_found(s3_storage):
|
||||
"""Retrieve a non-existent key, expect FileNotFoundError."""
|
||||
with pytest.raises(FileNotFoundError):
|
||||
await s3_storage.retrieve(f"nonexistent/{uuid.uuid4()}.txt")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_storage_get_download_url(s3_storage):
|
||||
"""Store a file, get a presigned URL, verify it's a valid URL string."""
|
||||
content = b"Presigned URL test content."
|
||||
key = f"test/{uuid.uuid4()}.txt"
|
||||
|
||||
await s3_storage.store(file_data=content, key=key)
|
||||
|
||||
url = await s3_storage.get_download_url(key, expires_in=300)
|
||||
assert isinstance(url, str)
|
||||
assert url.startswith("http")
|
||||
assert key in url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_file_retain_api_end_to_end(seaweedfs_container, memory_no_llm_verify):
|
||||
"""Full HTTP API flow: upload file via /files/retain with S3 storage backend."""
|
||||
from hindsight_api.api.http import create_app
|
||||
from hindsight_api.engine.storage.s3 import S3FileStorage
|
||||
|
||||
# Swap the engine's file storage to use the SeaweedFS-backed S3 storage
|
||||
original_storage = memory_no_llm_verify._file_storage
|
||||
s3_storage = S3FileStorage(
|
||||
bucket=seaweedfs_container["bucket"],
|
||||
region="us-east-1",
|
||||
endpoint=seaweedfs_container["endpoint"],
|
||||
access_key_id=seaweedfs_container["access_key"],
|
||||
secret_access_key=seaweedfs_container["secret_key"],
|
||||
)
|
||||
memory_no_llm_verify._file_storage = s3_storage
|
||||
|
||||
try:
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
bank_id = f"test-s3-bank-{uuid.uuid4().hex[:8]}"
|
||||
bank_response = await client.put(f"/v1/default/banks/{bank_id}", json={"name": "S3 Test Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
txt_content = b"Alice works at Acme Corp. She joined in 2024."
|
||||
request_data = {
|
||||
"document_tags": ["s3-test"],
|
||||
"async": True,
|
||||
}
|
||||
|
||||
files = {"files": ("notes.txt", txt_content, "text/plain")}
|
||||
data = {"request": json.dumps(request_data)}
|
||||
|
||||
response = await client.post(
|
||||
f"/v1/default/banks/{bank_id}/files/retain",
|
||||
files=files,
|
||||
data=data,
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "operation_ids" in result
|
||||
assert len(result["operation_ids"]) == 1
|
||||
finally:
|
||||
memory_no_llm_verify._file_storage = original_storage
|
||||
@@ -1,491 +0,0 @@
|
||||
"""
|
||||
Tests for hierarchical configuration system.
|
||||
|
||||
Tests config resolution hierarchy (global → tenant → bank),
|
||||
key normalization, API endpoints, validation, and caching.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.config import HindsightConfig, normalize_config_dict, normalize_config_key
|
||||
from hindsight_api.config_resolver import ConfigResolver
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
# Enable bank config API for all tests in this module
|
||||
os.environ["HINDSIGHT_API_ENABLE_BANK_CONFIG_API"] = "true"
|
||||
|
||||
|
||||
class MockTenantExtension(TenantExtension):
|
||||
"""Mock tenant extension for testing tenant-level config."""
|
||||
|
||||
def __init__(self, tenant_config: dict):
|
||||
self.tenant_config = tenant_config
|
||||
|
||||
async def authenticate(self, context):
|
||||
from hindsight_api.extensions.tenant import TenantContext
|
||||
|
||||
return TenantContext(schema_name="public")
|
||||
|
||||
async def list_tenants(self):
|
||||
from hindsight_api.extensions.tenant import Tenant
|
||||
|
||||
return [Tenant(schema="public")]
|
||||
|
||||
async def get_tenant_config(self, context):
|
||||
"""Return mock tenant config."""
|
||||
return self.tenant_config
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_key_normalization():
|
||||
"""Test that env var keys are normalized to Python field names."""
|
||||
# Test basic normalization
|
||||
assert normalize_config_key("HINDSIGHT_API_LLM_PROVIDER") == "llm_provider"
|
||||
assert normalize_config_key("HINDSIGHT_API_LLM_MODEL") == "llm_model"
|
||||
assert normalize_config_key("HINDSIGHT_API_RETAIN_LLM_PROVIDER") == "retain_llm_provider"
|
||||
|
||||
# Test already normalized keys
|
||||
assert normalize_config_key("llm_provider") == "llm_provider"
|
||||
assert normalize_config_key("llm_model") == "llm_model"
|
||||
|
||||
# Test dict normalization
|
||||
input_dict = {
|
||||
"HINDSIGHT_API_LLM_PROVIDER": "openai",
|
||||
"HINDSIGHT_API_LLM_MODEL": "gpt-4",
|
||||
"llm_base_url": "https://api.openai.com",
|
||||
}
|
||||
expected = {"llm_provider": "openai", "llm_model": "gpt-4", "llm_base_url": "https://api.openai.com"}
|
||||
assert normalize_config_dict(input_dict) == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hierarchical_fields_categorization():
|
||||
"""Test that fields are correctly categorized as configurable, credentials, or static."""
|
||||
configurable = HindsightConfig.get_configurable_fields()
|
||||
credentials = HindsightConfig.get_credential_fields()
|
||||
static = HindsightConfig.get_static_fields()
|
||||
|
||||
# Verify no overlap between configurable and credentials
|
||||
assert len(configurable & credentials) == 0, "Configurable fields should not include credentials"
|
||||
|
||||
# Verify configurable fields include behavioral settings (safe to modify)
|
||||
assert "retain_extraction_mode" in configurable
|
||||
assert "enable_observations" in configurable
|
||||
assert "retain_chunk_size" in configurable
|
||||
assert "retain_custom_instructions" in configurable
|
||||
|
||||
# Verify count is correct (only 4 fields)
|
||||
assert len(configurable) == 4
|
||||
|
||||
# Verify credential fields (NEVER exposed)
|
||||
assert "llm_api_key" in credentials
|
||||
assert "llm_base_url" in credentials
|
||||
assert "retain_llm_api_key" in credentials
|
||||
assert "reflect_llm_api_key" in credentials
|
||||
|
||||
# Verify static fields include server settings AND non-configurable LLM fields
|
||||
assert "database_url" in static
|
||||
assert "port" in static
|
||||
assert "host" in static
|
||||
assert "embeddings_provider" in static
|
||||
assert "reranker_provider" in static
|
||||
assert "worker_enabled" in static
|
||||
assert "llm_provider" in static # Not configurable (needs presets)
|
||||
assert "llm_model" in static # Not configurable (needs presets)
|
||||
assert "graph_retriever" in static # Performance tuning, not configurable
|
||||
assert "llm_max_concurrent" in static # Performance tuning, not configurable
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_hierarchy_resolution(memory, request_context):
|
||||
"""Test that config resolution follows global → tenant → bank hierarchy."""
|
||||
bank_id = "test-hierarchy-bank"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Set up mock tenant extension with tenant-level config (use configurable fields only)
|
||||
tenant_config = {"retain_chunk_size": 5000, "retain_extraction_mode": "tenant-mode"}
|
||||
mock_tenant = MockTenantExtension(tenant_config)
|
||||
|
||||
# Create config resolver with mock tenant extension
|
||||
resolver = ConfigResolver(pool=memory._pool, tenant_extension=mock_tenant)
|
||||
|
||||
# Test 1: Global config only (no overrides)
|
||||
context = RequestContext(api_key=None, api_key_id=None, tenant_id=None, internal=False)
|
||||
config = await resolver.get_bank_config(bank_id, context)
|
||||
|
||||
# Should have configurable fields from global config (NOT credentials or llm_provider/model)
|
||||
assert "retain_chunk_size" in config # Configurable field
|
||||
assert "llm_api_key" not in config # Credential - never exposed
|
||||
assert "llm_provider" not in config # Not configurable (needs presets)
|
||||
|
||||
# Test 2: Add tenant-level overrides
|
||||
config = await resolver.get_bank_config(bank_id, context)
|
||||
|
||||
# Should apply tenant overrides (only configurable fields)
|
||||
assert config["retain_chunk_size"] == 5000 # Tenant override
|
||||
assert config["retain_extraction_mode"] == "tenant-mode" # Tenant override
|
||||
|
||||
# Test 3: Add bank-level overrides (should take precedence)
|
||||
await resolver.update_bank_config(
|
||||
bank_id,
|
||||
{"retain_chunk_size": 2000, "retain_extraction_mode": "bank-mode"}, # Override tenant settings
|
||||
context,
|
||||
)
|
||||
|
||||
# Config should reflect changes immediately (no caching)
|
||||
config = await resolver.get_bank_config(bank_id, context)
|
||||
|
||||
# Bank overrides should take precedence over tenant
|
||||
assert config["retain_chunk_size"] == 2000 # Bank override wins
|
||||
assert config["retain_extraction_mode"] == "bank-mode" # Bank override wins
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_validation_rejects_static_fields(memory, request_context):
|
||||
"""Test that attempting to override static fields raises ValueError."""
|
||||
bank_id = "test-validation-bank"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
resolver = ConfigResolver(pool=memory._pool)
|
||||
|
||||
# Test 1: Configurable fields should work
|
||||
await resolver.update_bank_config(bank_id, {"retain_chunk_size": 4000, "retain_extraction_mode": "verbose"})
|
||||
|
||||
# Test 2: Static fields should raise ValueError
|
||||
with pytest.raises(ValueError, match="Cannot override static"):
|
||||
await resolver.update_bank_config(bank_id, {"port": 9000})
|
||||
|
||||
with pytest.raises(ValueError, match="Cannot override static"):
|
||||
await resolver.update_bank_config(bank_id, {"database_url": "postgresql://fake"})
|
||||
|
||||
with pytest.raises(ValueError, match="Cannot override static"):
|
||||
await resolver.update_bank_config(bank_id, {"embeddings_provider": "openai"})
|
||||
|
||||
# Test 3: Credential fields should raise ValueError
|
||||
with pytest.raises(ValueError, match="Cannot set credential fields"):
|
||||
await resolver.update_bank_config(bank_id, {"llm_api_key": "sk-fake"})
|
||||
|
||||
# Test 4: Non-configurable LLM fields should raise ValueError (need presets)
|
||||
with pytest.raises(ValueError, match="Cannot override static"):
|
||||
await resolver.update_bank_config(bank_id, {"llm_model": "gpt-4"})
|
||||
|
||||
# Test 5: Mix of configurable and static should fail
|
||||
with pytest.raises(ValueError, match="Cannot override static"):
|
||||
await resolver.update_bank_config(bank_id, {"retain_chunk_size": 4000, "port": 9000})
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_freshness_across_updates(memory, request_context):
|
||||
"""Test that config changes are immediately visible (no stale cache)."""
|
||||
bank1 = "freshness-test-1"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank1, request_context=request_context)
|
||||
|
||||
resolver = ConfigResolver(pool=memory._pool)
|
||||
|
||||
# Test 1: Initial config reflects global defaults
|
||||
config1 = await resolver.get_bank_config(bank1, None)
|
||||
initial_chunk_size = config1["retain_chunk_size"]
|
||||
|
||||
# Test 2: Update config
|
||||
await resolver.update_bank_config(bank1, {"retain_chunk_size": 4000})
|
||||
|
||||
# Test 3: Next call should see updated value immediately (no stale cache)
|
||||
config2 = await resolver.get_bank_config(bank1, None)
|
||||
assert config2["retain_chunk_size"] == 4000
|
||||
|
||||
# Test 4: Multiple updates are all immediately visible
|
||||
await resolver.update_bank_config(bank1, {"retain_chunk_size": 4500})
|
||||
config3 = await resolver.get_bank_config(bank1, None)
|
||||
assert config3["retain_chunk_size"] == 4500
|
||||
|
||||
# Test 5: Reset restores global defaults immediately
|
||||
await resolver.reset_bank_config(bank1)
|
||||
config4 = await resolver.get_bank_config(bank1, None)
|
||||
assert config4["retain_chunk_size"] == initial_chunk_size # Back to global default
|
||||
|
||||
# Test 6: Each call returns a fresh config dict (not a cached reference)
|
||||
config5 = await resolver.get_bank_config(bank1, None)
|
||||
config6 = await resolver.get_bank_config(bank1, None)
|
||||
assert config5 is not config6 # Different object instances
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank1, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_reset_to_defaults(memory, request_context):
|
||||
"""Test that resetting config removes all bank-specific overrides."""
|
||||
bank_id = "test-reset-bank"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
resolver = ConfigResolver(pool=memory._pool)
|
||||
|
||||
# Add bank-specific overrides
|
||||
await resolver.update_bank_config(
|
||||
bank_id,
|
||||
{
|
||||
"retain_chunk_size": 5500,
|
||||
"retain_extraction_mode": "custom",
|
||||
"retain_custom_instructions": "Custom instructions",
|
||||
},
|
||||
)
|
||||
|
||||
# Verify overrides applied
|
||||
config = await resolver.get_bank_config(bank_id, None)
|
||||
assert config["retain_chunk_size"] == 5500
|
||||
assert config["retain_extraction_mode"] == "custom"
|
||||
assert config["retain_custom_instructions"] == "Custom instructions"
|
||||
|
||||
# Reset to defaults
|
||||
await resolver.reset_bank_config(bank_id)
|
||||
|
||||
# Verify overrides removed (back to global defaults)
|
||||
config_reset = await resolver.get_bank_config(bank_id, None)
|
||||
assert config_reset["retain_chunk_size"] != 5500 # Should be global default
|
||||
assert config_reset["retain_extraction_mode"] != "custom" # Should be global default
|
||||
|
||||
# Verify bank_config is empty
|
||||
bank_overrides = await resolver._load_bank_config(bank_id)
|
||||
assert bank_overrides == {}
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_supports_both_key_formats(memory, request_context):
|
||||
"""Test that API accepts both env var and Python field formats."""
|
||||
bank_id = "test-key-format-bank"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
resolver = ConfigResolver(pool=memory._pool)
|
||||
|
||||
# Test 1: Python field format
|
||||
await resolver.update_bank_config(bank_id, {"retain_chunk_size": 7000})
|
||||
|
||||
config = await resolver.get_bank_config(bank_id, None)
|
||||
assert config["retain_chunk_size"] == 7000
|
||||
|
||||
# Test 2: Env var format (should be normalized)
|
||||
await resolver.update_bank_config(bank_id, {"HINDSIGHT_API_RETAIN_CHUNK_SIZE": 8000})
|
||||
|
||||
config = await resolver.get_bank_config(bank_id, None)
|
||||
assert config["retain_chunk_size"] == 8000
|
||||
|
||||
# Test 3: Mixed format in same request
|
||||
await resolver.update_bank_config(
|
||||
bank_id,
|
||||
{
|
||||
"retain_chunk_size": 9000, # Python format
|
||||
"HINDSIGHT_API_RETAIN_EXTRACTION_MODE": "verbose", # Env format
|
||||
},
|
||||
)
|
||||
|
||||
config = await resolver.get_bank_config(bank_id, None)
|
||||
assert config["retain_chunk_size"] == 9000
|
||||
assert config["retain_extraction_mode"] == "verbose"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_only_configurable_fields_stored(memory, request_context):
|
||||
"""Test that only configurable fields are stored in bank config."""
|
||||
bank_id = "test-filter-bank"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
resolver = ConfigResolver(pool=memory._pool)
|
||||
|
||||
# Add valid configurable field
|
||||
await resolver.update_bank_config(bank_id, {"retain_chunk_size": 3500})
|
||||
|
||||
# Load bank config and verify only configurable fields present
|
||||
bank_overrides = await resolver._load_bank_config(bank_id)
|
||||
|
||||
for key in bank_overrides.keys():
|
||||
assert key in HindsightConfig.get_configurable_fields(), f"Non-configurable field {key} in bank config"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_get_bank_config_no_static_or_credential_fields_leak(memory, request_context):
|
||||
"""
|
||||
SECURITY TEST: Verify get_bank_config() only returns configurable fields (no static/credentials).
|
||||
|
||||
This prevents leaking sensitive system configuration like database URLs,
|
||||
API keys, LLM providers/models, worker counts, etc. when retrieving bank configuration.
|
||||
"""
|
||||
bank_id = "test-security-bank"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
resolver = ConfigResolver(pool=memory._pool)
|
||||
|
||||
# Get bank config
|
||||
config = await resolver.get_bank_config(bank_id, None)
|
||||
|
||||
# Get field categorizations
|
||||
configurable_fields = HindsightConfig.get_configurable_fields()
|
||||
credential_fields = HindsightConfig.get_credential_fields()
|
||||
static_fields = HindsightConfig.get_static_fields()
|
||||
|
||||
# SECURITY: Verify ONLY configurable fields are returned (NO static, NO credentials)
|
||||
for key in config.keys():
|
||||
assert key in configurable_fields, (
|
||||
f"SECURITY VIOLATION: Non-configurable field '{key}' returned by get_bank_config(). "
|
||||
f"Only configurable fields should be returned to prevent leaking system config."
|
||||
)
|
||||
assert key not in credential_fields, (
|
||||
f"SECURITY VIOLATION: Credential field '{key}' returned by get_bank_config(). "
|
||||
f"Credentials must NEVER be exposed via API."
|
||||
)
|
||||
|
||||
# SECURITY: Verify specific sensitive fields are NOT present
|
||||
sensitive_fields = [
|
||||
"database_url", "api_port", "host", "worker_count", # Infrastructure
|
||||
"llm_api_key", "llm_base_url", # Credentials
|
||||
"retain_llm_api_key", "reflect_llm_api_key", # More credentials
|
||||
"llm_provider", "llm_model", # Not configurable (need presets)
|
||||
]
|
||||
for field in sensitive_fields:
|
||||
assert field not in config, (
|
||||
f"SECURITY VIOLATION: Sensitive field '{field}' returned by get_bank_config(). "
|
||||
f"Must not be exposed via bank config API."
|
||||
)
|
||||
|
||||
# Verify we have the expected configurable fields (small set)
|
||||
expected_configurable = ["retain_chunk_size", "retain_extraction_mode", "enable_observations"]
|
||||
for field in expected_configurable:
|
||||
assert field in config, f"Expected configurable field '{field}' missing from config"
|
||||
|
||||
# Should have a small number of configurable fields (not hundreds)
|
||||
assert len(config) < 20, f"Too many fields returned: {len(config)}"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_permissions_system(memory, request_context):
|
||||
"""
|
||||
Test that tenant extension can control which fields banks are allowed to modify.
|
||||
|
||||
Tests get_allowed_config_fields() permission system.
|
||||
"""
|
||||
bank_id = "test-permissions-bank"
|
||||
|
||||
class PermissionTenantExtension(TenantExtension):
|
||||
"""Mock tenant extension with configurable permissions."""
|
||||
|
||||
def __init__(self, allowed_fields: set[str] | None):
|
||||
self.allowed_fields = allowed_fields
|
||||
|
||||
async def authenticate(self, context):
|
||||
from hindsight_api.extensions.tenant import TenantContext
|
||||
|
||||
return TenantContext(schema_name="public")
|
||||
|
||||
async def list_tenants(self):
|
||||
from hindsight_api.extensions.tenant import Tenant
|
||||
|
||||
return [Tenant(schema="public")]
|
||||
|
||||
async def get_allowed_config_fields(self, context, bank_id):
|
||||
"""Return configured allowed fields."""
|
||||
return self.allowed_fields
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Test 1: None = allow all configurable fields
|
||||
extension = PermissionTenantExtension(allowed_fields=None)
|
||||
resolver = ConfigResolver(pool=memory._pool, tenant_extension=extension)
|
||||
|
||||
await resolver.update_bank_config(
|
||||
bank_id, {"retain_chunk_size": 4000, "retain_extraction_mode": "verbose"}, request_context
|
||||
)
|
||||
config = await resolver.get_bank_config(bank_id, request_context)
|
||||
assert config["retain_chunk_size"] == 4000
|
||||
assert config["retain_extraction_mode"] == "verbose"
|
||||
|
||||
# Reset for next test
|
||||
await resolver.reset_bank_config(bank_id)
|
||||
|
||||
# Test 2: Specific set = only those fields allowed
|
||||
extension = PermissionTenantExtension(allowed_fields={"retain_chunk_size"})
|
||||
resolver = ConfigResolver(pool=memory._pool, tenant_extension=extension)
|
||||
|
||||
# Should allow retain_chunk_size
|
||||
await resolver.update_bank_config(bank_id, {"retain_chunk_size": 5000}, request_context)
|
||||
config = await resolver.get_bank_config(bank_id, request_context)
|
||||
assert config["retain_chunk_size"] == 5000
|
||||
|
||||
# Should reject retain_extraction_mode (not in allowed list)
|
||||
with pytest.raises(ValueError, match="Not allowed to modify fields"):
|
||||
await resolver.update_bank_config(bank_id, {"retain_extraction_mode": "verbose"}, request_context)
|
||||
|
||||
# Should reject mix of allowed and disallowed
|
||||
with pytest.raises(ValueError, match="Not allowed to modify fields"):
|
||||
await resolver.update_bank_config(
|
||||
bank_id, {"retain_chunk_size": 6000, "retain_extraction_mode": "verbose"}, request_context
|
||||
)
|
||||
|
||||
# Reset for next test
|
||||
await resolver.reset_bank_config(bank_id)
|
||||
|
||||
# Test 3: Empty set = no modifications allowed (read-only)
|
||||
extension = PermissionTenantExtension(allowed_fields=set())
|
||||
resolver = ConfigResolver(pool=memory._pool, tenant_extension=extension)
|
||||
|
||||
with pytest.raises(ValueError, match="Not allowed to modify fields"):
|
||||
await resolver.update_bank_config(bank_id, {"retain_chunk_size": 7000}, request_context)
|
||||
|
||||
# Test 4: get_bank_config should filter response based on permissions
|
||||
extension = PermissionTenantExtension(allowed_fields={"retain_chunk_size", "enable_observations"})
|
||||
resolver = ConfigResolver(pool=memory._pool, tenant_extension=extension)
|
||||
|
||||
config = await resolver.get_bank_config(bank_id, request_context)
|
||||
|
||||
# Should only return allowed fields
|
||||
assert "retain_chunk_size" in config
|
||||
assert "enable_observations" in config
|
||||
# Other configurable fields should be filtered out
|
||||
assert "retain_extraction_mode" not in config
|
||||
assert "retain_custom_instructions" not in config
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -528,7 +528,7 @@ async def test_delete_bank(api_client):
|
||||
{
|
||||
"content": "Bob is the CTO and leads the engineering team.",
|
||||
"context": "team info",
|
||||
"document_id": "team-doc-2",
|
||||
"document_id": "team-doc-1",
|
||||
},
|
||||
]
|
||||
},
|
||||
|
||||
@@ -1,72 +0,0 @@
|
||||
"""
|
||||
Integration tests for the Iris file parser.
|
||||
|
||||
Tests are skipped automatically if HINDSIGHT_API_FILE_PARSER_IRIS_TOKEN
|
||||
and HINDSIGHT_API_FILE_PARSER_IRIS_ORG_ID are not set in the environment.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config import ENV_FILE_PARSER_IRIS_ORG_ID, ENV_FILE_PARSER_IRIS_TOKEN
|
||||
from hindsight_api.engine.parsers.iris import IrisParser
|
||||
|
||||
_token = os.getenv(ENV_FILE_PARSER_IRIS_TOKEN)
|
||||
_org_id = os.getenv(ENV_FILE_PARSER_IRIS_ORG_ID)
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not (_token and _org_id),
|
||||
reason="HINDSIGHT_API_FILE_PARSER_IRIS_TOKEN and HINDSIGHT_API_FILE_PARSER_IRIS_ORG_ID not set",
|
||||
)
|
||||
|
||||
# Minimal valid PDF with the text "Hello from Hindsight"
|
||||
_SAMPLE_PDF = b"""%PDF-1.4
|
||||
1 0 obj
|
||||
<< /Type /Catalog /Pages 2 0 R >>
|
||||
endobj
|
||||
2 0 obj
|
||||
<< /Type /Pages /Kids [3 0 R] /Count 1 >>
|
||||
endobj
|
||||
3 0 obj
|
||||
<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792]
|
||||
/Contents 4 0 R /Resources << /Font << /F1 << /Type /Font /Subtype /Type1 /BaseFont /Helvetica >> >> >> >>
|
||||
endobj
|
||||
4 0 obj
|
||||
<< /Length 44 >>
|
||||
stream
|
||||
BT /F1 12 Tf 100 700 Td (Hello from Hindsight) Tj ET
|
||||
endstream
|
||||
endobj
|
||||
xref
|
||||
0 5
|
||||
0000000000 65535 f
|
||||
0000000009 00000 n
|
||||
0000000058 00000 n
|
||||
0000000115 00000 n
|
||||
0000000274 00000 n
|
||||
trailer << /Size 5 /Root 1 0 R >>
|
||||
startxref
|
||||
369
|
||||
%%EOF"""
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def iris_parser() -> IrisParser:
|
||||
return IrisParser(token=_token, org_id=_org_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_iris_parser_converts_pdf(iris_parser: IrisParser):
|
||||
"""IrisParser should extract text from a valid PDF."""
|
||||
result = await iris_parser.convert(_SAMPLE_PDF, "sample.pdf")
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_iris_parser_name(iris_parser: IrisParser):
|
||||
"""IrisParser.name() should return 'iris'."""
|
||||
assert iris_parser.name() == "iris"
|
||||
|
||||
|
||||
@@ -12,9 +12,9 @@ import pytest
|
||||
@pytest.fixture(autouse=True)
|
||||
def enable_observations():
|
||||
"""Enable observations for all tests in this module."""
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
config = _get_raw_config()
|
||||
config = get_config()
|
||||
original_value = config.enable_observations
|
||||
config.enable_observations = True
|
||||
yield
|
||||
|
||||
@@ -1,392 +0,0 @@
|
||||
"""
|
||||
Tests for LiteLLMSDKCrossEncoder.
|
||||
|
||||
Tests the LiteLLM SDK-based cross-encoder implementation for reranking.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.cross_encoder import LiteLLMSDKCrossEncoder, create_cross_encoder_from_env
|
||||
|
||||
|
||||
class TestLiteLLMSDKCrossEncoder:
|
||||
"""Test suite for LiteLLMSDKCrossEncoder class."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialization_success(self):
|
||||
"""Test successful initialization with valid config."""
|
||||
encoder = LiteLLMSDKCrossEncoder(
|
||||
api_key="test_key",
|
||||
model="deepinfra/Qwen3-reranker-8B",
|
||||
)
|
||||
|
||||
assert encoder.provider_name == "litellm-sdk"
|
||||
assert encoder.api_key == "test_key"
|
||||
assert encoder.model == "deepinfra/Qwen3-reranker-8B"
|
||||
assert encoder._initialized is False
|
||||
|
||||
# Mock the litellm import
|
||||
mock_litellm = MagicMock()
|
||||
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
||||
await encoder.initialize()
|
||||
assert encoder._initialized is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialization_missing_package(self):
|
||||
"""Test initialization fails when litellm package is missing."""
|
||||
encoder = LiteLLMSDKCrossEncoder(
|
||||
api_key="test_key",
|
||||
model="cohere/rerank-english-v3.0",
|
||||
)
|
||||
|
||||
with patch.dict("sys.modules", {"litellm": None}):
|
||||
with pytest.raises(ImportError, match="litellm is required"):
|
||||
await encoder.initialize()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialization_idempotent(self):
|
||||
"""Test that calling initialize() multiple times is safe."""
|
||||
encoder = LiteLLMSDKCrossEncoder(
|
||||
api_key="test_key",
|
||||
model="cohere/rerank-english-v3.0",
|
||||
)
|
||||
|
||||
mock_litellm = MagicMock()
|
||||
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
||||
await encoder.initialize()
|
||||
assert encoder._initialized is True
|
||||
|
||||
# Second call should be no-op
|
||||
await encoder.initialize()
|
||||
assert encoder._initialized is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_predict_single_query(self):
|
||||
"""Test prediction with a single query and multiple documents."""
|
||||
encoder = LiteLLMSDKCrossEncoder(
|
||||
api_key="test_key",
|
||||
model="deepinfra/Qwen3-reranker-8B",
|
||||
)
|
||||
|
||||
# Create mock response with results as TypedDicts
|
||||
mock_response = MagicMock()
|
||||
mock_response.results = [
|
||||
{"index": 0, "relevance_score": 0.9},
|
||||
{"index": 1, "relevance_score": 0.7},
|
||||
{"index": 2, "relevance_score": 0.5},
|
||||
]
|
||||
|
||||
mock_litellm = MagicMock()
|
||||
mock_litellm.arerank = AsyncMock(return_value=mock_response)
|
||||
|
||||
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
||||
await encoder.initialize()
|
||||
|
||||
pairs = [
|
||||
("What is Python?", "Python is a programming language"),
|
||||
("What is Python?", "Python is a snake"),
|
||||
("What is Python?", "Python is a British comedy group"),
|
||||
]
|
||||
|
||||
scores = await encoder.predict(pairs)
|
||||
|
||||
assert len(scores) == 3
|
||||
assert scores == [0.9, 0.7, 0.5]
|
||||
|
||||
# Verify arerank was called correctly
|
||||
mock_litellm.arerank.assert_called_once()
|
||||
call_args = mock_litellm.arerank.call_args
|
||||
assert call_args.kwargs["model"] == "deepinfra/Qwen3-reranker-8B"
|
||||
assert call_args.kwargs["query"] == "What is Python?"
|
||||
assert len(call_args.kwargs["documents"]) == 3
|
||||
assert call_args.kwargs["api_key"] == "test_key"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_predict_multiple_queries(self):
|
||||
"""Test prediction with multiple different queries (grouped efficiently)."""
|
||||
encoder = LiteLLMSDKCrossEncoder(
|
||||
api_key="test_key",
|
||||
model="cohere/rerank-english-v3.0",
|
||||
)
|
||||
|
||||
# First query response
|
||||
mock_response1 = MagicMock()
|
||||
mock_response1.results = [
|
||||
{"index": 0, "relevance_score": 0.9},
|
||||
{"index": 1, "relevance_score": 0.7},
|
||||
]
|
||||
|
||||
# Second query response
|
||||
mock_response2 = MagicMock()
|
||||
mock_response2.results = [
|
||||
{"index": 0, "relevance_score": 0.8},
|
||||
]
|
||||
|
||||
mock_litellm = MagicMock()
|
||||
mock_litellm.arerank = AsyncMock(side_effect=[mock_response1, mock_response2])
|
||||
|
||||
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
||||
await encoder.initialize()
|
||||
|
||||
pairs = [
|
||||
("What is Python?", "Python is a programming language"),
|
||||
("What is Python?", "Python is a snake"),
|
||||
("What is Java?", "Java is a programming language"),
|
||||
]
|
||||
|
||||
scores = await encoder.predict(pairs)
|
||||
|
||||
assert len(scores) == 3
|
||||
assert scores[0] == 0.9 # First query, first doc
|
||||
assert scores[1] == 0.7 # First query, second doc
|
||||
assert scores[2] == 0.8 # Second query, first doc
|
||||
|
||||
# Verify arerank was called twice (once per unique query)
|
||||
assert mock_litellm.arerank.call_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_predict_empty_pairs(self):
|
||||
"""Test prediction with empty input."""
|
||||
encoder = LiteLLMSDKCrossEncoder(
|
||||
api_key="test_key",
|
||||
model="cohere/rerank-english-v3.0",
|
||||
)
|
||||
|
||||
mock_litellm = MagicMock()
|
||||
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
||||
await encoder.initialize()
|
||||
scores = await encoder.predict([])
|
||||
assert scores == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_predict_not_initialized(self):
|
||||
"""Test that predict fails if encoder not initialized."""
|
||||
encoder = LiteLLMSDKCrossEncoder(
|
||||
api_key="test_key",
|
||||
model="cohere/rerank-english-v3.0",
|
||||
)
|
||||
|
||||
pairs = [("query", "document")]
|
||||
|
||||
with pytest.raises(RuntimeError, match="not initialized"):
|
||||
await encoder.predict(pairs)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_predict_error_handling(self):
|
||||
"""Test that errors during prediction are raised."""
|
||||
encoder = LiteLLMSDKCrossEncoder(
|
||||
api_key="test_key",
|
||||
model="cohere/rerank-english-v3.0",
|
||||
)
|
||||
|
||||
# Mock litellm to raise an error
|
||||
mock_litellm = MagicMock()
|
||||
mock_litellm.arerank = AsyncMock(side_effect=Exception("API Error"))
|
||||
|
||||
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
||||
await encoder.initialize()
|
||||
|
||||
pairs = [
|
||||
("What is Python?", "Python is a programming language"),
|
||||
]
|
||||
|
||||
# Should raise the exception
|
||||
with pytest.raises(Exception, match="API Error"):
|
||||
await encoder.predict(pairs)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_api_base(self):
|
||||
"""Test that custom API base URL is passed to rerank calls."""
|
||||
encoder = LiteLLMSDKCrossEncoder(
|
||||
api_key="test_key",
|
||||
model="cohere/rerank-english-v3.0",
|
||||
api_base="https://custom.api.example.com",
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.results = [
|
||||
{"index": 0, "relevance_score": 0.9},
|
||||
]
|
||||
|
||||
mock_litellm = MagicMock()
|
||||
mock_litellm.arerank = AsyncMock(return_value=mock_response)
|
||||
|
||||
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
||||
await encoder.initialize()
|
||||
|
||||
# Test that api_base is passed to arerank
|
||||
pairs = [("query", "document")]
|
||||
scores = await encoder.predict(pairs)
|
||||
|
||||
assert scores == [0.9]
|
||||
mock_litellm.arerank.assert_called_once()
|
||||
call_args = mock_litellm.arerank.call_args
|
||||
assert call_args.kwargs["api_base"] == "https://custom.api.example.com"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_with_direct_score_list(self):
|
||||
"""Test handling of response format with direct score list."""
|
||||
encoder = LiteLLMSDKCrossEncoder(
|
||||
api_key="test_key",
|
||||
model="some-provider/model",
|
||||
)
|
||||
|
||||
# Mock litellm to return direct list of scores
|
||||
mock_litellm = MagicMock()
|
||||
mock_litellm.arerank = AsyncMock(return_value=[0.9, 0.7, 0.5])
|
||||
|
||||
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
||||
await encoder.initialize()
|
||||
|
||||
pairs = [
|
||||
("query", "doc1"),
|
||||
("query", "doc2"),
|
||||
("query", "doc3"),
|
||||
]
|
||||
|
||||
scores = await encoder.predict(pairs)
|
||||
|
||||
assert scores == [0.9, 0.7, 0.5]
|
||||
|
||||
|
||||
class TestFactoryFunction:
|
||||
"""Test suite for create_cross_encoder_from_env factory function."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_litellm_sdk_from_env(self):
|
||||
"""Test creating LiteLLM SDK cross-encoder from environment variables."""
|
||||
env_vars = {
|
||||
"HINDSIGHT_API_RERANKER_PROVIDER": "litellm-sdk",
|
||||
"HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY": "test_key",
|
||||
"HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL": "deepinfra/Qwen3-reranker-8B",
|
||||
}
|
||||
|
||||
with patch.dict(os.environ, env_vars, clear=False):
|
||||
# Need to reload config to pick up env vars
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
with patch("hindsight_api.config.get_config", return_value=config):
|
||||
encoder = create_cross_encoder_from_env()
|
||||
|
||||
assert isinstance(encoder, LiteLLMSDKCrossEncoder)
|
||||
assert encoder.api_key == "test_key"
|
||||
assert encoder.model == "deepinfra/Qwen3-reranker-8B"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_litellm_sdk_missing_api_key(self):
|
||||
"""Test that factory raises error when API key is missing."""
|
||||
env_vars = {
|
||||
"HINDSIGHT_API_RERANKER_PROVIDER": "litellm-sdk",
|
||||
"HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL": "deepinfra/Qwen3-reranker-8B",
|
||||
}
|
||||
|
||||
with patch.dict(os.environ, env_vars, clear=False):
|
||||
# Remove API key if set
|
||||
if "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY" in os.environ:
|
||||
del os.environ["HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY"]
|
||||
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
with patch("hindsight_api.config.get_config", return_value=config):
|
||||
with pytest.raises(ValueError, match="HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY is required"):
|
||||
create_cross_encoder_from_env()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_litellm_sdk_with_custom_api_base(self):
|
||||
"""Test creating LiteLLM SDK cross-encoder with custom API base."""
|
||||
env_vars = {
|
||||
"HINDSIGHT_API_RERANKER_PROVIDER": "litellm-sdk",
|
||||
"HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY": "test_key",
|
||||
"HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL": "cohere/rerank-english-v3.0",
|
||||
"HINDSIGHT_API_RERANKER_LITELLM_SDK_API_BASE": "https://custom.api.example.com",
|
||||
}
|
||||
|
||||
with patch.dict(os.environ, env_vars, clear=False):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
with patch("hindsight_api.config.get_config", return_value=config):
|
||||
encoder = create_cross_encoder_from_env()
|
||||
|
||||
assert isinstance(encoder, LiteLLMSDKCrossEncoder)
|
||||
assert encoder.api_base == "https://custom.api.example.com"
|
||||
|
||||
|
||||
class TestLiteLLMSDKCohereCrossEncoder:
|
||||
"""Tests for LiteLLM SDK calling Cohere (runs in CI with COHERE_API_KEY)."""
|
||||
|
||||
@pytest.fixture
|
||||
async def litellm_cohere_cross_encoder(self):
|
||||
"""Create LiteLLM SDK cross-encoder instance for Cohere."""
|
||||
if not os.environ.get("COHERE_API_KEY"):
|
||||
pytest.skip("Cohere API key not available (set COHERE_API_KEY)")
|
||||
|
||||
encoder = LiteLLMSDKCrossEncoder(
|
||||
api_key=os.environ["COHERE_API_KEY"],
|
||||
model="cohere/rerank-english-v3.0",
|
||||
)
|
||||
await encoder.initialize()
|
||||
return encoder
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_sdk_cohere_initialization(self, litellm_cohere_cross_encoder):
|
||||
"""Test that LiteLLM SDK Cohere cross-encoder initializes correctly."""
|
||||
assert litellm_cohere_cross_encoder.provider_name == "litellm-sdk"
|
||||
assert litellm_cohere_cross_encoder.model == "cohere/rerank-english-v3.0"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_sdk_cohere_predict(self, litellm_cohere_cross_encoder):
|
||||
"""Test that LiteLLM SDK can call Cohere rerank API."""
|
||||
pairs = [
|
||||
("What is the capital of France?", "Paris is the capital of France."),
|
||||
("What is the capital of France?", "The Eiffel Tower is in Paris."),
|
||||
("What is the capital of France?", "Python is a programming language."),
|
||||
]
|
||||
scores = await litellm_cohere_cross_encoder.predict(pairs)
|
||||
|
||||
assert len(scores) == 3
|
||||
assert all(isinstance(s, float) for s in scores)
|
||||
# The first result should be most relevant
|
||||
assert scores[0] > scores[2], "Direct answer should score higher than unrelated text"
|
||||
# All scores should be in valid range
|
||||
assert all(0.0 <= score <= 1.0 for score in scores)
|
||||
|
||||
|
||||
class TestIntegration:
|
||||
"""Integration tests with real API (optional - requires API keys)."""
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("DEEPINFRA_API_KEY"),
|
||||
reason="DEEPINFRA_API_KEY not set - skipping integration test",
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_deepinfra_api(self):
|
||||
"""Test with real DeepInfra API (requires DEEPINFRA_API_KEY env var)."""
|
||||
encoder = LiteLLMSDKCrossEncoder(
|
||||
api_key=os.environ["DEEPINFRA_API_KEY"],
|
||||
model="deepinfra/Qwen3-reranker-8B",
|
||||
)
|
||||
|
||||
await encoder.initialize()
|
||||
|
||||
pairs = [
|
||||
("What is Python?", "Python is a high-level programming language"),
|
||||
("What is Python?", "Python is a species of snake"),
|
||||
("What is Python?", "Python is unrelated text about cars"),
|
||||
]
|
||||
|
||||
scores = await encoder.predict(pairs)
|
||||
|
||||
# First doc should have highest score (most relevant)
|
||||
assert len(scores) == 3
|
||||
assert scores[0] > scores[1]
|
||||
assert scores[1] > scores[2]
|
||||
assert all(0.0 <= score <= 1.0 for score in scores)
|
||||
@@ -1,387 +0,0 @@
|
||||
"""
|
||||
Tests for LiteLLM SDK embeddings implementation.
|
||||
|
||||
These tests cover:
|
||||
1. Initialization (success, missing package, missing API key, idempotent)
|
||||
2. Encode (single text, multiple texts, batching, error handling)
|
||||
3. Provider-specific configuration (Cohere, OpenAI, etc.)
|
||||
4. Factory function (create from env, validation errors)
|
||||
5. Dimension detection
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config import (
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY,
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
ENV_EMBEDDINGS_PROVIDER,
|
||||
HindsightConfig,
|
||||
)
|
||||
from hindsight_api.engine.embeddings import LiteLLMSDKEmbeddings, create_embeddings_from_env
|
||||
|
||||
|
||||
class TestLiteLLMSDKEmbeddings:
|
||||
"""Unit tests for LiteLLMSDKEmbeddings with mocked litellm responses."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_litellm(self):
|
||||
"""Mock litellm module."""
|
||||
mock = MagicMock()
|
||||
|
||||
# Mock aembedding (async) for initialization
|
||||
mock_response = MagicMock()
|
||||
mock_response.data = [{"embedding": [0.1] * 768, "index": 0}]
|
||||
mock.aembedding = AsyncMock(return_value=mock_response)
|
||||
|
||||
# Mock embedding (sync) for encode
|
||||
mock_sync_response = MagicMock()
|
||||
mock_sync_response.data = [
|
||||
{"embedding": [0.1] * 768, "index": 0},
|
||||
{"embedding": [0.2] * 768, "index": 1},
|
||||
]
|
||||
mock.embedding = MagicMock(return_value=mock_sync_response)
|
||||
|
||||
return mock
|
||||
|
||||
@pytest.fixture
|
||||
async def embeddings(self, mock_litellm):
|
||||
"""Create initialized LiteLLMSDKEmbeddings instance."""
|
||||
emb = LiteLLMSDKEmbeddings(
|
||||
api_key="test_key",
|
||||
model="cohere/embed-english-v3.0",
|
||||
api_base=None,
|
||||
batch_size=100,
|
||||
timeout=60.0,
|
||||
)
|
||||
# Manually set the mock (simulating successful initialization)
|
||||
emb._litellm = mock_litellm
|
||||
emb._dimension = 768
|
||||
return emb
|
||||
|
||||
async def test_initialization_success(self, mock_litellm):
|
||||
"""Test successful initialization."""
|
||||
with patch("builtins.__import__", side_effect=lambda name, *args: mock_litellm if name == "litellm" else __import__(name, *args)):
|
||||
emb = LiteLLMSDKEmbeddings(
|
||||
api_key="test_key",
|
||||
model="cohere/embed-english-v3.0",
|
||||
api_base=None,
|
||||
batch_size=100,
|
||||
timeout=60.0,
|
||||
)
|
||||
|
||||
assert emb._litellm is None
|
||||
assert emb._dimension is None
|
||||
|
||||
await emb.initialize()
|
||||
|
||||
assert emb._litellm is not None
|
||||
assert emb._dimension == 768
|
||||
|
||||
# Verify test embedding was called
|
||||
mock_litellm.aembedding.assert_called_once_with(
|
||||
model="cohere/embed-english-v3.0",
|
||||
input=["test"],
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
async def test_initialization_missing_package(self):
|
||||
"""Test initialization fails gracefully when litellm is not installed."""
|
||||
def mock_import(name, *args):
|
||||
if name == "litellm":
|
||||
raise ImportError("No module named 'litellm'")
|
||||
return __import__(name, *args)
|
||||
|
||||
with patch("builtins.__import__", side_effect=mock_import):
|
||||
emb = LiteLLMSDKEmbeddings(
|
||||
api_key="test_key",
|
||||
model="cohere/embed-english-v3.0",
|
||||
api_base=None,
|
||||
batch_size=100,
|
||||
timeout=60.0,
|
||||
)
|
||||
|
||||
with pytest.raises(ImportError, match="litellm is required"):
|
||||
await emb.initialize()
|
||||
|
||||
async def test_initialization_idempotent(self, embeddings, mock_litellm):
|
||||
"""Test that calling initialize() multiple times is safe."""
|
||||
# embeddings._litellm is already set in fixture
|
||||
assert embeddings._litellm is not None
|
||||
|
||||
# Call again
|
||||
await embeddings.initialize()
|
||||
|
||||
# Should still have same litellm instance
|
||||
assert embeddings._litellm is not None
|
||||
|
||||
async def test_encode_single_text(self, embeddings, mock_litellm):
|
||||
"""Test encoding a single text."""
|
||||
# Set up mock response
|
||||
mock_litellm.embedding.return_value.data = [
|
||||
{"embedding": [0.5] * 768, "index": 0},
|
||||
]
|
||||
|
||||
result = embeddings.encode(["Hello world"])
|
||||
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 1
|
||||
assert len(result[0]) == 768
|
||||
assert all(isinstance(x, float) for x in result[0])
|
||||
assert all(abs(x - 0.5) < 0.001 for x in result[0])
|
||||
|
||||
# Verify call
|
||||
mock_litellm.embedding.assert_called_once_with(
|
||||
model="cohere/embed-english-v3.0",
|
||||
input=["Hello world"],
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
async def test_encode_multiple_texts(self, embeddings, mock_litellm):
|
||||
"""Test encoding multiple texts."""
|
||||
# Set up mock response
|
||||
mock_litellm.embedding.return_value.data = [
|
||||
{"embedding": [0.1] * 768, "index": 0},
|
||||
{"embedding": [0.2] * 768, "index": 1},
|
||||
{"embedding": [0.3] * 768, "index": 2},
|
||||
]
|
||||
|
||||
texts = ["First text", "Second text", "Third text"]
|
||||
result = embeddings.encode(texts)
|
||||
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 3
|
||||
assert len(result[0]) == 768
|
||||
assert len(result[1]) == 768
|
||||
assert len(result[2]) == 768
|
||||
assert all(abs(x - 0.1) < 0.001 for x in result[0])
|
||||
assert all(abs(x - 0.2) < 0.001 for x in result[1])
|
||||
assert all(abs(x - 0.3) < 0.001 for x in result[2])
|
||||
|
||||
async def test_encode_batching(self, embeddings, mock_litellm):
|
||||
"""Test that large inputs are batched correctly."""
|
||||
# Create embeddings with small batch size
|
||||
emb = LiteLLMSDKEmbeddings(
|
||||
api_key="test_key",
|
||||
model="cohere/embed-english-v3.0",
|
||||
api_base=None,
|
||||
batch_size=2, # Small batch for testing
|
||||
timeout=60.0,
|
||||
)
|
||||
emb._litellm = mock_litellm
|
||||
emb._initialized = True
|
||||
emb._dimension = 768
|
||||
|
||||
# Mock responses for each batch
|
||||
def mock_embedding_side_effect(model, input, **kwargs):
|
||||
mock_response = MagicMock()
|
||||
mock_response.data = [
|
||||
{"embedding": [float(i)] * 768, "index": i} for i in range(len(input))
|
||||
]
|
||||
return mock_response
|
||||
|
||||
mock_litellm.embedding.side_effect = mock_embedding_side_effect
|
||||
|
||||
# Encode 5 texts (should create 3 batches: 2, 2, 1)
|
||||
texts = [f"Text {i}" for i in range(5)]
|
||||
result = emb.encode(texts)
|
||||
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 5
|
||||
assert all(len(embedding) == 768 for embedding in result)
|
||||
|
||||
# Verify batching: should be called 3 times
|
||||
assert mock_litellm.embedding.call_count == 3
|
||||
|
||||
# Verify batch sizes
|
||||
calls = mock_litellm.embedding.call_args_list
|
||||
assert len(calls[0][1]["input"]) == 2 # First batch
|
||||
assert len(calls[1][1]["input"]) == 2 # Second batch
|
||||
assert len(calls[2][1]["input"]) == 1 # Third batch
|
||||
|
||||
async def test_encode_empty_list(self, embeddings):
|
||||
"""Test encoding empty list returns empty list."""
|
||||
result = embeddings.encode([])
|
||||
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 0
|
||||
|
||||
async def test_encode_before_initialization(self, mock_litellm):
|
||||
"""Test that encode raises error if not initialized."""
|
||||
emb = LiteLLMSDKEmbeddings(
|
||||
api_key="test_key",
|
||||
model="cohere/embed-english-v3.0",
|
||||
api_base=None,
|
||||
batch_size=100,
|
||||
timeout=60.0,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="not initialized"):
|
||||
emb.encode(["test"])
|
||||
|
||||
async def test_encode_error_handling(self, embeddings, mock_litellm):
|
||||
"""Test error handling during encoding."""
|
||||
# Make embedding raise an error
|
||||
mock_litellm.embedding.side_effect = Exception("API Error")
|
||||
|
||||
with pytest.raises(Exception, match="API Error"):
|
||||
embeddings.encode(["test"])
|
||||
|
||||
async def test_dimension_property(self, embeddings):
|
||||
"""Test dimension property."""
|
||||
assert embeddings.dimension == 768
|
||||
|
||||
async def test_dimension_before_initialization(self, mock_litellm):
|
||||
"""Test dimension raises error if not initialized."""
|
||||
emb = LiteLLMSDKEmbeddings(
|
||||
api_key="test_key",
|
||||
model="cohere/embed-english-v3.0",
|
||||
api_base=None,
|
||||
batch_size=100,
|
||||
timeout=60.0,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="not initialized"):
|
||||
_ = emb.dimension
|
||||
|
||||
async def test_custom_api_base(self, mock_litellm):
|
||||
"""Test custom API base URL is passed to embedding calls."""
|
||||
with patch("builtins.__import__", side_effect=lambda name, *args: mock_litellm if name == "litellm" else __import__(name, *args)):
|
||||
emb = LiteLLMSDKEmbeddings(
|
||||
api_key="test_key",
|
||||
model="cohere/embed-english-v3.0",
|
||||
api_base="https://custom.api.com",
|
||||
batch_size=100,
|
||||
timeout=60.0,
|
||||
)
|
||||
|
||||
await emb.initialize()
|
||||
|
||||
# Verify api_base is set
|
||||
assert emb.api_base == "https://custom.api.com"
|
||||
|
||||
# Verify api_base is passed to aembedding
|
||||
mock_litellm.aembedding.assert_called_once()
|
||||
call_args = mock_litellm.aembedding.call_args
|
||||
assert call_args.kwargs["api_base"] == "https://custom.api.com"
|
||||
|
||||
# Test encode also passes api_base
|
||||
mock_litellm.embedding.return_value.data = [{"embedding": [0.1] * 768, "index": 0}]
|
||||
emb.encode(["test"])
|
||||
|
||||
mock_litellm.embedding.assert_called_once()
|
||||
call_args = mock_litellm.embedding.call_args
|
||||
assert call_args.kwargs["api_base"] == "https://custom.api.com"
|
||||
|
||||
|
||||
class TestLiteLLMSDKEmbeddingsFactory:
|
||||
"""Test the factory function for creating LiteLLM SDK embeddings."""
|
||||
|
||||
def test_create_from_env_success(self, monkeypatch):
|
||||
"""Test creating embeddings from environment variables."""
|
||||
# Mock get_config() to return configured HindsightConfig
|
||||
mock_config = MagicMock()
|
||||
mock_config.embeddings_provider = "litellm-sdk"
|
||||
mock_config.embeddings_litellm_sdk_api_key = "test_key"
|
||||
mock_config.embeddings_litellm_sdk_model = "cohere/embed-english-v3.0"
|
||||
mock_config.embeddings_litellm_sdk_api_base = None
|
||||
|
||||
with patch("hindsight_api.config.get_config", return_value=mock_config):
|
||||
embeddings = create_embeddings_from_env()
|
||||
|
||||
assert isinstance(embeddings, LiteLLMSDKEmbeddings)
|
||||
assert embeddings.api_key == "test_key"
|
||||
assert embeddings.model == "cohere/embed-english-v3.0"
|
||||
|
||||
def test_create_from_env_missing_api_key(self, monkeypatch):
|
||||
"""Test that missing API key raises error."""
|
||||
# Mock get_config() with missing API key
|
||||
mock_config = MagicMock()
|
||||
mock_config.embeddings_provider = "litellm-sdk"
|
||||
mock_config.embeddings_litellm_sdk_api_key = None # Missing key
|
||||
mock_config.embeddings_litellm_sdk_model = "cohere/embed-english-v3.0"
|
||||
|
||||
with patch("hindsight_api.config.get_config", return_value=mock_config):
|
||||
with pytest.raises(ValueError, match=ENV_EMBEDDINGS_LITELLM_SDK_API_KEY):
|
||||
create_embeddings_from_env()
|
||||
|
||||
def test_create_from_env_with_api_base(self, monkeypatch):
|
||||
"""Test creating embeddings with custom API base."""
|
||||
# Mock get_config() with custom API base
|
||||
mock_config = MagicMock()
|
||||
mock_config.embeddings_provider = "litellm-sdk"
|
||||
mock_config.embeddings_litellm_sdk_api_key = "test_key"
|
||||
mock_config.embeddings_litellm_sdk_model = "cohere/embed-english-v3.0"
|
||||
mock_config.embeddings_litellm_sdk_api_base = "https://custom.api.com"
|
||||
|
||||
with patch("hindsight_api.config.get_config", return_value=mock_config):
|
||||
embeddings = create_embeddings_from_env()
|
||||
|
||||
assert isinstance(embeddings, LiteLLMSDKEmbeddings)
|
||||
assert embeddings.api_base == "https://custom.api.com"
|
||||
|
||||
|
||||
class TestLiteLLMSDKCohereEmbeddings:
|
||||
"""Integration tests calling real Cohere API (matches CI pattern)."""
|
||||
|
||||
@pytest.fixture
|
||||
async def litellm_cohere_embeddings(self):
|
||||
"""Create embeddings instance with real Cohere API key."""
|
||||
if not os.environ.get("COHERE_API_KEY"):
|
||||
pytest.skip("Cohere API key not available")
|
||||
|
||||
emb = LiteLLMSDKEmbeddings(
|
||||
api_key=os.environ["COHERE_API_KEY"],
|
||||
model="cohere/embed-english-v3.0",
|
||||
api_base=None,
|
||||
batch_size=100,
|
||||
timeout=60.0,
|
||||
)
|
||||
await emb.initialize()
|
||||
return emb
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_sdk_cohere_encode(self, litellm_cohere_embeddings):
|
||||
"""Test real Cohere API call for embeddings."""
|
||||
texts = [
|
||||
"The quick brown fox jumps over the lazy dog",
|
||||
"Machine learning is a subset of artificial intelligence",
|
||||
"Python is a popular programming language",
|
||||
]
|
||||
|
||||
result = litellm_cohere_embeddings.encode(texts)
|
||||
|
||||
# Verify result type and shape
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 3
|
||||
assert all(len(embedding) > 0 for embedding in result)
|
||||
assert all(isinstance(x, float) for x in result[0])
|
||||
|
||||
# Verify embeddings are not zeros (common API failure mode)
|
||||
for i, embedding in enumerate(result):
|
||||
assert not all(abs(x) < 0.0001 for x in embedding), f"Embedding {i} is all zeros"
|
||||
|
||||
# Verify embeddings are normalized (Cohere returns normalized vectors)
|
||||
for i, embedding in enumerate(result):
|
||||
norm = sum(x * x for x in embedding) ** 0.5
|
||||
assert 0.9 < norm < 1.1, f"Embedding {i} norm {norm} is not close to 1.0"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_sdk_cohere_dimension(self, litellm_cohere_embeddings):
|
||||
"""Test dimension detection with real Cohere API."""
|
||||
dimension = litellm_cohere_embeddings.dimension
|
||||
|
||||
# Cohere embed-english-v3.0 has 1024 dimensions
|
||||
assert dimension == 1024
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_sdk_cohere_single_text(self, litellm_cohere_embeddings):
|
||||
"""Test encoding single text with real Cohere API."""
|
||||
result = litellm_cohere_embeddings.encode(["Hello world"])
|
||||
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 1
|
||||
assert len(result[0]) == 1024
|
||||
assert not all(abs(x) < 0.0001 for x in result[0])
|
||||
@@ -209,7 +209,7 @@ class TestLargeBatchRetain:
|
||||
raise
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(240) # Increased timeout for VectorChord BM25 tokenization
|
||||
@pytest.mark.timeout(120)
|
||||
async def test_batch_chunking_behavior(self, memory_with_mock_llm, request_context):
|
||||
"""
|
||||
Test that large batches are properly chunked into sub-batches.
|
||||
|
||||
@@ -45,7 +45,7 @@ class TestMainModuleExtensionLoading:
|
||||
|
||||
with patch("hindsight_api.main.MemoryEngine") as mock_engine, \
|
||||
patch("hindsight_api.main.create_app") as mock_create_app, \
|
||||
patch("hindsight_api.main._get_raw_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.get_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.load_extension", side_effect=tracking_load_extension), \
|
||||
patch("hindsight_api.main.DefaultExtensionContext"), \
|
||||
patch("hindsight_api.main.print_banner"), \
|
||||
@@ -96,7 +96,7 @@ class TestMainModuleExtensionLoading:
|
||||
|
||||
with patch("hindsight_api.main.MemoryEngine") as mock_engine, \
|
||||
patch("hindsight_api.main.create_app") as mock_create_app, \
|
||||
patch("hindsight_api.main._get_raw_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.get_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.load_extension", side_effect=tracking_load_extension), \
|
||||
patch("hindsight_api.main.DefaultExtensionContext"), \
|
||||
patch("hindsight_api.main.print_banner"), \
|
||||
@@ -143,7 +143,7 @@ class TestMainModuleExtensionLoading:
|
||||
|
||||
with patch("hindsight_api.main.MemoryEngine", side_effect=capture_memory_engine), \
|
||||
patch("hindsight_api.main.create_app") as mock_create_app, \
|
||||
patch("hindsight_api.main._get_raw_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.get_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.DefaultExtensionContext"), \
|
||||
patch("hindsight_api.main.print_banner"), \
|
||||
patch("uvicorn.run"):
|
||||
@@ -200,7 +200,7 @@ class TestMainModuleExtensionLoading:
|
||||
|
||||
with patch("hindsight_api.main.MemoryEngine", side_effect=capture_memory_engine), \
|
||||
patch("hindsight_api.main.create_app") as mock_create_app, \
|
||||
patch("hindsight_api.main._get_raw_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.get_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.DefaultExtensionContext", side_effect=capture_context), \
|
||||
patch("hindsight_api.main.print_banner"), \
|
||||
patch("uvicorn.run"):
|
||||
@@ -242,7 +242,7 @@ class TestMainModuleExtensionLoading:
|
||||
|
||||
with patch("hindsight_api.main.MemoryEngine", side_effect=capture_memory_engine), \
|
||||
patch("hindsight_api.main.create_app") as mock_create_app, \
|
||||
patch("hindsight_api.main._get_raw_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.get_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.print_banner"), \
|
||||
patch("uvicorn.run"):
|
||||
|
||||
@@ -287,7 +287,7 @@ class TestMainModuleExtensionLoading:
|
||||
|
||||
with patch("hindsight_api.main.MemoryEngine") as mock_engine, \
|
||||
patch("hindsight_api.main.create_app", return_value=mock_app), \
|
||||
patch("hindsight_api.main._get_raw_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.get_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.print_banner"), \
|
||||
patch("uvicorn.run", side_effect=capture_uvicorn_run):
|
||||
|
||||
@@ -327,7 +327,7 @@ class TestMainModuleExtensionLoading:
|
||||
|
||||
with patch("hindsight_api.main.MemoryEngine") as mock_engine, \
|
||||
patch("hindsight_api.main.create_app") as mock_create_app, \
|
||||
patch("hindsight_api.main._get_raw_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.get_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.print_banner"), \
|
||||
patch("uvicorn.run", side_effect=capture_uvicorn_run):
|
||||
|
||||
@@ -352,47 +352,6 @@ class TestMainModuleExtensionLoading:
|
||||
"main.py should use import string when workers > 1"
|
||||
assert uvicorn_calls[0]["workers"] == 2
|
||||
|
||||
def test_main_sets_keepalive_timeout(self, monkeypatch):
|
||||
"""
|
||||
Verify that uvicorn is configured with timeout_keep_alive > aiohttp's
|
||||
default client keepalive timeout (15s), so the server never closes
|
||||
connections before the client does.
|
||||
"""
|
||||
monkeypatch.setenv("HINDSIGHT_API_WORKERS", "1")
|
||||
monkeypatch.delenv("HINDSIGHT_API_TENANT_EXTENSION", raising=False)
|
||||
|
||||
uvicorn_calls = []
|
||||
|
||||
def capture_uvicorn_run(**kwargs):
|
||||
uvicorn_calls.append(kwargs)
|
||||
|
||||
with patch("hindsight_api.main.MemoryEngine") as mock_engine, \
|
||||
patch("hindsight_api.main.create_app") as mock_create_app, \
|
||||
patch("hindsight_api.main._get_raw_config") as mock_get_config, \
|
||||
patch("hindsight_api.main.print_banner"), \
|
||||
patch("uvicorn.run", side_effect=capture_uvicorn_run):
|
||||
|
||||
mock_config = MagicMock()
|
||||
mock_config.host = "0.0.0.0"
|
||||
mock_config.port = 8888
|
||||
mock_config.log_level = "info"
|
||||
mock_config.mcp_enabled = False
|
||||
mock_config.run_migrations_on_startup = False
|
||||
mock_config.database_url = "postgresql://test:test@localhost/test"
|
||||
mock_get_config.return_value = mock_config
|
||||
mock_engine.return_value = MagicMock()
|
||||
mock_create_app.return_value = MagicMock()
|
||||
|
||||
with patch.object(sys, 'argv', ['hindsight-api']):
|
||||
from hindsight_api.main import main
|
||||
main()
|
||||
|
||||
assert len(uvicorn_calls) == 1
|
||||
assert "timeout_keep_alive" in uvicorn_calls[0], \
|
||||
"uvicorn config must set timeout_keep_alive"
|
||||
assert uvicorn_calls[0]["timeout_keep_alive"] > 15, \
|
||||
"timeout_keep_alive must exceed aiohttp's 15s client default"
|
||||
|
||||
|
||||
# Mock extensions for testing
|
||||
from hindsight_api.extensions import (
|
||||
|
||||
@@ -8,14 +8,14 @@ populated from the summary for backwards compatibility.
|
||||
import pytest
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.config import get_config
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def disable_observations():
|
||||
"""Disable observations for a specific test."""
|
||||
config = _get_raw_config()
|
||||
config = get_config()
|
||||
original_value = config.enable_observations
|
||||
config.enable_observations = False
|
||||
yield
|
||||
|
||||
@@ -1,274 +0,0 @@
|
||||
"""
|
||||
Test that recall chunks are fetched independently of max_tokens filtering.
|
||||
|
||||
This test verifies the new behavior where:
|
||||
1. Chunks are fetched BEFORE max_tokens filtering
|
||||
2. max_tokens=0 returns 0 facts but can still return chunks
|
||||
3. Chunks are fetched in batches to handle varying chunk sizes
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_chunks_independent_of_max_tokens(memory, request_context):
|
||||
"""
|
||||
Test that chunks are fetched independently of max_tokens.
|
||||
|
||||
When max_tokens=0, recall should:
|
||||
- Return 0 memory facts
|
||||
- Still return chunks (up to max_chunk_tokens)
|
||||
- Chunks should come from top-scored results before token filtering
|
||||
"""
|
||||
bank_id = "test-chunks-independence"
|
||||
|
||||
try:
|
||||
|
||||
# Retain some test content with substantial size to generate chunks
|
||||
test_content = """
|
||||
The quantum computing research team at MIT has made significant breakthroughs.
|
||||
Dr. Sarah Chen leads the team and focuses on quantum error correction.
|
||||
The team published three papers in Nature Physics this year.
|
||||
Their work on topological qubits shows promise for scalable quantum computers.
|
||||
Collaborators include IBM Research and Google Quantum AI.
|
||||
The research is funded by a $5M NSF grant running through 2026.
|
||||
""" * 10 # Repeat to ensure we get multiple chunks
|
||||
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=test_content,
|
||||
context="research notes",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Test 1: Normal recall with both facts and chunks
|
||||
result_normal = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="quantum computing",
|
||||
max_tokens=4096, # Normal token budget
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=2000,
|
||||
budget=Budget.MID,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result_normal.results) > 0, "Should return memory facts with normal max_tokens"
|
||||
assert result_normal.chunks is not None, "Should include chunks when requested"
|
||||
assert len(result_normal.chunks) > 0, "Should return at least one chunk"
|
||||
|
||||
# Test 2: Recall with max_tokens=0 but chunks enabled
|
||||
result_chunks_only = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="quantum computing",
|
||||
max_tokens=0, # Zero token budget for facts
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=2000, # But allow chunks
|
||||
budget=Budget.MID,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Key assertions for new behavior
|
||||
assert len(result_chunks_only.results) == 0, "max_tokens=0 should return 0 facts"
|
||||
assert result_chunks_only.chunks is not None, "Should still include chunks dict"
|
||||
assert len(result_chunks_only.chunks) > 0, "Should return chunks even with max_tokens=0"
|
||||
|
||||
# Verify chunks are from the same content (non-empty text)
|
||||
for chunk_id, chunk_info in result_chunks_only.chunks.items():
|
||||
assert len(chunk_info.chunk_text) > 0, "Chunks should contain text"
|
||||
assert chunk_info.chunk_index >= 0, "Chunk should have valid index"
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_chunks_batching_with_varying_sizes(memory, request_context):
|
||||
"""
|
||||
Test that chunk batching works correctly with varying chunk sizes.
|
||||
|
||||
This verifies that:
|
||||
1. Chunks are fetched in batches until token budget is exhausted
|
||||
2. The system handles varying chunk sizes across documents
|
||||
3. Token budget is respected across multiple batch fetches
|
||||
"""
|
||||
bank_id = "test-chunks-batching"
|
||||
|
||||
try:
|
||||
|
||||
# Retain multiple documents with different content sizes
|
||||
# Document 1: Short content (small chunks)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice is a software engineer who specializes in Python programming and machine learning.",
|
||||
context="doc1",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Document 2: Medium content
|
||||
content_bob = """
|
||||
Bob works as a data scientist at a tech startup in San Francisco.
|
||||
He has expertise in natural language processing and computer vision.
|
||||
Bob completed his PhD at Stanford University in 2020.
|
||||
He leads a team of five engineers working on AI-powered recommendation systems.
|
||||
""" * 5
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content_bob,
|
||||
context="doc2",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Document 3: Long content (large chunks)
|
||||
content_charlie = """
|
||||
Charlie is the CTO of a growing AI company focused on healthcare applications.
|
||||
He has over 15 years of experience in software architecture and distributed systems.
|
||||
Charlie's team builds machine learning models for medical image analysis and diagnosis.
|
||||
The company recently raised $50 million in Series B funding.
|
||||
They have partnerships with major hospitals in the United States and Europe.
|
||||
Charlie holds several patents in medical imaging and deep learning.
|
||||
""" * 20
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content_charlie,
|
||||
context="doc3",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Recall with modest chunk token budget
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="Alice Bob Charlie",
|
||||
max_tokens=0, # No facts, only chunks
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=1000, # Limited chunk budget
|
||||
budget=Budget.MID,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) == 0, "Should return 0 facts with max_tokens=0"
|
||||
assert result.chunks is not None, "Should include chunks"
|
||||
|
||||
# Verify we got chunks and respected the token budget
|
||||
if len(result.chunks) > 0:
|
||||
# Count total tokens (approximate)
|
||||
total_chunk_chars = sum(len(chunk.chunk_text) for chunk in result.chunks.values())
|
||||
# Very rough estimate: 1 token ≈ 4 characters
|
||||
estimated_tokens = total_chunk_chars // 4
|
||||
|
||||
# Should be reasonably close to budget (within 2x due to estimation and batching)
|
||||
assert estimated_tokens <= 1000 * 2, f"Should respect chunk token budget (got ~{estimated_tokens} tokens)"
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_chunks_ordering_by_relevance(memory, request_context):
|
||||
"""
|
||||
Test that chunks are returned in order of fact relevance.
|
||||
|
||||
Chunks should be ordered based on the top-scored (reranked) results,
|
||||
not in document order or random order.
|
||||
"""
|
||||
bank_id = "test-chunks-ordering"
|
||||
|
||||
try:
|
||||
|
||||
# Retain content with different relevance to query
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="The Python programming language is widely used for machine learning and data science applications.",
|
||||
context="topic: Python",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="JavaScript is commonly used for web development and frontend applications.",
|
||||
context="topic: JavaScript",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Python's scikit-learn library is excellent for traditional machine learning tasks and model training.",
|
||||
context="topic: Python ML",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Query specifically about Python - should rank Python facts higher
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="Python machine learning",
|
||||
max_tokens=0, # No facts
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=5000, # Enough for all chunks
|
||||
budget=Budget.HIGH, # Use high budget for better recall
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) == 0, "Should return 0 facts with max_tokens=0"
|
||||
assert result.chunks is not None, "Should include chunks"
|
||||
|
||||
# We should get chunks, and they should be ordered by relevance
|
||||
# The exact ordering depends on the reranker, but we should have chunks
|
||||
assert len(result.chunks) > 0, "Should return chunks from relevant facts"
|
||||
|
||||
# Verify chunks contain relevant content
|
||||
all_chunk_text = " ".join(chunk.chunk_text for chunk in result.chunks.values())
|
||||
# At least some chunks should mention Python (higher relevance)
|
||||
# This is a soft check since exact ordering depends on scoring
|
||||
assert "Python" in all_chunk_text or "python" in all_chunk_text.lower(), \
|
||||
"Chunks should include content about Python (relevant to query)"
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_chunks_without_include_flag(memory, request_context):
|
||||
"""
|
||||
Test that chunks are NOT returned when include_chunks=False (default).
|
||||
|
||||
This ensures backward compatibility - chunks are only fetched when explicitly requested.
|
||||
"""
|
||||
bank_id = "test-chunks-no-include"
|
||||
|
||||
try:
|
||||
|
||||
# Retain content
|
||||
test_content = """
|
||||
Sarah is a product manager at a fintech company in New York.
|
||||
She specializes in user experience design and agile methodologies.
|
||||
Sarah graduated from MIT with a degree in computer science.
|
||||
She has led the development of three successful mobile banking applications.
|
||||
"""
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=test_content,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Recall without include_chunks flag (default is False)
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="Sarah product manager",
|
||||
max_tokens=4096,
|
||||
request_context=request_context,
|
||||
# include_chunks=False is the default
|
||||
)
|
||||
|
||||
# Should have facts but no chunks
|
||||
assert len(result.results) > 0, "Should return facts"
|
||||
assert result.chunks is None or len(result.chunks) == 0, \
|
||||
"Should NOT return chunks when include_chunks=False"
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -1,103 +0,0 @@
|
||||
"""
|
||||
Test reflect endpoint with empty based_on (no memories scenario).
|
||||
|
||||
This test verifies that the API returns the correct based_on format:
|
||||
- v0.3.0 (old): returned based_on as list []
|
||||
- v0.4.0+ (current): returns based_on as object {"memories": [], "mental_models": [], "directives": []}
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import httpx
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client(memory):
|
||||
"""Create an async test client for the FastAPI app."""
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_with_no_memories_empty_bank(api_client):
|
||||
"""Test reflect on an empty bank (no memories) with include.facts enabled."""
|
||||
bank_id = "test_empty_bank"
|
||||
|
||||
# Reflect on empty bank with facts requested
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/reflect",
|
||||
json={
|
||||
"query": "What do you know about machine learning?",
|
||||
"budget": "low",
|
||||
"include": {
|
||||
"facts": {} # Request facts but bank is empty
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# DEBUG: Print what the API actually returned
|
||||
import json
|
||||
print("\n" + "="*80)
|
||||
print("API Response:")
|
||||
print(json.dumps(data, indent=2))
|
||||
print("="*80 + "\n")
|
||||
|
||||
# Verify response structure
|
||||
assert "text" in data
|
||||
assert "based_on" in data
|
||||
|
||||
# The API should return based_on as either:
|
||||
# 1. null/None (if include.facts not set)
|
||||
# 2. {"memories": [], "mental_models": [], "directives": []} (if include.facts set but empty)
|
||||
# It should NEVER return based_on: []
|
||||
|
||||
based_on = data.get("based_on")
|
||||
if based_on is not None:
|
||||
assert isinstance(based_on, dict), f"based_on should be dict or null, got {type(based_on)}: {based_on}"
|
||||
assert not isinstance(based_on, list), f"based_on should NEVER be a list! Got: {based_on}"
|
||||
assert "memories" in based_on
|
||||
assert "mental_models" in based_on
|
||||
assert "directives" in based_on
|
||||
# All should be empty lists
|
||||
assert based_on["memories"] == []
|
||||
assert based_on["mental_models"] == []
|
||||
assert based_on["directives"] == []
|
||||
|
||||
# Verify the structure is parseable as proper types
|
||||
assert isinstance(data["text"], str)
|
||||
if based_on is not None:
|
||||
# Verify it's the v0.4.0+ format (object with arrays)
|
||||
assert isinstance(based_on["memories"], list)
|
||||
assert isinstance(based_on["mental_models"], list)
|
||||
assert isinstance(based_on["directives"], list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_without_include_facts(api_client):
|
||||
"""Test reflect without requesting facts (based_on should be None)."""
|
||||
bank_id = "test_no_facts"
|
||||
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/reflect",
|
||||
json={
|
||||
"query": "Hello world",
|
||||
"budget": "low"
|
||||
# No include.facts
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# When include.facts is not set, based_on should not be in response (or be null)
|
||||
based_on = data.get("based_on")
|
||||
assert based_on is None, f"based_on should be None when not requested, got {type(based_on)}: {based_on}"
|
||||
|
||||
# Verify structure
|
||||
assert isinstance(data["text"], str)
|
||||
@@ -2093,7 +2093,7 @@ async def test_custom_extraction_mode():
|
||||
import os
|
||||
from hindsight_api import LLMConfig
|
||||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
|
||||
from hindsight_api.config import clear_config_cache, _get_raw_config
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
# Save original env vars
|
||||
original_mode = os.getenv("HINDSIGHT_API_RETAIN_EXTRACTION_MODE")
|
||||
@@ -2135,8 +2135,7 @@ If the text contains both Italian and English content, extract ONLY the Italian
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
context="team meeting notes",
|
||||
llm_config=llm_config,
|
||||
agent_name="TestUser",
|
||||
config=_get_raw_config(),
|
||||
agent_name="TestUser"
|
||||
)
|
||||
|
||||
logger.info(f"\nExtracted {len(facts)} facts with custom mode (Italian only):")
|
||||
|
||||
@@ -206,12 +206,8 @@ class TestWorkerPoller:
|
||||
assert len(claimed) == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_task_executor_marks_completed(self, pool, clean_operations):
|
||||
"""Test that executor's status marking is preserved by the poller.
|
||||
|
||||
The executor (MemoryEngine.execute_task) handles marking operations as completed/failed.
|
||||
The poller should NOT override those status updates.
|
||||
"""
|
||||
async def test_execute_task_marks_completed(self, pool, clean_operations):
|
||||
"""Test that successful task execution marks task as completed."""
|
||||
from hindsight_api.worker import WorkerPoller
|
||||
from hindsight_api.worker.poller import ClaimedTask
|
||||
|
||||
@@ -232,16 +228,7 @@ class TestWorkerPoller:
|
||||
executed = []
|
||||
|
||||
async def mock_executor(task_dict):
|
||||
"""Executor that marks its own status as completed (like MemoryEngine.execute_task)."""
|
||||
executed.append(task_dict)
|
||||
await pool.execute(
|
||||
"""
|
||||
UPDATE async_operations
|
||||
SET status = 'completed', completed_at = now(), updated_at = now()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
op_id,
|
||||
)
|
||||
|
||||
poller = WorkerPoller(
|
||||
pool=pool,
|
||||
@@ -259,7 +246,7 @@ class TestWorkerPoller:
|
||||
assert completed, "Task did not complete within timeout"
|
||||
assert len(executed) == 1
|
||||
|
||||
# Verify task is marked as completed (by executor, not overridden by poller)
|
||||
# Verify task is marked as completed
|
||||
row = await pool.fetchrow(
|
||||
"SELECT status, completed_at FROM async_operations WHERE operation_id = $1",
|
||||
op_id,
|
||||
@@ -268,24 +255,19 @@ class TestWorkerPoller:
|
||||
assert row["completed_at"] is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_executor_exception_does_not_crash_poller(self, pool, clean_operations):
|
||||
"""Test that unexpected exceptions from executor are caught and don't crash the poller.
|
||||
|
||||
If the executor raises an unexpected exception (which MemoryEngine.execute_task should NOT do,
|
||||
but could happen from schema setup or other infrastructure issues), the poller should catch it
|
||||
gracefully. Status remains 'processing' since neither executor nor poller handled it.
|
||||
"""
|
||||
async def test_execute_task_retries_on_failure(self, pool, clean_operations):
|
||||
"""Test that failed task execution triggers retry mechanism."""
|
||||
from hindsight_api.worker import WorkerPoller
|
||||
from hindsight_api.worker.poller import ClaimedTask
|
||||
|
||||
# Create a pending task
|
||||
# Create a pending task with retry_count=0
|
||||
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
|
||||
op_id = uuid.uuid4()
|
||||
payload = json.dumps({"type": "test_task", "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, worker_id)
|
||||
VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'test-worker-1')
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, retry_count)
|
||||
VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'test-worker-1', 0)
|
||||
""",
|
||||
op_id,
|
||||
bank_id,
|
||||
@@ -293,15 +275,16 @@ class TestWorkerPoller:
|
||||
)
|
||||
|
||||
async def failing_executor(task_dict):
|
||||
raise ValueError("Unexpected infrastructure failure")
|
||||
raise ValueError("Simulated failure")
|
||||
|
||||
poller = WorkerPoller(
|
||||
pool=pool,
|
||||
worker_id="test-worker-1",
|
||||
executor=failing_executor,
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
# Execute - should catch exception without crashing
|
||||
# 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)
|
||||
@@ -310,84 +293,61 @@ class TestWorkerPoller:
|
||||
completed = await poller.wait_for_active_tasks(timeout=5.0)
|
||||
assert completed, "Task did not complete within timeout"
|
||||
|
||||
# Status stays 'processing' since the poller no longer manages status
|
||||
# Verify task is back to pending with incremented retry_count
|
||||
row = await pool.fetchrow(
|
||||
"SELECT status FROM async_operations WHERE operation_id = $1",
|
||||
"SELECT status, retry_count, worker_id FROM async_operations WHERE operation_id = $1",
|
||||
op_id,
|
||||
)
|
||||
assert row["status"] == "processing"
|
||||
assert row["status"] == "pending"
|
||||
assert row["retry_count"] == 1
|
||||
assert row["worker_id"] is None # Worker ID cleared for retry
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_executor_failed_status_not_overridden(self, pool, clean_operations):
|
||||
"""REGRESSION TEST: Verify poller does NOT overwrite executor's 'failed' status to 'completed'.
|
||||
|
||||
This test catches the bug where the poller always called _mark_completed() after executor
|
||||
returned, overwriting the 'failed' status that the executor had already set.
|
||||
|
||||
Scenario:
|
||||
1. Executor catches an internal error and marks the operation as 'failed' in the DB
|
||||
2. Executor returns normally (does NOT re-raise) - this is how MemoryEngine.execute_task works
|
||||
3. The poller must NOT overwrite the 'failed' status to 'completed'
|
||||
|
||||
With the old buggy code, this test would FAIL (status would be 'completed').
|
||||
"""
|
||||
async def test_execute_task_fails_after_max_retries(self, pool, clean_operations):
|
||||
"""Test that task is marked failed after exceeding max retries."""
|
||||
from hindsight_api.worker import WorkerPoller
|
||||
from hindsight_api.worker.poller import ClaimedTask
|
||||
|
||||
# Create a task that has already used all retries
|
||||
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
|
||||
op_id = uuid.uuid4()
|
||||
payload = json.dumps({"type": "test_task", "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, worker_id)
|
||||
VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'test-worker-1')
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, retry_count)
|
||||
VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'test-worker-1', 3)
|
||||
""",
|
||||
op_id,
|
||||
bank_id,
|
||||
payload,
|
||||
)
|
||||
|
||||
async def executor_that_marks_failed(task_dict):
|
||||
"""Simulates MemoryEngine.execute_task behavior on internal error.
|
||||
|
||||
The executor catches the error, marks the operation as 'failed',
|
||||
and returns normally (does NOT re-raise the exception).
|
||||
"""
|
||||
# Simulate internal failure handling (like MemoryEngine._mark_operation_failed)
|
||||
await pool.execute(
|
||||
"""
|
||||
UPDATE async_operations
|
||||
SET status = 'failed', error_message = $2, completed_at = now(), updated_at = now()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
op_id,
|
||||
"Simulated conversion error: file format not supported",
|
||||
)
|
||||
# Returns normally - this is the key: executor does NOT re-raise
|
||||
async def failing_executor(task_dict):
|
||||
raise ValueError("Simulated failure")
|
||||
|
||||
poller = WorkerPoller(
|
||||
pool=pool,
|
||||
worker_id="test-worker-1",
|
||||
executor=executor_that_marks_failed,
|
||||
executor=failing_executor,
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
# 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"
|
||||
|
||||
# THE KEY ASSERTION: Status must be 'failed', NOT 'completed'
|
||||
# Verify task is marked as failed
|
||||
row = await pool.fetchrow(
|
||||
"SELECT status, error_message FROM async_operations WHERE operation_id = $1",
|
||||
op_id,
|
||||
)
|
||||
assert row["status"] == "failed", (
|
||||
f"REGRESSION: Poller overwrote executor's 'failed' status to '{row['status']}'. "
|
||||
"The poller must not override status set by the executor."
|
||||
)
|
||||
assert "Simulated conversion error" in row["error_message"]
|
||||
assert row["status"] == "failed"
|
||||
assert "Max retries" in row["error_message"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claim_batch_skips_consolidation_when_same_bank_processing(self, pool, clean_operations):
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "hindsight-cli"
|
||||
version = "0.4.11"
|
||||
version = "0.4.10"
|
||||
edition = "2021"
|
||||
authors = ["Hindsight Team"]
|
||||
description = "A beautiful CLI for Hindsight - semantic memory system"
|
||||
@@ -20,8 +20,8 @@ clap = { version = "4.5", features = ["derive", "env"] }
|
||||
# Async runtime
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
|
||||
# HTTP client (for timeout configuration and multipart file uploads)
|
||||
reqwest = { version = "0.12", features = ["multipart"] }
|
||||
# HTTP client (for timeout configuration)
|
||||
reqwest = "0.12"
|
||||
|
||||
# Serialization (for config and output formatting)
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
|
||||
+2
-107
@@ -58,16 +58,9 @@ pub struct MemoryPutResult {
|
||||
pub operation_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct FileRetainResult {
|
||||
pub operation_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ApiClient {
|
||||
client: AsyncClient,
|
||||
http_client: reqwest::Client,
|
||||
base_url: String,
|
||||
runtime: std::sync::Arc<tokio::runtime::Runtime>,
|
||||
}
|
||||
|
||||
@@ -91,8 +84,8 @@ impl ApiClient {
|
||||
|
||||
let http_client = client_builder.build()?;
|
||||
|
||||
let client = AsyncClient::new_with_client(&base_url, http_client.clone());
|
||||
Ok(ApiClient { client, http_client, base_url, runtime })
|
||||
let client = AsyncClient::new_with_client(&base_url, http_client);
|
||||
Ok(ApiClient { client, runtime })
|
||||
}
|
||||
|
||||
pub fn list_agents(&self, _verbose: bool) -> Result<Vec<types::BankListItem>> {
|
||||
@@ -175,67 +168,6 @@ impl ApiClient {
|
||||
})
|
||||
}
|
||||
|
||||
/// Upload files to the file retain endpoint (multipart/form-data).
|
||||
/// Returns a list of operation IDs for tracking. Always async server-side.
|
||||
pub fn file_retain(
|
||||
&self,
|
||||
bank_id: &str,
|
||||
files: Vec<(String, Vec<u8>)>,
|
||||
context: Option<String>,
|
||||
verbose: bool,
|
||||
) -> Result<FileRetainResult> {
|
||||
self.runtime.block_on(async {
|
||||
let url = format!("{}/v1/default/banks/{}/files/retain", self.base_url, bank_id);
|
||||
|
||||
let files_metadata: Vec<serde_json::Value> = files
|
||||
.iter()
|
||||
.map(|(name, _)| {
|
||||
let mut meta = serde_json::json!({});
|
||||
if let Some(ctx) = &context {
|
||||
meta["context"] = serde_json::Value::String(ctx.clone());
|
||||
}
|
||||
// Use filename stem as document_id for deduplication
|
||||
if let Some(stem) = std::path::Path::new(name)
|
||||
.file_stem()
|
||||
.and_then(|s| s.to_str())
|
||||
{
|
||||
meta["document_id"] = serde_json::Value::String(stem.to_string());
|
||||
}
|
||||
meta
|
||||
})
|
||||
.collect();
|
||||
|
||||
let request_json = serde_json::json!({
|
||||
"files_metadata": files_metadata,
|
||||
});
|
||||
|
||||
let mut form = reqwest::multipart::Form::new()
|
||||
.text("request", request_json.to_string());
|
||||
|
||||
for (filename, content) in files {
|
||||
let part = reqwest::multipart::Part::bytes(content)
|
||||
.file_name(filename)
|
||||
.mime_str("application/octet-stream")?;
|
||||
form = form.part("files", part);
|
||||
}
|
||||
|
||||
if verbose {
|
||||
eprintln!("POST {}", url);
|
||||
}
|
||||
|
||||
let response = self.http_client.post(&url).multipart(form).send().await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let text = response.text().await.unwrap_or_default();
|
||||
anyhow::bail!("File retain failed ({}): {}", status, text);
|
||||
}
|
||||
|
||||
let result: FileRetainResult = response.json().await?;
|
||||
Ok(result)
|
||||
})
|
||||
}
|
||||
|
||||
/// Poll an operation until it completes or fails.
|
||||
/// Returns Ok(true) if completed successfully, Ok(false) if failed, Err if polling error.
|
||||
pub fn poll_operation(&self, agent_id: &str, operation_id: &str, verbose: bool) -> Result<(bool, Option<String>)> {
|
||||
@@ -455,43 +387,6 @@ impl ApiClient {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn get_bank_config(
|
||||
&self,
|
||||
bank_id: &str,
|
||||
_verbose: bool,
|
||||
) -> Result<types::BankConfigResponse> {
|
||||
self.runtime.block_on(async {
|
||||
let response = self.client.get_bank_config(bank_id, None).await?;
|
||||
Ok(response.into_inner())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn update_bank_config(
|
||||
&self,
|
||||
bank_id: &str,
|
||||
updates: std::collections::HashMap<String, serde_json::Value>,
|
||||
_verbose: bool,
|
||||
) -> Result<types::BankConfigResponse> {
|
||||
self.runtime.block_on(async {
|
||||
// Convert HashMap to serde_json::Map
|
||||
let updates_map: serde_json::Map<String, serde_json::Value> = updates.into_iter().collect();
|
||||
let request = types::BankConfigUpdate { updates: updates_map };
|
||||
let response = self.client.update_bank_config(bank_id, None, &request).await?;
|
||||
Ok(response.into_inner())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn reset_bank_config(
|
||||
&self,
|
||||
bank_id: &str,
|
||||
_verbose: bool,
|
||||
) -> Result<types::BankConfigResponse> {
|
||||
self.runtime.block_on(async {
|
||||
let response = self.client.reset_bank_config(bank_id, None).await?;
|
||||
Ok(response.into_inner())
|
||||
})
|
||||
}
|
||||
|
||||
// --- Tag Methods ---
|
||||
|
||||
pub fn list_tags(
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use anyhow::{anyhow, Result};
|
||||
use anyhow::Result;
|
||||
use crate::api::ApiClient;
|
||||
use crate::output::{self, OutputFormat};
|
||||
use crate::ui;
|
||||
@@ -655,159 +655,3 @@ pub fn clear_observations(
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn config(
|
||||
client: &ApiClient,
|
||||
bank_id: &str,
|
||||
overrides_only: bool,
|
||||
verbose: bool,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
let spinner = if output_format == OutputFormat::Pretty {
|
||||
Some(ui::create_spinner("Fetching bank configuration..."))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let response = client.get_bank_config(bank_id, verbose);
|
||||
|
||||
if let Some(mut sp) = spinner {
|
||||
sp.finish();
|
||||
}
|
||||
|
||||
match response {
|
||||
Ok(result) => {
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_success(&format!("Configuration for bank '{}'", bank_id));
|
||||
println!();
|
||||
if overrides_only {
|
||||
println!("Bank-specific overrides:");
|
||||
if result.overrides.is_empty() {
|
||||
println!(" (none - using defaults)");
|
||||
} else {
|
||||
for (key, value) in result.overrides.iter() {
|
||||
println!(" {}: {:?}", key, value);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
println!("Resolved configuration (with all overrides applied):");
|
||||
for (key, value) in result.config.iter() {
|
||||
println!(" {}: {:?}", key, value);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if overrides_only {
|
||||
output::print_output(&result.overrides, output_format)?;
|
||||
} else {
|
||||
output::print_output(&result, output_format)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_config(
|
||||
client: &ApiClient,
|
||||
bank_id: &str,
|
||||
llm_provider: Option<String>,
|
||||
llm_model: Option<String>,
|
||||
llm_api_key: Option<String>,
|
||||
llm_base_url: Option<String>,
|
||||
verbose: bool,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
use std::collections::HashMap;
|
||||
|
||||
let mut updates: HashMap<String, serde_json::Value> = HashMap::new();
|
||||
|
||||
if let Some(provider) = llm_provider {
|
||||
updates.insert("llm_provider".to_string(), serde_json::Value::String(provider));
|
||||
}
|
||||
if let Some(model) = llm_model {
|
||||
updates.insert("llm_model".to_string(), serde_json::Value::String(model));
|
||||
}
|
||||
if let Some(api_key) = llm_api_key {
|
||||
updates.insert("llm_api_key".to_string(), serde_json::Value::String(api_key));
|
||||
}
|
||||
if let Some(base_url) = llm_base_url {
|
||||
updates.insert("llm_base_url".to_string(), serde_json::Value::String(base_url));
|
||||
}
|
||||
|
||||
if updates.is_empty() {
|
||||
return Err(anyhow!("No config updates provided. Use --llm-provider, --llm-model, --llm-api-key, or --llm-base-url".to_string()));
|
||||
}
|
||||
|
||||
let spinner = if output_format == OutputFormat::Pretty {
|
||||
Some(ui::create_spinner("Updating bank configuration..."))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let response = client.update_bank_config(bank_id, updates, verbose);
|
||||
|
||||
if let Some(mut sp) = spinner {
|
||||
sp.finish();
|
||||
}
|
||||
|
||||
match response {
|
||||
Ok(result) => {
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_success(&format!("Configuration updated for bank '{}'", bank_id));
|
||||
println!("\nUpdated overrides:");
|
||||
for (key, value) in result.overrides.iter() {
|
||||
println!(" {}: {:?}", key, value);
|
||||
}
|
||||
} else {
|
||||
output::print_output(&result, output_format)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn reset_config(
|
||||
client: &ApiClient,
|
||||
bank_id: &str,
|
||||
yes: bool,
|
||||
verbose: bool,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
if !yes && output_format == OutputFormat::Pretty {
|
||||
let confirmed = ui::prompt_confirmation(&format!(
|
||||
"Reset all configuration overrides for bank '{}'?",
|
||||
bank_id
|
||||
))?;
|
||||
|
||||
if !confirmed {
|
||||
ui::print_info("Operation cancelled");
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
let spinner = if output_format == OutputFormat::Pretty {
|
||||
Some(ui::create_spinner("Resetting bank configuration..."))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let response = client.reset_bank_config(bank_id, verbose);
|
||||
|
||||
if let Some(mut sp) = spinner {
|
||||
sp.finish();
|
||||
}
|
||||
|
||||
match response {
|
||||
Ok(result) => {
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_success(&format!("Configuration reset to defaults for bank '{}'", bank_id));
|
||||
} else {
|
||||
output::print_output(&result, output_format)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -220,23 +220,14 @@ pub fn get(
|
||||
}
|
||||
}
|
||||
|
||||
// Helper function to check if a file is supported by the file converter (markitdown)
|
||||
fn is_supported_file(path: &std::path::Path) -> bool {
|
||||
const SUPPORTED_EXTENSIONS: &[&str] = &[
|
||||
// Documents
|
||||
"pdf", "docx", "doc", "pptx", "ppt", "xlsx", "xls",
|
||||
// Images (OCR)
|
||||
"jpg", "jpeg", "png", "gif", "bmp", "webp", "tiff",
|
||||
// Web / markup
|
||||
"html", "htm",
|
||||
// Text / data
|
||||
"txt", "md", "csv", "json", "yaml", "yml", "toml", "xml", "rst", "adoc", "log",
|
||||
// Audio (transcription)
|
||||
"mp3", "wav", "ogg", "flac",
|
||||
// Helper function to check if a file has a text-based extension
|
||||
fn is_text_file(path: &std::path::Path) -> bool {
|
||||
const TEXT_EXTENSIONS: &[&str] = &[
|
||||
"txt", "md", "json", "yaml", "yml", "toml", "xml", "csv", "log", "rst", "adoc",
|
||||
];
|
||||
path.extension()
|
||||
.and_then(|ext| ext.to_str())
|
||||
.map(|ext| SUPPORTED_EXTENSIONS.contains(&ext.to_lowercase().as_str()))
|
||||
.map(|ext| TEXT_EXTENSIONS.contains(&ext.to_lowercase().as_str()))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
@@ -436,10 +427,10 @@ pub fn retain_files(
|
||||
anyhow::bail!("Path does not exist: {}", path.display());
|
||||
}
|
||||
|
||||
let mut file_paths = Vec::new();
|
||||
let mut files = Vec::new();
|
||||
|
||||
if path.is_file() {
|
||||
file_paths.push(path);
|
||||
files.push(path);
|
||||
} else if path.is_dir() {
|
||||
if recursive {
|
||||
for entry in WalkDir::new(&path)
|
||||
@@ -447,110 +438,133 @@ pub fn retain_files(
|
||||
.filter_map(|e| e.ok())
|
||||
.filter(|e| e.file_type().is_file())
|
||||
{
|
||||
let file_path = entry.path();
|
||||
if is_supported_file(file_path) {
|
||||
file_paths.push(file_path.to_path_buf());
|
||||
let path = entry.path();
|
||||
if is_text_file(&path) {
|
||||
files.push(path.to_path_buf());
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for entry in fs::read_dir(&path)? {
|
||||
let entry = entry?;
|
||||
let file_path = entry.path();
|
||||
if file_path.is_file() && is_supported_file(&file_path) {
|
||||
file_paths.push(file_path);
|
||||
let path = entry.path();
|
||||
if path.is_file() && is_text_file(&path) {
|
||||
files.push(path);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if file_paths.is_empty() {
|
||||
ui::print_warning("No supported files found. Supported formats: pdf, docx, pptx, xlsx, jpg, png, html, txt, md, csv, mp3, wav, and more.");
|
||||
if files.is_empty() {
|
||||
ui::print_warning("No text files found (supported: txt, md, json, yaml, yml, toml, xml, csv, log, rst, adoc)");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
ui::print_info(&format!("Found {} file(s) to import", file_paths.len()));
|
||||
ui::print_info(&format!("Found {} files to import", files.len()));
|
||||
|
||||
// Batch files (max 10 per request)
|
||||
const BATCH_SIZE: usize = 10;
|
||||
let batches: Vec<&[PathBuf]> = file_paths.chunks(BATCH_SIZE).collect();
|
||||
let mut all_operation_ids: Vec<String> = Vec::new();
|
||||
let pb = ui::create_progress_bar(files.len() as u64, "Processing files");
|
||||
|
||||
let pb = ui::create_progress_bar(file_paths.len() as u64, "Uploading files");
|
||||
let mut items = Vec::new();
|
||||
|
||||
for batch in &batches {
|
||||
let mut file_data: Vec<(String, Vec<u8>)> = Vec::new();
|
||||
for file_path in *batch {
|
||||
let filename = file_path
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_else(|| "file".to_string());
|
||||
let content = fs::read(file_path)
|
||||
.with_context(|| format!("Failed to read file: {}", file_path.display()))?;
|
||||
file_data.push((filename, content));
|
||||
pb.inc(1);
|
||||
}
|
||||
for file_path in &files {
|
||||
let content = fs::read_to_string(file_path)
|
||||
.with_context(|| format!("Failed to read file: {}", file_path.display()))?;
|
||||
|
||||
let result = client.file_retain(agent_id, file_data, context.clone(), verbose)?;
|
||||
all_operation_ids.extend(result.operation_ids);
|
||||
let doc_id = file_path
|
||||
.file_stem()
|
||||
.and_then(|s| s.to_str())
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_else(config::generate_doc_id);
|
||||
|
||||
items.push(MemoryItem {
|
||||
content,
|
||||
context: context.clone(),
|
||||
metadata: None,
|
||||
timestamp: None,
|
||||
document_id: Some(doc_id),
|
||||
entities: None,
|
||||
tags: None,
|
||||
});
|
||||
|
||||
pb.inc(1);
|
||||
}
|
||||
|
||||
pb.finish_with_message("Files uploaded");
|
||||
pb.finish_with_message("Files processed");
|
||||
|
||||
if r#async {
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_success("Files queued for processing");
|
||||
println!(" Files: {}", file_paths.len());
|
||||
for op_id in &all_operation_ids {
|
||||
println!(" Operation ID: {}", op_id);
|
||||
}
|
||||
} else {
|
||||
let result = serde_json::json!({ "operation_ids": all_operation_ids });
|
||||
output::print_output(&result, output_format)?;
|
||||
}
|
||||
// Always use async mode for the API call
|
||||
let request = RetainRequest {
|
||||
items,
|
||||
async_: true,
|
||||
document_tags: None,
|
||||
};
|
||||
|
||||
let spinner = if output_format == OutputFormat::Pretty {
|
||||
Some(ui::create_spinner("Submitting retain request..."))
|
||||
} else {
|
||||
// Poll all operations until they complete
|
||||
let poll_spinner = if output_format == OutputFormat::Pretty {
|
||||
Some(ui::create_spinner("Processing files..."))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
None
|
||||
};
|
||||
|
||||
let mut failed = Vec::new();
|
||||
for op_id in &all_operation_ids {
|
||||
let (success, error_msg) = client.poll_operation(agent_id, op_id, verbose)?;
|
||||
if !success {
|
||||
failed.push(error_msg.unwrap_or_else(|| "Unknown error".to_string()));
|
||||
}
|
||||
}
|
||||
let response = client.retain(agent_id, &request, true, verbose);
|
||||
|
||||
if let Some(mut sp) = poll_spinner {
|
||||
sp.finish();
|
||||
}
|
||||
if let Some(mut sp) = spinner {
|
||||
sp.finish();
|
||||
}
|
||||
|
||||
if failed.is_empty() {
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_success("Files retained successfully");
|
||||
println!(" Files processed: {}", file_paths.len());
|
||||
} else {
|
||||
let result = serde_json::json!({
|
||||
"success": true,
|
||||
"files_count": file_paths.len(),
|
||||
"operation_ids": all_operation_ids,
|
||||
});
|
||||
output::print_output(&result, output_format)?;
|
||||
}
|
||||
} else {
|
||||
for msg in &failed {
|
||||
match response {
|
||||
Ok(result) => {
|
||||
if r#async {
|
||||
// User requested async mode - return immediately
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_error(&format!("Retain operation failed: {}", msg));
|
||||
ui::print_success("Files queued for processing");
|
||||
println!(" Items: {}", result.items_count);
|
||||
if let Some(op_id) = &result.operation_id {
|
||||
println!(" Operation ID: {}", op_id);
|
||||
}
|
||||
} else {
|
||||
output::print_output(&result, output_format)?;
|
||||
}
|
||||
} else {
|
||||
// Poll until completion
|
||||
if let Some(operation_id) = &result.operation_id {
|
||||
let poll_spinner = if output_format == OutputFormat::Pretty {
|
||||
Some(ui::create_spinner("Processing memories..."))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let (success, error_msg) = client.poll_operation(agent_id, operation_id, verbose)?;
|
||||
|
||||
if let Some(mut sp) = poll_spinner {
|
||||
sp.finish();
|
||||
}
|
||||
|
||||
if success {
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_success("Files retained successfully");
|
||||
println!(" Items processed: {}", result.items_count);
|
||||
} else {
|
||||
output::print_output(&result, output_format)?;
|
||||
}
|
||||
} else {
|
||||
let msg = error_msg.unwrap_or_else(|| "Unknown error".to_string());
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_error(&format!("Retain operation failed: {}", msg));
|
||||
}
|
||||
anyhow::bail!("Retain operation failed: {}", msg);
|
||||
}
|
||||
} else {
|
||||
// No operation ID returned, shouldn't happen with async=true
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_success("Files retained successfully");
|
||||
println!(" Items processed: {}", result.items_count);
|
||||
} else {
|
||||
output::print_output(&result, output_format)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
anyhow::bail!("{} operation(s) failed", failed.len());
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => Err(e)
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn delete(
|
||||
@@ -665,71 +679,55 @@ mod tests {
|
||||
use std::path::Path;
|
||||
|
||||
#[test]
|
||||
fn test_is_supported_file_text_extensions() {
|
||||
fn test_is_text_file_supported_extensions() {
|
||||
let supported = [
|
||||
"file.txt", "file.md", "file.json", "file.yaml", "file.yml",
|
||||
"file.toml", "file.xml", "file.csv", "file.log", "file.rst", "file.adoc",
|
||||
];
|
||||
for filename in supported {
|
||||
assert!(
|
||||
is_supported_file(Path::new(filename)),
|
||||
"{} should be recognized as a supported file",
|
||||
is_text_file(Path::new(filename)),
|
||||
"{} should be recognized as a text file",
|
||||
filename
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_supported_file_binary_extensions() {
|
||||
let supported = [
|
||||
"file.pdf", "file.docx", "file.pptx", "file.xlsx",
|
||||
"file.png", "file.jpg", "file.jpeg", "file.gif",
|
||||
"file.mp3", "file.wav",
|
||||
];
|
||||
for filename in supported {
|
||||
assert!(
|
||||
is_supported_file(Path::new(filename)),
|
||||
"{} should be recognized as a supported file",
|
||||
filename
|
||||
);
|
||||
}
|
||||
fn test_is_text_file_case_insensitive() {
|
||||
assert!(is_text_file(Path::new("file.JSON")));
|
||||
assert!(is_text_file(Path::new("file.TXT")));
|
||||
assert!(is_text_file(Path::new("file.Md")));
|
||||
assert!(is_text_file(Path::new("file.YAML")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_supported_file_case_insensitive() {
|
||||
assert!(is_supported_file(Path::new("file.JSON")));
|
||||
assert!(is_supported_file(Path::new("file.TXT")));
|
||||
assert!(is_supported_file(Path::new("file.Md")));
|
||||
assert!(is_supported_file(Path::new("file.YAML")));
|
||||
assert!(is_supported_file(Path::new("file.PDF")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_supported_file_unsupported_extensions() {
|
||||
fn test_is_text_file_unsupported_extensions() {
|
||||
let unsupported = [
|
||||
"file.pdf", "file.doc", "file.docx", "file.png", "file.jpg",
|
||||
"file.exe", "file.bin", "file.zip", "file.tar", "file.gz",
|
||||
];
|
||||
for filename in unsupported {
|
||||
assert!(
|
||||
!is_supported_file(Path::new(filename)),
|
||||
"{} should NOT be recognized as a supported file",
|
||||
!is_text_file(Path::new(filename)),
|
||||
"{} should NOT be recognized as a text file",
|
||||
filename
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_supported_file_no_extension() {
|
||||
assert!(!is_supported_file(Path::new("README")));
|
||||
assert!(!is_supported_file(Path::new("Makefile")));
|
||||
assert!(!is_supported_file(Path::new(".gitignore")));
|
||||
fn test_is_text_file_no_extension() {
|
||||
assert!(!is_text_file(Path::new("README")));
|
||||
assert!(!is_text_file(Path::new("Makefile")));
|
||||
assert!(!is_text_file(Path::new(".gitignore")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_supported_file_with_path() {
|
||||
assert!(is_supported_file(Path::new("/some/path/to/file.json")));
|
||||
assert!(is_supported_file(Path::new("../relative/path/file.md")));
|
||||
assert!(is_supported_file(Path::new("/path/to/image.png")));
|
||||
fn test_is_text_file_with_path() {
|
||||
assert!(is_text_file(Path::new("/some/path/to/file.json")));
|
||||
assert!(is_text_file(Path::new("../relative/path/file.md")));
|
||||
assert!(!is_text_file(Path::new("/path/to/image.png")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -58,22 +58,8 @@ fn format_error_message(err: &anyhow::Error, api_url: &str) -> String {
|
||||
);
|
||||
}
|
||||
|
||||
// 404 Not Found - check for disabled features first
|
||||
// 404 Not Found
|
||||
if err_str.contains("404") {
|
||||
if err_str.contains("Bank configuration API is disabled") {
|
||||
return format!(
|
||||
"{} {}\n\n{}\n {}\n\n{}\n {}\n\n{}\n {}",
|
||||
"✗".bright_red().bold(),
|
||||
"Bank configuration API is disabled".bright_red().bold(),
|
||||
"API URL:".bright_yellow(),
|
||||
api_url.bright_white(),
|
||||
"This feature is disabled by default for security.".bright_yellow(),
|
||||
"To enable, set HINDSIGHT_API_ENABLE_BANK_CONFIG_API=true on the API server".bright_white(),
|
||||
"Note:".bright_cyan(),
|
||||
"This allows per-bank LLM configuration overrides via API".bright_white()
|
||||
);
|
||||
}
|
||||
|
||||
return format!(
|
||||
"{} {}\n\n{}\n {}\n\n{}\n • {}\n • {}\n\n{}\n {}",
|
||||
"✗".bright_red().bold(),
|
||||
@@ -88,8 +74,8 @@ fn format_error_message(err: &anyhow::Error, api_url: &str) -> String {
|
||||
);
|
||||
}
|
||||
|
||||
// 401 Authentication failed
|
||||
if err_str.contains("401") {
|
||||
// 401/403 Authentication
|
||||
if err_str.contains("401") || err_str.contains("403") {
|
||||
return format!(
|
||||
"{} {}\n\n{}\n {}\n\n{}\n • {}\n • {}\n\n{}\n {}",
|
||||
"✗".bright_red().bold(),
|
||||
@@ -104,22 +90,6 @@ fn format_error_message(err: &anyhow::Error, api_url: &str) -> String {
|
||||
);
|
||||
}
|
||||
|
||||
// 403 Forbidden
|
||||
if err_str.contains("403") {
|
||||
return format!(
|
||||
"{} {}\n\n{}\n {}\n\n{}\n • {}\n • {}\n\n{}\n {}",
|
||||
"✗".bright_red().bold(),
|
||||
"Permission denied (403)".bright_red().bold(),
|
||||
"API URL:".bright_yellow(),
|
||||
api_url.bright_white(),
|
||||
"Possible causes:".bright_yellow(),
|
||||
"This operation is not allowed".bright_white(),
|
||||
"The feature may be disabled on the server".bright_white(),
|
||||
"Try:".bright_green(),
|
||||
"Check server configuration or contact your administrator".bright_white()
|
||||
);
|
||||
}
|
||||
|
||||
// 500 Server Error
|
||||
if err_str.contains("500") || err_str.contains("502") || err_str.contains("503") {
|
||||
return format!(
|
||||
|
||||
@@ -279,48 +279,6 @@ enum BankCommands {
|
||||
#[arg(short = 'y', long)]
|
||||
yes: bool,
|
||||
},
|
||||
|
||||
/// Get bank configuration (hierarchical overrides)
|
||||
Config {
|
||||
/// Bank ID
|
||||
bank_id: String,
|
||||
|
||||
/// Show only bank-specific overrides (not full resolved config)
|
||||
#[arg(long)]
|
||||
overrides_only: bool,
|
||||
},
|
||||
|
||||
/// Update bank configuration (set hierarchical overrides)
|
||||
SetConfig {
|
||||
/// Bank ID
|
||||
bank_id: String,
|
||||
|
||||
/// LLM provider override
|
||||
#[arg(long)]
|
||||
llm_provider: Option<String>,
|
||||
|
||||
/// LLM model override
|
||||
#[arg(long)]
|
||||
llm_model: Option<String>,
|
||||
|
||||
/// LLM API key override
|
||||
#[arg(long)]
|
||||
llm_api_key: Option<String>,
|
||||
|
||||
/// LLM base URL override
|
||||
#[arg(long)]
|
||||
llm_base_url: Option<String>,
|
||||
},
|
||||
|
||||
/// Reset bank configuration to defaults (remove all overrides)
|
||||
ResetConfig {
|
||||
/// Bank ID
|
||||
bank_id: String,
|
||||
|
||||
/// Skip confirmation prompt
|
||||
#[arg(short = 'y', long)]
|
||||
yes: bool,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
@@ -818,15 +776,6 @@ fn run() -> Result<()> {
|
||||
BankCommands::ClearObservations { bank_id, yes } => {
|
||||
commands::bank::clear_observations(&client, &bank_id, yes, verbose, output_format)
|
||||
}
|
||||
BankCommands::Config { bank_id, overrides_only } => {
|
||||
commands::bank::config(&client, &bank_id, overrides_only, verbose, output_format)
|
||||
}
|
||||
BankCommands::SetConfig { bank_id, llm_provider, llm_model, llm_api_key, llm_base_url } => {
|
||||
commands::bank::set_config(&client, &bank_id, llm_provider, llm_model, llm_api_key, llm_base_url, verbose, output_format)
|
||||
}
|
||||
BankCommands::ResetConfig { bank_id, yes } => {
|
||||
commands::bank::reset_config(&client, &bank_id, yes, verbose, output_format)
|
||||
}
|
||||
},
|
||||
|
||||
// Memory commands
|
||||
|
||||
@@ -1,24 +0,0 @@
|
||||
# Compiled Object files, Static and Dynamic libs (Shared Objects)
|
||||
*.o
|
||||
*.a
|
||||
*.so
|
||||
|
||||
# Folders
|
||||
_obj
|
||||
_test
|
||||
|
||||
# Architecture specific extensions/prefixes
|
||||
*.[568vq]
|
||||
[568vq].out
|
||||
|
||||
*.cgo1.go
|
||||
*.cgo2.c
|
||||
_cgo_defun.c
|
||||
_cgo_gotypes.go
|
||||
_cgo_export.*
|
||||
|
||||
_testmain.go
|
||||
|
||||
*.exe
|
||||
*.test
|
||||
*.prof
|
||||
@@ -1,226 +0,0 @@
|
||||
# Go API client for hindsight
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
## Overview
|
||||
This API client was generated by the [OpenAPI Generator](https://openapi-generator.tech) project. By using the [OpenAPI-spec](https://www.openapis.org/) from a remote server, you can easily generate an API client.
|
||||
|
||||
- API version: 0.4.11
|
||||
- Package version: 1.0.0
|
||||
- Generator version: 7.10.0
|
||||
- Build package: org.openapitools.codegen.languages.GoClientCodegen
|
||||
|
||||
## Installation
|
||||
|
||||
Install the following dependencies:
|
||||
|
||||
```sh
|
||||
go get github.com/stretchr/testify/assert
|
||||
go get golang.org/x/net/context
|
||||
```
|
||||
|
||||
Put the package under your project folder and add the following in import:
|
||||
|
||||
```go
|
||||
import hindsight "github.com/vectorize-io/hindsight-client-go"
|
||||
```
|
||||
|
||||
To use a proxy, set the environment variable `HTTP_PROXY`:
|
||||
|
||||
```go
|
||||
os.Setenv("HTTP_PROXY", "http://proxy_name:proxy_port")
|
||||
```
|
||||
|
||||
## Configuration of Server URL
|
||||
|
||||
Default configuration comes with `Servers` field that contains server objects as defined in the OpenAPI specification.
|
||||
|
||||
### Select Server Configuration
|
||||
|
||||
For using other server than the one defined on index 0 set context value `hindsight.ContextServerIndex` of type `int`.
|
||||
|
||||
```go
|
||||
ctx := context.WithValue(context.Background(), hindsight.ContextServerIndex, 1)
|
||||
```
|
||||
|
||||
### Templated Server URL
|
||||
|
||||
Templated server URL is formatted using default variables from configuration or from context value `hindsight.ContextServerVariables` of type `map[string]string`.
|
||||
|
||||
```go
|
||||
ctx := context.WithValue(context.Background(), hindsight.ContextServerVariables, map[string]string{
|
||||
"basePath": "v2",
|
||||
})
|
||||
```
|
||||
|
||||
Note, enum values are always validated and all unused variables are silently ignored.
|
||||
|
||||
### URLs Configuration per Operation
|
||||
|
||||
Each operation can use different server URL defined using `OperationServers` map in the `Configuration`.
|
||||
An operation is uniquely identified by `"{classname}Service.{nickname}"` string.
|
||||
Similar rules for overriding default operation server index and variables applies by using `hindsight.ContextOperationServerIndices` and `hindsight.ContextOperationServerVariables` context maps.
|
||||
|
||||
```go
|
||||
ctx := context.WithValue(context.Background(), hindsight.ContextOperationServerIndices, map[string]int{
|
||||
"{classname}Service.{nickname}": 2,
|
||||
})
|
||||
ctx = context.WithValue(context.Background(), hindsight.ContextOperationServerVariables, map[string]map[string]string{
|
||||
"{classname}Service.{nickname}": {
|
||||
"port": "8443",
|
||||
},
|
||||
})
|
||||
```
|
||||
|
||||
## Documentation for API Endpoints
|
||||
|
||||
All URIs are relative to *http://localhost*
|
||||
|
||||
Class | Method | HTTP request | Description
|
||||
------------ | ------------- | ------------- | -------------
|
||||
*BanksAPI* | [**AddBankBackground**](docs/BanksAPI.md#addbankbackground) | **Post** /v1/default/banks/{bank_id}/background | Add/merge memory bank background (deprecated)
|
||||
*BanksAPI* | [**ClearObservations**](docs/BanksAPI.md#clearobservations) | **Delete** /v1/default/banks/{bank_id}/observations | Clear all observations
|
||||
*BanksAPI* | [**CreateOrUpdateBank**](docs/BanksAPI.md#createorupdatebank) | **Put** /v1/default/banks/{bank_id} | Create or update memory bank
|
||||
*BanksAPI* | [**DeleteBank**](docs/BanksAPI.md#deletebank) | **Delete** /v1/default/banks/{bank_id} | Delete memory bank
|
||||
*BanksAPI* | [**GetAgentStats**](docs/BanksAPI.md#getagentstats) | **Get** /v1/default/banks/{bank_id}/stats | Get statistics for memory bank
|
||||
*BanksAPI* | [**GetBankConfig**](docs/BanksAPI.md#getbankconfig) | **Get** /v1/default/banks/{bank_id}/config | Get bank configuration
|
||||
*BanksAPI* | [**GetBankProfile**](docs/BanksAPI.md#getbankprofile) | **Get** /v1/default/banks/{bank_id}/profile | Get memory bank profile
|
||||
*BanksAPI* | [**ListBanks**](docs/BanksAPI.md#listbanks) | **Get** /v1/default/banks | List all memory banks
|
||||
*BanksAPI* | [**ResetBankConfig**](docs/BanksAPI.md#resetbankconfig) | **Delete** /v1/default/banks/{bank_id}/config | Reset bank configuration
|
||||
*BanksAPI* | [**TriggerConsolidation**](docs/BanksAPI.md#triggerconsolidation) | **Post** /v1/default/banks/{bank_id}/consolidate | Trigger consolidation
|
||||
*BanksAPI* | [**UpdateBank**](docs/BanksAPI.md#updatebank) | **Patch** /v1/default/banks/{bank_id} | Partial update memory bank
|
||||
*BanksAPI* | [**UpdateBankConfig**](docs/BanksAPI.md#updatebankconfig) | **Patch** /v1/default/banks/{bank_id}/config | Update bank configuration
|
||||
*BanksAPI* | [**UpdateBankDisposition**](docs/BanksAPI.md#updatebankdisposition) | **Put** /v1/default/banks/{bank_id}/profile | Update memory bank disposition
|
||||
*DirectivesAPI* | [**CreateDirective**](docs/DirectivesAPI.md#createdirective) | **Post** /v1/default/banks/{bank_id}/directives | Create directive
|
||||
*DirectivesAPI* | [**DeleteDirective**](docs/DirectivesAPI.md#deletedirective) | **Delete** /v1/default/banks/{bank_id}/directives/{directive_id} | Delete directive
|
||||
*DirectivesAPI* | [**GetDirective**](docs/DirectivesAPI.md#getdirective) | **Get** /v1/default/banks/{bank_id}/directives/{directive_id} | Get directive
|
||||
*DirectivesAPI* | [**ListDirectives**](docs/DirectivesAPI.md#listdirectives) | **Get** /v1/default/banks/{bank_id}/directives | List directives
|
||||
*DirectivesAPI* | [**UpdateDirective**](docs/DirectivesAPI.md#updatedirective) | **Patch** /v1/default/banks/{bank_id}/directives/{directive_id} | Update directive
|
||||
*DocumentsAPI* | [**DeleteDocument**](docs/DocumentsAPI.md#deletedocument) | **Delete** /v1/default/banks/{bank_id}/documents/{document_id} | Delete a document
|
||||
*DocumentsAPI* | [**GetChunk**](docs/DocumentsAPI.md#getchunk) | **Get** /v1/default/chunks/{chunk_id} | Get chunk details
|
||||
*DocumentsAPI* | [**GetDocument**](docs/DocumentsAPI.md#getdocument) | **Get** /v1/default/banks/{bank_id}/documents/{document_id} | Get document details
|
||||
*DocumentsAPI* | [**ListDocuments**](docs/DocumentsAPI.md#listdocuments) | **Get** /v1/default/banks/{bank_id}/documents | List documents
|
||||
*EntitiesAPI* | [**GetEntity**](docs/EntitiesAPI.md#getentity) | **Get** /v1/default/banks/{bank_id}/entities/{entity_id} | Get entity details
|
||||
*EntitiesAPI* | [**ListEntities**](docs/EntitiesAPI.md#listentities) | **Get** /v1/default/banks/{bank_id}/entities | List entities
|
||||
*EntitiesAPI* | [**RegenerateEntityObservations**](docs/EntitiesAPI.md#regenerateentityobservations) | **Post** /v1/default/banks/{bank_id}/entities/{entity_id}/regenerate | Regenerate entity observations (deprecated)
|
||||
*MemoryAPI* | [**ClearBankMemories**](docs/MemoryAPI.md#clearbankmemories) | **Delete** /v1/default/banks/{bank_id}/memories | Clear memory bank memories
|
||||
*MemoryAPI* | [**GetGraph**](docs/MemoryAPI.md#getgraph) | **Get** /v1/default/banks/{bank_id}/graph | Get memory graph data
|
||||
*MemoryAPI* | [**GetMemory**](docs/MemoryAPI.md#getmemory) | **Get** /v1/default/banks/{bank_id}/memories/{memory_id} | Get memory unit
|
||||
*MemoryAPI* | [**ListMemories**](docs/MemoryAPI.md#listmemories) | **Get** /v1/default/banks/{bank_id}/memories/list | List memory units
|
||||
*MemoryAPI* | [**ListTags**](docs/MemoryAPI.md#listtags) | **Get** /v1/default/banks/{bank_id}/tags | List tags
|
||||
*MemoryAPI* | [**RecallMemories**](docs/MemoryAPI.md#recallmemories) | **Post** /v1/default/banks/{bank_id}/memories/recall | Recall memory
|
||||
*MemoryAPI* | [**Reflect**](docs/MemoryAPI.md#reflect) | **Post** /v1/default/banks/{bank_id}/reflect | Reflect and generate answer
|
||||
*MemoryAPI* | [**RetainMemories**](docs/MemoryAPI.md#retainmemories) | **Post** /v1/default/banks/{bank_id}/memories | Retain memories
|
||||
*MentalModelsAPI* | [**CreateMentalModel**](docs/MentalModelsAPI.md#creatementalmodel) | **Post** /v1/default/banks/{bank_id}/mental-models | Create mental model
|
||||
*MentalModelsAPI* | [**DeleteMentalModel**](docs/MentalModelsAPI.md#deletementalmodel) | **Delete** /v1/default/banks/{bank_id}/mental-models/{mental_model_id} | Delete mental model
|
||||
*MentalModelsAPI* | [**GetMentalModel**](docs/MentalModelsAPI.md#getmentalmodel) | **Get** /v1/default/banks/{bank_id}/mental-models/{mental_model_id} | Get mental model
|
||||
*MentalModelsAPI* | [**ListMentalModels**](docs/MentalModelsAPI.md#listmentalmodels) | **Get** /v1/default/banks/{bank_id}/mental-models | List mental models
|
||||
*MentalModelsAPI* | [**RefreshMentalModel**](docs/MentalModelsAPI.md#refreshmentalmodel) | **Post** /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/refresh | Refresh mental model
|
||||
*MentalModelsAPI* | [**UpdateMentalModel**](docs/MentalModelsAPI.md#updatementalmodel) | **Patch** /v1/default/banks/{bank_id}/mental-models/{mental_model_id} | Update mental model
|
||||
*MonitoringAPI* | [**GetVersion**](docs/MonitoringAPI.md#getversion) | **Get** /version | Get API version and feature flags
|
||||
*MonitoringAPI* | [**HealthEndpointHealthGet**](docs/MonitoringAPI.md#healthendpointhealthget) | **Get** /health | Health check endpoint
|
||||
*MonitoringAPI* | [**MetricsEndpointMetricsGet**](docs/MonitoringAPI.md#metricsendpointmetricsget) | **Get** /metrics | Prometheus metrics endpoint
|
||||
*OperationsAPI* | [**CancelOperation**](docs/OperationsAPI.md#canceloperation) | **Delete** /v1/default/banks/{bank_id}/operations/{operation_id} | Cancel a pending async operation
|
||||
*OperationsAPI* | [**GetOperationStatus**](docs/OperationsAPI.md#getoperationstatus) | **Get** /v1/default/banks/{bank_id}/operations/{operation_id} | Get operation status
|
||||
*OperationsAPI* | [**ListOperations**](docs/OperationsAPI.md#listoperations) | **Get** /v1/default/banks/{bank_id}/operations | List async operations
|
||||
|
||||
|
||||
## Documentation For Models
|
||||
|
||||
- [AddBackgroundRequest](docs/AddBackgroundRequest.md)
|
||||
- [AsyncOperationSubmitResponse](docs/AsyncOperationSubmitResponse.md)
|
||||
- [BackgroundResponse](docs/BackgroundResponse.md)
|
||||
- [BankConfigResponse](docs/BankConfigResponse.md)
|
||||
- [BankConfigUpdate](docs/BankConfigUpdate.md)
|
||||
- [BankListItem](docs/BankListItem.md)
|
||||
- [BankListResponse](docs/BankListResponse.md)
|
||||
- [BankProfileResponse](docs/BankProfileResponse.md)
|
||||
- [BankStatsResponse](docs/BankStatsResponse.md)
|
||||
- [Budget](docs/Budget.md)
|
||||
- [CancelOperationResponse](docs/CancelOperationResponse.md)
|
||||
- [ChunkData](docs/ChunkData.md)
|
||||
- [ChunkIncludeOptions](docs/ChunkIncludeOptions.md)
|
||||
- [ChunkResponse](docs/ChunkResponse.md)
|
||||
- [ConsolidationResponse](docs/ConsolidationResponse.md)
|
||||
- [CreateBankRequest](docs/CreateBankRequest.md)
|
||||
- [CreateDirectiveRequest](docs/CreateDirectiveRequest.md)
|
||||
- [CreateMentalModelRequest](docs/CreateMentalModelRequest.md)
|
||||
- [CreateMentalModelResponse](docs/CreateMentalModelResponse.md)
|
||||
- [DeleteDocumentResponse](docs/DeleteDocumentResponse.md)
|
||||
- [DeleteResponse](docs/DeleteResponse.md)
|
||||
- [DirectiveListResponse](docs/DirectiveListResponse.md)
|
||||
- [DirectiveResponse](docs/DirectiveResponse.md)
|
||||
- [DispositionTraits](docs/DispositionTraits.md)
|
||||
- [DocumentResponse](docs/DocumentResponse.md)
|
||||
- [EntityDetailResponse](docs/EntityDetailResponse.md)
|
||||
- [EntityIncludeOptions](docs/EntityIncludeOptions.md)
|
||||
- [EntityInput](docs/EntityInput.md)
|
||||
- [EntityListItem](docs/EntityListItem.md)
|
||||
- [EntityListResponse](docs/EntityListResponse.md)
|
||||
- [EntityObservationResponse](docs/EntityObservationResponse.md)
|
||||
- [EntityStateResponse](docs/EntityStateResponse.md)
|
||||
- [FeaturesInfo](docs/FeaturesInfo.md)
|
||||
- [GraphDataResponse](docs/GraphDataResponse.md)
|
||||
- [HTTPValidationError](docs/HTTPValidationError.md)
|
||||
- [IncludeOptions](docs/IncludeOptions.md)
|
||||
- [ListDocumentsResponse](docs/ListDocumentsResponse.md)
|
||||
- [ListMemoryUnitsResponse](docs/ListMemoryUnitsResponse.md)
|
||||
- [ListTagsResponse](docs/ListTagsResponse.md)
|
||||
- [MemoryItem](docs/MemoryItem.md)
|
||||
- [MentalModelListResponse](docs/MentalModelListResponse.md)
|
||||
- [MentalModelResponse](docs/MentalModelResponse.md)
|
||||
- [MentalModelTrigger](docs/MentalModelTrigger.md)
|
||||
- [OperationResponse](docs/OperationResponse.md)
|
||||
- [OperationStatusResponse](docs/OperationStatusResponse.md)
|
||||
- [OperationsListResponse](docs/OperationsListResponse.md)
|
||||
- [RecallRequest](docs/RecallRequest.md)
|
||||
- [RecallResponse](docs/RecallResponse.md)
|
||||
- [RecallResult](docs/RecallResult.md)
|
||||
- [ReflectBasedOn](docs/ReflectBasedOn.md)
|
||||
- [ReflectDirective](docs/ReflectDirective.md)
|
||||
- [ReflectFact](docs/ReflectFact.md)
|
||||
- [ReflectIncludeOptions](docs/ReflectIncludeOptions.md)
|
||||
- [ReflectLLMCall](docs/ReflectLLMCall.md)
|
||||
- [ReflectMentalModel](docs/ReflectMentalModel.md)
|
||||
- [ReflectRequest](docs/ReflectRequest.md)
|
||||
- [ReflectResponse](docs/ReflectResponse.md)
|
||||
- [ReflectToolCall](docs/ReflectToolCall.md)
|
||||
- [ReflectTrace](docs/ReflectTrace.md)
|
||||
- [RetainRequest](docs/RetainRequest.md)
|
||||
- [RetainResponse](docs/RetainResponse.md)
|
||||
- [TagItem](docs/TagItem.md)
|
||||
- [TokenUsage](docs/TokenUsage.md)
|
||||
- [ToolCallsIncludeOptions](docs/ToolCallsIncludeOptions.md)
|
||||
- [UpdateDirectiveRequest](docs/UpdateDirectiveRequest.md)
|
||||
- [UpdateDispositionRequest](docs/UpdateDispositionRequest.md)
|
||||
- [UpdateMentalModelRequest](docs/UpdateMentalModelRequest.md)
|
||||
- [ValidationError](docs/ValidationError.md)
|
||||
- [ValidationErrorLocInner](docs/ValidationErrorLocInner.md)
|
||||
- [VersionResponse](docs/VersionResponse.md)
|
||||
|
||||
|
||||
## Documentation For Authorization
|
||||
|
||||
Endpoints do not require authorization.
|
||||
|
||||
|
||||
## Documentation for Utility Methods
|
||||
|
||||
Due to the fact that model structure members are all pointers, this package contains
|
||||
a number of utility functions to easily obtain pointers to values of basic types.
|
||||
Each of these functions takes a value of the given basic type and returns a pointer to it:
|
||||
|
||||
* `PtrBool`
|
||||
* `PtrInt`
|
||||
* `PtrInt32`
|
||||
* `PtrInt64`
|
||||
* `PtrFloat`
|
||||
* `PtrFloat32`
|
||||
* `PtrFloat64`
|
||||
* `PtrString`
|
||||
* `PtrTime`
|
||||
|
||||
## Author
|
||||
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user