Compare commits
3
Commits
ci-gcp
..
deadc-oder
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7972fd3906 | ||
|
|
86b698460e | ||
|
|
6f9cef674b |
+1
-7
@@ -5,7 +5,7 @@
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
HINDSIGHT_API_LLM_MODEL=o3-mini
|
||||
HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
|
||||
# Example: Anthropic Claude configuration
|
||||
@@ -41,12 +41,6 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# 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
|
||||
|
||||
@@ -46,10 +46,6 @@ jobs:
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-crewai
|
||||
working-directory: ./hindsight-integrations/crewai
|
||||
run: uv build --out-dir dist
|
||||
|
||||
# Publish in order (client and api first, then hindsight-all which depends on them)
|
||||
- name: Publish hindsight-client to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
@@ -81,12 +77,6 @@ jobs:
|
||||
packages-dir: ./hindsight-embed/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-crewai to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-integrations/crewai/dist
|
||||
skip-existing: true
|
||||
|
||||
# Upload artifacts for GitHub release
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
@@ -98,7 +88,6 @@ jobs:
|
||||
hindsight/dist/*
|
||||
hindsight-integrations/litellm/dist/*
|
||||
hindsight-embed/dist/*
|
||||
hindsight-integrations/crewai/dist/*
|
||||
retention-days: 1
|
||||
|
||||
release-typescript-client:
|
||||
@@ -279,14 +268,11 @@ jobs:
|
||||
- name: Build
|
||||
run: npm run build --workspace=hindsight-control-plane
|
||||
|
||||
- name: Verify standalone build
|
||||
run: test -f hindsight-control-plane/standalone/server.js || (echo 'standalone/server.js missing - build failed' && exit 1)
|
||||
|
||||
- name: Publish to npm
|
||||
working-directory: ./hindsight-control-plane
|
||||
run: |
|
||||
set +e
|
||||
OUTPUT=$(npm publish --access public --ignore-scripts 2>&1)
|
||||
OUTPUT=$(npm publish --access public 2>&1)
|
||||
EXIT_CODE=$?
|
||||
echo "$OUTPUT"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
|
||||
+82
-479
@@ -171,9 +171,9 @@ jobs:
|
||||
test-rust-cli:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_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 }}
|
||||
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
|
||||
@@ -181,12 +181,6 @@ jobs:
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
|
||||
@@ -233,46 +227,25 @@ jobs:
|
||||
working-directory: ./hindsight-api
|
||||
run: uv sync --frozen --no-install-project --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api
|
||||
run: |
|
||||
uv run python -c "
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Downloading cross-encoder model...')
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
|
||||
print('Models downloaded successfully')
|
||||
"
|
||||
|
||||
- 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 }}
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY=/tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID
|
||||
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..120}; do
|
||||
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 120 ]; then
|
||||
echo "API server failed to start after 120s"
|
||||
if [ $i -eq 60 ]; then
|
||||
echo "API server failed to start after 60s"
|
||||
cat /tmp/api-server.log
|
||||
exit 1
|
||||
fi
|
||||
@@ -367,21 +340,12 @@ jobs:
|
||||
|
||||
# Only test slim variants to save disk space (they're much smaller)
|
||||
# Slim variants require external embedding providers
|
||||
- name: Setup GCP credentials for smoke test
|
||||
if: matrix.variant == 'slim'
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Smoke test - verify container starts
|
||||
if: matrix.variant == 'slim'
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID: ${{ env.HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID }}
|
||||
HINDSIGHT_API_EMBEDDINGS_PROVIDER: cohere
|
||||
GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
HINDSIGHT_API_EMBEDDINGS_PROVIDER: openai
|
||||
HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
|
||||
HINDSIGHT_API_RERANKER_PROVIDER: cohere
|
||||
HINDSIGHT_API_COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
|
||||
run: ./docker/test-image.sh "hindsight-${{ matrix.name }}:test" "${{ matrix.target }}"
|
||||
@@ -389,13 +353,14 @@ jobs:
|
||||
test-api:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_LLM_PROVIDER: groq
|
||||
HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
|
||||
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
|
||||
COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
|
||||
HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
|
||||
HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
|
||||
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
|
||||
@@ -403,12 +368,6 @@ jobs:
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
@@ -455,9 +414,9 @@ jobs:
|
||||
test-python-client:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_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)
|
||||
@@ -466,12 +425,6 @@ jobs:
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
@@ -499,46 +452,25 @@ jobs:
|
||||
working-directory: ./hindsight-api
|
||||
run: uv sync --frozen --no-install-project --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api
|
||||
run: |
|
||||
uv run python -c "
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Downloading cross-encoder model...')
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
|
||||
print('Models downloaded successfully')
|
||||
"
|
||||
|
||||
- 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 }}
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY=/tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID
|
||||
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..120}; do
|
||||
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 120 ]; then
|
||||
echo "API server failed to start after 120s"
|
||||
if [ $i -eq 60 ]; then
|
||||
echo "API server failed to start after 60s"
|
||||
cat /tmp/api-server.log
|
||||
exit 1
|
||||
fi
|
||||
@@ -558,9 +490,9 @@ jobs:
|
||||
test-typescript-client:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_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)
|
||||
@@ -569,12 +501,6 @@ jobs:
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
@@ -607,46 +533,25 @@ jobs:
|
||||
working-directory: ./hindsight-clients/typescript
|
||||
run: npm run build
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api
|
||||
run: |
|
||||
uv run python -c "
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Downloading cross-encoder model...')
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
|
||||
print('Models downloaded successfully')
|
||||
"
|
||||
|
||||
- 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 }}
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY=/tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID
|
||||
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..120}; do
|
||||
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 120 ]; then
|
||||
echo "API server failed to start after 120s"
|
||||
if [ $i -eq 60 ]; then
|
||||
echo "API server failed to start after 60s"
|
||||
cat /tmp/api-server.log
|
||||
exit 1
|
||||
fi
|
||||
@@ -666,9 +571,9 @@ jobs:
|
||||
test-rust-client:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_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)
|
||||
@@ -677,12 +582,6 @@ jobs:
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
@@ -714,46 +613,25 @@ jobs:
|
||||
working-directory: ./hindsight-api
|
||||
run: uv sync --frozen --no-install-project --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api
|
||||
run: |
|
||||
uv run python -c "
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Downloading cross-encoder model...')
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
|
||||
print('Models downloaded successfully')
|
||||
"
|
||||
|
||||
- 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 }}
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY=/tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID
|
||||
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..120}; do
|
||||
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 120 ]; then
|
||||
echo "API server failed to start after 120s"
|
||||
if [ $i -eq 60 ]; then
|
||||
echo "API server failed to start after 60s"
|
||||
cat /tmp/api-server.log
|
||||
exit 1
|
||||
fi
|
||||
@@ -770,225 +648,12 @@ 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: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
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: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@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: Cache HuggingFace models
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api
|
||||
run: |
|
||||
uv run python -c "
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Downloading cross-encoder model...')
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
|
||||
print('Models downloaded successfully')
|
||||
"
|
||||
|
||||
- name: Create .env file
|
||||
run: |
|
||||
cat > .env << EOF
|
||||
HINDSIGHT_API_LLM_PROVIDER=${{ env.HINDSIGHT_API_LLM_PROVIDER }}
|
||||
HINDSIGHT_API_LLM_MODEL=${{ env.HINDSIGHT_API_LLM_MODEL }}
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY=/tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID
|
||||
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..120}; 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 120 ]; then
|
||||
echo "API server failed to start after 120s"
|
||||
cat /tmp/api-server.log
|
||||
exit 1
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
|
||||
- name: Build Go client
|
||||
working-directory: ./hindsight-clients/go
|
||||
run: go build ./...
|
||||
|
||||
- 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-openclaw-integration:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_URL: http://localhost:8888
|
||||
HINDSIGHT_EMBED_PACKAGE_PATH: ${{ github.workspace }}/hindsight-embed
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@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 Node.js
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '22'
|
||||
|
||||
- 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: Install embed dependencies
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv sync --frozen --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api
|
||||
run: |
|
||||
uv run python -c "
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Downloading cross-encoder model...')
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
|
||||
print('Models downloaded successfully')
|
||||
"
|
||||
|
||||
- name: Install openclaw integration dependencies
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm ci
|
||||
|
||||
- name: Create .env file
|
||||
run: |
|
||||
cat > .env << EOF
|
||||
HINDSIGHT_API_LLM_PROVIDER=${{ env.HINDSIGHT_API_LLM_PROVIDER }}
|
||||
HINDSIGHT_API_LLM_MODEL=${{ env.HINDSIGHT_API_LLM_MODEL }}
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY=/tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID
|
||||
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..120}; 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 120 ]; then
|
||||
echo "API server failed to start after 120s"
|
||||
cat /tmp/api-server.log
|
||||
exit 1
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
|
||||
- name: Run openclaw integration tests
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm run test: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:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_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 }}
|
||||
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
|
||||
@@ -996,12 +661,6 @@ jobs:
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
@@ -1049,22 +708,21 @@ jobs:
|
||||
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 }}
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY=/tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID
|
||||
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..120}; do
|
||||
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 120 ]; then
|
||||
echo "API server failed to start after 120s"
|
||||
if [ $i -eq 60 ]; then
|
||||
echo "API server failed to start after 60s"
|
||||
cat /tmp/api-server.log
|
||||
exit 1
|
||||
fi
|
||||
@@ -1081,35 +739,6 @@ jobs:
|
||||
echo "=== API Server Logs ==="
|
||||
cat /tmp/api-server.log || echo "No API server log found"
|
||||
|
||||
test-crewai-integration:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build crewai integration
|
||||
working-directory: ./hindsight-integrations/crewai
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/crewai
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/crewai
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-litellm-integration:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
@@ -1142,21 +771,15 @@ jobs:
|
||||
test-embed:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_LLM_PROVIDER: groq
|
||||
HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
|
||||
# Prefer CPU-only PyTorch in CI
|
||||
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
@@ -1192,25 +815,19 @@ jobs:
|
||||
test-hindsight-all:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_LLM_PROVIDER: groq
|
||||
HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
|
||||
# For test_server_integration.py compatibility
|
||||
HINDSIGHT_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_LLM_PROVIDER: groq
|
||||
HINDSIGHT_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
HINDSIGHT_LLM_MODEL: openai/gpt-oss-20b
|
||||
# Prefer CPU-only PyTorch in CI
|
||||
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
@@ -1247,9 +864,9 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
needs: test-rust-cli
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_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 }}
|
||||
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
|
||||
@@ -1257,12 +874,6 @@ jobs:
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Download CLI artifact
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
@@ -1305,57 +916,55 @@ jobs:
|
||||
npm ci --workspace=hindsight-clients/typescript
|
||||
npm run build --workspace=hindsight-clients/typescript
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api
|
||||
run: |
|
||||
uv run python -c "
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Downloading reranker model...')
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
|
||||
print('Models downloaded successfully')
|
||||
"
|
||||
|
||||
- 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 }}
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY=/tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID
|
||||
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..120}; do
|
||||
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 120 ]; then
|
||||
echo "API server failed to start after 120s"
|
||||
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 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()
|
||||
@@ -1366,9 +975,9 @@ jobs:
|
||||
test-upgrade:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_LLM_PROVIDER: groq
|
||||
HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
|
||||
|
||||
@@ -1377,12 +986,6 @@ jobs:
|
||||
with:
|
||||
fetch-depth: 0 # Full history needed for git clone of tags
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Fetch tags
|
||||
run: git fetch --tags
|
||||
|
||||
|
||||
@@ -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
|
||||
```
|
||||
|
||||
@@ -317,7 +310,7 @@ npm install
|
||||
Required env vars:
|
||||
- `HINDSIGHT_API_LLM_PROVIDER`: openai, anthropic, gemini, groq, ollama, lmstudio
|
||||
- `HINDSIGHT_API_LLM_API_KEY`: Your API key
|
||||
- `HINDSIGHT_API_LLM_MODEL`: Model name (e.g., gpt-4o-mini, claude-sonnet-4-20250514)
|
||||
- `HINDSIGHT_API_LLM_MODEL`: Model name (e.g., o3-mini, claude-sonnet-4-20250514)
|
||||
|
||||
Optional (uses local models by default):
|
||||
- `HINDSIGHT_API_EMBEDDINGS_PROVIDER`: local (default) or tei
|
||||
|
||||
@@ -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:
|
||||
@@ -170,11 +170,6 @@ RUN chown -R hindsight:hindsight /app
|
||||
|
||||
USER hindsight
|
||||
|
||||
# Create pg0 data directory as hindsight user so that Docker seeds new named
|
||||
# volumes with correct ownership (UID 1000) on first use, avoiding the
|
||||
# "Permission denied" error when mounting a fresh root-owned volume.
|
||||
RUN mkdir -p /home/hindsight/.pg0
|
||||
|
||||
ENV PATH="/app/api/.venv/bin:${PATH}"
|
||||
|
||||
# Pre-download tiktoken encoding (ALWAYS - required for token counting even in air-gapped envs)
|
||||
@@ -326,11 +321,6 @@ RUN chown -R hindsight:hindsight /app
|
||||
|
||||
USER hindsight
|
||||
|
||||
# Create pg0 data directory as hindsight user so that Docker seeds new named
|
||||
# volumes with correct ownership (UID 1000) on first use, avoiding the
|
||||
# "Permission denied" error when mounting a fresh root-owned volume.
|
||||
RUN mkdir -p /home/hindsight/.pg0
|
||||
|
||||
ENV PATH="/app/api/.venv/bin:${PATH}"
|
||||
|
||||
# Pre-download tiktoken encoding (ALWAYS - required for token counting even in air-gapped envs)
|
||||
|
||||
+10
-26
@@ -13,9 +13,9 @@
|
||||
# target - Optional: 'cp-only' for control plane, otherwise assumes API image (default: api)
|
||||
#
|
||||
# Environment variables:
|
||||
# HINDSIGHT_API_LLM_API_KEY - Required for API/standalone images (LLM verification)
|
||||
# HINDSIGHT_API_LLM_PROVIDER - LLM provider (default: openai)
|
||||
# HINDSIGHT_API_LLM_MODEL - LLM model (default: gpt-4o-mini)
|
||||
# GROQ_API_KEY - Required for API/standalone images (LLM verification)
|
||||
# HINDSIGHT_API_LLM_PROVIDER - LLM provider (default: groq)
|
||||
# HINDSIGHT_API_LLM_MODEL - LLM model (default: llama-3.3-70b-versatile)
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER - Embeddings provider (optional, for slim images: openai, cohere, tei)
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY - OpenAI API key for embeddings (optional)
|
||||
# HINDSIGHT_API_RERANKER_PROVIDER - Reranker provider (optional, for slim images: cohere, tei)
|
||||
@@ -34,7 +34,7 @@
|
||||
# ./docker/test-image.sh hindsight-control-plane:test cp-only
|
||||
#
|
||||
# # Test slim image with external providers
|
||||
# export HINDSIGHT_API_LLM_API_KEY=sk_xxx
|
||||
# export GROQ_API_KEY=gsk_xxx
|
||||
# export HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai
|
||||
# export HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=sk-xxx
|
||||
# export HINDSIGHT_API_RERANKER_PROVIDER=cohere
|
||||
@@ -60,8 +60,8 @@ IMAGE="${1:-}"
|
||||
TARGET="${2:-api}"
|
||||
TIMEOUT="${SMOKE_TEST_TIMEOUT:-120}"
|
||||
CONTAINER_NAME="${SMOKE_TEST_CONTAINER_NAME:-hindsight-smoke-test}"
|
||||
LLM_PROVIDER="${HINDSIGHT_API_LLM_PROVIDER:-openai}"
|
||||
LLM_MODEL="${HINDSIGHT_API_LLM_MODEL:-gpt-4o-mini}"
|
||||
LLM_PROVIDER="${HINDSIGHT_API_LLM_PROVIDER:-groq}"
|
||||
LLM_MODEL="${HINDSIGHT_API_LLM_MODEL:-llama-3.3-70b-versatile}"
|
||||
|
||||
# Validate arguments
|
||||
if [ -z "$IMAGE" ]; then
|
||||
@@ -88,9 +88,9 @@ else
|
||||
fi
|
||||
|
||||
# Check for required environment variables
|
||||
if [ "$NEEDS_LLM" = true ] && [ "$LLM_PROVIDER" != "vertexai" ] && [ -z "${HINDSIGHT_API_LLM_API_KEY:-}" ]; then
|
||||
echo -e "${RED}Error: HINDSIGHT_API_LLM_API_KEY environment variable is required for API/standalone images${NC}"
|
||||
echo "Set it with: export HINDSIGHT_API_LLM_API_KEY=your-api-key"
|
||||
if [ "$NEEDS_LLM" = true ] && [ -z "${GROQ_API_KEY:-}" ]; then
|
||||
echo -e "${RED}Error: GROQ_API_KEY environment variable is required for API/standalone images${NC}"
|
||||
echo "Set it with: export GROQ_API_KEY=your-api-key"
|
||||
exit 2
|
||||
fi
|
||||
|
||||
@@ -123,25 +123,9 @@ else
|
||||
# Build docker run command with required and optional env vars
|
||||
DOCKER_CMD="docker run -d --name $CONTAINER_NAME"
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_LLM_PROVIDER=$LLM_PROVIDER"
|
||||
if [ -n "${HINDSIGHT_API_LLM_API_KEY:-}" ]; then
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_LLM_API_KEY=${HINDSIGHT_API_LLM_API_KEY}"
|
||||
fi
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_LLM_API_KEY=${GROQ_API_KEY}"
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_LLM_MODEL=$LLM_MODEL"
|
||||
|
||||
# Add Vertex AI config if provider is vertexai
|
||||
if [ "$LLM_PROVIDER" = "vertexai" ]; then
|
||||
if [ -n "${HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY:-}" ]; then
|
||||
DOCKER_CMD="$DOCKER_CMD -v ${HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY}:/tmp/gcp-credentials.json:ro"
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY=/tmp/gcp-credentials.json"
|
||||
fi
|
||||
if [ -n "${HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID:-}" ]; then
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=${HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID}"
|
||||
fi
|
||||
if [ -n "${HINDSIGHT_API_LLM_VERTEXAI_REGION:-}" ]; then
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_LLM_VERTEXAI_REGION=${HINDSIGHT_API_LLM_VERTEXAI_REGION}"
|
||||
fi
|
||||
fi
|
||||
|
||||
# Add optional embeddings provider config
|
||||
if [ -n "${HINDSIGHT_API_EMBEDDINGS_PROVIDER:-}" ]; then
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_EMBEDDINGS_PROVIDER=${HINDSIGHT_API_EMBEDDINGS_PROVIDER}"
|
||||
|
||||
@@ -6,17 +6,24 @@
|
||||
# It expects API keys to be set in environment variables.
|
||||
#
|
||||
# Usage:
|
||||
# export GROQ_API_KEY=gsk_xxx
|
||||
# export OPENAI_API_KEY=sk-xxx
|
||||
# export COHERE_API_KEY=xxx
|
||||
# ./docker/test-slim-local.sh
|
||||
#
|
||||
# Or inline:
|
||||
# OPENAI_API_KEY=sk_xxx COHERE_API_KEY=xxx ./docker/test-slim-local.sh
|
||||
# GROQ_API_KEY=gsk_xxx OPENAI_API_KEY=sk_xxx COHERE_API_KEY=xxx ./docker/test-slim-local.sh
|
||||
#
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# Check for required API keys
|
||||
if [ -z "${GROQ_API_KEY:-}" ]; then
|
||||
echo "❌ Error: GROQ_API_KEY environment variable is required"
|
||||
echo "Set it with: export GROQ_API_KEY=gsk_xxx"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [ -z "${OPENAI_API_KEY:-}" ]; then
|
||||
echo "❌ Error: OPENAI_API_KEY environment variable is required"
|
||||
echo "Set it with: export OPENAI_API_KEY=sk-xxx"
|
||||
@@ -34,10 +41,7 @@ IMAGE="${1:-hindsight-slim:test}"
|
||||
echo "Testing image: $IMAGE"
|
||||
echo ""
|
||||
|
||||
# Set up LLM and external providers
|
||||
export HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
export HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY
|
||||
export HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
# Set up external providers
|
||||
export HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai
|
||||
export HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=$OPENAI_API_KEY
|
||||
export HINDSIGHT_API_RERANKER_PROVIDER=cohere
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.4.13
|
||||
appVersion: "0.4.13"
|
||||
version: 0.4.10
|
||||
appVersion: "0.4.10"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -46,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.4.13"
|
||||
__version__ = "0.4.10"
|
||||
|
||||
@@ -24,35 +24,14 @@ depends_on: str | Sequence[str] | None = None
|
||||
|
||||
def _detect_vector_extension() -> str:
|
||||
"""
|
||||
Detect or validate vector extension: 'pgvector', 'vchord', or 'pgvectorscale'.
|
||||
Detect or validate vector extension: 'vchord' or 'pgvector'.
|
||||
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":
|
||||
if vector_extension == "vchord":
|
||||
vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar()
|
||||
if not vchord_check:
|
||||
raise RuntimeError(
|
||||
@@ -67,14 +46,12 @@ def _detect_vector_extension() -> str:
|
||||
)
|
||||
return "pgvector"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {vector_extension}. Must be 'pgvector', 'vchord', or 'pgvectorscale'"
|
||||
)
|
||||
raise ValueError(f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {vector_extension}. Must be 'pgvector' or 'vchord'")
|
||||
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Detect or validate text search extension: 'native' or 'vchord'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
"""
|
||||
@@ -92,23 +69,11 @@ def _detect_text_search_extension() -> str:
|
||||
# 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'"
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native' or 'vchord'"
|
||||
)
|
||||
|
||||
|
||||
@@ -267,12 +232,6 @@ def upgrade() -> None:
|
||||
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("""
|
||||
@@ -312,21 +271,7 @@ def upgrade() -> None:
|
||||
# Create vector index - conditional based on available extension
|
||||
vector_ext = _detect_vector_extension()
|
||||
|
||||
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":
|
||||
if vector_ext == "vchord":
|
||||
# Use vchordrq index for vchord (supports high-dimensional embeddings)
|
||||
op.execute("""
|
||||
CREATE INDEX idx_memory_units_embedding ON memory_units
|
||||
@@ -350,14 +295,6 @@ def upgrade() -> None:
|
||||
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("""
|
||||
|
||||
@@ -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")
|
||||
+7
-85
@@ -31,35 +31,14 @@ def _get_schema_prefix() -> str:
|
||||
|
||||
def _detect_vector_extension() -> str:
|
||||
"""
|
||||
Detect or validate vector extension: 'pgvector', 'vchord', or 'pgvectorscale'.
|
||||
Detect or validate vector extension: 'vchord' or 'pgvector'.
|
||||
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":
|
||||
if vector_extension == "vchord":
|
||||
vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar()
|
||||
if not vchord_check:
|
||||
raise RuntimeError(
|
||||
@@ -74,14 +53,12 @@ def _detect_vector_extension() -> str:
|
||||
)
|
||||
return "pgvector"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {vector_extension}. Must be 'pgvector', 'vchord', or 'pgvectorscale'"
|
||||
)
|
||||
raise ValueError(f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {vector_extension}. Must be 'pgvector' or 'vchord'")
|
||||
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Detect or validate text search extension: 'native' or 'vchord'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
"""
|
||||
@@ -99,23 +76,11 @@ def _detect_text_search_extension() -> str:
|
||||
# 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'"
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native' or 'vchord'"
|
||||
)
|
||||
|
||||
|
||||
@@ -157,19 +122,7 @@ def upgrade() -> None:
|
||||
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":
|
||||
if vector_ext == "vchord":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_embedding ON {schema}learnings
|
||||
USING vchordrq (embedding vector_l2_ops)
|
||||
@@ -193,15 +146,6 @@ def upgrade() -> None:
|
||||
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"""
|
||||
@@ -236,19 +180,7 @@ def upgrade() -> None:
|
||||
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":
|
||||
if vector_ext == "vchord":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_embedding ON {schema}pinned_reflections
|
||||
USING vchordrq (embedding vector_l2_ops)
|
||||
@@ -272,16 +204,6 @@ def upgrade() -> None:
|
||||
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"""
|
||||
|
||||
-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
|
||||
|
||||
@@ -74,7 +74,7 @@ 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
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES, MemoryFact, TokenUsage
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES, TokenUsage
|
||||
from hindsight_api.engine.search.tags import TagsMatch
|
||||
from hindsight_api.extensions import HttpExtension, OperationValidationError, load_extension
|
||||
from hindsight_api.metrics import create_metrics_collector, get_metrics_collector, initialize_metrics
|
||||
@@ -97,12 +97,6 @@ class ChunkIncludeOptions(BaseModel):
|
||||
max_tokens: int = Field(default=8192, description="Maximum tokens for chunks (chunks may be truncated)")
|
||||
|
||||
|
||||
class SourceFactsIncludeOptions(BaseModel):
|
||||
"""Options for including source facts for observation-type results."""
|
||||
|
||||
max_tokens: int = Field(default=4096, description="Maximum tokens for source facts")
|
||||
|
||||
|
||||
class IncludeOptions(BaseModel):
|
||||
"""Options for including additional data in recall results."""
|
||||
|
||||
@@ -113,10 +107,6 @@ class IncludeOptions(BaseModel):
|
||||
chunks: ChunkIncludeOptions | None = Field(
|
||||
default=None, description="Include raw chunks. Set to {} to enable, null to disable (default: disabled)."
|
||||
)
|
||||
source_facts: SourceFactsIncludeOptions | None = Field(
|
||||
default=None,
|
||||
description="Include source facts for observation-type results. Set to {} to enable, null to disable (default: disabled).",
|
||||
)
|
||||
|
||||
|
||||
class RecallRequest(BaseModel):
|
||||
@@ -199,9 +189,6 @@ class RecallResult(BaseModel):
|
||||
metadata: dict[str, str] | None = None # User-defined metadata
|
||||
chunk_id: str | None = None # Chunk this fact was extracted from
|
||||
tags: list[str] | None = None # Visibility scope tags
|
||||
source_fact_ids: list[str] | None = (
|
||||
None # IDs of source facts (observation type only, when source_facts is enabled)
|
||||
)
|
||||
|
||||
|
||||
class EntityObservationResponse(BaseModel):
|
||||
@@ -353,9 +340,6 @@ class RecallResponse(BaseModel):
|
||||
default=None, description="Entity states for entities mentioned in results"
|
||||
)
|
||||
chunks: dict[str, ChunkData] | None = Field(default=None, description="Chunks for facts, keyed by chunk_id")
|
||||
source_facts: dict[str, RecallResult] | None = Field(
|
||||
default=None, description="Source facts for observation-type results, keyed by fact ID"
|
||||
)
|
||||
|
||||
|
||||
class EntityInput(BaseModel):
|
||||
@@ -429,6 +413,7 @@ class RetainRequest(BaseModel):
|
||||
},
|
||||
],
|
||||
"async": False,
|
||||
"document_tags": ["user_a", "user_b"],
|
||||
}
|
||||
}
|
||||
)
|
||||
@@ -441,38 +426,7 @@ class RetainRequest(BaseModel):
|
||||
)
|
||||
document_tags: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Deprecated. Use item-level tags instead.",
|
||||
deprecated=True,
|
||||
)
|
||||
|
||||
|
||||
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)",
|
||||
description="Tags applied to all items in this request. These are merged with any item-level tags.",
|
||||
)
|
||||
|
||||
|
||||
@@ -500,7 +454,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,
|
||||
@@ -508,26 +462,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."""
|
||||
|
||||
@@ -1423,16 +1357,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."""
|
||||
|
||||
@@ -1457,13 +1381,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):
|
||||
@@ -1489,7 +1406,6 @@ class FeaturesInfo(BaseModel):
|
||||
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):
|
||||
@@ -1504,7 +1420,6 @@ class VersionResponse(BaseModel):
|
||||
"mcp": True,
|
||||
"worker": True,
|
||||
"bank_config_api": False,
|
||||
"file_upload_api": True,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -1799,7 +1714,6 @@ def _register_routes(app: FastAPI):
|
||||
mcp=config.mcp_enabled,
|
||||
worker=config.worker_enabled,
|
||||
bank_config_api=config.enable_bank_config_api,
|
||||
file_upload_api=config.enable_file_upload_api,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -1975,10 +1889,6 @@ def _register_routes(app: FastAPI):
|
||||
include_chunks = request.include.chunks is not None
|
||||
max_chunk_tokens = request.include.chunks.max_tokens if include_chunks else 8192
|
||||
|
||||
# Determine source facts inclusion settings
|
||||
include_source_facts = request.include.source_facts is not None
|
||||
max_source_facts_tokens = request.include.source_facts.max_tokens if include_source_facts else 4096
|
||||
|
||||
pre_recall = time.time() - handler_start
|
||||
# Run recall with tracing (record metrics)
|
||||
with metrics.record_operation(
|
||||
@@ -1997,16 +1907,14 @@ def _register_routes(app: FastAPI):
|
||||
max_entity_tokens=max_entity_tokens,
|
||||
include_chunks=include_chunks,
|
||||
max_chunk_tokens=max_chunk_tokens,
|
||||
include_source_facts=include_source_facts,
|
||||
max_source_facts_tokens=max_source_facts_tokens,
|
||||
request_context=request_context,
|
||||
tags=request.tags,
|
||||
tags_match=request.tags_match,
|
||||
)
|
||||
|
||||
# Convert core MemoryFact objects to API RecallResult objects (excluding internal metrics)
|
||||
def _fact_to_result(fact: "MemoryFact") -> RecallResult:
|
||||
return RecallResult(
|
||||
recall_results = [
|
||||
RecallResult(
|
||||
id=fact.id,
|
||||
text=fact.text,
|
||||
type=fact.fact_type,
|
||||
@@ -2018,10 +1926,9 @@ def _register_routes(app: FastAPI):
|
||||
document_id=fact.document_id,
|
||||
chunk_id=fact.chunk_id,
|
||||
tags=fact.tags,
|
||||
source_fact_ids=fact.source_fact_ids,
|
||||
)
|
||||
|
||||
recall_results = [_fact_to_result(fact) for fact in core_result.results]
|
||||
for fact in core_result.results
|
||||
]
|
||||
|
||||
# Convert chunks from engine to HTTP API format
|
||||
chunks_response = None
|
||||
@@ -2049,19 +1956,11 @@ def _register_routes(app: FastAPI):
|
||||
],
|
||||
)
|
||||
|
||||
# Convert source facts dict to API format
|
||||
source_facts_response = None
|
||||
if core_result.source_facts:
|
||||
source_facts_response = {
|
||||
fact_id: _fact_to_result(fact) for fact_id, fact in core_result.source_facts.items()
|
||||
}
|
||||
|
||||
response = RecallResponse(
|
||||
results=recall_results,
|
||||
trace=core_result.trace,
|
||||
entities=entities_response,
|
||||
chunks=chunks_response,
|
||||
source_facts=source_facts_response,
|
||||
)
|
||||
|
||||
handler_duration = time.time() - handler_start
|
||||
@@ -3664,21 +3563,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(
|
||||
@@ -3716,147 +3600,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,
|
||||
|
||||
@@ -78,9 +78,10 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
If False, only expose bank-scoped tools without bank_id parameters.
|
||||
|
||||
Returns:
|
||||
Configured FastMCP server instance
|
||||
Configured FastMCP server instance with stateless_http enabled
|
||||
"""
|
||||
mcp = FastMCP("hindsight-mcp-server")
|
||||
# Use stateless_http=True for Claude Code compatibility
|
||||
mcp = FastMCP("hindsight-mcp-server", stateless_http=True)
|
||||
|
||||
# Configure and register tools using shared module
|
||||
config = MCPToolsConfig(
|
||||
@@ -210,9 +211,9 @@ class MCPMiddleware:
|
||||
else:
|
||||
# Create servers internally (for direct construction / tests)
|
||||
self.multi_bank_server = create_mcp_server(memory, multi_bank=True)
|
||||
self.multi_bank_app = self.multi_bank_server.http_app(path="/", stateless_http=True)
|
||||
self.multi_bank_app = self.multi_bank_server.http_app(path="/")
|
||||
self.single_bank_server = create_mcp_server(memory, multi_bank=False)
|
||||
self.single_bank_app = self.single_bank_server.http_app(path="/", stateless_http=True)
|
||||
self.single_bank_app = self.single_bank_server.http_app(path="/")
|
||||
|
||||
def _get_header(self, scope: dict, name: str) -> str | None:
|
||||
"""Extract a header value from ASGI scope."""
|
||||
@@ -378,9 +379,9 @@ def create_mcp_servers(memory: MemoryEngine):
|
||||
Tuple of (multi_bank_server, single_bank_server, multi_bank_app, single_bank_app)
|
||||
"""
|
||||
multi_bank_server = create_mcp_server(memory, multi_bank=True)
|
||||
multi_bank_app = multi_bank_server.http_app(path="/", stateless_http=True)
|
||||
multi_bank_app = multi_bank_server.http_app(path="/")
|
||||
|
||||
single_bank_server = create_mcp_server(memory, multi_bank=False)
|
||||
single_bank_app = single_bank_server.http_app(path="/", stateless_http=True)
|
||||
single_bank_app = single_bank_server.http_app(path="/")
|
||||
|
||||
return multi_bank_server, single_bank_server, multi_bank_app, single_bank_app
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -129,11 +129,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 +189,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"
|
||||
@@ -233,6 +220,8 @@ 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"
|
||||
ENV_RECALL_CONNECTION_BUDGET = "HINDSIGHT_API_RECALL_CONNECTION_BUDGET"
|
||||
ENV_MCP_LOCAL_BANK_ID = "HINDSIGHT_API_MCP_LOCAL_BANK_ID"
|
||||
ENV_MCP_INSTRUCTIONS = "HINDSIGHT_API_MCP_INSTRUCTIONS"
|
||||
ENV_MENTAL_MODEL_REFRESH_CONCURRENCY = "HINDSIGHT_API_MENTAL_MODEL_REFRESH_CONCURRENCY"
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
@@ -253,29 +242,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"
|
||||
@@ -314,18 +280,18 @@ DEFAULT_LLM_PROVIDER = "openai"
|
||||
|
||||
# Provider-specific default models
|
||||
PROVIDER_DEFAULT_MODELS = {
|
||||
"openai": "gpt-4o-mini",
|
||||
"openai": "o3-mini",
|
||||
"anthropic": "claude-haiku-4-5-20251001",
|
||||
"gemini": "gemini-2.5-flash",
|
||||
"groq": "openai/gpt-oss-120b",
|
||||
"ollama": "gemma3:12b",
|
||||
"lmstudio": "local-model",
|
||||
"vertexai": "google/gemini-2.5-flash-lite",
|
||||
"vertexai": "gemini-2.0-flash-001",
|
||||
"openai-codex": "gpt-5.2-codex",
|
||||
"claude-code": "claude-sonnet-4-5-20250929",
|
||||
"mock": "mock-model",
|
||||
}
|
||||
DEFAULT_LLM_MODEL = "gpt-4o-mini" # Fallback if provider not in table
|
||||
DEFAULT_LLM_MODEL = "o3-mini" # Fallback if provider not in table
|
||||
DEFAULT_LLM_MAX_CONCURRENT = 32
|
||||
DEFAULT_LLM_MAX_RETRIES = 10 # Max retry attempts for LLM API calls
|
||||
DEFAULT_LLM_INITIAL_BACKOFF = 1.0 # Initial backoff in seconds for retry exponential backoff
|
||||
@@ -360,21 +326,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"
|
||||
# Vector extension (pgvector vs vchord)
|
||||
DEFAULT_VECTOR_EXTENSION = "pgvector" # Options: "pgvector", "vchord"
|
||||
|
||||
# Text search extension (native PostgreSQL, vchord BM25, or Timescale pg_textsearch)
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord", "pg_textsearch"
|
||||
# Text search extension (native PostgreSQL vs vchord BM25)
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord"
|
||||
|
||||
# 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
|
||||
@@ -387,6 +349,7 @@ DEFAULT_GRAPH_RETRIEVER = "link_expansion" # Options: "link_expansion", "mpfp",
|
||||
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
|
||||
DEFAULT_RECALL_CONNECTION_BUDGET = 4 # Max concurrent DB connections per recall operation
|
||||
DEFAULT_MCP_LOCAL_BANK_ID = "mcp"
|
||||
DEFAULT_MENTAL_MODEL_REFRESH_CONCURRENCY = 8 # Max concurrent mental model refreshes
|
||||
|
||||
# Retain settings
|
||||
@@ -396,17 +359,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
|
||||
@@ -530,8 +482,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
|
||||
@@ -582,9 +532,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
|
||||
@@ -602,9 +549,6 @@ 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
|
||||
@@ -628,29 +572,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
|
||||
@@ -708,13 +629,6 @@ class HindsightConfig:
|
||||
"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
|
||||
@@ -729,11 +643,6 @@ class HindsightConfig:
|
||||
"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]:
|
||||
"""
|
||||
@@ -790,14 +699,14 @@ class HindsightConfig:
|
||||
def validate(self) -> None:
|
||||
"""Validate configuration values and raise errors for invalid combinations."""
|
||||
# Validate vector_extension
|
||||
valid_extensions = ("pgvector", "vchord", "pgvectorscale")
|
||||
valid_extensions = ("pgvector", "vchord")
|
||||
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")
|
||||
valid_text_search = ("native", "vchord")
|
||||
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)}"
|
||||
@@ -840,8 +749,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),
|
||||
@@ -940,12 +847,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),
|
||||
@@ -975,10 +876,6 @@ 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)),
|
||||
@@ -1014,39 +911,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(
|
||||
|
||||
@@ -18,8 +18,6 @@ import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ...config import get_config
|
||||
from ..memory_engine import fq_table
|
||||
from ..retain import embedding_utils
|
||||
@@ -33,22 +31,10 @@ if TYPE_CHECKING:
|
||||
|
||||
from ...api.http import RequestContext
|
||||
from ..memory_engine import MemoryEngine
|
||||
from ..response_models import MemoryFact, RecallResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class _ConsolidationAction(BaseModel):
|
||||
action: str # "update" | "create"
|
||||
text: str
|
||||
reason: str = ""
|
||||
learning_id: str | None = None # required for "update" actions
|
||||
|
||||
|
||||
class _ConsolidationResponse(BaseModel):
|
||||
actions: list[_ConsolidationAction]
|
||||
|
||||
|
||||
class ConsolidationPerfLog:
|
||||
"""Performance logging for consolidation operations."""
|
||||
|
||||
@@ -459,7 +445,8 @@ async def _process_memory(
|
||||
# Find related observations using the full recall system
|
||||
# SECURITY: Pass tags to ensure observations don't leak across security boundaries
|
||||
t0 = time.time()
|
||||
recall_result = await _find_related_observations(
|
||||
related_observations = await _find_related_observations(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
query=fact_text,
|
||||
@@ -475,7 +462,7 @@ async def _process_memory(
|
||||
actions = await _consolidate_with_llm(
|
||||
memory_engine=memory_engine,
|
||||
fact_text=fact_text,
|
||||
recall_result=recall_result,
|
||||
observations=related_observations, # Can be empty list
|
||||
mission=mission,
|
||||
)
|
||||
if perf:
|
||||
@@ -496,7 +483,7 @@ async def _process_memory(
|
||||
bank_id=bank_id,
|
||||
memory_id=memory_id,
|
||||
action=action,
|
||||
observations=recall_result.results,
|
||||
observations=related_observations,
|
||||
source_fact_tags=fact_tags, # Pass source fact's tags for security
|
||||
source_occurred_start=memory.get("occurred_start"),
|
||||
source_occurred_end=memory.get("occurred_end"),
|
||||
@@ -550,7 +537,7 @@ async def _execute_update_action(
|
||||
bank_id: str,
|
||||
memory_id: uuid.UUID,
|
||||
action: dict[str, Any],
|
||||
observations: list["MemoryFact"],
|
||||
observations: list[dict[str, Any]],
|
||||
source_fact_tags: list[str] | None = None,
|
||||
source_occurred_start: datetime | None = None,
|
||||
source_occurred_end: datetime | None = None,
|
||||
@@ -579,27 +566,28 @@ async def _execute_update_action(
|
||||
return {"action": "skipped", "reason": "missing_learning_id_or_text"}
|
||||
|
||||
# Find the observation
|
||||
model = next((m for m in observations if m.id == learning_id), None)
|
||||
model = next((m for m in observations if str(m["id"]) == learning_id), None)
|
||||
if not model:
|
||||
return {"action": "skipped", "reason": "learning_not_found"}
|
||||
|
||||
# Build history entry (history is fetched fresh from DB on update to avoid stale state)
|
||||
history = [
|
||||
# Build history entry
|
||||
history = list(model.get("history", []))
|
||||
history.append(
|
||||
{
|
||||
"previous_text": model.text,
|
||||
"previous_text": model["text"],
|
||||
"changed_at": datetime.now(timezone.utc).isoformat(),
|
||||
"reason": reason,
|
||||
"source_memory_id": str(memory_id),
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
# Update source_memory_ids
|
||||
source_ids = list(model.source_fact_ids or [])
|
||||
source_ids = list(model.get("source_memory_ids", []))
|
||||
source_ids.append(memory_id)
|
||||
|
||||
# SECURITY: Merge source fact's tags into existing observation tags
|
||||
# This ensures all contributors can see the observation they contributed to
|
||||
existing_tags = set(model.tags or [])
|
||||
existing_tags = set(model.get("tags", []) or [])
|
||||
source_tags = set(source_fact_tags or [])
|
||||
merged_tags = list(existing_tags | source_tags) # Union of both tag sets
|
||||
if source_tags and source_tags != existing_tags:
|
||||
@@ -735,12 +723,13 @@ async def _create_memory_links(
|
||||
|
||||
|
||||
async def _find_related_observations(
|
||||
conn: "Connection",
|
||||
memory_engine: "MemoryEngine",
|
||||
bank_id: str,
|
||||
query: str,
|
||||
request_context: "RequestContext",
|
||||
tags: list[str] | None = None,
|
||||
) -> "RecallResult":
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Find observations related to the given query using optimized recall.
|
||||
|
||||
@@ -785,51 +774,96 @@ async def _find_related_observations(
|
||||
request_context=request_context,
|
||||
tags=tags, # Filter by source memory's tags
|
||||
tags_match=tags_match, # Use strict matching for security
|
||||
include_source_facts=True, # Embed source facts so we avoid a separate DB fetch
|
||||
max_source_facts_tokens=-1, # No token limit — we need all source facts for consolidation
|
||||
_quiet=True, # Suppress logging
|
||||
)
|
||||
finally:
|
||||
if recall_span:
|
||||
recall_span.end()
|
||||
|
||||
return recall_result
|
||||
# If no observations returned, return empty list
|
||||
if not recall_result.results:
|
||||
return []
|
||||
|
||||
# Batch fetch all observations in a single query (no artificial limit)
|
||||
observation_ids = [uuid.UUID(obs.id) for obs in recall_result.results]
|
||||
|
||||
def _build_observations_for_llm(
|
||||
observations: "list[MemoryFact]",
|
||||
source_facts: "dict[str, MemoryFact]",
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Serialize MemoryFact observations into dicts for the consolidation LLM prompt."""
|
||||
obs_list = []
|
||||
for obs in observations:
|
||||
obs_data: dict[str, Any] = {
|
||||
"id": obs.id,
|
||||
"text": obs.text,
|
||||
"proof_count": len(obs.source_fact_ids or []) or 1,
|
||||
"tags": obs.tags or [],
|
||||
}
|
||||
if obs.occurred_start:
|
||||
obs_data["occurred_start"] = obs.occurred_start
|
||||
if obs.occurred_end:
|
||||
obs_data["occurred_end"] = obs.occurred_end
|
||||
if obs.mentioned_at:
|
||||
obs_data["mentioned_at"] = obs.mentioned_at
|
||||
source_memories = [
|
||||
{"text": sf.text, "occurred_start": sf.occurred_start}
|
||||
for sid in (obs.source_fact_ids or [])[:3]
|
||||
if (sf := source_facts.get(sid)) is not None
|
||||
]
|
||||
if source_memories:
|
||||
obs_data["source_memories"] = source_memories
|
||||
obs_list.append(obs_data)
|
||||
return obs_list
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, proof_count, history, tags, source_memory_ids, created_at, updated_at,
|
||||
occurred_start, occurred_end, mentioned_at
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1) AND bank_id = $2 AND fact_type = 'observation'
|
||||
""",
|
||||
observation_ids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
# Build results list preserving recall order
|
||||
id_to_row = {row["id"]: row for row in rows}
|
||||
results = []
|
||||
|
||||
for obs in recall_result.results:
|
||||
obs_id = uuid.UUID(obs.id)
|
||||
if obs_id not in id_to_row:
|
||||
continue
|
||||
|
||||
row = id_to_row[obs_id]
|
||||
history = row["history"]
|
||||
if isinstance(history, str):
|
||||
history = json.loads(history)
|
||||
elif history is None:
|
||||
history = []
|
||||
|
||||
# Fetch source memories to include their text and dates
|
||||
source_memory_ids = row["source_memory_ids"] or []
|
||||
source_memories = []
|
||||
|
||||
if source_memory_ids:
|
||||
source_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT text, occurred_start, occurred_end, mentioned_at, event_date
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1) AND bank_id = $2
|
||||
ORDER BY created_at ASC
|
||||
LIMIT 5
|
||||
""",
|
||||
source_memory_ids[:5], # Limit to first 5 source memories for token efficiency
|
||||
bank_id,
|
||||
)
|
||||
|
||||
for src_row in source_rows:
|
||||
source_memories.append(
|
||||
{
|
||||
"text": src_row["text"],
|
||||
"occurred_start": src_row["occurred_start"],
|
||||
"occurred_end": src_row["occurred_end"],
|
||||
"mentioned_at": src_row["mentioned_at"],
|
||||
"event_date": src_row["event_date"],
|
||||
}
|
||||
)
|
||||
|
||||
results.append(
|
||||
{
|
||||
"id": row["id"],
|
||||
"text": row["text"],
|
||||
"proof_count": row["proof_count"] or 1,
|
||||
"tags": row["tags"] or [],
|
||||
"source_memories": source_memories,
|
||||
"occurred_start": row["occurred_start"],
|
||||
"occurred_end": row["occurred_end"],
|
||||
"mentioned_at": row["mentioned_at"],
|
||||
"created_at": row["created_at"],
|
||||
"updated_at": row["updated_at"],
|
||||
}
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
async def _consolidate_with_llm(
|
||||
memory_engine: "MemoryEngine",
|
||||
fact_text: str,
|
||||
recall_result: "RecallResult",
|
||||
observations: list[dict[str, Any]],
|
||||
mission: str,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
@@ -850,11 +884,40 @@ async def _consolidate_with_llm(
|
||||
- {"action": "create", "text": "...", "reason": "..."}
|
||||
- [] if fact is purely ephemeral (no durable knowledge)
|
||||
"""
|
||||
observations = recall_result.results
|
||||
source_facts = recall_result.source_facts or {}
|
||||
|
||||
# Format observations as JSON with source memories and dates
|
||||
if observations:
|
||||
obs_list = _build_observations_for_llm(observations, source_facts)
|
||||
obs_list = []
|
||||
for obs in observations:
|
||||
obs_data = {
|
||||
"id": str(obs["id"]),
|
||||
"text": obs["text"],
|
||||
"proof_count": obs["proof_count"],
|
||||
"tags": obs["tags"],
|
||||
"created_at": obs["created_at"].isoformat() if obs.get("created_at") else None,
|
||||
"updated_at": obs["updated_at"].isoformat() if obs.get("updated_at") else None,
|
||||
}
|
||||
|
||||
# Include temporal info if available
|
||||
if obs.get("occurred_start"):
|
||||
obs_data["occurred_start"] = obs["occurred_start"].isoformat()
|
||||
if obs.get("occurred_end"):
|
||||
obs_data["occurred_end"] = obs["occurred_end"].isoformat()
|
||||
if obs.get("mentioned_at"):
|
||||
obs_data["mentioned_at"] = obs["mentioned_at"].isoformat()
|
||||
|
||||
# Include source memories (up to 3 for brevity)
|
||||
if obs.get("source_memories"):
|
||||
obs_data["source_memories"] = [
|
||||
{
|
||||
"text": sm["text"],
|
||||
"event_date": sm["event_date"].isoformat() if sm.get("event_date") else None,
|
||||
"occurred_start": sm["occurred_start"].isoformat() if sm.get("occurred_start") else None,
|
||||
}
|
||||
for sm in obs["source_memories"][:3] # Limit to 3 for token efficiency
|
||||
]
|
||||
|
||||
obs_list.append(obs_data)
|
||||
|
||||
observations_text = json.dumps(obs_list, indent=2)
|
||||
else:
|
||||
observations_text = "[]"
|
||||
@@ -879,12 +942,42 @@ Focus on DURABLE knowledge that serves this mission, not ephemeral state.
|
||||
{"role": "user", "content": user_prompt},
|
||||
]
|
||||
|
||||
response: _ConsolidationResponse = await memory_engine._consolidation_llm_config.call(
|
||||
messages=messages,
|
||||
response_format=_ConsolidationResponse,
|
||||
scope="consolidation",
|
||||
)
|
||||
return [a.model_dump() for a in response.actions]
|
||||
try:
|
||||
result = await memory_engine._consolidation_llm_config.call(
|
||||
messages=messages,
|
||||
skip_validation=True, # Raw JSON response
|
||||
scope="consolidation",
|
||||
)
|
||||
# Parse JSON response - should be an array
|
||||
if isinstance(result, str):
|
||||
# Strip markdown code fences (some models wrap JSON in ```json ... ```)
|
||||
clean = result.strip()
|
||||
if clean.startswith("```"):
|
||||
clean = clean.split("\n", 1)[1] if "\n" in clean else clean[3:]
|
||||
if clean.endswith("```"):
|
||||
clean = clean[:-3]
|
||||
clean = clean.strip()
|
||||
result = json.loads(clean)
|
||||
# Ensure result is a list
|
||||
if isinstance(result, list):
|
||||
return result
|
||||
# Handle legacy single-action format for backward compatibility
|
||||
if isinstance(result, dict):
|
||||
if result.get("related_ids") and result.get("consolidated_text"):
|
||||
# Convert old format to new format
|
||||
return [
|
||||
{
|
||||
"action": "update",
|
||||
"learning_id": result["related_ids"][0],
|
||||
"text": result["consolidated_text"],
|
||||
"reason": result.get("reason", ""),
|
||||
}
|
||||
]
|
||||
return []
|
||||
return []
|
||||
except Exception as e:
|
||||
logger.warning(f"Error in consolidation LLM call: {e}")
|
||||
return []
|
||||
|
||||
|
||||
async def _create_observation_directly(
|
||||
@@ -937,9 +1030,8 @@ async def _create_observation_directly(
|
||||
tokenize($3, 'llmlingua2')::bm25_catalog.bm25vector)
|
||||
RETURNING id
|
||||
"""
|
||||
else: # native or pg_textsearch
|
||||
else: # native
|
||||
# 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,
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
CONSOLIDATION_SYSTEM_PROMPT = """You are a memory consolidation system. Your job is to convert facts into durable knowledge (observations) and merge with existing knowledge when appropriate.
|
||||
|
||||
You must output a JSON object with an "actions" array. The "text" field within each action should use markdown formatting (headers, lists, bold, etc.) for clarity and readability.
|
||||
You must output ONLY valid JSON with no markdown code blocks or additional text. However, the "text" field within each observation should use markdown formatting (headers, lists, bold, etc.) for clarity and readability.
|
||||
|
||||
## EXTRACT DURABLE KNOWLEDGE, NOT EPHEMERAL STATE
|
||||
Facts often describe events or actions. Extract the DURABLE KNOWLEDGE implied by the fact, not the transient state.
|
||||
@@ -58,6 +58,7 @@ Each observation includes:
|
||||
- text: the observation content
|
||||
- proof_count: number of supporting memories
|
||||
- tags: visibility scope (handled automatically)
|
||||
- created_at/updated_at: when observation was created/modified
|
||||
- occurred_start/occurred_end: temporal range of source facts
|
||||
- source_memories: array of supporting facts with their text and dates
|
||||
|
||||
@@ -68,15 +69,15 @@ Instructions:
|
||||
4. Compare with observations:
|
||||
- Same topic → UPDATE with learning_id
|
||||
- New topic → CREATE new observation
|
||||
- Purely ephemeral → return empty actions list
|
||||
- Purely ephemeral → return []
|
||||
|
||||
Output a JSON object with an "actions" array (the "text" field should use markdown formatting for structure):
|
||||
{{"actions": [
|
||||
Output JSON array of actions (the "text" field should use markdown formatting for structure):
|
||||
[
|
||||
{{"action": "update", "learning_id": "uuid-from-observations", "text": "## Updated Knowledge\n\n**Key point**: details here\n\n- Supporting detail 1\n- Supporting detail 2", "reason": "..."}},
|
||||
{{"action": "create", "text": "## New Durable Knowledge\n\nDescription with **emphasis** and proper structure", "reason": "..."}}
|
||||
]}}
|
||||
]
|
||||
|
||||
Return {{"actions": []}} if fact contains no durable knowledge.
|
||||
Return [] if fact contains no durable knowledge.
|
||||
|
||||
IMPORTANT: Format the "text" field with markdown for better readability:
|
||||
- Use headers, lists, bold/italic, tables where appropriate
|
||||
|
||||
@@ -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.)."""
|
||||
|
||||
@@ -60,59 +60,6 @@ class OutputTooLongError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def parse_llm_json(raw: str) -> Any:
|
||||
"""
|
||||
Robustly parse JSON returned by an LLM.
|
||||
|
||||
Handles common LLM output quirks:
|
||||
1. Markdown code fences (```json ... ```) — strip them before parsing.
|
||||
2. Embedded control characters (\\x00-\\x1f, \\x7f) — replace with space
|
||||
and retry if the initial parse fails.
|
||||
|
||||
Args:
|
||||
raw: Raw text returned by the LLM.
|
||||
|
||||
Returns:
|
||||
Parsed Python object (dict, list, etc.).
|
||||
|
||||
Raises:
|
||||
json.JSONDecodeError: If the text cannot be parsed even after cleanup.
|
||||
"""
|
||||
text = raw.strip()
|
||||
|
||||
# Strip markdown code fences (some models wrap JSON in ```json ... ```)
|
||||
if text.startswith("```"):
|
||||
text = text.split("\n", 1)[1] if "\n" in text else text[3:]
|
||||
if text.endswith("```"):
|
||||
text = text[:-3]
|
||||
text = text.strip()
|
||||
|
||||
try:
|
||||
return json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
# Some models (e.g. Gemini) embed raw control characters inside JSON
|
||||
# string values. Replacing them with a space usually produces valid JSON.
|
||||
cleaned = re.sub(r"[\x00-\x1f\x7f]", " ", text)
|
||||
return json.loads(cleaned)
|
||||
|
||||
|
||||
_PROVIDERS_WITHOUT_API_KEY = frozenset(
|
||||
{
|
||||
"ollama",
|
||||
"lmstudio",
|
||||
"openai-codex",
|
||||
"claude-code",
|
||||
"mock",
|
||||
"vertexai",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def requires_api_key(provider: str) -> bool:
|
||||
"""Return True if the given provider requires an API key to operate."""
|
||||
return provider.lower() not in _PROVIDERS_WITHOUT_API_KEY
|
||||
|
||||
|
||||
def create_llm_provider(
|
||||
provider: str,
|
||||
api_key: str,
|
||||
@@ -120,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,
|
||||
@@ -134,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).
|
||||
@@ -211,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:
|
||||
@@ -233,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.
|
||||
@@ -244,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 = [
|
||||
@@ -331,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,
|
||||
@@ -605,9 +545,8 @@ class LLMProvider:
|
||||
provider = os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq")
|
||||
api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY", "")
|
||||
|
||||
# API key not needed for openai-codex (uses OAuth), claude-code (uses Keychain OAuth),
|
||||
# ollama (local), or vertexai (uses GCP service account credentials)
|
||||
if not api_key and provider not in ("openai-codex", "claude-code", "ollama", "vertexai"):
|
||||
# API key not needed for openai-codex (uses OAuth) or claude-code (uses Keychain OAuth)
|
||||
if not api_key and provider not in ("openai-codex", "claude-code"):
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_API_KEY environment variable is required (unless using openai-codex or claude-code)"
|
||||
)
|
||||
@@ -623,9 +562,8 @@ class LLMProvider:
|
||||
provider = os.getenv("HINDSIGHT_API_ANSWER_LLM_PROVIDER", os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"))
|
||||
api_key = os.getenv("HINDSIGHT_API_ANSWER_LLM_API_KEY", os.getenv("HINDSIGHT_API_LLM_API_KEY", ""))
|
||||
|
||||
# API key not needed for openai-codex (uses OAuth), claude-code (uses Keychain OAuth),
|
||||
# ollama (local), or vertexai (uses GCP service account credentials)
|
||||
if not api_key and provider not in ("openai-codex", "claude-code", "ollama", "vertexai"):
|
||||
# API key not needed for openai-codex (uses OAuth) or claude-code (uses Keychain OAuth)
|
||||
if not api_key and provider not in ("openai-codex", "claude-code"):
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_API_ANSWER_LLM_API_KEY environment variable is required "
|
||||
"(unless using openai-codex or claude-code)"
|
||||
@@ -642,9 +580,8 @@ class LLMProvider:
|
||||
provider = os.getenv("HINDSIGHT_API_JUDGE_LLM_PROVIDER", os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"))
|
||||
api_key = os.getenv("HINDSIGHT_API_JUDGE_LLM_API_KEY", os.getenv("HINDSIGHT_API_LLM_API_KEY", ""))
|
||||
|
||||
# API key not needed for openai-codex (uses OAuth), claude-code (uses Keychain OAuth),
|
||||
# ollama (local), or vertexai (uses GCP service account credentials)
|
||||
if not api_key and provider not in ("openai-codex", "claude-code", "ollama", "vertexai"):
|
||||
# API key not needed for openai-codex (uses OAuth) or claude-code (uses Keychain OAuth)
|
||||
if not api_key and provider not in ("openai-codex", "claude-code"):
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_API_JUDGE_LLM_API_KEY environment variable is required "
|
||||
"(unless using openai-codex or claude-code)"
|
||||
|
||||
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"
|
||||
@@ -18,7 +18,6 @@ from google.genai import errors as genai_errors
|
||||
from google.genai import types as genai_types
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.llm_wrapper import parse_llm_json
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
|
||||
@@ -222,13 +221,10 @@ class GeminiLLM(LLMInterface):
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=gemini_contents,
|
||||
config=generation_config,
|
||||
),
|
||||
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
|
||||
response = await self._client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=gemini_contents,
|
||||
config=generation_config,
|
||||
)
|
||||
|
||||
content = response.text
|
||||
@@ -251,7 +247,7 @@ class GeminiLLM(LLMInterface):
|
||||
|
||||
# Parse structured output if requested
|
||||
if response_format is not None:
|
||||
json_data = parse_llm_json(content)
|
||||
json_data = json.loads(content)
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
else:
|
||||
@@ -409,57 +405,31 @@ class GeminiLLM(LLMInterface):
|
||||
# Convert messages
|
||||
system_instruction = None
|
||||
gemini_contents = []
|
||||
msg_list = list(messages)
|
||||
i = 0
|
||||
while i < len(msg_list):
|
||||
msg = msg_list[i]
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
system_instruction = (system_instruction + "\n\n" + content) if system_instruction else content
|
||||
i += 1
|
||||
elif role == "tool":
|
||||
# Gemini requires ALL tool responses for a given model turn to be grouped
|
||||
# into a single Content with multiple FunctionResponse parts.
|
||||
# Consecutive role="tool" messages correspond to one model turn's tool calls.
|
||||
parts = []
|
||||
while i < len(msg_list) and msg_list[i].get("role") == "tool":
|
||||
tool_msg = msg_list[i]
|
||||
tool_content = tool_msg.get("content", "")
|
||||
parts.append(
|
||||
genai_types.Part(
|
||||
function_response=genai_types.FunctionResponse(
|
||||
name=tool_msg.get("name", ""),
|
||||
response={"result": tool_content},
|
||||
# Gemini uses function_response
|
||||
gemini_contents.append(
|
||||
genai_types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
genai_types.Part(
|
||||
function_response=genai_types.FunctionResponse(
|
||||
name=msg.get("name", ""),
|
||||
response={"result": content},
|
||||
)
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
i += 1
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=parts))
|
||||
)
|
||||
elif role == "assistant":
|
||||
tool_calls_in_msg = msg.get("tool_calls", [])
|
||||
if tool_calls_in_msg:
|
||||
# Convert OpenAI-style tool_calls to Gemini function_call parts
|
||||
# This is required for proper multi-turn conversation history
|
||||
parts = []
|
||||
if content:
|
||||
parts.append(genai_types.Part(text=content))
|
||||
for tc in tool_calls_in_msg:
|
||||
fn = tc.get("function", {})
|
||||
fn_name = fn.get("name", "")
|
||||
fn_args_str = fn.get("arguments", "{}")
|
||||
fn_args = parse_llm_json(fn_args_str)
|
||||
parts.append(
|
||||
genai_types.Part(function_call=genai_types.FunctionCall(name=fn_name, args=fn_args))
|
||||
)
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=parts))
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=[genai_types.Part(text=content)]))
|
||||
i += 1
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=[genai_types.Part(text=content)]))
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
|
||||
i += 1
|
||||
|
||||
config_kwargs: dict[str, Any] = {"tools": gemini_tools}
|
||||
if system_instruction:
|
||||
@@ -467,40 +437,15 @@ class GeminiLLM(LLMInterface):
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
|
||||
# Map OpenAI-style tool_choice to Gemini FunctionCallingConfig
|
||||
if tool_choice == "required":
|
||||
config_kwargs["tool_config"] = genai_types.ToolConfig(
|
||||
function_calling_config=genai_types.FunctionCallingConfig(
|
||||
mode="ANY",
|
||||
)
|
||||
)
|
||||
elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
|
||||
fn_name = tool_choice.get("function", {}).get("name")
|
||||
if fn_name:
|
||||
config_kwargs["tool_config"] = genai_types.ToolConfig(
|
||||
function_calling_config=genai_types.FunctionCallingConfig(
|
||||
mode="ANY",
|
||||
allowed_function_names=[fn_name],
|
||||
)
|
||||
)
|
||||
elif tool_choice == "none":
|
||||
config_kwargs["tool_config"] = genai_types.ToolConfig(
|
||||
function_calling_config=genai_types.FunctionCallingConfig(mode="NONE")
|
||||
)
|
||||
# "auto" is the default (no tool_config needed)
|
||||
|
||||
config = genai_types.GenerateContentConfig(**config_kwargs)
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=gemini_contents,
|
||||
config=config,
|
||||
),
|
||||
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
|
||||
response = await self._client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=gemini_contents,
|
||||
config=config,
|
||||
)
|
||||
|
||||
# Extract content and tool calls
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -20,18 +20,26 @@ from .tools_schema import get_reflect_tools
|
||||
|
||||
|
||||
def _build_directives_applied(directives: list[dict[str, Any]] | None) -> list[DirectiveInfo]:
|
||||
"""Build list of DirectiveInfo from directives."""
|
||||
"""Build list of DirectiveInfo from directive mental models.
|
||||
|
||||
Handles multiple directive formats:
|
||||
1. New format: directives have direct 'content' field
|
||||
2. Fallback: directives have 'description' field
|
||||
"""
|
||||
if not directives:
|
||||
return []
|
||||
|
||||
return [
|
||||
DirectiveInfo(
|
||||
id=directive.get("id", ""),
|
||||
name=directive.get("name", ""),
|
||||
content=directive.get("content", ""),
|
||||
)
|
||||
for directive in directives
|
||||
]
|
||||
result = []
|
||||
for directive in directives:
|
||||
directive_id = directive.get("id", "")
|
||||
directive_name = directive.get("name", "")
|
||||
|
||||
# Get content from 'content' field or fallback to 'description'
|
||||
content = directive.get("content", "") or directive.get("description", "")
|
||||
|
||||
result.append(DirectiveInfo(id=directive_id, name=directive_name, content=content))
|
||||
|
||||
return result
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -382,7 +390,6 @@ async def run_reflect_agent(
|
||||
f"total={elapsed_ms}ms"
|
||||
)
|
||||
|
||||
consecutive_errors = 0
|
||||
for iteration in range(max_iterations):
|
||||
is_last = iteration == max_iterations - 1
|
||||
|
||||
@@ -436,29 +443,14 @@ async def run_reflect_agent(
|
||||
# Call LLM with tools
|
||||
llm_start = time.time()
|
||||
|
||||
# Determine tool_choice for this iteration.
|
||||
# With mental models:
|
||||
# 0 → search_mental_models, 1+ → auto
|
||||
# Without mental models, enforce a minimum retrieval path:
|
||||
# 0 → search_observations, 1 → recall, 2+ → auto
|
||||
if iteration == 0 and has_mental_models:
|
||||
iter_tool_choice: str | dict = {"type": "function", "function": {"name": "search_mental_models"}}
|
||||
elif iteration == 0:
|
||||
iter_tool_choice = {"type": "function", "function": {"name": "search_observations"}}
|
||||
elif iteration == 1 and not has_mental_models:
|
||||
iter_tool_choice = {"type": "function", "function": {"name": "recall"}}
|
||||
else:
|
||||
iter_tool_choice = "auto"
|
||||
|
||||
try:
|
||||
result = await llm_config.call_with_tools(
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
scope="reflect_tool_call",
|
||||
tool_choice=iter_tool_choice,
|
||||
tool_choice="required" if iteration == 0 else "auto", # Force tool use on first iteration
|
||||
)
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
consecutive_errors = 0
|
||||
total_input_tokens += result.input_tokens
|
||||
total_output_tokens += result.output_tokens
|
||||
llm_trace.append(
|
||||
@@ -472,14 +464,13 @@ async def run_reflect_agent(
|
||||
|
||||
except Exception as e:
|
||||
err_duration = int((time.time() - llm_start) * 1000)
|
||||
consecutive_errors += 1
|
||||
logger.warning(f"[REFLECT {reflect_id}] LLM error on iteration {iteration + 1}: {e} ({err_duration}ms)")
|
||||
llm_trace.append({"scope": f"agent_{iteration + 1}_err", "duration_ms": err_duration})
|
||||
# Guardrail: If no evidence gathered yet, retry (but cap consecutive errors to avoid long hangs)
|
||||
# Guardrail: If no evidence gathered yet, retry
|
||||
has_gathered_evidence = (
|
||||
bool(available_memory_ids) or bool(available_mental_model_ids) or bool(available_observation_ids)
|
||||
)
|
||||
if not has_gathered_evidence and iteration < max_iterations - 1 and consecutive_errors < 2:
|
||||
if not has_gathered_evidence and iteration < max_iterations - 1:
|
||||
continue
|
||||
prompt = build_final_prompt(query, context_history, bank_profile, context)
|
||||
llm_start = time.time()
|
||||
|
||||
@@ -12,20 +12,57 @@ from typing import Any
|
||||
|
||||
|
||||
def _extract_directive_rules(directives: list[dict[str, Any]]) -> list[str]:
|
||||
"""Extract directive rules as a list of strings."""
|
||||
"""
|
||||
Extract directive rules as a list of strings.
|
||||
|
||||
Args:
|
||||
directives: List of directives with name and content
|
||||
|
||||
Returns:
|
||||
List of directive rule strings
|
||||
"""
|
||||
rules = []
|
||||
for directive in directives:
|
||||
name = directive.get("name", "")
|
||||
directive_name = directive.get("name", "")
|
||||
# New format: directives have direct content field
|
||||
content = directive.get("content", "")
|
||||
if content:
|
||||
rules.append(f"**{name}**: {content}" if name else content)
|
||||
if directive_name:
|
||||
rules.append(f"**{directive_name}**: {content}")
|
||||
else:
|
||||
rules.append(content)
|
||||
else:
|
||||
# Legacy format: check for observations
|
||||
observations = directive.get("observations", [])
|
||||
if observations:
|
||||
for obs in observations:
|
||||
# Support both Pydantic Observation objects and dicts
|
||||
if hasattr(obs, "title"):
|
||||
title = obs.title
|
||||
obs_content = obs.content
|
||||
else:
|
||||
title = obs.get("title", "")
|
||||
obs_content = obs.get("content", "")
|
||||
if title and obs_content:
|
||||
rules.append(f"**{title}**: {obs_content}")
|
||||
elif obs_content:
|
||||
rules.append(obs_content)
|
||||
elif directive_name:
|
||||
# Fallback to description
|
||||
desc = directive.get("description", "")
|
||||
if desc:
|
||||
rules.append(f"**{directive_name}**: {desc}")
|
||||
return rules
|
||||
|
||||
|
||||
def build_directives_section(directives: list[dict[str, Any]]) -> str:
|
||||
"""Build the directives section for the system prompt.
|
||||
"""
|
||||
Build the directives section for the system prompt.
|
||||
|
||||
Directives are hard rules that MUST be followed in all responses.
|
||||
|
||||
Args:
|
||||
directives: List of directive mental models with observations
|
||||
"""
|
||||
if not directives:
|
||||
return ""
|
||||
@@ -132,12 +169,6 @@ def build_system_prompt_for_tools(
|
||||
|
||||
parts.extend(
|
||||
[
|
||||
"## LANGUAGE RULE (default - directives take precedence)",
|
||||
"- By default, detect the language of the user's question and respond in that SAME language.",
|
||||
"- If the question is in Chinese, respond in Chinese. If in Japanese, respond in Japanese.",
|
||||
"- IMPORTANT: The DIRECTIVES section above has HIGHER PRIORITY than this rule.",
|
||||
" If a directive specifies a language (e.g. 'Always respond in French'), follow the directive.",
|
||||
"",
|
||||
"## CRITICAL RULES",
|
||||
"- ONLY use information from tool results - no external knowledge or guessing",
|
||||
"- You SHOULD synthesize, infer, and reason from the retrieved memories",
|
||||
@@ -174,7 +205,6 @@ def build_system_prompt_for_tools(
|
||||
"### 3. RAW FACTS (recall) - Ground Truth",
|
||||
"- Individual memories (world facts and experiences)",
|
||||
"- Use when: no mental models/observations exist, they're stale, or you need specific details",
|
||||
"- MANDATORY: If search_mental_models and search_observations both return 0 results, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that other levels are built from",
|
||||
"",
|
||||
]
|
||||
@@ -192,7 +222,6 @@ def build_system_prompt_for_tools(
|
||||
"### 2. RAW FACTS (recall) - Ground Truth",
|
||||
"- Individual memories (world facts and experiences)",
|
||||
"- Use when: no observations exist, they're stale, or you need specific details",
|
||||
"- MANDATORY: If search_observations returns 0 results or count=0, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that observations are built from",
|
||||
"",
|
||||
]
|
||||
@@ -270,7 +299,7 @@ def build_system_prompt_for_tools(
|
||||
parts.extend(
|
||||
[
|
||||
"1. First, try search_observations() - check for consolidated knowledge",
|
||||
"2. If search_observations returns 0 results OR observations are stale, you MUST call recall() for raw facts",
|
||||
"2. If observations are stale OR you need specific details, use recall() for raw facts",
|
||||
"3. Use expand() if you need more context on specific memories",
|
||||
"4. When ready, call done() with your answer and supporting IDs",
|
||||
]
|
||||
|
||||
@@ -9,7 +9,7 @@ Implements hierarchical retrieval:
|
||||
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -20,6 +20,9 @@ if TYPE_CHECKING:
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Observation is considered stale if not updated in this many days
|
||||
STALE_THRESHOLD_DAYS = 7
|
||||
|
||||
|
||||
async def tool_search_mental_models(
|
||||
conn: "Connection",
|
||||
@@ -30,7 +33,6 @@ async def tool_search_mental_models(
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
exclude_ids: list[str] | None = None,
|
||||
pending_consolidation: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Search user-curated mental models by semantic similarity.
|
||||
@@ -85,6 +87,7 @@ async def tool_search_mental_models(
|
||||
*params,
|
||||
)
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
mental_models = []
|
||||
|
||||
for row in rows:
|
||||
@@ -92,10 +95,11 @@ async def tool_search_mental_models(
|
||||
if last_refreshed_at and last_refreshed_at.tzinfo is None:
|
||||
last_refreshed_at = last_refreshed_at.replace(tzinfo=timezone.utc)
|
||||
|
||||
# A mental model is stale when there are memories that haven't been consolidated yet —
|
||||
# the same signal used for observations staleness.
|
||||
is_stale = pending_consolidation > 0
|
||||
staleness_reason = f"{pending_consolidation} memories pending consolidation" if is_stale else None
|
||||
# Calculate freshness
|
||||
is_stale = False
|
||||
if last_refreshed_at:
|
||||
age = now - last_refreshed_at
|
||||
is_stale = age > timedelta(days=STALE_THRESHOLD_DAYS)
|
||||
|
||||
mental_models.append(
|
||||
{
|
||||
@@ -106,7 +110,6 @@ async def tool_search_mental_models(
|
||||
"relevance": round(row["relevance"], 4),
|
||||
"updated_at": last_refreshed_at.isoformat() if last_refreshed_at else None,
|
||||
"is_stale": is_stale,
|
||||
"staleness_reason": staleness_reason,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -47,8 +47,7 @@ TOOL_SEARCH_OBSERVATIONS = {
|
||||
"description": (
|
||||
"Search consolidated observations (auto-generated knowledge). These are automatically "
|
||||
"synthesized from memories. Returns observations with freshness info (updated_at, is_stale). "
|
||||
"If an observation is STALE, you should ALSO use recall() to verify with current facts. "
|
||||
"IMPORTANT: If search_mental_models is available, you MUST call it FIRST before using this tool."
|
||||
"If an observation is STALE, you should ALSO use recall() to verify with current facts."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
@@ -140,7 +139,7 @@ TOOL_DONE_ANSWER = {
|
||||
"properties": {
|
||||
"answer": {
|
||||
"type": "string",
|
||||
"description": "Your response as well-formatted markdown. Use headers, lists, bold/italic, and code blocks for clarity. NEVER include memory IDs, UUIDs, or 'Memory references' in this text - put IDs only in memory_ids array. LANGUAGE: By default, write in the SAME language as the user's question. However, if a language directive in the system prompt specifies a different language, follow that directive instead.",
|
||||
"description": "Your response as well-formatted markdown. Use headers, lists, bold/italic, and code blocks for clarity. NEVER include memory IDs, UUIDs, or 'Memory references' in this text - put IDs only in memory_ids array.",
|
||||
},
|
||||
"memory_ids": {
|
||||
"type": "array",
|
||||
@@ -191,11 +190,7 @@ def _build_done_tool_with_directives(directive_rules: list[str]) -> dict:
|
||||
"properties": {
|
||||
"answer": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"Your response as well-formatted markdown. Use headers, lists, bold/italic, and code blocks for clarity. "
|
||||
"NEVER include memory IDs, UUIDs, or 'Memory references' in this text - put IDs only in memory_ids array. "
|
||||
f"MANDATORY: Your answer MUST comply with ALL directives:\n{rules_list}"
|
||||
),
|
||||
"description": "Your response as well-formatted markdown. Use headers, lists, bold/italic, and code blocks for clarity. NEVER include memory IDs, UUIDs, or 'Memory references' in this text - put IDs only in memory_ids array.",
|
||||
},
|
||||
"memory_ids": {
|
||||
"type": "array",
|
||||
|
||||
@@ -159,10 +159,6 @@ class MemoryFact(BaseModel):
|
||||
None, description="ID of the chunk this fact was extracted from (format: bank_id_document_id_chunk_index)"
|
||||
)
|
||||
tags: list[str] | None = Field(None, description="Visibility scope tags associated with this fact")
|
||||
source_fact_ids: list[str] | None = Field(
|
||||
None,
|
||||
description="IDs of source facts this observation was derived from (observation type only, when source_facts is enabled)",
|
||||
)
|
||||
|
||||
|
||||
class ChunkInfo(BaseModel):
|
||||
@@ -230,9 +226,6 @@ class RecallResult(BaseModel):
|
||||
chunks: dict[str, ChunkInfo] | None = Field(
|
||||
None, description="Chunks for facts, keyed by '{document_id}_{chunk_index}'"
|
||||
)
|
||||
source_facts: dict[str, MemoryFact] | None = Field(
|
||||
None, description="Source facts for observation-type results, keyed by fact ID"
|
||||
)
|
||||
|
||||
|
||||
class ReflectResult(BaseModel):
|
||||
|
||||
@@ -26,6 +26,8 @@ def _infer_temporal_date(fact_text: str, event_date: datetime) -> str | None:
|
||||
This is a fallback for when the LLM fails to extract temporal information
|
||||
from relative time expressions like "last night", "yesterday", etc.
|
||||
"""
|
||||
import re
|
||||
|
||||
fact_lower = fact_text.lower()
|
||||
|
||||
# Map relative time expressions to day offsets
|
||||
@@ -438,7 +440,7 @@ def _chunk_conversation(turns: list[dict], max_chars: int) -> list[str]:
|
||||
# Uses {extraction_guidelines} placeholder for mode-specific instructions
|
||||
_BASE_FACT_EXTRACTION_PROMPT = """Extract SIGNIFICANT facts from text. Be SELECTIVE - only extract facts worth remembering long-term.
|
||||
|
||||
LANGUAGE: MANDATORY — Detect the language of the input text and produce ALL output in that EXACT same language. You are STRICTLY FORBIDDEN from translating or switching to any other language. Every single word of your output must be in the same language as the input. Do NOT output in a different language under any circumstance.
|
||||
LANGUAGE REQUIREMENT: Detect the language of the input text. All extracted facts, entity names, descriptions, and other output MUST be in the SAME language as the input. Do not translate to another language.
|
||||
|
||||
{fact_types_instruction}
|
||||
|
||||
@@ -481,9 +483,7 @@ TEMPORAL HANDLING
|
||||
══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
Use "Event Date" from input as reference for relative dates.
|
||||
- CRITICAL: Convert ALL relative temporal expressions to absolute dates in the fact text itself.
|
||||
"yesterday" → write the resolved date (e.g. "on November 12, 2024"), NOT the word "yesterday"
|
||||
"last night", "this morning", "today", "tonight" → convert to the resolved absolute date
|
||||
- "yesterday" relative to Event Date, not today
|
||||
- For events: set occurred_start AND occurred_end (same for point events)
|
||||
- For conversation facts: NO occurred dates
|
||||
|
||||
@@ -521,7 +521,7 @@ CONSOLIDATE related statements into ONE fact when possible."""
|
||||
_CONCISE_EXAMPLES = """
|
||||
|
||||
══════════════════════════════════════════════════════════════════════════
|
||||
EXAMPLES (shown in English for illustration; for non-English input, ALL output values MUST be in the input language)
|
||||
EXAMPLES
|
||||
══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
Example 1 - Selective extraction (Event Date: June 10, 2024):
|
||||
@@ -567,7 +567,8 @@ CUSTOM_FACT_EXTRACTION_PROMPT = _BASE_FACT_EXTRACTION_PROMPT.format(
|
||||
# Verbose extraction prompt - detailed, comprehensive facts (legacy mode)
|
||||
VERBOSE_FACT_EXTRACTION_PROMPT = """Extract facts from text into structured format with FIVE required dimensions - BE EXTREMELY DETAILED.
|
||||
|
||||
LANGUAGE: MANDATORY — Detect the language of the input text and produce ALL output in that EXACT same language. You are STRICTLY FORBIDDEN from translating or switching to any other language. Every single word of your output must be in the same language as the input. Do NOT output in a different language under any circumstance.
|
||||
LANGUAGE REQUIREMENT: Detect the language of the input text. All extracted facts, entity names, descriptions,
|
||||
and other output MUST be in the SAME language as the input. Do not translate to English if the input is in another language.
|
||||
|
||||
{fact_types_instruction}
|
||||
|
||||
@@ -694,91 +695,6 @@ 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]:
|
||||
"""
|
||||
Build extraction prompt and response schema based on config.
|
||||
|
||||
Returns:
|
||||
Tuple of (prompt, response_schema)
|
||||
"""
|
||||
fact_types_instruction = "Extract ONLY 'world' and 'assistant' type facts."
|
||||
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":
|
||||
if not config.retain_custom_instructions:
|
||||
base_prompt = CONCISE_FACT_EXTRACTION_PROMPT
|
||||
prompt = base_prompt.format(fact_types_instruction=fact_types_instruction)
|
||||
else:
|
||||
base_prompt = CUSTOM_FACT_EXTRACTION_PROMPT
|
||||
prompt = base_prompt.format(
|
||||
fact_types_instruction=fact_types_instruction,
|
||||
custom_instructions=config.retain_custom_instructions,
|
||||
)
|
||||
elif extraction_mode == "verbose":
|
||||
base_prompt = VERBOSE_FACT_EXTRACTION_PROMPT
|
||||
prompt = base_prompt.format(fact_types_instruction=fact_types_instruction)
|
||||
else:
|
||||
base_prompt = CONCISE_FACT_EXTRACTION_PROMPT
|
||||
prompt = base_prompt.format(fact_types_instruction=fact_types_instruction)
|
||||
|
||||
# Add causal relationships section if enabled
|
||||
if extract_causal_links:
|
||||
prompt = prompt + CAUSAL_RELATIONSHIPS_SECTION
|
||||
response_schema = FactExtractionResponseVerbose if extraction_mode == "verbose" else 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
|
||||
|
||||
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.
|
||||
|
||||
Chunk: {chunk_index + 1}/{total_chunks}
|
||||
Event Date: {event_date_formatted} ({event_date.isoformat()})
|
||||
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,
|
||||
@@ -801,20 +717,72 @@ async def _extract_facts_from_chunk(
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Build prompt and schema using helper function
|
||||
prompt, response_schema = _build_extraction_prompt_and_schema(config)
|
||||
# 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
|
||||
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)
|
||||
# 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:
|
||||
base_prompt = CUSTOM_FACT_EXTRACTION_PROMPT
|
||||
prompt = base_prompt.format(
|
||||
fact_types_instruction=fact_types_instruction,
|
||||
custom_instructions=config.retain_custom_instructions,
|
||||
)
|
||||
elif extraction_mode == "verbose":
|
||||
base_prompt = VERBOSE_FACT_EXTRACTION_PROMPT
|
||||
prompt = base_prompt.format(fact_types_instruction=fact_types_instruction)
|
||||
else:
|
||||
base_prompt = CONCISE_FACT_EXTRACTION_PROMPT
|
||||
prompt = base_prompt.format(fact_types_instruction=fact_types_instruction)
|
||||
|
||||
# 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
|
||||
if extraction_mode == "verbose":
|
||||
response_schema = FactExtractionResponseVerbose
|
||||
else:
|
||||
response_schema = FactExtractionResponse
|
||||
else:
|
||||
response_schema = FactExtractionResponseNoCausal
|
||||
|
||||
# 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"
|
||||
|
||||
# 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()})
|
||||
Context: {sanitized_context}
|
||||
|
||||
Text:
|
||||
{sanitized_chunk}"""
|
||||
|
||||
usage = TokenUsage() # Track cumulative usage across retries
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
@@ -1273,424 +1241,12 @@ from .types import ExtractedFact as ExtractedFactType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Each fact gets 10ms offset to preserve ordering within a document
|
||||
SECONDS_PER_FACT = 0.01
|
||||
|
||||
|
||||
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
|
||||
# Each fact gets 10 seconds offset to preserve ordering within a document
|
||||
SECONDS_PER_FACT = 10
|
||||
|
||||
|
||||
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, config
|
||||
) -> tuple[list[ExtractedFactType], list[ChunkMetadata], TokenUsage]:
|
||||
"""
|
||||
Extract facts from multiple content items in parallel.
|
||||
@@ -1701,16 +1257,11 @@ 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)
|
||||
@@ -1718,12 +1269,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:
|
||||
|
||||
@@ -97,9 +97,8 @@ async def insert_facts_batch(
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
else: # native or pg_textsearch
|
||||
else: # native
|
||||
# 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(
|
||||
|
||||
@@ -82,8 +82,6 @@ async def retain_batch(
|
||||
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.
|
||||
@@ -149,29 +147,20 @@ async def retain_batch(
|
||||
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
|
||||
contents, llm_config, agent_name, config
|
||||
)
|
||||
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"
|
||||
)
|
||||
|
||||
if not extracted_facts:
|
||||
# Still need to create document if document_id was provided or chunks exist
|
||||
from collections import defaultdict
|
||||
|
||||
docs_tracked = 0
|
||||
# Still need to create document if document_id was provided
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
await fact_storage.ensure_bank_exists(conn, bank_id)
|
||||
|
||||
# Group contents by document_id (consistent with normal path)
|
||||
contents_by_doc_early = defaultdict(list)
|
||||
for idx, content_dict in enumerate(contents_dicts):
|
||||
doc_id = content_dict.get("document_id")
|
||||
contents_by_doc_early[doc_id].append((idx, content_dict))
|
||||
|
||||
# Handle document tracking even with no facts
|
||||
if document_id:
|
||||
# Legacy: single document_id parameter
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
# Collect tags from all content items and merge with document_tags
|
||||
all_tags = set(document_tags or [])
|
||||
@@ -196,57 +185,45 @@ async def retain_batch(
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags
|
||||
)
|
||||
docs_tracked += 1
|
||||
else:
|
||||
# Handle per-item document_ids and/or chunks (mirrors normal path logic)
|
||||
has_any_doc_ids = any(item.get("document_id") for item in contents_dicts)
|
||||
# Check for per-item document_ids
|
||||
from collections import defaultdict
|
||||
|
||||
if has_any_doc_ids or chunks:
|
||||
for original_doc_id, doc_contents in contents_by_doc_early.items():
|
||||
should_create_doc = (original_doc_id is not None) or chunks
|
||||
if not should_create_doc:
|
||||
continue
|
||||
contents_by_doc = defaultdict(list)
|
||||
for idx, content_dict in enumerate(contents_dicts):
|
||||
doc_id = content_dict.get("document_id")
|
||||
if doc_id:
|
||||
contents_by_doc[doc_id].append((idx, content_dict))
|
||||
|
||||
actual_doc_id = original_doc_id
|
||||
if actual_doc_id is None:
|
||||
# No document_id but have chunks - generate one
|
||||
actual_doc_id = str(uuid.uuid4())
|
||||
for doc_id, doc_contents in contents_by_doc.items():
|
||||
combined_content = "\n".join([c.get("content", "") for _, c in doc_contents])
|
||||
# Collect tags from all content items for this document and merge with document_tags
|
||||
all_tags = set(document_tags or [])
|
||||
for _, item in doc_contents:
|
||||
item_tags = item.get("tags", []) or []
|
||||
all_tags.update(item_tags)
|
||||
merged_tags = list(all_tags)
|
||||
|
||||
combined_content = "\n".join([c.get("content", "") for _, c in doc_contents])
|
||||
all_tags = set(document_tags or [])
|
||||
for _, item in doc_contents:
|
||||
item_tags = item.get("tags", []) or []
|
||||
all_tags.update(item_tags)
|
||||
merged_tags = list(all_tags)
|
||||
|
||||
retain_params = {}
|
||||
if doc_contents:
|
||||
first_item = doc_contents[0][1]
|
||||
if first_item.get("context"):
|
||||
retain_params["context"] = first_item["context"]
|
||||
if first_item.get("event_date"):
|
||||
retain_params["event_date"] = (
|
||||
first_item["event_date"].isoformat()
|
||||
if hasattr(first_item["event_date"], "isoformat")
|
||||
else str(first_item["event_date"])
|
||||
)
|
||||
if first_item.get("metadata"):
|
||||
retain_params["metadata"] = first_item["metadata"]
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn,
|
||||
bank_id,
|
||||
actual_doc_id,
|
||||
combined_content,
|
||||
is_first_batch,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
docs_tracked += 1
|
||||
retain_params = {}
|
||||
if doc_contents:
|
||||
first_item = doc_contents[0][1]
|
||||
if first_item.get("context"):
|
||||
retain_params["context"] = first_item["context"]
|
||||
if first_item.get("event_date"):
|
||||
retain_params["event_date"] = (
|
||||
first_item["event_date"].isoformat()
|
||||
if hasattr(first_item["event_date"], "isoformat")
|
||||
else str(first_item["event_date"])
|
||||
)
|
||||
if first_item.get("metadata"):
|
||||
retain_params["metadata"] = first_item["metadata"]
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn, bank_id, doc_id, combined_content, is_first_batch, retain_params, merged_tags
|
||||
)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
doc_status = f"{docs_tracked} document(s) tracked" if docs_tracked > 0 else "no document tracked"
|
||||
logger.info(
|
||||
f"RETAIN_BATCH COMPLETE: 0 facts extracted from {len(contents)} contents in {total_time:.3f}s ({doc_status}, no facts)"
|
||||
f"RETAIN_BATCH COMPLETE: 0 facts extracted from {len(contents)} contents in {total_time:.3f}s (document tracked, no facts)"
|
||||
)
|
||||
return [[] for _ in contents], usage
|
||||
|
||||
|
||||
@@ -164,74 +164,99 @@ async def retrieve_semantic_bm25_combined(
|
||||
# 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]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
|
||||
query = 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,
|
||||
NULL::float AS bm25_score,
|
||||
'semantic' AS source,
|
||||
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY embedding <=> $1::vector) AS rn
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND fact_type = ANY($3)
|
||||
AND (1 - (embedding <=> $1::vector)) >= 0.3
|
||||
{tags_clause}
|
||||
),
|
||||
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,
|
||||
search_vector <&> to_bm25query('idx_memory_units_text_search', tokenize($5, 'llmlingua2')) AS bm25_score,
|
||||
'bm25' AS source,
|
||||
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY search_vector <&> to_bm25query('idx_memory_units_text_search', tokenize($5, 'llmlingua2')) DESC) AS rn
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = ANY($3)
|
||||
{tags_clause}
|
||||
),
|
||||
semantic AS (
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
|
||||
similarity, bm25_score, source
|
||||
FROM semantic_ranked WHERE rn <= $4
|
||||
),
|
||||
bm25 AS (
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
|
||||
similarity, bm25_score, source
|
||||
FROM bm25_ranked WHERE rn <= $4
|
||||
)
|
||||
SELECT * FROM semantic
|
||||
UNION ALL
|
||||
SELECT * FROM bm25
|
||||
"""
|
||||
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]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
|
||||
if tags:
|
||||
params.append(tags)
|
||||
|
||||
# Single query template with backend-specific parts injected
|
||||
query = 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,
|
||||
NULL::float AS bm25_score,
|
||||
'semantic' AS source,
|
||||
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY embedding <=> $1::vector) AS rn
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND fact_type = ANY($3)
|
||||
AND (1 - (embedding <=> $1::vector)) >= 0.3
|
||||
{tags_clause}
|
||||
),
|
||||
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,
|
||||
'bm25' AS source,
|
||||
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY {bm25_order_by}) AS rn
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = ANY($3)
|
||||
{bm25_where_filter}
|
||||
{tags_clause}
|
||||
),
|
||||
semantic AS (
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
|
||||
similarity, bm25_score, source
|
||||
FROM semantic_ranked WHERE rn <= $4
|
||||
),
|
||||
bm25 AS (
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
|
||||
similarity, bm25_score, source
|
||||
FROM bm25_ranked WHERE rn <= $4
|
||||
)
|
||||
SELECT * FROM semantic
|
||||
UNION ALL
|
||||
SELECT * FROM bm25
|
||||
"""
|
||||
query = 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,
|
||||
NULL::float AS bm25_score,
|
||||
'semantic' AS source,
|
||||
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY embedding <=> $1::vector) AS rn
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND fact_type = ANY($3)
|
||||
AND (1 - (embedding <=> $1::vector)) >= 0.3
|
||||
{tags_clause}
|
||||
),
|
||||
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,
|
||||
ts_rank_cd(search_vector, to_tsquery('english', $5)) AS bm25_score,
|
||||
'bm25' AS source,
|
||||
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)
|
||||
AND search_vector @@ to_tsquery('english', $5)
|
||||
{tags_clause}
|
||||
),
|
||||
semantic AS (
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
|
||||
similarity, bm25_score, source
|
||||
FROM semantic_ranked WHERE rn <= $4
|
||||
),
|
||||
bm25 AS (
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
|
||||
similarity, bm25_score, source
|
||||
FROM bm25_ranked WHERE rn <= $4
|
||||
)
|
||||
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
|
||||
|
||||
@@ -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))
|
||||
@@ -166,8 +166,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 +208,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,9 +223,6 @@ 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,
|
||||
@@ -247,27 +239,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 +340,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 +368,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
|
||||
|
||||
@@ -1,14 +1,8 @@
|
||||
"""
|
||||
Local MCP server entry point for use with Claude Code (HTTP transport).
|
||||
Local MCP server for use with Claude Code (stdio transport).
|
||||
|
||||
This is a thin wrapper around the main hindsight-api server that pre-configures
|
||||
sensible defaults for local use (embedded PostgreSQL via pg0, warning log level).
|
||||
|
||||
The full API runs on localhost:8888. Configure Claude Code's MCP settings:
|
||||
claude mcp add --transport http hindsight http://localhost:8888/mcp/
|
||||
|
||||
Or pinned to a specific bank (single-bank mode):
|
||||
claude mcp add --transport http hindsight http://localhost:8888/mcp/default/
|
||||
This runs a fully local Hindsight instance with embedded PostgreSQL (pg0).
|
||||
No external database or server required.
|
||||
|
||||
Run with:
|
||||
hindsight-local-mcp
|
||||
@@ -16,24 +10,148 @@ Run with:
|
||||
Or with uvx:
|
||||
uvx hindsight-api@latest hindsight-local-mcp
|
||||
|
||||
Configure in Claude Code's MCP settings:
|
||||
{
|
||||
"mcpServers": {
|
||||
"hindsight": {
|
||||
"command": "uvx",
|
||||
"args": ["hindsight-api@latest", "hindsight-local-mcp"],
|
||||
"env": {
|
||||
"HINDSIGHT_API_LLM_API_KEY": "your-openai-key"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Environment variables:
|
||||
HINDSIGHT_API_LLM_API_KEY: Required. API key for LLM provider.
|
||||
HINDSIGHT_API_LLM_PROVIDER: Optional. LLM provider (default: "openai").
|
||||
HINDSIGHT_API_LLM_MODEL: Optional. LLM model (default: "gpt-4o-mini").
|
||||
HINDSIGHT_API_DATABASE_URL: Optional. Override database URL (default: pg0://hindsight-mcp).
|
||||
HINDSIGHT_API_MCP_LOCAL_BANK_ID: Optional. Memory bank ID (default: "mcp").
|
||||
HINDSIGHT_API_LOG_LEVEL: Optional. Log level (default: "warning").
|
||||
HINDSIGHT_API_MCP_INSTRUCTIONS: Optional. Additional instructions appended to both retain and recall tools.
|
||||
|
||||
Example custom instructions (these are ADDED to the default behavior):
|
||||
To also store assistant actions:
|
||||
HINDSIGHT_API_MCP_INSTRUCTIONS="Also store every action you take, including tool calls, code written, and decisions made."
|
||||
|
||||
To also store conversation summaries:
|
||||
HINDSIGHT_API_MCP_INSTRUCTIONS="Also store summaries of important conversations and their outcomes."
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
from hindsight_api.config import (
|
||||
DEFAULT_MCP_LOCAL_BANK_ID,
|
||||
DEFAULT_MCP_RECALL_DESCRIPTION,
|
||||
DEFAULT_MCP_RETAIN_DESCRIPTION,
|
||||
ENV_MCP_INSTRUCTIONS,
|
||||
ENV_MCP_LOCAL_BANK_ID,
|
||||
)
|
||||
from hindsight_api.mcp_tools import MCPToolsConfig, register_mcp_tools
|
||||
|
||||
# Configure logging - default to warning to avoid polluting stderr during MCP init
|
||||
# MCP clients interpret stderr output as errors, so we suppress INFO logs by default
|
||||
_log_level_str = os.environ.get("HINDSIGHT_API_LOG_LEVEL", "warning").lower()
|
||||
_log_level_map = {
|
||||
"critical": logging.CRITICAL,
|
||||
"error": logging.ERROR,
|
||||
"warning": logging.WARNING,
|
||||
"info": logging.INFO,
|
||||
"debug": logging.DEBUG,
|
||||
}
|
||||
logging.basicConfig(
|
||||
level=_log_level_map.get(_log_level_str, logging.WARNING),
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
stream=sys.stderr, # MCP uses stdout for protocol, logs go to stderr
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Start the Hindsight API server with local defaults."""
|
||||
# Set local defaults (only if not already configured by the user)
|
||||
os.environ.setdefault("HINDSIGHT_API_DATABASE_URL", "pg0://hindsight-mcp")
|
||||
def create_local_mcp_server(bank_id: str, memory=None) -> FastMCP:
|
||||
"""
|
||||
Create a stdio MCP server with retain/recall tools.
|
||||
|
||||
from hindsight_api.main import main as api_main
|
||||
Args:
|
||||
bank_id: The memory bank ID to use for all operations.
|
||||
memory: Optional MemoryEngine instance. If not provided, creates one with pg0.
|
||||
|
||||
api_main()
|
||||
Returns:
|
||||
Configured FastMCP server instance.
|
||||
"""
|
||||
# Import here to avoid slow startup if just checking --help
|
||||
from hindsight_api import MemoryEngine
|
||||
|
||||
# Create memory engine with pg0 embedded database if not provided
|
||||
if memory is None:
|
||||
memory = MemoryEngine(db_url="pg0://hindsight-mcp")
|
||||
|
||||
# Get custom instructions from environment variable (appended to both tools)
|
||||
extra_instructions = os.environ.get(ENV_MCP_INSTRUCTIONS, "")
|
||||
|
||||
retain_description = DEFAULT_MCP_RETAIN_DESCRIPTION
|
||||
recall_description = DEFAULT_MCP_RECALL_DESCRIPTION
|
||||
|
||||
if extra_instructions:
|
||||
retain_description = f"{DEFAULT_MCP_RETAIN_DESCRIPTION}\n\nAdditional instructions: {extra_instructions}"
|
||||
recall_description = f"{DEFAULT_MCP_RECALL_DESCRIPTION}\n\nAdditional instructions: {extra_instructions}"
|
||||
|
||||
mcp = FastMCP("hindsight")
|
||||
|
||||
# Configure and register tools using shared module
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=lambda: bank_id,
|
||||
include_bank_id_param=False, # Local MCP uses fixed bank_id
|
||||
tools={"retain", "recall"}, # Local MCP only has retain and recall
|
||||
retain_description=retain_description,
|
||||
recall_description=recall_description,
|
||||
retain_fire_and_forget=True, # Local MCP uses fire-and-forget pattern
|
||||
)
|
||||
|
||||
register_mcp_tools(mcp, memory, config)
|
||||
|
||||
return mcp
|
||||
|
||||
|
||||
async def _initialize_and_run(bank_id: str):
|
||||
"""Initialize memory and run the MCP server."""
|
||||
from hindsight_api import MemoryEngine
|
||||
|
||||
# Create and initialize memory engine with pg0 embedded database
|
||||
# Note: We avoid printing to stderr during init as MCP clients show it as "errors"
|
||||
memory = MemoryEngine(db_url="pg0://hindsight-mcp")
|
||||
await memory.initialize()
|
||||
|
||||
# Create and run the server
|
||||
mcp = create_local_mcp_server(bank_id, memory=memory)
|
||||
await mcp.run_stdio_async()
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point for the stdio MCP server."""
|
||||
import asyncio
|
||||
|
||||
from hindsight_api.config import ENV_LLM_API_KEY, get_config
|
||||
|
||||
# Check for required environment variables
|
||||
config = get_config()
|
||||
if not config.llm_api_key:
|
||||
print(f"Error: {ENV_LLM_API_KEY} environment variable is required", file=sys.stderr)
|
||||
print("Set it in your MCP configuration or shell environment", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# Get bank ID from environment, default to "mcp"
|
||||
bank_id = os.environ.get(ENV_MCP_LOCAL_BANK_ID, DEFAULT_MCP_LOCAL_BANK_ID)
|
||||
|
||||
# Note: We don't print to stderr as MCP clients display it as "error output"
|
||||
# Use HINDSIGHT_API_LOG_LEVEL=debug for verbose startup logging
|
||||
|
||||
# Run the async initialization and server
|
||||
asyncio.run(_initialize_and_run(bank_id))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -35,46 +35,20 @@ MIGRATION_LOCK_ID = 123456789
|
||||
|
||||
def _detect_vector_extension(conn, vector_extension: str = "pgvector") -> str:
|
||||
"""
|
||||
Validate vector extension: 'pgvector', 'vchord', or 'pgvectorscale'.
|
||||
Validate vector extension: 'vchord' or 'pgvector'.
|
||||
|
||||
Args:
|
||||
conn: SQLAlchemy connection object
|
||||
vector_extension: Configured extension ("pgvector", "vchord", or "pgvectorscale")
|
||||
vector_extension: Configured extension ("pgvector" or "vchord")
|
||||
|
||||
Returns:
|
||||
"pgvector", "vchord", "pgvectorscale", or "pg_diskann"
|
||||
"vchord" or "pgvector"
|
||||
|
||||
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":
|
||||
if vector_extension == "vchord":
|
||||
vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar()
|
||||
if not vchord_check:
|
||||
raise RuntimeError(
|
||||
@@ -91,9 +65,7 @@ def _detect_vector_extension(conn, vector_extension: str = "pgvector") -> str:
|
||||
logger.debug("Using configured vector extension: pgvector")
|
||||
return "pgvector"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid vector_extension: {vector_extension}. Must be 'pgvector', 'vchord', or 'pgvectorscale'"
|
||||
)
|
||||
raise ValueError(f"Invalid vector_extension: {vector_extension}. Must be 'pgvector' or 'vchord'")
|
||||
|
||||
|
||||
def _get_schema_lock_id(schema: str) -> int:
|
||||
@@ -305,48 +277,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:
|
||||
@@ -545,17 +475,7 @@ 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":
|
||||
if vector_ext == "vchord":
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_embedding_vchordrq
|
||||
@@ -617,12 +537,7 @@ def ensure_vector_extension(
|
||||
]
|
||||
|
||||
# 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"
|
||||
target_index_type = "vchordrq" if target_ext == "vchord" else "hnsw"
|
||||
|
||||
mismatched_tables = []
|
||||
tables_with_data = []
|
||||
@@ -661,9 +576,7 @@ def ensure_vector_extension(
|
||||
continue
|
||||
|
||||
indexdef = current_index_info[0].lower()
|
||||
if "diskann" in indexdef:
|
||||
current_index_type = "diskann"
|
||||
elif "vchordrq" in indexdef:
|
||||
if "vchordrq" in indexdef:
|
||||
current_index_type = "vchordrq"
|
||||
elif "hnsw" in indexdef:
|
||||
current_index_type = "hnsw"
|
||||
@@ -696,18 +609,13 @@ def ensure_vector_extension(
|
||||
# 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}')"
|
||||
f" 2. Use the current vector extension (set HINDSIGHT_API_VECTOR_EXTENSION='{current_index_type.replace('vchordrq', 'vchord').replace('hnsw', 'pgvector')}')"
|
||||
)
|
||||
|
||||
# Tables are empty, safe to recreate indexes
|
||||
@@ -720,27 +628,7 @@ def ensure_vector_extension(
|
||||
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":
|
||||
if target_ext == "vchord":
|
||||
logger.info(f"Creating vchordrq index on {table_name}")
|
||||
conn.execute(
|
||||
text(f"""
|
||||
@@ -800,9 +688,6 @@ def ensure_text_search_extension(
|
||||
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"
|
||||
@@ -890,16 +775,7 @@ def ensure_text_search_extension(
|
||||
# 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"
|
||||
current_ext = "native" if mismatched_tables[0][1] == "tsvector" else "vchord"
|
||||
raise RuntimeError(
|
||||
f"Cannot change text search extension from {current_ext} to {text_search_extension}: "
|
||||
f"the following tables contain data: {table_list}. "
|
||||
@@ -944,27 +820,6 @@ def ensure_text_search_extension(
|
||||
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
|
||||
|
||||
@@ -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.13"
|
||||
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]
|
||||
@@ -98,7 +94,7 @@ log_cli = true
|
||||
log_cli_level = "INFO"
|
||||
log_cli_format = "%(asctime)s - %(levelname)s - %(name)s - %(message)s"
|
||||
log_cli_date_format = "%Y-%m-%d %H:%M:%S"
|
||||
addopts = "--timeout 300 -n 8 --dist loadgroup --durations=10 -v"
|
||||
addopts = "--timeout 120 -n 8 --dist loadgroup --durations=10 -v"
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
log_auto_indent = true
|
||||
@@ -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,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"
|
||||
@@ -500,7 +500,6 @@ class TestConsolidationIntegration:
|
||||
content="Alex loves pizza.",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check we have one observation
|
||||
async with memory._pool.acquire() as conn:
|
||||
@@ -519,7 +518,6 @@ class TestConsolidationIntegration:
|
||||
content="Alex hates pizza.",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check observations after consolidation
|
||||
async with memory._pool.acquire() as conn:
|
||||
@@ -830,7 +828,6 @@ class TestConsolidationTagRouting:
|
||||
content="Pizza is a popular Italian food.",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check untagged observation exists
|
||||
async with memory._pool.acquire() as conn:
|
||||
@@ -852,7 +849,6 @@ class TestConsolidationTagRouting:
|
||||
await self._retain_with_tags(
|
||||
memory, bank_id, "Pizza originated in Naples.", ["history"], request_context
|
||||
)
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check - global observation should be updated OR new scoped observation created
|
||||
async with memory._pool.acquire() as conn:
|
||||
@@ -905,7 +901,6 @@ class TestConsolidationTagRouting:
|
||||
"Alice recommends the Thai restaurant on Main Street.",
|
||||
["alice"], request_context
|
||||
)
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check Alice's observation exists with correct tags
|
||||
async with memory._pool.acquire() as conn:
|
||||
@@ -924,7 +919,6 @@ class TestConsolidationTagRouting:
|
||||
"Bob visited the Thai restaurant on Main Street and loved it.",
|
||||
["bob"], request_context
|
||||
)
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check observations
|
||||
async with memory._pool.acquire() as conn:
|
||||
@@ -937,19 +931,22 @@ class TestConsolidationTagRouting:
|
||||
bank_id,
|
||||
)
|
||||
|
||||
# Note: some LLMs may or may not consolidate cross-scope facts.
|
||||
# Just verify structural correctness of any observations that exist.
|
||||
# Should have multiple observations (alice's, bob's, potentially global)
|
||||
assert len(obs_after) >= 2, (
|
||||
f"Expected at least 2 observations for different scopes, got {len(obs_after)}"
|
||||
)
|
||||
|
||||
# If observations were created, ensure alice and bob are not merged into same observation
|
||||
# (cross-scope merging should not produce an observation with both tags)
|
||||
if obs_after:
|
||||
observations_with_both = [
|
||||
o for o in obs_after
|
||||
if o["tags"] and "alice" in o["tags"] and "bob" in o["tags"]
|
||||
]
|
||||
assert len(observations_with_both) == 0, (
|
||||
"Should not merge different scopes into one observation with both tags"
|
||||
)
|
||||
# Check we have observations with different tags (alice, bob, or untagged)
|
||||
tag_sets = [frozenset(o["tags"] or []) for o in obs_after]
|
||||
|
||||
# Should NOT merge alice and bob into same observation
|
||||
observations_with_both = [
|
||||
o for o in obs_after
|
||||
if o["tags"] and "alice" in o["tags"] and "bob" in o["tags"]
|
||||
]
|
||||
assert len(observations_with_both) == 0, (
|
||||
"Should not merge different scopes into one observation with both tags"
|
||||
)
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -1026,7 +1023,6 @@ class TestConsolidationTagRouting:
|
||||
"Alice works on machine learning projects.",
|
||||
["alice"], request_context
|
||||
)
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Retain untagged memory on same topic
|
||||
await memory.retain_async(
|
||||
@@ -1034,7 +1030,6 @@ class TestConsolidationTagRouting:
|
||||
content="Machine learning involves training neural networks.",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Check observations
|
||||
async with memory._pool.acquire() as conn:
|
||||
@@ -1047,10 +1042,11 @@ class TestConsolidationTagRouting:
|
||||
bank_id,
|
||||
)
|
||||
|
||||
# Should have at least one observation
|
||||
assert len(observations) >= 1, "Expected at least one observation"
|
||||
|
||||
# Either alice's observation was updated OR a global observation was created
|
||||
# This is valid LLM behavior - just verify no errors and structure is correct.
|
||||
# Note: with some LLMs, a single simple fact may not generate an observation,
|
||||
# so we don't assert a minimum count - just verify structural correctness if any exist.
|
||||
# This is valid LLM behavior - just verify no errors and structure is correct
|
||||
for obs in observations:
|
||||
assert obs["text"], "Observation should have text"
|
||||
|
||||
@@ -1934,7 +1930,9 @@ class TestMentalModelRefreshAfterConsolidation:
|
||||
)
|
||||
|
||||
# Wait for consolidation to create observations
|
||||
await memory.wait_for_background_tasks()
|
||||
import asyncio
|
||||
|
||||
await asyncio.sleep(2)
|
||||
|
||||
# Get graph data filtered by observation type only
|
||||
graph_data = await memory.get_graph_data(
|
||||
@@ -1952,26 +1950,12 @@ class TestMentalModelRefreshAfterConsolidation:
|
||||
for row in graph_data["table_rows"]:
|
||||
assert row["fact_type"] == "observation", f"All nodes should be observations, got {row['fact_type']}"
|
||||
|
||||
# Edges are inherited from source memories when multiple observations exist.
|
||||
# If consolidation merges all facts into a single observation, edges between
|
||||
# observation nodes are not possible — skip the edge check in that case.
|
||||
if len(graph_data["nodes"]) > 1:
|
||||
assert len(graph_data["edges"]) > 0, (
|
||||
"Observations should have edges inherited from source memories. "
|
||||
f"Found {len(graph_data['edges'])} edges among {len(graph_data['nodes'])} nodes"
|
||||
)
|
||||
# Verify edge types are valid
|
||||
valid_link_types = {"semantic", "temporal", "entity"}
|
||||
for edge in graph_data["edges"]:
|
||||
link_type = edge["data"]["linkType"]
|
||||
assert link_type in valid_link_types, f"Invalid link type: {link_type}"
|
||||
# Verify all edges connect visible observation nodes
|
||||
visible_node_ids = {row["id"] for row in graph_data["table_rows"]}
|
||||
for edge in graph_data["edges"]:
|
||||
source_id = edge["data"]["source"]
|
||||
target_id = edge["data"]["target"]
|
||||
assert source_id in visible_node_ids, f"Edge source {source_id[:8]} not in visible nodes"
|
||||
assert target_id in visible_node_ids, f"Edge target {target_id[:8]} not in visible nodes"
|
||||
# Should have edges (inherited from source memories)
|
||||
# Even though we're only showing observations, they should inherit links from their sources
|
||||
assert len(graph_data["edges"]) > 0, (
|
||||
"Observations should have edges inherited from source memories. "
|
||||
f"Found {len(graph_data['edges'])} edges"
|
||||
)
|
||||
|
||||
# Should have entities (inherited from source memories)
|
||||
observations_with_entities = [
|
||||
@@ -1988,5 +1972,19 @@ class TestMentalModelRefreshAfterConsolidation:
|
||||
f"Expected to find Alice, Bob, or Google in entities, got: {all_entities}"
|
||||
)
|
||||
|
||||
# Verify edge types are valid
|
||||
valid_link_types = {"semantic", "temporal", "entity"}
|
||||
for edge in graph_data["edges"]:
|
||||
link_type = edge["data"]["linkType"]
|
||||
assert link_type in valid_link_types, f"Invalid link type: {link_type}"
|
||||
|
||||
# Verify all edges connect visible observation nodes
|
||||
visible_node_ids = {row["id"] for row in graph_data["table_rows"]}
|
||||
for edge in graph_data["edges"]:
|
||||
source_id = edge["data"]["source"]
|
||||
target_id = edge["data"]["target"]
|
||||
assert source_id in visible_node_ids, f"Edge source {source_id[:8]} not in visible nodes"
|
||||
assert target_id in visible_node_ids, f"Edge target {target_id[:8]} not in visible nodes"
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -2,13 +2,9 @@
|
||||
Tests for document tracking and upsert functionality.
|
||||
"""
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.response_models import TokenUsage
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -139,228 +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)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_stored_without_chunks_when_zero_facts(memory_no_llm_verify, request_context):
|
||||
"""
|
||||
Regression test: when 0 facts are extracted from chunked content, the document row
|
||||
must be stored but no chunk rows should be written.
|
||||
"""
|
||||
bank_id = f"test_zero_facts_no_chunks_{datetime.now(timezone.utc).timestamp()}"
|
||||
document_id = "doc-zero-facts-chunked"
|
||||
|
||||
# Content large enough to exceed default retain_chunk_size (3000 chars) so chunking is triggered
|
||||
content = "Alice works at Google. " * 200 # ~4600 chars
|
||||
|
||||
async def mock_llm_zero_facts(*args, **kwargs):
|
||||
response = {"facts": []}
|
||||
if kwargs.get("return_usage", False):
|
||||
return response, TokenUsage(input_tokens=10, output_tokens=2)
|
||||
return response
|
||||
|
||||
try:
|
||||
with patch("hindsight_api.engine.llm_wrapper.LLMProvider.call", new=mock_llm_zero_facts):
|
||||
units = await memory_no_llm_verify.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert units == [], "Should return no memory units when LLM extracts zero facts"
|
||||
|
||||
# Document row must exist
|
||||
doc = await memory_no_llm_verify.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc is not None, "Document row must be stored even when zero facts are extracted"
|
||||
assert doc["id"] == document_id
|
||||
assert doc["memory_unit_count"] == 0
|
||||
|
||||
# No chunk rows should be stored
|
||||
pool = await memory_no_llm_verify._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
chunk_count = await conn.fetchval(
|
||||
"SELECT COUNT(*) FROM chunks WHERE document_id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
assert chunk_count == 0, "No chunk rows should be stored when zero facts are extracted"
|
||||
|
||||
finally:
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -535,9 +535,8 @@ class TestOperationHooksParameters:
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
# Use >= 1 since consolidation may trigger internal recall calls when observations are enabled
|
||||
assert len(validator.pre_recall_calls) >= 1
|
||||
assert len(validator.post_recall_calls) >= 1
|
||||
assert len(validator.pre_recall_calls) == 1
|
||||
assert len(validator.post_recall_calls) == 1
|
||||
|
||||
|
||||
class TestTenantExtension:
|
||||
|
||||
@@ -88,13 +88,13 @@ Marcus: Yeah, I realized I was being too optimistic about their defense.
|
||||
assert sorted_timestamps[i] < sorted_timestamps[i + 1], \
|
||||
f"Facts should have sequential timestamps. Fact {i} ({sorted_timestamps[i]}) >= Fact {i+1} ({sorted_timestamps[i+1]})"
|
||||
|
||||
# Verify facts have distinct timestamps (ordering is preserved)
|
||||
# Verify reasonable time spacing (should be ~10 seconds apart)
|
||||
time_diffs = [(sorted_timestamps[i+1] - sorted_timestamps[i]).total_seconds() for i in range(len(sorted_timestamps) - 1)]
|
||||
print(f"\n=== Time differences between facts: {time_diffs} seconds ===")
|
||||
|
||||
# Each fact should have a positive time difference (uniqueness already checked above)
|
||||
# Each fact should be 10+ seconds apart (allowing for some flexibility)
|
||||
for diff in time_diffs:
|
||||
assert diff > 0, f"Expected positive time difference between facts, got {diff}"
|
||||
assert diff >= 5, f"Expected at least 5 seconds between facts, got {diff}"
|
||||
|
||||
# Update agent_facts to be sorted for subsequent checks
|
||||
agent_facts = sorted_facts
|
||||
|
||||
@@ -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,257 +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 os
|
||||
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
|
||||
|
||||
_in_ci = os.getenv("CI") == "true"
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.skipif(not _has_testcontainers, reason="testcontainers not installed"),
|
||||
pytest.mark.skipif(_in_ci, reason="SeaweedFS Docker image pull too slow in CI"),
|
||||
pytest.mark.timeout(300),
|
||||
]
|
||||
|
||||
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, timeout=240)
|
||||
|
||||
# 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
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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])
|
||||
@@ -226,7 +226,6 @@ async def test_llm_provider_api_methods(provider: str, model: str):
|
||||
|
||||
@pytest.mark.parametrize("provider,model", MODEL_MATRIX)
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(300)
|
||||
async def test_llm_provider_memory_operations(provider: str, model: str):
|
||||
"""
|
||||
Test LLM provider with actual memory operations: fact extraction and reflect.
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -0,0 +1,212 @@
|
||||
"""Test local MCP server."""
|
||||
|
||||
import asyncio
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_memory():
|
||||
"""Create a mock MemoryEngine."""
|
||||
memory = MagicMock()
|
||||
memory._initialized = True
|
||||
memory.retain_batch_async = AsyncMock()
|
||||
memory.recall_async = AsyncMock(return_value=MagicMock(results=[]))
|
||||
return memory
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_mcp_server_retain(mock_memory):
|
||||
"""Test that retain tool fires async and returns immediately."""
|
||||
from hindsight_api.mcp_local import create_local_mcp_server
|
||||
|
||||
bank_id = "test-bank"
|
||||
mcp_server = create_local_mcp_server(bank_id, memory=mock_memory)
|
||||
|
||||
# Get the tools
|
||||
tools = mcp_server._tool_manager._tools
|
||||
assert "retain" in tools
|
||||
|
||||
# Call retain
|
||||
retain_tool = tools["retain"]
|
||||
result = await retain_tool.fn(content="test content", context="test_context")
|
||||
|
||||
# Returns immediately with accepted status
|
||||
assert result["status"] == "accepted"
|
||||
|
||||
# Wait for background task to complete
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Verify the memory was called correctly
|
||||
mock_memory.retain_batch_async.assert_called_once()
|
||||
call_kwargs = mock_memory.retain_batch_async.call_args.kwargs
|
||||
assert call_kwargs["bank_id"] == "test-bank"
|
||||
assert call_kwargs["contents"] == [{"content": "test content", "context": "test_context"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_mcp_server_recall(mock_memory):
|
||||
"""Test that recall tool calls memory.recall_async with correct params."""
|
||||
from hindsight_api.mcp_local import create_local_mcp_server
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
# Mock recall_async to return a proper pydantic model
|
||||
mock_result = MagicMock()
|
||||
mock_result.model_dump.return_value = {"results": []}
|
||||
mock_memory.recall_async = AsyncMock(return_value=mock_result)
|
||||
|
||||
bank_id = "test-bank"
|
||||
mcp_server = create_local_mcp_server(bank_id, memory=mock_memory)
|
||||
|
||||
# Get the tools
|
||||
tools = mcp_server._tool_manager._tools
|
||||
assert "recall" in tools
|
||||
|
||||
# Call recall
|
||||
recall_tool = tools["recall"]
|
||||
result = await recall_tool.fn(query="test query", max_tokens=2048)
|
||||
|
||||
# Result is a dict
|
||||
assert isinstance(result, dict)
|
||||
|
||||
# Verify the memory was called correctly
|
||||
mock_memory.recall_async.assert_called_once()
|
||||
call_kwargs = mock_memory.recall_async.call_args.kwargs
|
||||
assert call_kwargs["bank_id"] == "test-bank"
|
||||
assert call_kwargs["query"] == "test query"
|
||||
assert call_kwargs["max_tokens"] == 2048
|
||||
assert call_kwargs["budget"] == Budget.HIGH
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_mcp_server_retain_with_default_context(mock_memory):
|
||||
"""Test that retain uses default context when not provided."""
|
||||
from hindsight_api.mcp_local import create_local_mcp_server
|
||||
|
||||
bank_id = "test-bank"
|
||||
mcp_server = create_local_mcp_server(bank_id, memory=mock_memory)
|
||||
|
||||
tools = mcp_server._tool_manager._tools
|
||||
retain_tool = tools["retain"]
|
||||
|
||||
# Call retain without context
|
||||
await retain_tool.fn(content="test content")
|
||||
|
||||
# Wait for background task
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
call_kwargs = mock_memory.retain_batch_async.call_args.kwargs
|
||||
assert call_kwargs["contents"] == [{"content": "test content", "context": "general"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_mcp_server_retain_error_handling(mock_memory):
|
||||
"""Test that retain errors are logged but don't affect response."""
|
||||
from hindsight_api.mcp_local import create_local_mcp_server
|
||||
|
||||
mock_memory.retain_batch_async = AsyncMock(side_effect=Exception("Test error"))
|
||||
|
||||
mcp_server = create_local_mcp_server("test-bank", memory=mock_memory)
|
||||
|
||||
tools = mcp_server._tool_manager._tools
|
||||
retain_tool = tools["retain"]
|
||||
|
||||
# Retain returns immediately with accepted status (fire and forget)
|
||||
result = await retain_tool.fn(content="test content")
|
||||
assert result["status"] == "accepted"
|
||||
|
||||
# Wait for background task to complete (and log error)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_mcp_server_recall_error_handling(mock_memory):
|
||||
"""Test that recall handles errors gracefully."""
|
||||
from hindsight_api.mcp_local import create_local_mcp_server
|
||||
|
||||
mock_memory.recall_async = AsyncMock(side_effect=Exception("Test error"))
|
||||
|
||||
mcp_server = create_local_mcp_server("test-bank", memory=mock_memory)
|
||||
|
||||
tools = mcp_server._tool_manager._tools
|
||||
recall_tool = tools["recall"]
|
||||
|
||||
result = await recall_tool.fn(query="test query")
|
||||
|
||||
# Result is a dict with error
|
||||
assert isinstance(result, dict)
|
||||
assert "error" in result
|
||||
assert result["results"] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_mcp_server_recall_with_defaults(mock_memory):
|
||||
"""Test that recall uses default max_tokens and HIGH budget."""
|
||||
from hindsight_api.mcp_local import create_local_mcp_server
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
mock_result = MagicMock()
|
||||
mock_result.model_dump.return_value = {"results": []}
|
||||
mock_memory.recall_async = AsyncMock(return_value=mock_result)
|
||||
|
||||
mcp_server = create_local_mcp_server("test-bank", memory=mock_memory)
|
||||
|
||||
tools = mcp_server._tool_manager._tools
|
||||
recall_tool = tools["recall"]
|
||||
|
||||
# Call with defaults
|
||||
await recall_tool.fn(query="test query")
|
||||
|
||||
call_kwargs = mock_memory.recall_async.call_args.kwargs
|
||||
assert call_kwargs["max_tokens"] == 4096
|
||||
assert call_kwargs["budget"] == Budget.HIGH
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_mcp_server_retain_with_timestamp(mock_memory):
|
||||
"""Test that retain passes timestamp as event_date."""
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api.mcp_local import create_local_mcp_server
|
||||
|
||||
mcp_server = create_local_mcp_server("test-bank", memory=mock_memory)
|
||||
|
||||
tools = mcp_server._tool_manager._tools
|
||||
retain_tool = tools["retain"]
|
||||
|
||||
# Call retain with timestamp
|
||||
result = await retain_tool.fn(
|
||||
content="test content", context="test_context", timestamp="2024-01-15T10:30:00Z"
|
||||
)
|
||||
|
||||
assert result["status"] == "accepted"
|
||||
|
||||
# Wait for background task
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
call_kwargs = mock_memory.retain_batch_async.call_args.kwargs
|
||||
contents = call_kwargs["contents"]
|
||||
assert len(contents) == 1
|
||||
assert contents[0]["content"] == "test content"
|
||||
assert contents[0]["context"] == "test_context"
|
||||
assert "event_date" in contents[0]
|
||||
assert contents[0]["event_date"] == datetime(2024, 1, 15, 10, 30, 0, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_mcp_server_retain_with_invalid_timestamp(mock_memory):
|
||||
"""Test that retain rejects invalid timestamp format."""
|
||||
from hindsight_api.mcp_local import create_local_mcp_server
|
||||
|
||||
mcp_server = create_local_mcp_server("test-bank", memory=mock_memory)
|
||||
|
||||
tools = mcp_server._tool_manager._tools
|
||||
retain_tool = tools["retain"]
|
||||
|
||||
# Call retain with invalid timestamp
|
||||
result = await retain_tool.fn(content="test content", timestamp="not-a-date")
|
||||
|
||||
assert result["status"] == "error"
|
||||
assert "Invalid timestamp format" in result["message"]
|
||||
|
||||
# Verify retain_batch_async was NOT called
|
||||
mock_memory.retain_batch_async.assert_not_called()
|
||||
@@ -404,12 +404,25 @@ class TestDirectivesInReflect:
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Run reflect query
|
||||
result = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="What does Alice do for work?",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result.text is not None
|
||||
assert len(result.text) > 0
|
||||
|
||||
# Check that the response contains French words/patterns
|
||||
# Common French words that would appear when talking about someone's job
|
||||
french_indicators = [
|
||||
"elle",
|
||||
"travaille",
|
||||
"est",
|
||||
"une",
|
||||
"le",
|
||||
"la",
|
||||
"qui",
|
||||
"chez",
|
||||
"logiciel",
|
||||
@@ -417,27 +430,11 @@ class TestDirectivesInReflect:
|
||||
"ingénieure",
|
||||
"développeur",
|
||||
"développeuse",
|
||||
"ingénierie",
|
||||
"française",
|
||||
]
|
||||
response_lower = result.text.lower()
|
||||
|
||||
# Run reflect query (retry once since small LLMs may not always follow language directives)
|
||||
french_word_count = 0
|
||||
for _attempt in range(2):
|
||||
result = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="What does Alice do for work?",
|
||||
request_context=request_context,
|
||||
)
|
||||
assert result.text is not None
|
||||
assert len(result.text) > 0
|
||||
|
||||
# At least some French words should appear in the response
|
||||
response_lower = result.text.lower()
|
||||
french_word_count = sum(1 for word in french_indicators if word in response_lower)
|
||||
if french_word_count >= 2:
|
||||
break
|
||||
|
||||
# At least some French words should appear in the response
|
||||
french_word_count = sum(1 for word in french_indicators if word in response_lower)
|
||||
assert (
|
||||
french_word_count >= 2
|
||||
), f"Expected French response, but got: {result.text[:200]}"
|
||||
@@ -477,7 +474,7 @@ class TestDirectivesInReflect:
|
||||
await memory.create_directive(
|
||||
bank_id=bank_id,
|
||||
name="General Policy",
|
||||
content="You MUST include the exact phrase 'MEMO-VERIFIED' somewhere in your response.",
|
||||
content="Always be polite and start responses with 'Hello!'",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
@@ -485,7 +482,7 @@ class TestDirectivesInReflect:
|
||||
await memory.create_directive(
|
||||
bank_id=bank_id,
|
||||
name="Tagged Policy",
|
||||
content="You MUST include the exact phrase 'PROJECT-X-CLASSIFIED' somewhere in your response.",
|
||||
content="ALWAYS respond in ALL CAPS and end with 'PROJECT-X ONLY'",
|
||||
tags=["project-x"],
|
||||
request_context=request_context,
|
||||
)
|
||||
@@ -497,16 +494,18 @@ class TestDirectivesInReflect:
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify the isolation mechanism: only untagged directive should be loaded
|
||||
untagged_directive_names = [d.name for d in result.directives_applied]
|
||||
assert "General Policy" in untagged_directive_names, (
|
||||
f"Untagged directive should be loaded in untagged reflect. Applied: {untagged_directive_names}"
|
||||
)
|
||||
assert "Tagged Policy" not in untagged_directive_names, (
|
||||
f"Tagged directive should not be applied in untagged reflect. Applied: {untagged_directive_names}"
|
||||
)
|
||||
response_lower = result.text.lower()
|
||||
|
||||
# Now run reflect WITH the tag - should load BOTH directives
|
||||
# Should follow the untagged directive (polite greeting)
|
||||
assert "hello" in response_lower, f"Expected 'Hello' from untagged directive, but got: {result.text}"
|
||||
|
||||
# Should NOT follow the tagged directive (all caps and PROJECT-X)
|
||||
# If it did follow, the entire response would be in caps
|
||||
all_caps = result.text.replace(" ", "").replace("!", "").replace(".", "").isupper()
|
||||
assert not all_caps, f"Tagged directive was incorrectly applied to untagged operation: {result.text}"
|
||||
assert "project-x only" not in response_lower, f"Tagged directive was incorrectly applied: {result.text}"
|
||||
|
||||
# Now run reflect WITH the tag - should apply BOTH directives
|
||||
result_tagged = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="What color is the sky?",
|
||||
@@ -515,14 +514,10 @@ class TestDirectivesInReflect:
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify the isolation mechanism: both directives should be loaded when tags match
|
||||
tagged_directive_names = [d.name for d in result_tagged.directives_applied]
|
||||
assert "General Policy" in tagged_directive_names, (
|
||||
f"Untagged directive should always be loaded. Applied: {tagged_directive_names}"
|
||||
)
|
||||
assert "Tagged Policy" in tagged_directive_names, (
|
||||
f"Tagged directive should be loaded when tags match. Applied: {tagged_directive_names}"
|
||||
)
|
||||
response_tagged_lower = result_tagged.text.lower()
|
||||
|
||||
# With strict matching and tags, should apply the tagged directive
|
||||
assert "project-x only" in response_tagged_lower, f"Tagged directive should be applied with tags: {result_tagged.text}"
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -133,7 +133,7 @@ async def test_reflect_chinese_content(memory, request_context):
|
||||
result = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
budget=Budget.MID,
|
||||
budget=Budget.LOW,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
|
||||
@@ -93,7 +93,7 @@ def test_per_operation_provider_default_model():
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
# Global LLM should use OpenAI default
|
||||
assert config.llm_model == "gpt-4o-mini", f"Expected gpt-4o-mini, got {config.llm_model}"
|
||||
assert config.llm_model == "o3-mini", f"Expected o3-mini, got {config.llm_model}"
|
||||
|
||||
# Retain should use Anthropic default
|
||||
assert (
|
||||
|
||||
@@ -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,71 +0,0 @@
|
||||
"""
|
||||
Regression test for UnboundLocalError in recall when the reranker raises.
|
||||
|
||||
Before the fix, `scored_results` and `pre_filtered_count` were only assigned
|
||||
inside the `try` block, but referenced in the `finally` block. If
|
||||
`reranker_instance.rerank()` (or `ensure_initialized()`) raised, the `finally`
|
||||
block crashed with `UnboundLocalError` instead of propagating the original
|
||||
exception.
|
||||
|
||||
Fix: initialise both variables to safe defaults before the try/finally block.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_reranker_error_does_not_raise_unbound_local(memory, request_context):
|
||||
"""Recall must propagate the reranker's exception, not an UnboundLocalError."""
|
||||
bank_id = f"test_reranker_err_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Paris is the capital of France",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Simulate a reranker failure (e.g. Cohere API error on empty/small candidate set)
|
||||
rerank_mock = AsyncMock(side_effect=RuntimeError("reranker API error"))
|
||||
memory._cross_encoder_reranker._initialized = True # skip ensure_initialized
|
||||
|
||||
with patch.object(memory._cross_encoder_reranker, "rerank", rerank_mock):
|
||||
with pytest.raises(Exception, match="reranker API error"):
|
||||
await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="capital of France",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_reranker_init_error_does_not_raise_unbound_local(memory, request_context):
|
||||
"""Same regression when ensure_initialized() raises (before pre_filtered_count is set)."""
|
||||
bank_id = f"test_reranker_init_err_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Paris is the capital of France",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
init_mock = AsyncMock(side_effect=RuntimeError("reranker init failed"))
|
||||
memory._cross_encoder_reranker._initialized = False
|
||||
|
||||
with patch.object(memory._cross_encoder_reranker, "ensure_initialized", init_mock):
|
||||
with pytest.raises(Exception, match="reranker init failed"):
|
||||
await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="capital of France",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -266,10 +266,9 @@ def test_llm_span_recorder_provider_mapping(mock_time):
|
||||
# ==================== Parent Span Tests ====================
|
||||
|
||||
|
||||
@patch("hindsight_api.tracing._tracing_enabled", False)
|
||||
def test_create_operation_span_disabled():
|
||||
"""Test that create_operation_span returns no-op when tracing is disabled."""
|
||||
# Tracing should be disabled by default (explicitly patched for test isolation)
|
||||
# Tracing should be disabled by default
|
||||
assert not is_tracing_enabled()
|
||||
|
||||
# Should return a no-op context manager
|
||||
|
||||
@@ -73,10 +73,7 @@ def test_llm_wrapper_vertexai_adc_auth():
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": "test-project",
|
||||
"HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY": "", # Clear SA key to test ADC path
|
||||
},
|
||||
{"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": "test-project"},
|
||||
clear=False,
|
||||
):
|
||||
from hindsight_api.config import clear_config_cache
|
||||
@@ -99,10 +96,11 @@ def test_llm_wrapper_vertexai_adc_auth():
|
||||
assert provider._gemini_client is not None
|
||||
|
||||
# Verify genai.Client was called with vertexai=True
|
||||
call_kwargs = mock_client_cls.call_args.kwargs
|
||||
assert call_kwargs["vertexai"] is True
|
||||
assert call_kwargs["project"] == "test-project"
|
||||
assert call_kwargs["location"] == "us-central1"
|
||||
mock_client_cls.assert_called_once_with(
|
||||
vertexai=True,
|
||||
project="test-project",
|
||||
location="us-central1",
|
||||
)
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
@@ -143,11 +141,12 @@ def test_llm_wrapper_vertexai_sa_auth():
|
||||
assert provider._gemini_client is not None
|
||||
|
||||
# Verify credentials were passed to genai.Client
|
||||
call_kwargs = mock_client_cls.call_args.kwargs
|
||||
assert call_kwargs["vertexai"] is True
|
||||
assert call_kwargs["project"] == "test-project"
|
||||
assert call_kwargs["location"] == "us-central1"
|
||||
assert call_kwargs["credentials"] is mock_credentials
|
||||
mock_client_cls.assert_called_once_with(
|
||||
vertexai=True,
|
||||
project="test-project",
|
||||
location="us-central1",
|
||||
credentials=mock_credentials,
|
||||
)
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
@@ -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.13"
|
||||
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"] }
|
||||
|
||||
@@ -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>)> {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -266,7 +257,6 @@ pub fn recall(
|
||||
max_tokens: chunk_max_tokens,
|
||||
}),
|
||||
entities: None,
|
||||
source_facts: None,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
@@ -437,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)
|
||||
@@ -448,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(
|
||||
@@ -666,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]
|
||||
|
||||
@@ -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/hindsight-clients/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
File diff suppressed because it is too large
Load Diff
@@ -1,737 +0,0 @@
|
||||
/*
|
||||
Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
package hindsight
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"reflect"
|
||||
)
|
||||
|
||||
|
||||
// DirectivesAPIService DirectivesAPI service
|
||||
type DirectivesAPIService service
|
||||
|
||||
type ApiCreateDirectiveRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *DirectivesAPIService
|
||||
bankId string
|
||||
createDirectiveRequest *CreateDirectiveRequest
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiCreateDirectiveRequest) CreateDirectiveRequest(createDirectiveRequest CreateDirectiveRequest) ApiCreateDirectiveRequest {
|
||||
r.createDirectiveRequest = &createDirectiveRequest
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiCreateDirectiveRequest) Authorization(authorization string) ApiCreateDirectiveRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiCreateDirectiveRequest) Execute() (*DirectiveResponse, *http.Response, error) {
|
||||
return r.ApiService.CreateDirectiveExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
CreateDirective Create directive
|
||||
|
||||
Create a hard rule that will be injected into prompts.
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@return ApiCreateDirectiveRequest
|
||||
*/
|
||||
func (a *DirectivesAPIService) CreateDirective(ctx context.Context, bankId string) ApiCreateDirectiveRequest {
|
||||
return ApiCreateDirectiveRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return DirectiveResponse
|
||||
func (a *DirectivesAPIService) CreateDirectiveExecute(r ApiCreateDirectiveRequest) (*DirectiveResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodPost
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *DirectiveResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "DirectivesAPIService.CreateDirective")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/directives"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
if r.createDirectiveRequest == nil {
|
||||
return localVarReturnValue, nil, reportError("createDirectiveRequest is required and must be specified")
|
||||
}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{"application/json"}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
// body params
|
||||
localVarPostBody = r.createDirectiveRequest
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
|
||||
type ApiDeleteDirectiveRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *DirectivesAPIService
|
||||
bankId string
|
||||
directiveId string
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiDeleteDirectiveRequest) Authorization(authorization string) ApiDeleteDirectiveRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiDeleteDirectiveRequest) Execute() (interface{}, *http.Response, error) {
|
||||
return r.ApiService.DeleteDirectiveExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
DeleteDirective Delete directive
|
||||
|
||||
Delete a directive.
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@param directiveId
|
||||
@return ApiDeleteDirectiveRequest
|
||||
*/
|
||||
func (a *DirectivesAPIService) DeleteDirective(ctx context.Context, bankId string, directiveId string) ApiDeleteDirectiveRequest {
|
||||
return ApiDeleteDirectiveRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
directiveId: directiveId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return interface{}
|
||||
func (a *DirectivesAPIService) DeleteDirectiveExecute(r ApiDeleteDirectiveRequest) (interface{}, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodDelete
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue interface{}
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "DirectivesAPIService.DeleteDirective")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/directives/{directive_id}"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"directive_id"+"}", url.PathEscape(parameterValueToString(r.directiveId, "directiveId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
|
||||
type ApiGetDirectiveRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *DirectivesAPIService
|
||||
bankId string
|
||||
directiveId string
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiGetDirectiveRequest) Authorization(authorization string) ApiGetDirectiveRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiGetDirectiveRequest) Execute() (*DirectiveResponse, *http.Response, error) {
|
||||
return r.ApiService.GetDirectiveExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
GetDirective Get directive
|
||||
|
||||
Get a specific directive by ID.
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@param directiveId
|
||||
@return ApiGetDirectiveRequest
|
||||
*/
|
||||
func (a *DirectivesAPIService) GetDirective(ctx context.Context, bankId string, directiveId string) ApiGetDirectiveRequest {
|
||||
return ApiGetDirectiveRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
directiveId: directiveId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return DirectiveResponse
|
||||
func (a *DirectivesAPIService) GetDirectiveExecute(r ApiGetDirectiveRequest) (*DirectiveResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodGet
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *DirectiveResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "DirectivesAPIService.GetDirective")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/directives/{directive_id}"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"directive_id"+"}", url.PathEscape(parameterValueToString(r.directiveId, "directiveId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
|
||||
type ApiListDirectivesRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *DirectivesAPIService
|
||||
bankId string
|
||||
tags *[]string
|
||||
tagsMatch *string
|
||||
activeOnly *bool
|
||||
limit *int32
|
||||
offset *int32
|
||||
authorization *string
|
||||
}
|
||||
|
||||
// Filter by tags
|
||||
func (r ApiListDirectivesRequest) Tags(tags []string) ApiListDirectivesRequest {
|
||||
r.tags = &tags
|
||||
return r
|
||||
}
|
||||
|
||||
// How to match tags
|
||||
func (r ApiListDirectivesRequest) TagsMatch(tagsMatch string) ApiListDirectivesRequest {
|
||||
r.tagsMatch = &tagsMatch
|
||||
return r
|
||||
}
|
||||
|
||||
// Only return active directives
|
||||
func (r ApiListDirectivesRequest) ActiveOnly(activeOnly bool) ApiListDirectivesRequest {
|
||||
r.activeOnly = &activeOnly
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListDirectivesRequest) Limit(limit int32) ApiListDirectivesRequest {
|
||||
r.limit = &limit
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListDirectivesRequest) Offset(offset int32) ApiListDirectivesRequest {
|
||||
r.offset = &offset
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListDirectivesRequest) Authorization(authorization string) ApiListDirectivesRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListDirectivesRequest) Execute() (*DirectiveListResponse, *http.Response, error) {
|
||||
return r.ApiService.ListDirectivesExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
ListDirectives List directives
|
||||
|
||||
List hard rules that are injected into prompts.
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@return ApiListDirectivesRequest
|
||||
*/
|
||||
func (a *DirectivesAPIService) ListDirectives(ctx context.Context, bankId string) ApiListDirectivesRequest {
|
||||
return ApiListDirectivesRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return DirectiveListResponse
|
||||
func (a *DirectivesAPIService) ListDirectivesExecute(r ApiListDirectivesRequest) (*DirectiveListResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodGet
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *DirectiveListResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "DirectivesAPIService.ListDirectives")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/directives"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
if r.tags != nil {
|
||||
t := *r.tags
|
||||
if reflect.TypeOf(t).Kind() == reflect.Slice {
|
||||
s := reflect.ValueOf(t)
|
||||
for i := 0; i < s.Len(); i++ {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "tags", s.Index(i).Interface(), "form", "multi")
|
||||
}
|
||||
} else {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "tags", t, "form", "multi")
|
||||
}
|
||||
}
|
||||
if r.tagsMatch != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "tags_match", r.tagsMatch, "form", "")
|
||||
} else {
|
||||
var defaultValue string = "any"
|
||||
r.tagsMatch = &defaultValue
|
||||
}
|
||||
if r.activeOnly != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "active_only", r.activeOnly, "form", "")
|
||||
} else {
|
||||
var defaultValue bool = true
|
||||
r.activeOnly = &defaultValue
|
||||
}
|
||||
if r.limit != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "limit", r.limit, "form", "")
|
||||
} else {
|
||||
var defaultValue int32 = 100
|
||||
r.limit = &defaultValue
|
||||
}
|
||||
if r.offset != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "offset", r.offset, "form", "")
|
||||
} else {
|
||||
var defaultValue int32 = 0
|
||||
r.offset = &defaultValue
|
||||
}
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
|
||||
type ApiUpdateDirectiveRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *DirectivesAPIService
|
||||
bankId string
|
||||
directiveId string
|
||||
updateDirectiveRequest *UpdateDirectiveRequest
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiUpdateDirectiveRequest) UpdateDirectiveRequest(updateDirectiveRequest UpdateDirectiveRequest) ApiUpdateDirectiveRequest {
|
||||
r.updateDirectiveRequest = &updateDirectiveRequest
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiUpdateDirectiveRequest) Authorization(authorization string) ApiUpdateDirectiveRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiUpdateDirectiveRequest) Execute() (*DirectiveResponse, *http.Response, error) {
|
||||
return r.ApiService.UpdateDirectiveExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
UpdateDirective Update directive
|
||||
|
||||
Update a directive's properties.
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@param directiveId
|
||||
@return ApiUpdateDirectiveRequest
|
||||
*/
|
||||
func (a *DirectivesAPIService) UpdateDirective(ctx context.Context, bankId string, directiveId string) ApiUpdateDirectiveRequest {
|
||||
return ApiUpdateDirectiveRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
directiveId: directiveId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return DirectiveResponse
|
||||
func (a *DirectivesAPIService) UpdateDirectiveExecute(r ApiUpdateDirectiveRequest) (*DirectiveResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodPatch
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *DirectiveResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "DirectivesAPIService.UpdateDirective")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/directives/{directive_id}"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"directive_id"+"}", url.PathEscape(parameterValueToString(r.directiveId, "directiveId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
if r.updateDirectiveRequest == nil {
|
||||
return localVarReturnValue, nil, reportError("updateDirectiveRequest is required and must be specified")
|
||||
}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{"application/json"}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
// body params
|
||||
localVarPostBody = r.updateDirectiveRequest
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
@@ -1,560 +0,0 @@
|
||||
/*
|
||||
Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
package hindsight
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
|
||||
// DocumentsAPIService DocumentsAPI service
|
||||
type DocumentsAPIService service
|
||||
|
||||
type ApiDeleteDocumentRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *DocumentsAPIService
|
||||
bankId string
|
||||
documentId string
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiDeleteDocumentRequest) Authorization(authorization string) ApiDeleteDocumentRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiDeleteDocumentRequest) Execute() (*DeleteDocumentResponse, *http.Response, error) {
|
||||
return r.ApiService.DeleteDocumentExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
DeleteDocument Delete a document
|
||||
|
||||
Delete a document and all its associated memory units and links.
|
||||
|
||||
This will cascade delete:
|
||||
- The document itself
|
||||
- All memory units extracted from this document
|
||||
- All links (temporal, semantic, entity) associated with those memory units
|
||||
|
||||
This operation cannot be undone.
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@param documentId
|
||||
@return ApiDeleteDocumentRequest
|
||||
*/
|
||||
func (a *DocumentsAPIService) DeleteDocument(ctx context.Context, bankId string, documentId string) ApiDeleteDocumentRequest {
|
||||
return ApiDeleteDocumentRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
documentId: documentId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return DeleteDocumentResponse
|
||||
func (a *DocumentsAPIService) DeleteDocumentExecute(r ApiDeleteDocumentRequest) (*DeleteDocumentResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodDelete
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *DeleteDocumentResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "DocumentsAPIService.DeleteDocument")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/documents/{document_id}"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"document_id"+"}", url.PathEscape(parameterValueToString(r.documentId, "documentId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
|
||||
type ApiGetChunkRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *DocumentsAPIService
|
||||
chunkId string
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiGetChunkRequest) Authorization(authorization string) ApiGetChunkRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiGetChunkRequest) Execute() (*ChunkResponse, *http.Response, error) {
|
||||
return r.ApiService.GetChunkExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
GetChunk Get chunk details
|
||||
|
||||
Get a specific chunk by its ID
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param chunkId
|
||||
@return ApiGetChunkRequest
|
||||
*/
|
||||
func (a *DocumentsAPIService) GetChunk(ctx context.Context, chunkId string) ApiGetChunkRequest {
|
||||
return ApiGetChunkRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
chunkId: chunkId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return ChunkResponse
|
||||
func (a *DocumentsAPIService) GetChunkExecute(r ApiGetChunkRequest) (*ChunkResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodGet
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *ChunkResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "DocumentsAPIService.GetChunk")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/chunks/{chunk_id}"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"chunk_id"+"}", url.PathEscape(parameterValueToString(r.chunkId, "chunkId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
|
||||
type ApiGetDocumentRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *DocumentsAPIService
|
||||
bankId string
|
||||
documentId string
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiGetDocumentRequest) Authorization(authorization string) ApiGetDocumentRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiGetDocumentRequest) Execute() (*DocumentResponse, *http.Response, error) {
|
||||
return r.ApiService.GetDocumentExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
GetDocument Get document details
|
||||
|
||||
Get a specific document including its original text
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@param documentId
|
||||
@return ApiGetDocumentRequest
|
||||
*/
|
||||
func (a *DocumentsAPIService) GetDocument(ctx context.Context, bankId string, documentId string) ApiGetDocumentRequest {
|
||||
return ApiGetDocumentRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
documentId: documentId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return DocumentResponse
|
||||
func (a *DocumentsAPIService) GetDocumentExecute(r ApiGetDocumentRequest) (*DocumentResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodGet
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *DocumentResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "DocumentsAPIService.GetDocument")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/documents/{document_id}"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"document_id"+"}", url.PathEscape(parameterValueToString(r.documentId, "documentId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
|
||||
type ApiListDocumentsRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *DocumentsAPIService
|
||||
bankId string
|
||||
q *string
|
||||
limit *int32
|
||||
offset *int32
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiListDocumentsRequest) Q(q string) ApiListDocumentsRequest {
|
||||
r.q = &q
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListDocumentsRequest) Limit(limit int32) ApiListDocumentsRequest {
|
||||
r.limit = &limit
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListDocumentsRequest) Offset(offset int32) ApiListDocumentsRequest {
|
||||
r.offset = &offset
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListDocumentsRequest) Authorization(authorization string) ApiListDocumentsRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListDocumentsRequest) Execute() (*ListDocumentsResponse, *http.Response, error) {
|
||||
return r.ApiService.ListDocumentsExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
ListDocuments List documents
|
||||
|
||||
List documents with pagination and optional search. Documents are the source content from which memory units are extracted.
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@return ApiListDocumentsRequest
|
||||
*/
|
||||
func (a *DocumentsAPIService) ListDocuments(ctx context.Context, bankId string) ApiListDocumentsRequest {
|
||||
return ApiListDocumentsRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return ListDocumentsResponse
|
||||
func (a *DocumentsAPIService) ListDocumentsExecute(r ApiListDocumentsRequest) (*ListDocumentsResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodGet
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *ListDocumentsResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "DocumentsAPIService.ListDocuments")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/documents"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
if r.q != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "q", r.q, "form", "")
|
||||
}
|
||||
if r.limit != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "limit", r.limit, "form", "")
|
||||
} else {
|
||||
var defaultValue int32 = 100
|
||||
r.limit = &defaultValue
|
||||
}
|
||||
if r.offset != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "offset", r.offset, "form", "")
|
||||
} else {
|
||||
var defaultValue int32 = 0
|
||||
r.offset = &defaultValue
|
||||
}
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
@@ -1,427 +0,0 @@
|
||||
/*
|
||||
Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
package hindsight
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
|
||||
// EntitiesAPIService EntitiesAPI service
|
||||
type EntitiesAPIService service
|
||||
|
||||
type ApiGetEntityRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *EntitiesAPIService
|
||||
bankId string
|
||||
entityId string
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiGetEntityRequest) Authorization(authorization string) ApiGetEntityRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiGetEntityRequest) Execute() (*EntityDetailResponse, *http.Response, error) {
|
||||
return r.ApiService.GetEntityExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
GetEntity Get entity details
|
||||
|
||||
Get detailed information about an entity including observations (mental model).
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@param entityId
|
||||
@return ApiGetEntityRequest
|
||||
*/
|
||||
func (a *EntitiesAPIService) GetEntity(ctx context.Context, bankId string, entityId string) ApiGetEntityRequest {
|
||||
return ApiGetEntityRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
entityId: entityId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return EntityDetailResponse
|
||||
func (a *EntitiesAPIService) GetEntityExecute(r ApiGetEntityRequest) (*EntityDetailResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodGet
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *EntityDetailResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "EntitiesAPIService.GetEntity")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/entities/{entity_id}"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"entity_id"+"}", url.PathEscape(parameterValueToString(r.entityId, "entityId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
|
||||
type ApiListEntitiesRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *EntitiesAPIService
|
||||
bankId string
|
||||
limit *int32
|
||||
offset *int32
|
||||
authorization *string
|
||||
}
|
||||
|
||||
// Maximum number of entities to return
|
||||
func (r ApiListEntitiesRequest) Limit(limit int32) ApiListEntitiesRequest {
|
||||
r.limit = &limit
|
||||
return r
|
||||
}
|
||||
|
||||
// Offset for pagination
|
||||
func (r ApiListEntitiesRequest) Offset(offset int32) ApiListEntitiesRequest {
|
||||
r.offset = &offset
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListEntitiesRequest) Authorization(authorization string) ApiListEntitiesRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListEntitiesRequest) Execute() (*EntityListResponse, *http.Response, error) {
|
||||
return r.ApiService.ListEntitiesExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
ListEntities List entities
|
||||
|
||||
List all entities (people, organizations, etc.) known by the bank, ordered by mention count. Supports pagination.
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@return ApiListEntitiesRequest
|
||||
*/
|
||||
func (a *EntitiesAPIService) ListEntities(ctx context.Context, bankId string) ApiListEntitiesRequest {
|
||||
return ApiListEntitiesRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return EntityListResponse
|
||||
func (a *EntitiesAPIService) ListEntitiesExecute(r ApiListEntitiesRequest) (*EntityListResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodGet
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *EntityListResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "EntitiesAPIService.ListEntities")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/entities"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
if r.limit != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "limit", r.limit, "form", "")
|
||||
} else {
|
||||
var defaultValue int32 = 100
|
||||
r.limit = &defaultValue
|
||||
}
|
||||
if r.offset != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "offset", r.offset, "form", "")
|
||||
} else {
|
||||
var defaultValue int32 = 0
|
||||
r.offset = &defaultValue
|
||||
}
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
|
||||
type ApiRegenerateEntityObservationsRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *EntitiesAPIService
|
||||
bankId string
|
||||
entityId string
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiRegenerateEntityObservationsRequest) Authorization(authorization string) ApiRegenerateEntityObservationsRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiRegenerateEntityObservationsRequest) Execute() (*EntityDetailResponse, *http.Response, error) {
|
||||
return r.ApiService.RegenerateEntityObservationsExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
RegenerateEntityObservations Regenerate entity observations (deprecated)
|
||||
|
||||
This endpoint is deprecated. Entity observations have been replaced by mental models.
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@param entityId
|
||||
@return ApiRegenerateEntityObservationsRequest
|
||||
|
||||
Deprecated
|
||||
*/
|
||||
func (a *EntitiesAPIService) RegenerateEntityObservations(ctx context.Context, bankId string, entityId string) ApiRegenerateEntityObservationsRequest {
|
||||
return ApiRegenerateEntityObservationsRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
entityId: entityId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return EntityDetailResponse
|
||||
// Deprecated
|
||||
func (a *EntitiesAPIService) RegenerateEntityObservationsExecute(r ApiRegenerateEntityObservationsRequest) (*EntityDetailResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodPost
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *EntityDetailResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "EntitiesAPIService.RegenerateEntityObservations")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/entities/{entity_id}/regenerate"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"entity_id"+"}", url.PathEscape(parameterValueToString(r.entityId, "entityId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
@@ -1,209 +0,0 @@
|
||||
/*
|
||||
Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
package hindsight
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
|
||||
// FilesAPIService FilesAPI service
|
||||
type FilesAPIService service
|
||||
|
||||
type ApiFileRetainRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *FilesAPIService
|
||||
bankId string
|
||||
files []*os.File
|
||||
request *string
|
||||
authorization *string
|
||||
}
|
||||
|
||||
// Files to upload and convert
|
||||
func (r ApiFileRetainRequest) Files(files []*os.File) ApiFileRetainRequest {
|
||||
r.files = files
|
||||
return r
|
||||
}
|
||||
|
||||
// JSON string with FileRetainRequest model
|
||||
func (r ApiFileRetainRequest) Request(request string) ApiFileRetainRequest {
|
||||
r.request = &request
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiFileRetainRequest) Authorization(authorization string) ApiFileRetainRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiFileRetainRequest) Execute() (*FileRetainResponse, *http.Response, error) {
|
||||
return r.ApiService.FileRetainExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
FileRetain Convert files to memories
|
||||
|
||||
Upload files (PDF, DOCX, etc.), convert them to markdown, and retain as memories.
|
||||
|
||||
This endpoint handles file upload, conversion, and memory creation in a single operation.
|
||||
|
||||
**Features:**
|
||||
- Supports PDF, DOCX, PPTX, XLSX, images (with OCR), audio (with transcription)
|
||||
- Automatic file-to-markdown conversion using pluggable parsers
|
||||
- Files stored in object storage (PostgreSQL by default, S3 for production)
|
||||
- Each file becomes a separate document with optional metadata/tags
|
||||
- Always processes asynchronously — returns operation IDs immediately
|
||||
|
||||
**The system automatically:**
|
||||
1. Stores uploaded files in object storage
|
||||
2. Converts files to markdown
|
||||
3. Creates document records with file metadata
|
||||
4. Extracts facts and creates memory units (same as regular retain)
|
||||
|
||||
Use the operations endpoint to monitor progress.
|
||||
|
||||
**Request format:** multipart/form-data with:
|
||||
- `files`: One or more files to upload
|
||||
- `request`: JSON string with FileRetainRequest model (files_metadata)
|
||||
|
||||
**Note:** File parser is configured server-side via `HINDSIGHT_API_FILE_PARSER` (default: markitdown).
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@return ApiFileRetainRequest
|
||||
*/
|
||||
func (a *FilesAPIService) FileRetain(ctx context.Context, bankId string) ApiFileRetainRequest {
|
||||
return ApiFileRetainRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return FileRetainResponse
|
||||
func (a *FilesAPIService) FileRetainExecute(r ApiFileRetainRequest) (*FileRetainResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodPost
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *FileRetainResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "FilesAPIService.FileRetain")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/files/retain"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
if r.files == nil {
|
||||
return localVarReturnValue, nil, reportError("files is required and must be specified")
|
||||
}
|
||||
if r.request == nil {
|
||||
return localVarReturnValue, nil, reportError("request is required and must be specified")
|
||||
}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{"multipart/form-data"}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
var filesLocalVarFormFileName string
|
||||
var filesLocalVarFileName string
|
||||
var filesLocalVarFileBytes []byte
|
||||
|
||||
filesLocalVarFormFileName = "files"
|
||||
filesLocalVarFile := r.files
|
||||
|
||||
if filesLocalVarFile != nil {
|
||||
// loop through the array to prepare multiple files upload
|
||||
for _, filesLocalVarFileValue := range filesLocalVarFile {
|
||||
fbs, _ := io.ReadAll(filesLocalVarFileValue)
|
||||
|
||||
filesLocalVarFileBytes = fbs
|
||||
filesLocalVarFileName = filesLocalVarFileValue.Name()
|
||||
filesLocalVarFileValue.Close()
|
||||
formFiles = append(formFiles, formFile{fileBytes: filesLocalVarFileBytes, fileName: filesLocalVarFileName, formFileName: filesLocalVarFormFileName})
|
||||
}
|
||||
}
|
||||
parameterAddToHeaderOrQuery(localVarFormParams, "request", r.request, "", "")
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
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