Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bdef0d80f4 | ||
|
|
7db1eb1b3a | ||
|
|
c96ae58c28 | ||
|
|
d349376df7 | ||
|
|
b403b94795 | ||
|
|
fe5d3a999b | ||
|
|
6168a77846 | ||
|
|
da44a5e839 | ||
|
|
32bca12c6f | ||
|
|
26850a0156 | ||
|
|
2a0c490c9e |
@@ -42,6 +42,10 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/litellm
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-embed
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv build --out-dir dist
|
||||
|
||||
# Publish in order (client and api first, then hindsight-all which depends on them)
|
||||
- name: Publish hindsight-client to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
@@ -67,6 +71,12 @@ jobs:
|
||||
packages-dir: ./hindsight-integrations/litellm/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-embed to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-embed/dist
|
||||
skip-existing: true
|
||||
|
||||
# Upload artifacts for GitHub release
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
@@ -77,6 +87,7 @@ jobs:
|
||||
hindsight-api/dist/*
|
||||
hindsight/dist/*
|
||||
hindsight-integrations/litellm/dist/*
|
||||
hindsight-embed/dist/*
|
||||
retention-days: 1
|
||||
|
||||
release-typescript-client:
|
||||
@@ -416,6 +427,7 @@ jobs:
|
||||
cp artifacts/python-packages/hindsight-api/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-integrations/litellm/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-embed/dist/* release-assets/ || true
|
||||
# TypeScript client
|
||||
cp artifacts/typescript-client/*.tgz release-assets/ || true
|
||||
# Control Plane
|
||||
|
||||
@@ -20,6 +20,8 @@ jobs:
|
||||
path: hindsight-api
|
||||
- name: hindsight-client
|
||||
path: hindsight-clients/python
|
||||
- name: hindsight-embed
|
||||
path: hindsight-embed
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
@@ -563,6 +565,46 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/litellm
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-embed:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_EMBED_LLM_PROVIDER: groq
|
||||
HINDSIGHT_EMBED_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
HINDSIGHT_EMBED_LLM_MODEL: openai/gpt-oss-20b
|
||||
# Prefer CPU-only PyTorch in CI
|
||||
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv sync --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-embed-${{ hashFiles('hindsight-embed/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-embed-
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Run smoke test
|
||||
working-directory: ./hindsight-embed
|
||||
run: ./test.sh
|
||||
|
||||
test-doc-examples:
|
||||
runs-on: ubuntu-latest
|
||||
needs: build-rust-cli
|
||||
|
||||
@@ -14,6 +14,7 @@ This document captures architectural decisions and coding conventions for the Hi
|
||||
hindsight/ # Python package for embedded usage
|
||||
hindsight-api/ # FastAPI server (core memory engine)
|
||||
hindsight-cli/ # Rust CLI client
|
||||
hindsight-embed/ # Embedded CLI (no server needed)
|
||||
hindsight-control-plane/ # Next.js admin UI
|
||||
hindsight-docs/ # Docusaurus documentation site
|
||||
hindsight-dev/ # Development tools and benchmarks
|
||||
@@ -148,4 +149,5 @@ Note: The maintained wrapper `hindsight_client.py` and `README.md` are preserved
|
||||
|
||||
# Branding
|
||||
## Colors
|
||||
- Primary: gradient from #0074d9 to #009296
|
||||
- Primary: gradient from #0074d9 to #009296
|
||||
|
||||
|
||||
@@ -2,16 +2,19 @@
|
||||
# Supports building API-only, Control Plane-only, or both
|
||||
#
|
||||
# Build args:
|
||||
# INCLUDE_API=true/false - Include API (default: true)
|
||||
# INCLUDE_CP=true/false - Include Control Plane (default: true)
|
||||
# INCLUDE_API=true/false - Include API (default: true)
|
||||
# INCLUDE_CP=true/false - Include Control Plane (default: true)
|
||||
# PRELOAD_ML_MODELS=true/false - Pre-download ML models during build (default: true)
|
||||
#
|
||||
# Examples:
|
||||
# docker build -t hindsight . # Both (standalone)
|
||||
# docker build -t hindsight-api --build-arg INCLUDE_CP=false . # API only
|
||||
# docker build -t hindsight-cp --build-arg INCLUDE_API=false . # Control Plane only
|
||||
# docker build -t hindsight . # Both (standalone)
|
||||
# docker build -t hindsight-api --build-arg INCLUDE_CP=false . # API only
|
||||
# docker build -t hindsight-cp --build-arg INCLUDE_API=false . # Control Plane only
|
||||
# docker build -t hindsight --build-arg PRELOAD_ML_MODELS=false . # Skip ML model preload
|
||||
|
||||
ARG INCLUDE_API=true
|
||||
ARG INCLUDE_CP=true
|
||||
ARG PRELOAD_ML_MODELS=true
|
||||
|
||||
# =============================================================================
|
||||
# Stage: API Builder
|
||||
@@ -159,14 +162,17 @@ ENV PG0_HOME=/home/hindsight/.pg0-cache
|
||||
|
||||
ENV PG0_HOME=/home/hindsight/.pg0
|
||||
|
||||
# Pre-download ML models to avoid runtime download
|
||||
RUN /app/api/.venv/bin/python -c "\
|
||||
# Pre-download ML models to avoid runtime download (conditional)
|
||||
ARG PRELOAD_ML_MODELS
|
||||
RUN if [ "$PRELOAD_ML_MODELS" = "true" ]; then \
|
||||
/app/api/.venv/bin/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 cached successfully')"
|
||||
print('Models cached successfully')"; \
|
||||
else echo "Skipping ML model preload"; fi
|
||||
|
||||
EXPOSE 8888
|
||||
|
||||
@@ -277,14 +283,17 @@ print('PostgreSQL pre-cached to PG0_HOME')" || echo "Pre-download skipped"
|
||||
|
||||
ENV PG0_HOME=/home/hindsight/.pg0
|
||||
|
||||
# Pre-download ML models to avoid runtime download
|
||||
RUN /app/api/.venv/bin/python -c "\
|
||||
# Pre-download ML models to avoid runtime download (conditional)
|
||||
ARG PRELOAD_ML_MODELS
|
||||
RUN if [ "$PRELOAD_ML_MODELS" = "true" ]; then \
|
||||
/app/api/.venv/bin/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 cached successfully')"
|
||||
print('Models cached successfully')"; \
|
||||
else echo "Skipping ML model preload"; fi
|
||||
|
||||
EXPOSE 8888 9999
|
||||
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.1.11
|
||||
appVersion: "0.1.11"
|
||||
version: 0.1.12
|
||||
appVersion: "0.1.12"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -21,9 +21,11 @@ from .engine.search.trace import (
|
||||
WeightComponents,
|
||||
)
|
||||
from .engine.search.tracer import SearchTracer
|
||||
from .models import RequestContext
|
||||
|
||||
__all__ = [
|
||||
"MemoryEngine",
|
||||
"RequestContext",
|
||||
"HindsightConfig",
|
||||
"get_config",
|
||||
"SearchTrace",
|
||||
|
||||
@@ -109,6 +109,9 @@ def run_migrations_online() -> None:
|
||||
|
||||
get_database_url() # Process and set the database URL in config
|
||||
|
||||
# Check if we're targeting a specific schema (for multi-tenant isolation)
|
||||
target_schema = config.get_main_option("target_schema")
|
||||
|
||||
connectable = engine_from_config(
|
||||
config.get_section(config.config_ini_section, {}),
|
||||
prefix="sqlalchemy.",
|
||||
@@ -121,14 +124,34 @@ def run_migrations_online() -> None:
|
||||
def set_read_write_mode(dbapi_connection, connection_record):
|
||||
cursor = dbapi_connection.cursor()
|
||||
cursor.execute("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE")
|
||||
# If targeting a specific schema, set search_path
|
||||
# Include public in search_path for access to shared extensions (pgvector)
|
||||
if target_schema:
|
||||
cursor.execute(f'CREATE SCHEMA IF NOT EXISTS "{target_schema}"')
|
||||
cursor.execute(f'SET search_path TO "{target_schema}", public')
|
||||
cursor.close()
|
||||
|
||||
with connectable.connect() as connection:
|
||||
# Also explicitly set read-write mode on this connection
|
||||
connection.execute(text("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE"))
|
||||
|
||||
# If targeting a specific schema, set search_path
|
||||
# Include public in search_path for access to shared extensions (pgvector)
|
||||
if target_schema:
|
||||
connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{target_schema}"'))
|
||||
connection.execute(text(f'SET search_path TO "{target_schema}", public'))
|
||||
|
||||
connection.commit() # Commit the SET command
|
||||
|
||||
context.configure(connection=connection, target_metadata=target_metadata)
|
||||
# Configure context with version_table_schema if using a specific schema
|
||||
context_opts = {
|
||||
"connection": connection,
|
||||
"target_metadata": target_metadata,
|
||||
}
|
||||
if target_schema:
|
||||
context_opts["version_table_schema"] = target_schema
|
||||
|
||||
context.configure(**context_opts)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
+14
-4
@@ -6,7 +6,7 @@ Create Date: 2024-12-04 15:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
from alembic import context, op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "d9f6a3b4c5e2"
|
||||
@@ -15,14 +15,22 @@ branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (e.g., 'tenant_x.' or '' for public)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade():
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Drop old check constraint FIRST (before updating data)
|
||||
op.drop_constraint("memory_units_fact_type_check", "memory_units", type_="check")
|
||||
|
||||
# Update existing 'bank' values to 'experience'
|
||||
op.execute("UPDATE memory_units SET fact_type = 'experience' WHERE fact_type = 'bank'")
|
||||
op.execute(f"UPDATE {schema}memory_units SET fact_type = 'experience' WHERE fact_type = 'bank'")
|
||||
# Also update any 'interactions' values (in case of partial migration)
|
||||
op.execute("UPDATE memory_units SET fact_type = 'experience' WHERE fact_type = 'interactions'")
|
||||
op.execute(f"UPDATE {schema}memory_units SET fact_type = 'experience' WHERE fact_type = 'interactions'")
|
||||
|
||||
# Create new check constraint with 'experience' instead of 'bank'
|
||||
op.create_check_constraint(
|
||||
@@ -31,11 +39,13 @@ def upgrade():
|
||||
|
||||
|
||||
def downgrade():
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Drop new check constraint FIRST
|
||||
op.drop_constraint("memory_units_fact_type_check", "memory_units", type_="check")
|
||||
|
||||
# Update 'experience' back to 'bank'
|
||||
op.execute("UPDATE memory_units SET fact_type = 'bank' WHERE fact_type = 'experience'")
|
||||
op.execute(f"UPDATE {schema}memory_units SET fact_type = 'bank' WHERE fact_type = 'experience'")
|
||||
|
||||
# Recreate old check constraint
|
||||
op.create_check_constraint(
|
||||
|
||||
+54
-13
@@ -12,7 +12,7 @@ system (skepticism, literalism, empathy with 1-5 integer values).
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from alembic import context, op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "e0a1b2c3d4e5"
|
||||
@@ -21,9 +21,36 @@ 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 (e.g., 'tenant_x.' or '' for public)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _get_target_schema() -> str:
|
||||
"""Get the target schema name (tenant schema or 'public')."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return schema if schema else "public"
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Convert Big Five disposition to 3-trait disposition."""
|
||||
conn = op.get_bind()
|
||||
schema = _get_schema_prefix()
|
||||
target_schema = _get_target_schema()
|
||||
|
||||
# Check if disposition column exists (should have been created by previous migration)
|
||||
result = conn.execute(
|
||||
sa.text("""
|
||||
SELECT column_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
|
||||
"""),
|
||||
{"schema": target_schema},
|
||||
)
|
||||
if not result.fetchone():
|
||||
# Column doesn't exist yet (shouldn't happen but be safe)
|
||||
return
|
||||
|
||||
# Update all existing banks to use the new disposition format
|
||||
# Convert from old format to new format with reasonable mappings:
|
||||
@@ -32,18 +59,18 @@ def upgrade() -> None:
|
||||
# - empathy: derived from agreeableness + inverse of neuroticism
|
||||
# Default all to 3 (neutral) for simplicity
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
UPDATE banks
|
||||
SET disposition = '{"skepticism": 3, "literalism": 3, "empathy": 3}'::jsonb
|
||||
sa.text(f"""
|
||||
UPDATE {schema}banks
|
||||
SET disposition = '{{"skepticism": 3, "literalism": 3, "empathy": 3}}'::jsonb
|
||||
WHERE disposition IS NOT NULL
|
||||
""")
|
||||
)
|
||||
|
||||
# Update the default for new banks
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
ALTER TABLE banks
|
||||
ALTER COLUMN disposition SET DEFAULT '{"skepticism": 3, "literalism": 3, "empathy": 3}'::jsonb
|
||||
sa.text(f"""
|
||||
ALTER TABLE {schema}banks
|
||||
ALTER COLUMN disposition SET DEFAULT '{{"skepticism": 3, "literalism": 3, "empathy": 3}}'::jsonb
|
||||
""")
|
||||
)
|
||||
|
||||
@@ -51,20 +78,34 @@ def upgrade() -> None:
|
||||
def downgrade() -> None:
|
||||
"""Convert back to Big Five disposition."""
|
||||
conn = op.get_bind()
|
||||
schema = _get_schema_prefix()
|
||||
target_schema = _get_target_schema()
|
||||
|
||||
# Check if disposition column exists
|
||||
result = conn.execute(
|
||||
sa.text("""
|
||||
SELECT column_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
|
||||
"""),
|
||||
{"schema": target_schema},
|
||||
)
|
||||
if not result.fetchone():
|
||||
return
|
||||
|
||||
# Revert to Big Five format with default values
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
UPDATE banks
|
||||
SET disposition = '{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}'::jsonb
|
||||
sa.text(f"""
|
||||
UPDATE {schema}banks
|
||||
SET disposition = '{{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}}'::jsonb
|
||||
WHERE disposition IS NOT NULL
|
||||
""")
|
||||
)
|
||||
|
||||
# Update the default for new banks
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
ALTER TABLE banks
|
||||
ALTER COLUMN disposition SET DEFAULT '{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}'::jsonb
|
||||
sa.text(f"""
|
||||
ALTER TABLE {schema}banks
|
||||
ALTER COLUMN disposition SET DEFAULT '{{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}}'::jsonb
|
||||
""")
|
||||
)
|
||||
|
||||
@@ -9,7 +9,7 @@ Create Date: 2024-12-04
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from alembic import context, op
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -19,17 +19,25 @@ branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_target_schema() -> str:
|
||||
"""Get the target schema name (tenant schema or 'public')."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return schema if schema else "public"
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Rename personality column to disposition in banks table (if it exists)."""
|
||||
conn = op.get_bind()
|
||||
target_schema = _get_target_schema()
|
||||
|
||||
# Check if 'personality' column exists (old database)
|
||||
result = conn.execute(
|
||||
sa.text("""
|
||||
SELECT column_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_name = 'banks' AND column_name = 'personality'
|
||||
""")
|
||||
WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'personality'
|
||||
"""),
|
||||
{"schema": target_schema},
|
||||
)
|
||||
has_personality = result.fetchone() is not None
|
||||
|
||||
@@ -38,8 +46,9 @@ def upgrade() -> None:
|
||||
sa.text("""
|
||||
SELECT column_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_name = 'banks' AND column_name = 'disposition'
|
||||
""")
|
||||
WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
|
||||
"""),
|
||||
{"schema": target_schema},
|
||||
)
|
||||
has_disposition = result.fetchone() is not None
|
||||
|
||||
@@ -63,12 +72,14 @@ def upgrade() -> None:
|
||||
def downgrade() -> None:
|
||||
"""Revert disposition column back to personality."""
|
||||
conn = op.get_bind()
|
||||
target_schema = _get_target_schema()
|
||||
result = conn.execute(
|
||||
sa.text("""
|
||||
SELECT column_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_name = 'banks' AND column_name = 'disposition'
|
||||
""")
|
||||
WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
|
||||
"""),
|
||||
{"schema": target_schema},
|
||||
)
|
||||
if result.fetchone():
|
||||
op.alter_column("banks", "disposition", new_column_name="personality")
|
||||
|
||||
@@ -12,7 +12,7 @@ from contextlib import asynccontextmanager
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from fastapi import FastAPI, HTTPException, Query
|
||||
from fastapi import Depends, FastAPI, Header, HTTPException, Query
|
||||
|
||||
|
||||
def _parse_metadata(metadata: Any) -> dict[str, Any]:
|
||||
@@ -33,9 +33,11 @@ from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.engine.db_utils import acquire_with_retry
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.engine.memory_engine import Budget, fq_table
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
|
||||
from hindsight_api.extensions import HttpExtension, load_extension
|
||||
from hindsight_api.metrics import create_metrics_collector, get_metrics_collector, initialize_metrics
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -337,7 +339,7 @@ class RetainResponse(BaseModel):
|
||||
success: bool
|
||||
bank_id: str
|
||||
items_count: int
|
||||
async_: bool = Field(
|
||||
is_async: bool = Field(
|
||||
alias="async", serialization_alias="async", description="Whether the operation was processed asynchronously"
|
||||
)
|
||||
|
||||
@@ -706,7 +708,11 @@ class DeleteResponse(BaseModel):
|
||||
deleted_count: int | None = None
|
||||
|
||||
|
||||
def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
|
||||
def create_app(
|
||||
memory: MemoryEngine,
|
||||
initialize_memory: bool = True,
|
||||
http_extension: HttpExtension | None = None,
|
||||
) -> FastAPI:
|
||||
"""
|
||||
Create and configure the FastAPI application.
|
||||
|
||||
@@ -714,6 +720,8 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
|
||||
memory: MemoryEngine instance (already initialized with required parameters).
|
||||
Migrations are controlled by the MemoryEngine's run_migrations parameter.
|
||||
initialize_memory: Whether to initialize memory system on startup (default: True)
|
||||
http_extension: Optional HTTP extension to mount custom endpoints under /extension/.
|
||||
If None, attempts to load from HINDSIGHT_API_HTTP_EXTENSION env var.
|
||||
|
||||
Returns:
|
||||
Configured FastAPI application
|
||||
@@ -723,6 +731,11 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
|
||||
In that case, you should call memory.initialize() manually before starting the server
|
||||
and memory.close() when shutting down.
|
||||
"""
|
||||
# Load HTTP extension from environment if not provided
|
||||
if http_extension is None:
|
||||
http_extension = load_extension("HTTP", HttpExtension)
|
||||
if http_extension:
|
||||
logging.info(f"Loaded HTTP extension: {http_extension.__class__.__name__}")
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
@@ -746,8 +759,18 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
|
||||
await memory.initialize()
|
||||
logging.info("Memory system initialized")
|
||||
|
||||
# Call HTTP extension startup hook
|
||||
if http_extension:
|
||||
await http_extension.on_startup()
|
||||
logging.info("HTTP extension started")
|
||||
|
||||
yield
|
||||
|
||||
# Call HTTP extension shutdown hook
|
||||
if http_extension:
|
||||
await http_extension.on_shutdown()
|
||||
logging.info("HTTP extension stopped")
|
||||
|
||||
# Shutdown: Cleanup memory system
|
||||
await memory.close()
|
||||
logging.info("Memory system closed")
|
||||
@@ -775,12 +798,36 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
|
||||
# Register all routes
|
||||
_register_routes(app)
|
||||
|
||||
# Mount HTTP extension router if available
|
||||
if http_extension:
|
||||
extension_router = http_extension.get_router(memory)
|
||||
app.include_router(extension_router, prefix="/ext", tags=["Extension"])
|
||||
logging.info("HTTP extension router mounted at /ext/")
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def _register_routes(app: FastAPI):
|
||||
"""Register all API routes on the given app instance."""
|
||||
|
||||
def get_request_context(authorization: str | None = Header(default=None)) -> RequestContext:
|
||||
"""
|
||||
Extract request context from Authorization header.
|
||||
|
||||
Supports:
|
||||
- Bearer token: "Bearer <api_key>"
|
||||
- Direct API key: "<api_key>"
|
||||
|
||||
Returns RequestContext with extracted API key (may be None if no auth header).
|
||||
"""
|
||||
api_key = None
|
||||
if authorization:
|
||||
if authorization.lower().startswith("bearer "):
|
||||
api_key = authorization[7:].strip()
|
||||
else:
|
||||
api_key = authorization.strip()
|
||||
return RequestContext(api_key=api_key)
|
||||
|
||||
@app.get(
|
||||
"/health",
|
||||
summary="Health check endpoint",
|
||||
@@ -821,10 +868,12 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="get_graph",
|
||||
tags=["Memory"],
|
||||
)
|
||||
async def api_graph(bank_id: str, type: str | None = None):
|
||||
async def api_graph(
|
||||
bank_id: str, type: str | None = None, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Get graph data from database, filtered by bank_id and optionally by type."""
|
||||
try:
|
||||
data = await app.state.memory.get_graph_data(bank_id, type)
|
||||
data = await app.state.memory.get_graph_data(bank_id, type, request_context=request_context)
|
||||
return data
|
||||
except Exception as e:
|
||||
import traceback
|
||||
@@ -841,7 +890,14 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="list_memories",
|
||||
tags=["Memory"],
|
||||
)
|
||||
async def api_list(bank_id: str, type: str | None = None, q: str | None = None, limit: int = 100, offset: int = 0):
|
||||
async def api_list(
|
||||
bank_id: str,
|
||||
type: str | None = None,
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""
|
||||
List memory units for table view with optional full-text search.
|
||||
|
||||
@@ -857,7 +913,12 @@ def _register_routes(app: FastAPI):
|
||||
"""
|
||||
try:
|
||||
data = await app.state.memory.list_memory_units(
|
||||
bank_id=bank_id, fact_type=type, search_query=q, limit=limit, offset=offset
|
||||
bank_id=bank_id,
|
||||
fact_type=type,
|
||||
search_query=q,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
request_context=request_context,
|
||||
)
|
||||
return data
|
||||
except Exception as e:
|
||||
@@ -880,7 +941,9 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="recall_memories",
|
||||
tags=["Memory"],
|
||||
)
|
||||
async def api_recall(bank_id: str, request: RecallRequest):
|
||||
async def api_recall(
|
||||
bank_id: str, request: RecallRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Run a recall and return results with trace."""
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
@@ -923,6 +986,7 @@ def _register_routes(app: FastAPI):
|
||||
max_entity_tokens=max_entity_tokens,
|
||||
include_chunks=include_chunks,
|
||||
max_chunk_tokens=max_chunk_tokens,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Convert core MemoryFact objects to API RecallResult objects (excluding internal metrics)
|
||||
@@ -995,14 +1059,20 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="reflect",
|
||||
tags=["Memory"],
|
||||
)
|
||||
async def api_reflect(bank_id: str, request: ReflectRequest):
|
||||
async def api_reflect(
|
||||
bank_id: str, request: ReflectRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
try:
|
||||
# Use the memory system's reflect_async method (record metrics)
|
||||
with metrics.record_operation("reflect", bank_id=bank_id, budget=request.budget.value):
|
||||
core_result = await app.state.memory.reflect_async(
|
||||
bank_id=bank_id, query=request.query, budget=request.budget, context=request.context
|
||||
bank_id=bank_id,
|
||||
query=request.query,
|
||||
budget=request.budget,
|
||||
context=request.context,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Convert core MemoryFact objects to API ReflectFact objects if facts are requested
|
||||
@@ -1041,10 +1111,10 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="list_banks",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_list_banks():
|
||||
async def api_list_banks(request_context: RequestContext = Depends(get_request_context)):
|
||||
"""Get list of all banks with their profiles."""
|
||||
try:
|
||||
banks = await app.state.memory.list_banks()
|
||||
banks = await app.state.memory.list_banks(request_context=request_context)
|
||||
return BankListResponse(banks=banks)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
@@ -1067,9 +1137,9 @@ def _register_routes(app: FastAPI):
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Get node counts by fact_type
|
||||
node_stats = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT fact_type, COUNT(*) as count
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
GROUP BY fact_type
|
||||
""",
|
||||
@@ -1078,10 +1148,10 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
# Get link counts by link_type
|
||||
link_stats = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT ml.link_type, COUNT(*) as count
|
||||
FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.from_unit_id = mu.id
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
|
||||
WHERE mu.bank_id = $1
|
||||
GROUP BY ml.link_type
|
||||
""",
|
||||
@@ -1090,10 +1160,10 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
# Get link counts by fact_type (from nodes)
|
||||
link_fact_type_stats = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT mu.fact_type, COUNT(*) as count
|
||||
FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.from_unit_id = mu.id
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
|
||||
WHERE mu.bank_id = $1
|
||||
GROUP BY mu.fact_type
|
||||
""",
|
||||
@@ -1102,10 +1172,10 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
# Get link counts by fact_type AND link_type
|
||||
link_breakdown_stats = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT mu.fact_type, ml.link_type, COUNT(*) as count
|
||||
FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.from_unit_id = mu.id
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
|
||||
WHERE mu.bank_id = $1
|
||||
GROUP BY mu.fact_type, ml.link_type
|
||||
""",
|
||||
@@ -1114,9 +1184,9 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
# Get pending and failed operations counts
|
||||
ops_stats = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT status, COUNT(*) as count
|
||||
FROM async_operations
|
||||
FROM {fq_table("async_operations")}
|
||||
WHERE bank_id = $1
|
||||
GROUP BY status
|
||||
""",
|
||||
@@ -1128,9 +1198,9 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
# Get document count
|
||||
doc_count_result = await conn.fetchrow(
|
||||
"""
|
||||
f"""
|
||||
SELECT COUNT(*) as count
|
||||
FROM documents
|
||||
FROM {fq_table("documents")}
|
||||
WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
@@ -1184,11 +1254,13 @@ def _register_routes(app: FastAPI):
|
||||
tags=["Entities"],
|
||||
)
|
||||
async def api_list_entities(
|
||||
bank_id: str, limit: int = Query(default=100, description="Maximum number of entities to return")
|
||||
bank_id: str,
|
||||
limit: int = Query(default=100, description="Maximum number of entities to return"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""List entities for a memory bank."""
|
||||
try:
|
||||
entities = await app.state.memory.list_entities(bank_id, limit=limit)
|
||||
entities = await app.state.memory.list_entities(bank_id, limit=limit, request_context=request_context)
|
||||
return EntityListResponse(items=[EntityListItem(**e) for e in entities])
|
||||
except Exception as e:
|
||||
import traceback
|
||||
@@ -1205,37 +1277,26 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="get_entity",
|
||||
tags=["Entities"],
|
||||
)
|
||||
async def api_get_entity(bank_id: str, entity_id: str):
|
||||
async def api_get_entity(
|
||||
bank_id: str, entity_id: str, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Get entity details with observations."""
|
||||
try:
|
||||
# First get the entity metadata
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
entity_row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
|
||||
FROM entities
|
||||
WHERE bank_id = $1 AND id = $2
|
||||
""",
|
||||
bank_id,
|
||||
uuid.UUID(entity_id),
|
||||
)
|
||||
entity = await app.state.memory.get_entity(bank_id, entity_id, request_context=request_context)
|
||||
|
||||
if not entity_row:
|
||||
if entity is None:
|
||||
raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found")
|
||||
|
||||
# Get observations for the entity
|
||||
observations = await app.state.memory.get_entity_observations(bank_id, entity_id, limit=20)
|
||||
|
||||
return EntityDetailResponse(
|
||||
id=str(entity_row["id"]),
|
||||
canonical_name=entity_row["canonical_name"],
|
||||
mention_count=entity_row["mention_count"],
|
||||
first_seen=entity_row["first_seen"].isoformat() if entity_row["first_seen"] else None,
|
||||
last_seen=entity_row["last_seen"].isoformat() if entity_row["last_seen"] else None,
|
||||
metadata=_parse_metadata(entity_row["metadata"]),
|
||||
id=entity["id"],
|
||||
canonical_name=entity["canonical_name"],
|
||||
mention_count=entity["mention_count"],
|
||||
first_seen=entity["first_seen"],
|
||||
last_seen=entity["last_seen"],
|
||||
metadata=_parse_metadata(entity["metadata"]),
|
||||
observations=[
|
||||
EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at) for obs in observations
|
||||
EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at)
|
||||
for obs in entity["observations"]
|
||||
],
|
||||
)
|
||||
except HTTPException:
|
||||
@@ -1255,42 +1316,40 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="regenerate_entity_observations",
|
||||
tags=["Entities"],
|
||||
)
|
||||
async def api_regenerate_entity_observations(bank_id: str, entity_id: str):
|
||||
async def api_regenerate_entity_observations(
|
||||
bank_id: str,
|
||||
entity_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Regenerate observations for an entity."""
|
||||
try:
|
||||
# First get the entity metadata
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
entity_row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
|
||||
FROM entities
|
||||
WHERE bank_id = $1 AND id = $2
|
||||
""",
|
||||
bank_id,
|
||||
uuid.UUID(entity_id),
|
||||
)
|
||||
# Get the entity to verify it exists and get canonical_name
|
||||
entity = await app.state.memory.get_entity(bank_id, entity_id, request_context=request_context)
|
||||
|
||||
if not entity_row:
|
||||
if entity is None:
|
||||
raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found")
|
||||
|
||||
# Regenerate observations
|
||||
await app.state.memory.regenerate_entity_observations(
|
||||
bank_id=bank_id, entity_id=entity_id, entity_name=entity_row["canonical_name"]
|
||||
bank_id=bank_id,
|
||||
entity_id=entity_id,
|
||||
entity_name=entity["canonical_name"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Get updated observations
|
||||
observations = await app.state.memory.get_entity_observations(bank_id, entity_id, limit=20)
|
||||
# Get updated entity with new observations
|
||||
entity = await app.state.memory.get_entity(bank_id, entity_id, request_context=request_context)
|
||||
|
||||
return EntityDetailResponse(
|
||||
id=str(entity_row["id"]),
|
||||
canonical_name=entity_row["canonical_name"],
|
||||
mention_count=entity_row["mention_count"],
|
||||
first_seen=entity_row["first_seen"].isoformat() if entity_row["first_seen"] else None,
|
||||
last_seen=entity_row["last_seen"].isoformat() if entity_row["last_seen"] else None,
|
||||
metadata=_parse_metadata(entity_row["metadata"]),
|
||||
id=entity["id"],
|
||||
canonical_name=entity["canonical_name"],
|
||||
mention_count=entity["mention_count"],
|
||||
first_seen=entity["first_seen"],
|
||||
last_seen=entity["last_seen"],
|
||||
metadata=_parse_metadata(entity["metadata"]),
|
||||
observations=[
|
||||
EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at) for obs in observations
|
||||
EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at)
|
||||
for obs in entity["observations"]
|
||||
],
|
||||
)
|
||||
except HTTPException:
|
||||
@@ -1310,7 +1369,13 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="list_documents",
|
||||
tags=["Documents"],
|
||||
)
|
||||
async def api_list_documents(bank_id: str, q: str | None = None, limit: int = 100, offset: int = 0):
|
||||
async def api_list_documents(
|
||||
bank_id: str,
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""
|
||||
List documents for a memory bank with optional search.
|
||||
|
||||
@@ -1321,7 +1386,9 @@ def _register_routes(app: FastAPI):
|
||||
offset: Offset for pagination (default: 0)
|
||||
"""
|
||||
try:
|
||||
data = await app.state.memory.list_documents(bank_id=bank_id, search_query=q, limit=limit, offset=offset)
|
||||
data = await app.state.memory.list_documents(
|
||||
bank_id=bank_id, search_query=q, limit=limit, offset=offset, request_context=request_context
|
||||
)
|
||||
return data
|
||||
except Exception as e:
|
||||
import traceback
|
||||
@@ -1338,7 +1405,9 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="get_document",
|
||||
tags=["Documents"],
|
||||
)
|
||||
async def api_get_document(bank_id: str, document_id: str):
|
||||
async def api_get_document(
|
||||
bank_id: str, document_id: str, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""
|
||||
Get a specific document with its original text.
|
||||
|
||||
@@ -1347,7 +1416,7 @@ def _register_routes(app: FastAPI):
|
||||
document_id: Document ID (from path)
|
||||
"""
|
||||
try:
|
||||
document = await app.state.memory.get_document(document_id, bank_id)
|
||||
document = await app.state.memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
if not document:
|
||||
raise HTTPException(status_code=404, detail="Document not found")
|
||||
return document
|
||||
@@ -1368,7 +1437,7 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="get_chunk",
|
||||
tags=["Documents"],
|
||||
)
|
||||
async def api_get_chunk(chunk_id: str):
|
||||
async def api_get_chunk(chunk_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
"""
|
||||
Get a specific chunk with its text.
|
||||
|
||||
@@ -1376,7 +1445,7 @@ def _register_routes(app: FastAPI):
|
||||
chunk_id: Chunk ID (from path, format: bank_id_document_id_chunk_index)
|
||||
"""
|
||||
try:
|
||||
chunk = await app.state.memory.get_chunk(chunk_id)
|
||||
chunk = await app.state.memory.get_chunk(chunk_id, request_context=request_context)
|
||||
if not chunk:
|
||||
raise HTTPException(status_code=404, detail="Chunk not found")
|
||||
return chunk
|
||||
@@ -1401,7 +1470,9 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="delete_document",
|
||||
tags=["Documents"],
|
||||
)
|
||||
async def api_delete_document(bank_id: str, document_id: str):
|
||||
async def api_delete_document(
|
||||
bank_id: str, document_id: str, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""
|
||||
Delete a document and all its associated memory units and links.
|
||||
|
||||
@@ -1410,7 +1481,7 @@ def _register_routes(app: FastAPI):
|
||||
document_id: Document ID to delete (from path)
|
||||
"""
|
||||
try:
|
||||
result = await app.state.memory.delete_document(document_id, bank_id)
|
||||
result = await app.state.memory.delete_document(document_id, bank_id, request_context=request_context)
|
||||
|
||||
if result["document_deleted"] == 0:
|
||||
raise HTTPException(status_code=404, detail="Document not found")
|
||||
@@ -1437,45 +1508,14 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="list_operations",
|
||||
tags=["Operations"],
|
||||
)
|
||||
async def api_list_operations(bank_id: str):
|
||||
async def api_list_operations(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
"""List all async operations (pending and failed) for a memory bank."""
|
||||
try:
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
operations = await conn.fetch(
|
||||
"""
|
||||
SELECT operation_id, bank_id, operation_type, created_at, status, error_message, result_metadata
|
||||
FROM async_operations
|
||||
WHERE bank_id = $1
|
||||
ORDER BY created_at DESC
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
def parse_metadata(metadata):
|
||||
"""Parse result_metadata which may be a string or dict."""
|
||||
if metadata is None:
|
||||
return {}
|
||||
if isinstance(metadata, str):
|
||||
return json.loads(metadata)
|
||||
return metadata
|
||||
|
||||
return {
|
||||
"bank_id": bank_id,
|
||||
"operations": [
|
||||
{
|
||||
"id": str(row["operation_id"]),
|
||||
"task_type": row["operation_type"],
|
||||
"items_count": parse_metadata(row["result_metadata"]).get("items_count", 0),
|
||||
"document_id": parse_metadata(row["result_metadata"]).get("document_id"),
|
||||
"created_at": row["created_at"].isoformat(),
|
||||
"status": row["status"],
|
||||
"error_message": row["error_message"],
|
||||
}
|
||||
for row in operations
|
||||
],
|
||||
}
|
||||
|
||||
operations = await app.state.memory.list_operations(bank_id, request_context=request_context)
|
||||
return {
|
||||
"bank_id": bank_id,
|
||||
"operations": operations,
|
||||
}
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
@@ -1490,39 +1530,21 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="cancel_operation",
|
||||
tags=["Operations"],
|
||||
)
|
||||
async def api_cancel_operation(bank_id: str, operation_id: str):
|
||||
async def api_cancel_operation(
|
||||
bank_id: str, operation_id: str, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Cancel a pending async operation."""
|
||||
try:
|
||||
# Validate UUID format
|
||||
try:
|
||||
op_uuid = uuid.UUID(operation_id)
|
||||
uuid.UUID(operation_id)
|
||||
except ValueError:
|
||||
raise HTTPException(status_code=400, detail=f"Invalid operation_id format: {operation_id}")
|
||||
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Check if operation exists and belongs to this memory bank
|
||||
result = await conn.fetchrow(
|
||||
"SELECT bank_id FROM async_operations WHERE operation_id = $1 AND bank_id = $2", op_uuid, bank_id
|
||||
)
|
||||
|
||||
if not result:
|
||||
raise HTTPException(
|
||||
status_code=404, detail=f"Operation {operation_id} not found for memory bank {bank_id}"
|
||||
)
|
||||
|
||||
# Delete the operation
|
||||
await conn.execute("DELETE FROM async_operations WHERE operation_id = $1", op_uuid)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"message": f"Operation {operation_id} cancelled",
|
||||
"operation_id": operation_id,
|
||||
"bank_id": bank_id,
|
||||
}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
result = await app.state.memory.cancel_operation(bank_id, operation_id, request_context=request_context)
|
||||
return result
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
@@ -1538,10 +1560,10 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="get_bank_profile",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_get_bank_profile(bank_id: str):
|
||||
async def api_get_bank_profile(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
"""Get memory bank profile (disposition + background)."""
|
||||
try:
|
||||
profile = await app.state.memory.get_bank_profile(bank_id)
|
||||
profile = await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
# Convert DispositionTraits object to dict for Pydantic
|
||||
disposition_dict = (
|
||||
profile["disposition"].model_dump()
|
||||
@@ -1569,14 +1591,18 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="update_bank_disposition",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_update_bank_disposition(bank_id: str, request: UpdateDispositionRequest):
|
||||
async def api_update_bank_disposition(
|
||||
bank_id: str, request: UpdateDispositionRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Update bank disposition traits."""
|
||||
try:
|
||||
# Update disposition
|
||||
await app.state.memory.update_bank_disposition(bank_id, request.disposition.model_dump())
|
||||
await app.state.memory.update_bank_disposition(
|
||||
bank_id, request.disposition.model_dump(), request_context=request_context
|
||||
)
|
||||
|
||||
# Get updated profile
|
||||
profile = await app.state.memory.get_bank_profile(bank_id)
|
||||
profile = await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
disposition_dict = (
|
||||
profile["disposition"].model_dump()
|
||||
if hasattr(profile["disposition"], "model_dump")
|
||||
@@ -1603,11 +1629,13 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="add_bank_background",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_add_bank_background(bank_id: str, request: AddBackgroundRequest):
|
||||
async def api_add_bank_background(
|
||||
bank_id: str, request: AddBackgroundRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Add or merge bank background information. Optionally infer disposition traits."""
|
||||
try:
|
||||
result = await app.state.memory.merge_bank_background(
|
||||
bank_id, request.content, update_disposition=request.update_disposition
|
||||
bank_id, request.content, update_disposition=request.update_disposition, request_context=request_context
|
||||
)
|
||||
|
||||
response = BackgroundResponse(background=result["background"])
|
||||
@@ -1630,51 +1658,31 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="create_or_update_bank",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_create_or_update_bank(bank_id: str, request: CreateBankRequest):
|
||||
async def api_create_or_update_bank(
|
||||
bank_id: str, request: CreateBankRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Create or update an agent with disposition and background."""
|
||||
try:
|
||||
# Get existing profile or create with defaults
|
||||
profile = await app.state.memory.get_bank_profile(bank_id)
|
||||
# Ensure bank exists by getting profile (auto-creates with defaults)
|
||||
await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Update name if provided
|
||||
if request.name is not None:
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE banks
|
||||
SET name = $2,
|
||||
updated_at = NOW()
|
||||
WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
request.name,
|
||||
)
|
||||
profile["name"] = request.name
|
||||
# Update name and/or background if provided
|
||||
if request.name is not None or request.background is not None:
|
||||
await app.state.memory.update_bank(
|
||||
bank_id,
|
||||
name=request.name,
|
||||
background=request.background,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Update disposition if provided
|
||||
if request.disposition is not None:
|
||||
await app.state.memory.update_bank_disposition(bank_id, request.disposition.model_dump())
|
||||
profile["disposition"] = request.disposition.model_dump()
|
||||
|
||||
# Update background if provided (replace, not merge)
|
||||
if request.background is not None:
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE banks
|
||||
SET background = $2,
|
||||
updated_at = NOW()
|
||||
WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
request.background,
|
||||
)
|
||||
profile["background"] = request.background
|
||||
await app.state.memory.update_bank_disposition(
|
||||
bank_id, request.disposition.model_dump(), request_context=request_context
|
||||
)
|
||||
|
||||
# Get final profile
|
||||
final_profile = await app.state.memory.get_bank_profile(bank_id)
|
||||
final_profile = await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
disposition_dict = (
|
||||
final_profile["disposition"].model_dump()
|
||||
if hasattr(final_profile["disposition"], "model_dump")
|
||||
@@ -1702,10 +1710,10 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="delete_bank",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_delete_bank(bank_id: str):
|
||||
async def api_delete_bank(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
"""Delete an entire memory bank and all its data."""
|
||||
try:
|
||||
result = await app.state.memory.delete_bank(bank_id)
|
||||
result = await app.state.memory.delete_bank(bank_id, request_context=request_context)
|
||||
return DeleteResponse(
|
||||
success=True,
|
||||
message=f"Bank '{bank_id}' and all associated data deleted successfully",
|
||||
@@ -1745,7 +1753,9 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="retain_memories",
|
||||
tags=["Memory"],
|
||||
)
|
||||
async def api_retain(bank_id: str, request: RetainRequest):
|
||||
async def api_retain(
|
||||
bank_id: str, request: RetainRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Retain memories with optional async processing."""
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
@@ -1766,43 +1776,25 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
if request.async_:
|
||||
# Async processing: queue task and return immediately
|
||||
operation_id = uuid.uuid4()
|
||||
|
||||
# Insert operation record into database
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
""",
|
||||
operation_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps({"items_count": len(contents)}),
|
||||
)
|
||||
|
||||
# Submit task to background queue
|
||||
await app.state.memory._task_backend.submit_task(
|
||||
result = await app.state.memory.submit_async_retain(bank_id, contents, request_context=request_context)
|
||||
return RetainResponse.model_validate(
|
||||
{
|
||||
"type": "batch_retain",
|
||||
"operation_id": str(operation_id),
|
||||
"success": True,
|
||||
"bank_id": bank_id,
|
||||
"contents": contents,
|
||||
"items_count": result["items_count"],
|
||||
"async": True,
|
||||
}
|
||||
)
|
||||
|
||||
logging.info(
|
||||
f"Retain task queued for bank_id={bank_id}, {len(contents)} items, operation_id={operation_id}"
|
||||
)
|
||||
|
||||
return RetainResponse(success=True, bank_id=bank_id, items_count=len(contents), async_=True)
|
||||
else:
|
||||
# Synchronous processing: wait for completion (record metrics)
|
||||
with metrics.record_operation("retain", bank_id=bank_id):
|
||||
result = await app.state.memory.retain_batch_async(bank_id=bank_id, contents=contents)
|
||||
result = await app.state.memory.retain_batch_async(
|
||||
bank_id=bank_id, contents=contents, request_context=request_context
|
||||
)
|
||||
|
||||
return RetainResponse(success=True, bank_id=bank_id, items_count=len(contents), async_=False)
|
||||
return RetainResponse.model_validate(
|
||||
{"success": True, "bank_id": bank_id, "items_count": len(contents), "async": False}
|
||||
)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
@@ -1810,9 +1802,11 @@ def _register_routes(app: FastAPI):
|
||||
input_summary = []
|
||||
for i, item in enumerate(request.items):
|
||||
content_preview = item.content[:100] + "..." if len(item.content) > 100 else item.content
|
||||
input_summary.append(f" [{i}] content={content_preview!r}, context={item.context}, timestamp={item.timestamp}")
|
||||
input_summary.append(
|
||||
f" [{i}] content={content_preview!r}, context={item.context}, timestamp={item.timestamp}"
|
||||
)
|
||||
input_debug = "\n".join(input_summary)
|
||||
|
||||
|
||||
error_detail = (
|
||||
f"{str(e)}\n\n"
|
||||
f"Input ({len(request.items)} items):\n{input_debug}\n\n"
|
||||
@@ -1832,10 +1826,11 @@ def _register_routes(app: FastAPI):
|
||||
async def api_clear_bank_memories(
|
||||
bank_id: str,
|
||||
type: str | None = Query(None, description="Optional fact type filter (world, experience, opinion)"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Clear memories for a memory bank, optionally filtered by type."""
|
||||
try:
|
||||
await app.state.memory.delete_bank(bank_id, fact_type=type)
|
||||
await app.state.memory.delete_bank(bank_id, fact_type=type, request_context=request_context)
|
||||
|
||||
return DeleteResponse(success=True)
|
||||
except Exception as e:
|
||||
|
||||
@@ -9,6 +9,7 @@ from fastmcp import FastMCP
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
# Configure logging from HINDSIGHT_API_LOG_LEVEL environment variable
|
||||
_log_level_str = os.environ.get("HINDSIGHT_API_LOG_LEVEL", "info").lower()
|
||||
@@ -67,7 +68,11 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
|
||||
"""
|
||||
try:
|
||||
bank_id = get_current_bank_id()
|
||||
await memory.retain_batch_async(bank_id=bank_id, contents=[{"content": content, "context": context}])
|
||||
if bank_id is None:
|
||||
return "Error: No bank_id configured"
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id, contents=[{"content": content, "context": context}], request_context=RequestContext()
|
||||
)
|
||||
return "Memory stored successfully"
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing memory: {e}", exc_info=True)
|
||||
@@ -90,10 +95,16 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
|
||||
"""
|
||||
try:
|
||||
bank_id = get_current_bank_id()
|
||||
if bank_id is None:
|
||||
return "Error: No bank_id configured"
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
search_result = await memory.recall_async(
|
||||
bank_id=bank_id, query=query, fact_type=list(VALID_RECALL_FACT_TYPES), budget=Budget.LOW
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
fact_type=list(VALID_RECALL_FACT_TYPES),
|
||||
budget=Budget.LOW,
|
||||
request_context=RequestContext(),
|
||||
)
|
||||
|
||||
results = [
|
||||
@@ -102,7 +113,7 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
|
||||
"text": fact.text,
|
||||
"type": fact.fact_type,
|
||||
"context": fact.context,
|
||||
"event_date": fact.event_date,
|
||||
"occurred_start": fact.occurred_start,
|
||||
}
|
||||
for fact in search_result.results[:max_results]
|
||||
]
|
||||
|
||||
@@ -33,6 +33,10 @@ ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER"
|
||||
ENV_MCP_LOCAL_BANK_ID = "HINDSIGHT_API_MCP_LOCAL_BANK_ID"
|
||||
ENV_MCP_INSTRUCTIONS = "HINDSIGHT_API_MCP_INSTRUCTIONS"
|
||||
|
||||
# Optimization flags
|
||||
ENV_SKIP_LLM_VERIFICATION = "HINDSIGHT_API_SKIP_LLM_VERIFICATION"
|
||||
ENV_LAZY_RERANKER = "HINDSIGHT_API_LAZY_RERANKER"
|
||||
|
||||
# Default values
|
||||
DEFAULT_DATABASE_URL = "pg0"
|
||||
DEFAULT_LLM_PROVIDER = "openai"
|
||||
@@ -107,6 +111,10 @@ class HindsightConfig:
|
||||
# Recall
|
||||
graph_retriever: str
|
||||
|
||||
# Optimization flags
|
||||
skip_llm_verification: bool
|
||||
lazy_reranker: bool
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> "HindsightConfig":
|
||||
"""Create configuration from environment variables."""
|
||||
@@ -133,6 +141,9 @@ class HindsightConfig:
|
||||
mcp_enabled=os.getenv(ENV_MCP_ENABLED, str(DEFAULT_MCP_ENABLED)).lower() == "true",
|
||||
# Recall
|
||||
graph_retriever=os.getenv(ENV_GRAPH_RETRIEVER, DEFAULT_GRAPH_RETRIEVER),
|
||||
# Optimization flags
|
||||
skip_llm_verification=os.getenv(ENV_SKIP_LLM_VERIFICATION, "false").lower() == "true",
|
||||
lazy_reranker=os.getenv(ENV_LAZY_RERANKER, "false").lower() == "true",
|
||||
)
|
||||
|
||||
def get_llm_base_url(self) -> str:
|
||||
|
||||
@@ -0,0 +1,204 @@
|
||||
"""
|
||||
Daemon mode support for Hindsight API.
|
||||
|
||||
Provides idle timeout and lockfile management for running as a background daemon.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import fcntl
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default daemon configuration
|
||||
DEFAULT_DAEMON_PORT = 8889
|
||||
DEFAULT_IDLE_TIMEOUT = 0 # 0 = no auto-exit (hindsight-embed passes its own timeout)
|
||||
LOCKFILE_PATH = Path.home() / ".hindsight" / "daemon.lock"
|
||||
DAEMON_LOG_PATH = Path.home() / ".hindsight" / "daemon.log"
|
||||
|
||||
|
||||
class IdleTimeoutMiddleware:
|
||||
"""ASGI middleware that tracks activity and exits after idle timeout."""
|
||||
|
||||
def __init__(self, app, idle_timeout: int = DEFAULT_IDLE_TIMEOUT):
|
||||
self.app = app
|
||||
self.idle_timeout = idle_timeout
|
||||
self.last_activity = time.time()
|
||||
self._checker_task = None
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
# Update activity timestamp on each request
|
||||
self.last_activity = time.time()
|
||||
await self.app(scope, receive, send)
|
||||
|
||||
def start_idle_checker(self):
|
||||
"""Start the background task that checks for idle timeout."""
|
||||
self._checker_task = asyncio.create_task(self._check_idle())
|
||||
|
||||
async def _check_idle(self):
|
||||
"""Background task that exits the process after idle timeout."""
|
||||
# If idle_timeout is 0, don't auto-exit
|
||||
if self.idle_timeout <= 0:
|
||||
return
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(30) # Check every 30 seconds
|
||||
idle_time = time.time() - self.last_activity
|
||||
if idle_time > self.idle_timeout:
|
||||
logger.info(f"Idle timeout reached ({self.idle_timeout}s), shutting down daemon")
|
||||
# Give a moment for any in-flight requests
|
||||
await asyncio.sleep(1)
|
||||
os._exit(0)
|
||||
|
||||
|
||||
class DaemonLock:
|
||||
"""
|
||||
File-based lock to prevent multiple daemon instances.
|
||||
|
||||
Uses fcntl.flock for atomic locking on Unix systems.
|
||||
"""
|
||||
|
||||
def __init__(self, lockfile: Path = LOCKFILE_PATH):
|
||||
self.lockfile = lockfile
|
||||
self._fd = None
|
||||
|
||||
def acquire(self) -> bool:
|
||||
"""
|
||||
Try to acquire the daemon lock.
|
||||
|
||||
Returns True if lock acquired, False if another daemon is running.
|
||||
"""
|
||||
self.lockfile.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
self._fd = open(self.lockfile, "w")
|
||||
fcntl.flock(self._fd.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
# Write PID for debugging
|
||||
self._fd.write(str(os.getpid()))
|
||||
self._fd.flush()
|
||||
return True
|
||||
except (IOError, OSError):
|
||||
# Lock is held by another process
|
||||
if self._fd:
|
||||
self._fd.close()
|
||||
self._fd = None
|
||||
return False
|
||||
|
||||
def release(self):
|
||||
"""Release the daemon lock."""
|
||||
if self._fd:
|
||||
try:
|
||||
fcntl.flock(self._fd.fileno(), fcntl.LOCK_UN)
|
||||
self._fd.close()
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
self._fd = None
|
||||
# Remove lockfile
|
||||
try:
|
||||
self.lockfile.unlink()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def is_locked(self) -> bool:
|
||||
"""Check if the lock is held by another process."""
|
||||
if not self.lockfile.exists():
|
||||
return False
|
||||
|
||||
try:
|
||||
fd = open(self.lockfile, "r")
|
||||
fcntl.flock(fd.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
# We got the lock, so no one else has it
|
||||
fcntl.flock(fd.fileno(), fcntl.LOCK_UN)
|
||||
fd.close()
|
||||
return False
|
||||
except (IOError, OSError):
|
||||
return True
|
||||
|
||||
def get_pid(self) -> int | None:
|
||||
"""Get the PID of the daemon holding the lock."""
|
||||
if not self.lockfile.exists():
|
||||
return None
|
||||
try:
|
||||
with open(self.lockfile, "r") as f:
|
||||
return int(f.read().strip())
|
||||
except (ValueError, IOError):
|
||||
return None
|
||||
|
||||
|
||||
def daemonize():
|
||||
"""
|
||||
Fork the current process into a background daemon.
|
||||
|
||||
Uses double-fork technique to properly detach from terminal.
|
||||
"""
|
||||
# First fork
|
||||
pid = os.fork()
|
||||
if pid > 0:
|
||||
# Parent exits
|
||||
sys.exit(0)
|
||||
|
||||
# Create new session
|
||||
os.setsid()
|
||||
|
||||
# Second fork to prevent zombie processes
|
||||
pid = os.fork()
|
||||
if pid > 0:
|
||||
sys.exit(0)
|
||||
|
||||
# Redirect standard file descriptors to log file
|
||||
DAEMON_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
|
||||
# Redirect stdin to /dev/null
|
||||
with open("/dev/null", "r") as devnull:
|
||||
os.dup2(devnull.fileno(), sys.stdin.fileno())
|
||||
|
||||
# Redirect stdout/stderr to log file
|
||||
log_fd = open(DAEMON_LOG_PATH, "a")
|
||||
os.dup2(log_fd.fileno(), sys.stdout.fileno())
|
||||
os.dup2(log_fd.fileno(), sys.stderr.fileno())
|
||||
|
||||
|
||||
def check_daemon_running(port: int = DEFAULT_DAEMON_PORT) -> bool:
|
||||
"""Check if a daemon is running and responsive on the given port."""
|
||||
import socket
|
||||
|
||||
try:
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
sock.settimeout(1)
|
||||
result = sock.connect_ex(("127.0.0.1", port))
|
||||
sock.close()
|
||||
return result == 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def stop_daemon(port: int = DEFAULT_DAEMON_PORT) -> bool:
|
||||
"""Stop a running daemon by sending SIGTERM to the process."""
|
||||
lock = DaemonLock()
|
||||
pid = lock.get_pid()
|
||||
|
||||
if pid is None:
|
||||
return False
|
||||
|
||||
try:
|
||||
import signal
|
||||
|
||||
os.kill(pid, signal.SIGTERM)
|
||||
# Wait for process to exit
|
||||
for _ in range(50): # Wait up to 5 seconds
|
||||
time.sleep(0.1)
|
||||
try:
|
||||
os.kill(pid, 0) # Check if process exists
|
||||
except OSError:
|
||||
return True # Process exited
|
||||
return False
|
||||
except OSError:
|
||||
return False
|
||||
@@ -11,7 +11,13 @@ from .cross_encoder import CrossEncoderModel, LocalSTCrossEncoder, RemoteTEICros
|
||||
from .db_utils import acquire_with_retry
|
||||
from .embeddings import Embeddings, LocalSTEmbeddings, RemoteTEIEmbeddings
|
||||
from .llm_wrapper import LLMConfig
|
||||
from .memory_engine import MemoryEngine
|
||||
from .memory_engine import (
|
||||
MemoryEngine,
|
||||
UnqualifiedTableError,
|
||||
fq_table,
|
||||
get_current_schema,
|
||||
validate_sql_schema,
|
||||
)
|
||||
from .response_models import MemoryFact, RecallResult, ReflectResult
|
||||
from .search.trace import (
|
||||
EntryPoint,
|
||||
@@ -49,4 +55,9 @@ __all__ = [
|
||||
"RecallResult",
|
||||
"ReflectResult",
|
||||
"MemoryFact",
|
||||
# Schema safety utilities
|
||||
"fq_table",
|
||||
"get_current_schema",
|
||||
"validate_sql_schema",
|
||||
"UnqualifiedTableError",
|
||||
]
|
||||
|
||||
@@ -11,6 +11,7 @@ from difflib import SequenceMatcher
|
||||
import asyncpg
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
from .memory_engine import fq_table
|
||||
|
||||
# Load spaCy model (singleton)
|
||||
_nlp = None
|
||||
@@ -68,9 +69,9 @@ class EntityResolver:
|
||||
) -> list[str]:
|
||||
# Query ALL candidates for this bank
|
||||
all_entities = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT canonical_name, id, metadata, last_seen, mention_count
|
||||
FROM entities
|
||||
FROM {fq_table("entities")}
|
||||
WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
@@ -82,11 +83,11 @@ class EntityResolver:
|
||||
# Query ALL co-occurrences for this bank's entities in one query
|
||||
# This builds a map of entity_id -> set of co-occurring entity names
|
||||
all_cooccurrences = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT ec.entity_id_1, ec.entity_id_2, ec.cooccurrence_count
|
||||
FROM entity_cooccurrences ec
|
||||
WHERE ec.entity_id_1 IN (SELECT id FROM entities WHERE bank_id = $1)
|
||||
OR ec.entity_id_2 IN (SELECT id FROM entities WHERE bank_id = $1)
|
||||
FROM {fq_table("entity_cooccurrences")} ec
|
||||
WHERE ec.entity_id_1 IN (SELECT id FROM {fq_table("entities")} WHERE bank_id = $1)
|
||||
OR ec.entity_id_2 IN (SELECT id FROM {fq_table("entities")} WHERE bank_id = $1)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
@@ -195,8 +196,8 @@ class EntityResolver:
|
||||
# Batch update existing entities
|
||||
if entities_to_update:
|
||||
await conn.executemany(
|
||||
"""
|
||||
UPDATE entities SET
|
||||
f"""
|
||||
UPDATE {fq_table("entities")} SET
|
||||
mention_count = mention_count + 1,
|
||||
last_seen = $2
|
||||
WHERE id = $1::uuid
|
||||
@@ -232,13 +233,13 @@ class EntityResolver:
|
||||
# Batch INSERT ... ON CONFLICT with RETURNING
|
||||
# This is much faster than individual inserts
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
INSERT INTO entities (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
f"""
|
||||
INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
SELECT $1, name, event_date, event_date, 1
|
||||
FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date)
|
||||
ON CONFLICT (bank_id, LOWER(canonical_name))
|
||||
DO UPDATE SET
|
||||
mention_count = entities.mention_count + 1,
|
||||
mention_count = {fq_table("entities")}.mention_count + 1,
|
||||
last_seen = EXCLUDED.last_seen
|
||||
RETURNING id
|
||||
""",
|
||||
@@ -279,9 +280,9 @@ class EntityResolver:
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
# Find candidate entities with similar name
|
||||
candidates = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, canonical_name, metadata, last_seen
|
||||
FROM entities
|
||||
FROM {fq_table("entities")}
|
||||
WHERE bank_id = $1
|
||||
AND (
|
||||
canonical_name ILIKE $2
|
||||
@@ -326,10 +327,10 @@ class EntityResolver:
|
||||
# Get entities that co-occurred with this candidate before
|
||||
# Use the materialized co-occurrence cache for fast lookup
|
||||
co_entity_rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT e.canonical_name, ec.cooccurrence_count
|
||||
FROM entity_cooccurrences ec
|
||||
JOIN entities e ON (
|
||||
FROM {fq_table("entity_cooccurrences")} ec
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
CASE
|
||||
WHEN ec.entity_id_1 = $1 THEN ec.entity_id_2
|
||||
WHEN ec.entity_id_2 = $1 THEN ec.entity_id_1
|
||||
@@ -365,8 +366,8 @@ class EntityResolver:
|
||||
if best_score > threshold:
|
||||
# Update entity
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE entities
|
||||
f"""
|
||||
UPDATE {fq_table("entities")}
|
||||
SET mention_count = mention_count + 1,
|
||||
last_seen = $1
|
||||
WHERE id = $2
|
||||
@@ -402,12 +403,12 @@ class EntityResolver:
|
||||
Entity ID
|
||||
"""
|
||||
entity_id = await conn.fetchval(
|
||||
"""
|
||||
INSERT INTO entities (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
f"""
|
||||
INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
VALUES ($1, $2, $3, $4, 1)
|
||||
ON CONFLICT (bank_id, LOWER(canonical_name))
|
||||
DO UPDATE SET
|
||||
mention_count = entities.mention_count + 1,
|
||||
mention_count = {fq_table("entities")}.mention_count + 1,
|
||||
last_seen = EXCLUDED.last_seen
|
||||
RETURNING id
|
||||
""",
|
||||
@@ -430,8 +431,8 @@ class EntityResolver:
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
# Insert unit-entity link
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO unit_entities (unit_id, entity_id)
|
||||
f"""
|
||||
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
@@ -441,9 +442,9 @@ class EntityResolver:
|
||||
|
||||
# Update co-occurrence cache: find other entities in this unit
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT entity_id
|
||||
FROM unit_entities
|
||||
FROM {fq_table("unit_entities")}
|
||||
WHERE unit_id = $1 AND entity_id != $2
|
||||
""",
|
||||
unit_id,
|
||||
@@ -472,12 +473,12 @@ class EntityResolver:
|
||||
entity_id_1, entity_id_2 = entity_id_2, entity_id_1
|
||||
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO entity_cooccurrences (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
f"""
|
||||
INSERT INTO {fq_table("entity_cooccurrences")} (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
VALUES ($1, $2, 1, NOW())
|
||||
ON CONFLICT (entity_id_1, entity_id_2)
|
||||
DO UPDATE SET
|
||||
cooccurrence_count = entity_cooccurrences.cooccurrence_count + 1,
|
||||
cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + 1,
|
||||
last_cooccurred = NOW()
|
||||
""",
|
||||
entity_id_1,
|
||||
@@ -506,8 +507,8 @@ class EntityResolver:
|
||||
async def _link_units_to_entities_batch_impl(self, conn, unit_entity_pairs: list[tuple[str, str]]):
|
||||
# Batch insert all unit-entity links
|
||||
await conn.executemany(
|
||||
"""
|
||||
INSERT INTO unit_entities (unit_id, entity_id)
|
||||
f"""
|
||||
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
@@ -541,12 +542,12 @@ class EntityResolver:
|
||||
if cooccurrence_pairs:
|
||||
now = datetime.now(UTC)
|
||||
await conn.executemany(
|
||||
"""
|
||||
INSERT INTO entity_cooccurrences (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
f"""
|
||||
INSERT INTO {fq_table("entity_cooccurrences")} (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
ON CONFLICT (entity_id_1, entity_id_2)
|
||||
DO UPDATE SET
|
||||
cooccurrence_count = entity_cooccurrences.cooccurrence_count + 1,
|
||||
cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + 1,
|
||||
last_cooccurred = EXCLUDED.last_cooccurred
|
||||
""",
|
||||
[(e1, e2, 1, now) for e1, e2 in cooccurrence_pairs],
|
||||
@@ -565,9 +566,9 @@ class EntityResolver:
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT unit_id
|
||||
FROM unit_entities
|
||||
FROM {fq_table("unit_entities")}
|
||||
WHERE entity_id = $1
|
||||
ORDER BY unit_id
|
||||
LIMIT $2
|
||||
@@ -594,8 +595,8 @@ class EntityResolver:
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT id FROM entities
|
||||
f"""
|
||||
SELECT id FROM {fq_table("entities")}
|
||||
WHERE bank_id = $1
|
||||
AND canonical_name ILIKE $2
|
||||
ORDER BY mention_count DESC
|
||||
|
||||
@@ -0,0 +1,592 @@
|
||||
"""Abstract interface for MemoryEngine public methods.
|
||||
|
||||
This module defines the public API that HTTP endpoints and extensions should use
|
||||
to interact with the memory system. All methods require a RequestContext for
|
||||
authentication when a TenantExtension is configured.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.engine.response_models import RecallResult, ReflectResult
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
class MemoryEngineInterface(ABC):
|
||||
"""
|
||||
Abstract interface for the Memory Engine.
|
||||
|
||||
This defines the public API that should be used by HTTP endpoints and extensions.
|
||||
All methods require a RequestContext for authentication.
|
||||
"""
|
||||
|
||||
# =========================================================================
|
||||
# Health & Status
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def health_check(self) -> dict:
|
||||
"""
|
||||
Check the health of the memory system.
|
||||
|
||||
Returns:
|
||||
Dict with 'status' key ('healthy' or 'unhealthy') and additional info.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Core Memory Operations
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def retain_batch_async(
|
||||
self,
|
||||
bank_id: str,
|
||||
contents: list[dict[str, Any]],
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Retain a batch of memory items.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
contents: List of content dicts with 'content', optional 'event_date',
|
||||
'context', 'metadata', 'document_id'.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with processing results.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def recall_async(
|
||||
self,
|
||||
bank_id: str,
|
||||
query: str,
|
||||
*,
|
||||
budget: "Budget | None" = None,
|
||||
max_tokens: int = 4096,
|
||||
enable_trace: bool = False,
|
||||
fact_type: list[str] | None = None,
|
||||
question_date: datetime | None = None,
|
||||
include_entities: bool = False,
|
||||
max_entity_tokens: int = 500,
|
||||
include_chunks: bool = False,
|
||||
max_chunk_tokens: int = 8192,
|
||||
request_context: "RequestContext",
|
||||
) -> "RecallResult":
|
||||
"""
|
||||
Recall memories relevant to a query.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
query: The search query.
|
||||
budget: Search budget (LOW, MID, HIGH).
|
||||
max_tokens: Maximum tokens in response.
|
||||
enable_trace: Include trace information.
|
||||
fact_type: Filter by fact types.
|
||||
question_date: Context date for temporal relevance.
|
||||
include_entities: Include entity observations.
|
||||
max_entity_tokens: Max tokens for entity observations.
|
||||
include_chunks: Include raw chunks.
|
||||
max_chunk_tokens: Max tokens for chunks.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
RecallResult with matching memories.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def reflect_async(
|
||||
self,
|
||||
bank_id: str,
|
||||
query: str,
|
||||
*,
|
||||
budget: "Budget | None" = None,
|
||||
context: str | None = None,
|
||||
request_context: "RequestContext",
|
||||
) -> "ReflectResult":
|
||||
"""
|
||||
Reflect on a query and generate a thoughtful response.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
query: The question to reflect on.
|
||||
budget: Search budget for retrieving context.
|
||||
context: Additional context for the reflection.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
ReflectResult with generated response and supporting facts.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Bank Management
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def list_banks(
|
||||
self,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
List all memory banks.
|
||||
|
||||
Args:
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
List of bank info dicts.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_bank_profile(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Get bank profile including disposition and background.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Bank profile dict.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def update_bank_disposition(
|
||||
self,
|
||||
bank_id: str,
|
||||
disposition: dict[str, int],
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> None:
|
||||
"""
|
||||
Update bank disposition traits.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
disposition: Dict with trait values.
|
||||
request_context: Request context for authentication.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def merge_bank_background(
|
||||
self,
|
||||
bank_id: str,
|
||||
new_info: str,
|
||||
*,
|
||||
update_disposition: bool = True,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Merge new background information into bank profile.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
new_info: New background information to merge.
|
||||
update_disposition: Whether to infer disposition from background.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Updated background info.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def delete_bank(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
fact_type: str | None = None,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, int]:
|
||||
"""
|
||||
Delete a bank or its memories.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
fact_type: If specified, only delete memories of this type.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with deletion counts.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Memory Units
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def list_memory_units(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
List memory units with pagination.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
fact_type: Filter by fact type.
|
||||
search_query: Full-text search query.
|
||||
limit: Maximum results.
|
||||
offset: Pagination offset.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with 'items', 'total', 'limit', 'offset'.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def delete_memory_unit(
|
||||
self,
|
||||
unit_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Delete a specific memory unit.
|
||||
|
||||
Args:
|
||||
unit_id: The memory unit ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Deletion result.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_graph_data(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
fact_type: str | None = None,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Get graph data for visualization.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
fact_type: Filter by fact type.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with nodes, edges, table_rows, total_units.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Documents
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def list_documents(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
search_query: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
List documents with pagination.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
search_query: Search query.
|
||||
limit: Maximum results.
|
||||
offset: Pagination offset.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with 'items', 'total', 'limit', 'offset'.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_document(
|
||||
self,
|
||||
document_id: str,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get a specific document.
|
||||
|
||||
Args:
|
||||
document_id: The document ID.
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Document dict or None if not found.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def delete_document(
|
||||
self,
|
||||
document_id: str,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, int]:
|
||||
"""
|
||||
Delete a document and its memory units.
|
||||
|
||||
Args:
|
||||
document_id: The document ID.
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with deletion counts.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_chunk(
|
||||
self,
|
||||
chunk_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get a specific chunk.
|
||||
|
||||
Args:
|
||||
chunk_id: The chunk ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Chunk dict or None if not found.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Entities
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def list_entities(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
limit: int = 100,
|
||||
request_context: "RequestContext",
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
List entities for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
limit: Maximum results.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
List of entity dicts.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_entity_observations(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_id: str,
|
||||
*,
|
||||
limit: int = 10,
|
||||
request_context: "RequestContext",
|
||||
) -> list[Any]:
|
||||
"""
|
||||
Get observations for an entity.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
entity_id: The entity ID.
|
||||
limit: Maximum observations.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
List of EntityObservation objects.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def regenerate_entity_observations(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_id: str,
|
||||
entity_name: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> None:
|
||||
"""
|
||||
Regenerate observations for an entity.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
entity_id: The entity ID.
|
||||
entity_name: The entity's canonical name.
|
||||
request_context: Request context for authentication.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Statistics & Operations
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def get_bank_stats(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Get statistics about memory nodes and links for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with node_counts, link_counts, link_counts_by_fact_type,
|
||||
link_breakdown, and operations stats.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_entity(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get entity details including metadata and observations.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
entity_id: The entity ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Entity dict with id, canonical_name, mention_count, first_seen,
|
||||
last_seen, metadata, and observations. None if not found.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def list_operations(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
List async operations for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
List of operation dicts with id, task_type, status, etc.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def cancel_operation(
|
||||
self,
|
||||
bank_id: str,
|
||||
operation_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Cancel a pending async operation.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
operation_id: The operation ID to cancel.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with success status and message.
|
||||
|
||||
Raises:
|
||||
ValueError: If operation not found.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def update_bank(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
name: str | None = None,
|
||||
background: str | None = None,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Update bank name and/or background.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
name: New bank name (optional).
|
||||
background: New background text (optional, replaces existing).
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Updated bank profile dict.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def submit_async_retain(
|
||||
self,
|
||||
bank_id: str,
|
||||
contents: list[dict[str, Any]],
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Submit a batch retain operation to run asynchronously.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
contents: List of content dicts to retain.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with operation_id and items_count.
|
||||
"""
|
||||
...
|
||||
@@ -3,11 +3,13 @@ LLM wrapper for unified configuration across providers.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from google import genai
|
||||
from google.genai import errors as genai_errors
|
||||
from google.genai import types as genai_types
|
||||
@@ -96,7 +98,7 @@ class LLMProvider:
|
||||
client_kwargs = {"api_key": self.api_key, "max_retries": 0}
|
||||
if self.base_url:
|
||||
client_kwargs["base_url"] = self.base_url
|
||||
self._client = AsyncOpenAI(**client_kwargs)
|
||||
self._client = AsyncOpenAI(**client_kwargs) # type: ignore[invalid-argument-type] - dict kwargs
|
||||
self._gemini_client = None
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
@@ -157,7 +159,6 @@ class LLMProvider:
|
||||
"""
|
||||
async with _global_llm_semaphore:
|
||||
start_time = time.time()
|
||||
import json
|
||||
|
||||
# Handle Gemini provider separately
|
||||
if self.provider == "gemini":
|
||||
@@ -165,6 +166,20 @@ class LLMProvider:
|
||||
messages, response_format, max_retries, initial_backoff, max_backoff, skip_validation, start_time
|
||||
)
|
||||
|
||||
# Handle Ollama with native API for structured output (better schema enforcement)
|
||||
if self.provider == "ollama" and response_format is not None:
|
||||
return await self._call_ollama_native(
|
||||
messages,
|
||||
response_format,
|
||||
max_completion_tokens,
|
||||
temperature,
|
||||
max_retries,
|
||||
initial_backoff,
|
||||
max_backoff,
|
||||
skip_validation,
|
||||
start_time,
|
||||
)
|
||||
|
||||
call_params = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
@@ -227,7 +242,7 @@ class LLMProvider:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
|
||||
content = response.choices[0].message.content
|
||||
|
||||
|
||||
# Log raw LLM response for debugging JSON parse issues
|
||||
try:
|
||||
json_data = json.loads(content)
|
||||
@@ -324,6 +339,129 @@ class LLMProvider:
|
||||
raise last_exception
|
||||
raise RuntimeError("LLM call failed after all retries with no exception captured")
|
||||
|
||||
async def _call_ollama_native(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
response_format: Any,
|
||||
max_completion_tokens: int | None,
|
||||
temperature: float | None,
|
||||
max_retries: int,
|
||||
initial_backoff: float,
|
||||
max_backoff: float,
|
||||
skip_validation: bool,
|
||||
start_time: float,
|
||||
) -> Any:
|
||||
"""
|
||||
Call Ollama using native API with JSON schema enforcement.
|
||||
|
||||
Ollama's native API supports passing a full JSON schema in the 'format' parameter,
|
||||
which provides better structured output control than the OpenAI-compatible API.
|
||||
"""
|
||||
# Get the JSON schema from the Pydantic model
|
||||
schema = response_format.model_json_schema() if hasattr(response_format, "model_json_schema") else None
|
||||
|
||||
# Build the base URL for Ollama's native API
|
||||
# Default OpenAI-compatible URL is http://localhost:11434/v1
|
||||
# Native API is at http://localhost:11434/api/chat
|
||||
base_url = self.base_url or "http://localhost:11434/v1"
|
||||
if base_url.endswith("/v1"):
|
||||
native_url = base_url[:-3] + "/api/chat"
|
||||
else:
|
||||
native_url = base_url.rstrip("/") + "/api/chat"
|
||||
|
||||
# Build request payload
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"stream": False,
|
||||
}
|
||||
|
||||
# Add schema as format parameter for structured output
|
||||
if schema:
|
||||
payload["format"] = schema
|
||||
|
||||
# Add optional parameters with optimized defaults for Ollama
|
||||
# Benchmarking shows num_ctx=16384 + num_batch=512 is optimal
|
||||
options = {
|
||||
"num_ctx": 16384, # 16k context window for larger prompts
|
||||
"num_batch": 512, # Optimal batch size for prompt processing
|
||||
}
|
||||
if max_completion_tokens:
|
||||
options["num_predict"] = max_completion_tokens
|
||||
if temperature is not None:
|
||||
options["temperature"] = temperature
|
||||
payload["options"] = options
|
||||
|
||||
last_exception = None
|
||||
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await client.post(native_url, json=payload)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
content = result.get("message", {}).get("content", "")
|
||||
|
||||
# Parse JSON response
|
||||
try:
|
||||
json_data = json.loads(content)
|
||||
except json.JSONDecodeError as json_err:
|
||||
content_preview = content[:500] if content else "<empty>"
|
||||
if content and len(content) > 700:
|
||||
content_preview = f"{content[:500]}...TRUNCATED...{content[-200:]}"
|
||||
logger.warning(
|
||||
f"Ollama JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {json_err}\n"
|
||||
f" Model: ollama/{self.model}\n"
|
||||
f" Content length: {len(content) if content else 0} chars\n"
|
||||
f" Content preview: {content_preview!r}"
|
||||
)
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
last_exception = json_err
|
||||
continue
|
||||
else:
|
||||
raise
|
||||
|
||||
# Validate against Pydantic model or return raw JSON
|
||||
if skip_validation:
|
||||
return json_data
|
||||
else:
|
||||
return response_format.model_validate(json_data)
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"Ollama HTTP error (attempt {attempt + 1}/{max_retries + 1}): {e.response.status_code}"
|
||||
)
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"Ollama HTTP error after {max_retries + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
except httpx.RequestError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(f"Ollama connection error (attempt {attempt + 1}/{max_retries + 1}): {e}")
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"Ollama connection error after {max_retries + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error during Ollama call: {type(e).__name__}: {e}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("Ollama call failed after all retries")
|
||||
|
||||
async def _call_gemini(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
@@ -335,8 +473,6 @@ class LLMProvider:
|
||||
start_time: float,
|
||||
) -> Any:
|
||||
"""Handle Gemini-specific API calls."""
|
||||
import json
|
||||
|
||||
# Convert OpenAI-style messages to Gemini format
|
||||
system_instruction = None
|
||||
gemini_contents = []
|
||||
@@ -467,6 +603,8 @@ class LLMProvider:
|
||||
"""Create provider for memory operations from environment variables."""
|
||||
provider = os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq")
|
||||
api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("HINDSIGHT_API_LLM_API_KEY environment variable is required")
|
||||
base_url = os.getenv("HINDSIGHT_API_LLM_BASE_URL", "")
|
||||
model = os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b")
|
||||
|
||||
@@ -477,6 +615,10 @@ class LLMProvider:
|
||||
"""Create provider for answer generation. Falls back to memory config if not set."""
|
||||
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"))
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_API_ANSWER_LLM_API_KEY environment variable is required"
|
||||
)
|
||||
base_url = os.getenv("HINDSIGHT_API_ANSWER_LLM_BASE_URL", os.getenv("HINDSIGHT_API_LLM_BASE_URL", ""))
|
||||
model = os.getenv("HINDSIGHT_API_ANSWER_LLM_MODEL", os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"))
|
||||
|
||||
@@ -487,6 +629,10 @@ class LLMProvider:
|
||||
"""Create provider for judge/evaluator operations. Falls back to memory config if not set."""
|
||||
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"))
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_API_JUDGE_LLM_API_KEY environment variable is required"
|
||||
)
|
||||
base_url = os.getenv("HINDSIGHT_API_JUDGE_LLM_BASE_URL", os.getenv("HINDSIGHT_API_LLM_BASE_URL", ""))
|
||||
model = os.getenv("HINDSIGHT_API_JUDGE_LLM_MODEL", os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"))
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -10,6 +10,7 @@ from typing import TypedDict
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from ..response_models import DispositionTraits
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -51,9 +52,9 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Try to get existing bank
|
||||
row = await conn.fetchrow(
|
||||
"""
|
||||
f"""
|
||||
SELECT name, disposition, background
|
||||
FROM banks WHERE bank_id = $1
|
||||
FROM {fq_table("banks")} WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
@@ -70,8 +71,8 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
|
||||
|
||||
# Bank doesn't exist, create with defaults
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO banks (bank_id, name, disposition, background)
|
||||
f"""
|
||||
INSERT INTO {fq_table("banks")} (bank_id, name, disposition, background)
|
||||
VALUES ($1, $2, $3::jsonb, $4)
|
||||
ON CONFLICT (bank_id) DO NOTHING
|
||||
""",
|
||||
@@ -98,8 +99,8 @@ async def update_bank_disposition(pool, bank_id: str, disposition: dict[str, int
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE banks
|
||||
f"""
|
||||
UPDATE {fq_table("banks")}
|
||||
SET disposition = $2::jsonb,
|
||||
updated_at = NOW()
|
||||
WHERE bank_id = $1
|
||||
@@ -140,8 +141,8 @@ async def merge_bank_background(pool, llm_config, bank_id: str, new_info: str, u
|
||||
if inferred_disposition:
|
||||
# Update both background and disposition
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE banks
|
||||
f"""
|
||||
UPDATE {fq_table("banks")}
|
||||
SET background = $2,
|
||||
disposition = $3::jsonb,
|
||||
updated_at = NOW()
|
||||
@@ -154,8 +155,8 @@ async def merge_bank_background(pool, llm_config, bank_id: str, new_info: str, u
|
||||
else:
|
||||
# Update only background
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE banks
|
||||
f"""
|
||||
UPDATE {fq_table("banks")}
|
||||
SET background = $2,
|
||||
updated_at = NOW()
|
||||
WHERE bank_id = $1
|
||||
@@ -361,9 +362,9 @@ async def list_banks(pool) -> list:
|
||||
"""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT bank_id, name, disposition, background, created_at, updated_at
|
||||
FROM banks
|
||||
FROM {fq_table("banks")}
|
||||
ORDER BY updated_at DESC
|
||||
"""
|
||||
)
|
||||
|
||||
@@ -6,6 +6,7 @@ Handles storage of document chunks in the database.
|
||||
|
||||
import logging
|
||||
|
||||
from ..memory_engine import fq_table
|
||||
from .types import ChunkMetadata
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -42,8 +43,8 @@ async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[
|
||||
|
||||
# Batch insert all chunks
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO chunks (chunk_id, document_id, bank_id, chunk_text, chunk_index)
|
||||
f"""
|
||||
INSERT INTO {fq_table("chunks")} (chunk_id, document_id, bank_id, chunk_text, chunk_index)
|
||||
SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[], $5::integer[])
|
||||
""",
|
||||
chunk_ids,
|
||||
|
||||
@@ -7,6 +7,7 @@ Handles insertion of facts into the database.
|
||||
import json
|
||||
import logging
|
||||
|
||||
from ..memory_engine import fq_table
|
||||
from .types import ProcessedFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -67,8 +68,8 @@ async def insert_facts_batch(
|
||||
|
||||
# Batch insert all facts
|
||||
results = await conn.fetch(
|
||||
"""
|
||||
INSERT INTO memory_units (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
f"""
|
||||
INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, access_count, metadata, chunk_id, document_id)
|
||||
SELECT $1, * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
@@ -107,8 +108,8 @@ async def ensure_bank_exists(conn, bank_id: str) -> None:
|
||||
bank_id: Bank identifier
|
||||
"""
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO banks (bank_id, disposition, background)
|
||||
f"""
|
||||
INSERT INTO {fq_table("banks")} (bank_id, disposition, background)
|
||||
VALUES ($1, $2::jsonb, $3)
|
||||
ON CONFLICT (bank_id) DO UPDATE
|
||||
SET updated_at = NOW()
|
||||
@@ -141,12 +142,14 @@ async def handle_document_tracking(
|
||||
# Always delete old document first if it exists (cascades to units and links)
|
||||
# Only delete on the first batch to avoid deleting data we just inserted
|
||||
if is_first_batch:
|
||||
await conn.fetchval("DELETE FROM documents WHERE id = $1 AND bank_id = $2 RETURNING id", document_id, bank_id)
|
||||
await conn.fetchval(
|
||||
f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING id", document_id, bank_id
|
||||
)
|
||||
|
||||
# Insert document (or update if exists from concurrent operations)
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO documents (id, bank_id, original_text, content_hash, metadata, retain_params)
|
||||
f"""
|
||||
INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, metadata, retain_params)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
ON CONFLICT (id, bank_id) DO UPDATE
|
||||
SET original_text = EXCLUDED.original_text,
|
||||
|
||||
@@ -7,6 +7,7 @@ import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from uuid import UUID
|
||||
|
||||
from ..memory_engine import fq_table
|
||||
from .types import EntityLink
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -290,9 +291,9 @@ async def extract_entities_batch_optimized(
|
||||
|
||||
entity_id_list = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in all_entity_ids]
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT entity_id, unit_id
|
||||
FROM unit_entities
|
||||
FROM {fq_table("unit_entities")}
|
||||
WHERE entity_id = ANY($1::uuid[])
|
||||
""",
|
||||
entity_id_list,
|
||||
@@ -413,9 +414,9 @@ async def create_temporal_links_batch_per_fact(
|
||||
# Get the event_date for each new unit
|
||||
fetch_dates_start = time_mod.time()
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, event_date
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id::text = ANY($1)
|
||||
""",
|
||||
unit_ids,
|
||||
@@ -432,9 +433,9 @@ async def create_temporal_links_batch_per_fact(
|
||||
|
||||
fetch_neighbors_start = time_mod.time()
|
||||
all_candidates = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, event_date
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND event_date BETWEEN $2 AND $3
|
||||
AND id::text != ALL($4)
|
||||
@@ -479,8 +480,8 @@ async def create_temporal_links_batch_per_fact(
|
||||
if links:
|
||||
insert_start = time_mod.time()
|
||||
await conn.executemany(
|
||||
"""
|
||||
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
f"""
|
||||
INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
|
||||
""",
|
||||
@@ -535,9 +536,9 @@ async def create_semantic_links_batch(
|
||||
# Fetch ALL existing units with embeddings in ONE query
|
||||
fetch_start = time_mod.time()
|
||||
all_existing = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, embedding
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND embedding IS NOT NULL
|
||||
AND id::text != ALL($2)
|
||||
@@ -644,8 +645,8 @@ async def create_semantic_links_batch(
|
||||
if all_links:
|
||||
insert_start = time_mod.time()
|
||||
await conn.executemany(
|
||||
"""
|
||||
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
f"""
|
||||
INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
|
||||
""",
|
||||
@@ -721,8 +722,8 @@ async def insert_entity_links_batch(conn, links: list[EntityLink], chunk_size: i
|
||||
|
||||
# Insert from temp table with ON CONFLICT (single query for all rows)
|
||||
insert_start = time_mod.time()
|
||||
await conn.execute("""
|
||||
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
await conn.execute(f"""
|
||||
INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
SELECT from_unit_id, to_unit_id, link_type, weight, entity_id
|
||||
FROM _temp_entity_links
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
|
||||
@@ -808,8 +809,8 @@ async def create_causal_links_batch(
|
||||
insert_start = time_mod.time()
|
||||
try:
|
||||
await conn.executemany(
|
||||
"""
|
||||
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
f"""
|
||||
INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
|
||||
""",
|
||||
|
||||
@@ -9,6 +9,7 @@ import time
|
||||
import uuid
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from ..memory_engine import fq_table
|
||||
from ..search import observation_utils
|
||||
from . import embedding_utils
|
||||
from .types import EntityLink
|
||||
@@ -75,8 +76,8 @@ async def regenerate_observations_batch(
|
||||
|
||||
# Batch query for entity names
|
||||
entity_rows = await conn.fetch(
|
||||
"""
|
||||
SELECT id, canonical_name FROM entities
|
||||
f"""
|
||||
SELECT id, canonical_name FROM {fq_table("entities")}
|
||||
WHERE id = ANY($1) AND bank_id = $2
|
||||
""",
|
||||
entity_uuids,
|
||||
@@ -86,10 +87,10 @@ async def regenerate_observations_batch(
|
||||
|
||||
# Batch query for fact counts
|
||||
fact_counts = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT ue.entity_id, COUNT(*) as cnt
|
||||
FROM unit_entities ue
|
||||
JOIN memory_units mu ON ue.unit_id = mu.id
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("memory_units")} mu ON ue.unit_id = mu.id
|
||||
WHERE ue.entity_id = ANY($1) AND mu.bank_id = $2
|
||||
GROUP BY ue.entity_id
|
||||
""",
|
||||
@@ -154,10 +155,10 @@ async def _regenerate_entity_observations(
|
||||
|
||||
# Get all facts mentioning this entity (exclude observations themselves)
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.fact_type
|
||||
FROM memory_units mu
|
||||
JOIN unit_entities ue ON mu.id = ue.unit_id
|
||||
FROM {fq_table("memory_units")} mu
|
||||
JOIN {fq_table("unit_entities")} ue ON mu.id = ue.unit_id
|
||||
WHERE mu.bank_id = $1
|
||||
AND ue.entity_id = $2
|
||||
AND mu.fact_type IN ('world', 'experience')
|
||||
@@ -193,12 +194,12 @@ async def _regenerate_entity_observations(
|
||||
|
||||
# Delete old observations for this entity
|
||||
await conn.execute(
|
||||
"""
|
||||
DELETE FROM memory_units
|
||||
f"""
|
||||
DELETE FROM {fq_table("memory_units")}
|
||||
WHERE id IN (
|
||||
SELECT mu.id
|
||||
FROM memory_units mu
|
||||
JOIN unit_entities ue ON mu.id = ue.unit_id
|
||||
FROM {fq_table("memory_units")} mu
|
||||
JOIN {fq_table("unit_entities")} ue ON mu.id = ue.unit_id
|
||||
WHERE mu.bank_id = $1
|
||||
AND mu.fact_type = 'observation'
|
||||
AND ue.entity_id = $2
|
||||
@@ -217,8 +218,8 @@ async def _regenerate_entity_observations(
|
||||
|
||||
for obs_text, embedding in zip(observations, embeddings):
|
||||
result = await conn.fetchrow(
|
||||
"""
|
||||
INSERT INTO memory_units (
|
||||
f"""
|
||||
INSERT INTO {fq_table("memory_units")} (
|
||||
bank_id, text, embedding, context, event_date,
|
||||
occurred_start, occurred_end, mentioned_at,
|
||||
fact_type, access_count
|
||||
@@ -240,8 +241,8 @@ async def _regenerate_entity_observations(
|
||||
|
||||
# Link observation to entity
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO unit_entities (unit_id, entity_id)
|
||||
f"""
|
||||
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
||||
VALUES ($1, $2)
|
||||
""",
|
||||
uuid.UUID(obs_id),
|
||||
|
||||
@@ -8,7 +8,6 @@ import logging
|
||||
import time
|
||||
import uuid
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from . import bank_utils
|
||||
@@ -29,7 +28,7 @@ from . import (
|
||||
link_creation,
|
||||
observation_regeneration,
|
||||
)
|
||||
from .types import ExtractedFact, ProcessedFact, RetainContent
|
||||
from .types import ExtractedFact, ProcessedFact, RetainContent, RetainContentDict
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -43,7 +42,7 @@ async def retain_batch(
|
||||
format_date_fn,
|
||||
duplicate_checker_fn,
|
||||
bank_id: str,
|
||||
contents_dicts: list[dict[str, Any]],
|
||||
contents_dicts: list[RetainContentDict],
|
||||
document_id: str | None = None,
|
||||
is_first_batch: bool = True,
|
||||
fact_type_override: str | None = None,
|
||||
|
||||
@@ -7,9 +7,33 @@ from content input to fact storage.
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from typing import TypedDict
|
||||
from uuid import UUID
|
||||
|
||||
|
||||
class RetainContentDict(TypedDict, total=False):
|
||||
"""Type definition for content items in retain_batch_async.
|
||||
|
||||
Fields:
|
||||
content: Text content to store (required)
|
||||
context: Context about the content (optional)
|
||||
event_date: When the content occurred (optional, defaults to now)
|
||||
metadata: Custom key-value metadata (optional)
|
||||
document_id: Document ID for this content item (optional)
|
||||
"""
|
||||
|
||||
content: str # Required
|
||||
context: str
|
||||
event_date: datetime
|
||||
metadata: dict[str, str]
|
||||
document_id: str
|
||||
|
||||
|
||||
def _now_utc() -> datetime:
|
||||
"""Factory function for default event_date."""
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainContent:
|
||||
"""
|
||||
@@ -20,16 +44,9 @@ class RetainContent:
|
||||
|
||||
content: str
|
||||
context: str = ""
|
||||
event_date: datetime | None = None
|
||||
event_date: datetime = field(default_factory=_now_utc)
|
||||
metadata: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self):
|
||||
"""Ensure event_date is set."""
|
||||
if self.event_date is None:
|
||||
from datetime import datetime
|
||||
|
||||
self.event_date = datetime.now(UTC)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChunkMetadata:
|
||||
|
||||
@@ -10,6 +10,7 @@ import logging
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from .types import RetrievalResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -139,11 +140,11 @@ class BFSGraphRetriever(GraphRetriever):
|
||||
|
||||
# Step 1: Find entry points
|
||||
entry_points = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
||||
1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND fact_type = $3
|
||||
@@ -188,13 +189,13 @@ class BFSGraphRetriever(GraphRetriever):
|
||||
if batch_nodes and budget_remaining > 0:
|
||||
max_neighbors = len(batch_nodes) * 20
|
||||
neighbors = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.occurred_end,
|
||||
mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type,
|
||||
mu.document_id, mu.chunk_id,
|
||||
ml.weight, ml.link_type, ml.from_unit_id
|
||||
FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.to_unit_id = mu.id
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.weight >= $2
|
||||
AND mu.fact_type = $3
|
||||
|
||||
@@ -20,6 +20,7 @@ from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from .graph_retrieval import GraphRetriever
|
||||
from .types import RetrievalResult
|
||||
|
||||
@@ -217,10 +218,10 @@ async def load_typed_adjacency(pool, bank_id: str) -> TypedAdjacency:
|
||||
"""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT ml.from_unit_id, ml.to_unit_id, ml.link_type, ml.weight
|
||||
FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.from_unit_id = mu.id
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
|
||||
WHERE mu.bank_id = $1
|
||||
AND ml.weight >= 0.1
|
||||
ORDER BY ml.from_unit_id, ml.weight DESC
|
||||
@@ -252,10 +253,10 @@ async def fetch_memory_units_by_ids(
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, access_count, embedding, fact_type, document_id, chunk_id
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND fact_type = $2
|
||||
""",
|
||||
@@ -418,9 +419,9 @@ class MPFPGraphRetriever(GraphRetriever):
|
||||
"""Fallback: find semantic seeds via embedding search."""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, 1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND fact_type = $3
|
||||
|
||||
@@ -26,6 +26,23 @@ class CrossEncoderReranker:
|
||||
|
||||
cross_encoder = create_cross_encoder_from_env()
|
||||
self.cross_encoder = cross_encoder
|
||||
self._initialized = False
|
||||
|
||||
async def ensure_initialized(self):
|
||||
"""Ensure the cross-encoder model is initialized (for lazy initialization)."""
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
import asyncio
|
||||
|
||||
cross_encoder = self.cross_encoder
|
||||
# For local providers, run in thread pool to avoid blocking event loop
|
||||
if cross_encoder.provider_name == "local":
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, lambda: asyncio.run(cross_encoder.initialize()))
|
||||
else:
|
||||
await cross_encoder.initialize()
|
||||
self._initialized = True
|
||||
|
||||
def rerank(self, query: str, candidates: list[MergedCandidate]) -> list[ScoredResult]:
|
||||
"""
|
||||
|
||||
@@ -16,6 +16,7 @@ from typing import Optional
|
||||
|
||||
from ...config import get_config
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from .graph_retrieval import BFSGraphRetriever, GraphRetriever
|
||||
from .mpfp_retrieval import MPFPGraphRetriever
|
||||
from .types import RetrievalResult
|
||||
@@ -80,10 +81,10 @@ async def retrieve_semantic(
|
||||
List of RetrievalResult objects
|
||||
"""
|
||||
results = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
||||
1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND fact_type = $3
|
||||
@@ -131,10 +132,10 @@ async def retrieve_bm25(conn, query_text: str, bank_id: str, fact_type: str, lim
|
||||
query_tsquery = " | ".join(tokens)
|
||||
|
||||
results = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
||||
ts_rank_cd(search_vector, to_tsquery('english', $1)) AS bm25_score
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = $3
|
||||
AND search_vector @@ to_tsquery('english', $1)
|
||||
@@ -188,10 +189,10 @@ async def retrieve_temporal(
|
||||
end_date = end_date.replace(tzinfo=UTC)
|
||||
|
||||
entry_points = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
||||
1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = $3
|
||||
AND embedding IS NOT NULL
|
||||
@@ -272,12 +273,12 @@ async def retrieve_temporal(
|
||||
# Get neighbors via temporal and causal links
|
||||
if budget_remaining > 0:
|
||||
neighbors = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type, mu.document_id, mu.chunk_id,
|
||||
ml.weight, ml.link_type,
|
||||
1 - (mu.embedding <=> $1::vector) AS similarity
|
||||
FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.to_unit_id = mu.id
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = $2
|
||||
AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents')
|
||||
AND ml.weight >= 0.1
|
||||
@@ -546,11 +547,11 @@ async def _get_temporal_entry_points(
|
||||
end_date = end_date.replace(tzinfo=UTC)
|
||||
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
access_count, embedding, fact_type, document_id, chunk_id,
|
||||
1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = $3
|
||||
AND embedding IS NOT NULL
|
||||
|
||||
@@ -101,7 +101,7 @@ def build_think_prompt(
|
||||
name: str,
|
||||
disposition: DispositionTraits,
|
||||
background: str,
|
||||
context: str = None,
|
||||
context: str | None = None,
|
||||
) -> str:
|
||||
"""Build the think prompt for the LLM."""
|
||||
disposition_desc = build_disposition_description(disposition)
|
||||
|
||||
@@ -115,7 +115,7 @@ class SearchTracer:
|
||||
node_id: str,
|
||||
text: str,
|
||||
context: str,
|
||||
event_date: datetime,
|
||||
event_date: datetime | None,
|
||||
access_count: int,
|
||||
is_entry_point: bool,
|
||||
parent_node_id: str | None,
|
||||
|
||||
@@ -89,6 +89,38 @@ class TaskBackend(ABC):
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
class SyncTaskBackend(TaskBackend):
|
||||
"""
|
||||
Synchronous task backend that executes tasks immediately.
|
||||
|
||||
This is useful for embedded/CLI usage where we don't want background
|
||||
workers that prevent clean exit. Tasks are executed inline rather than
|
||||
being queued.
|
||||
"""
|
||||
|
||||
async def initialize(self):
|
||||
"""No-op for sync backend."""
|
||||
self._initialized = True
|
||||
logger.debug("SyncTaskBackend initialized")
|
||||
|
||||
async def submit_task(self, task_dict: dict[str, Any]):
|
||||
"""
|
||||
Execute the task immediately (synchronously).
|
||||
|
||||
Args:
|
||||
task_dict: Task dictionary to execute
|
||||
"""
|
||||
if not self._initialized:
|
||||
await self.initialize()
|
||||
|
||||
await self._execute_task(task_dict)
|
||||
|
||||
async def shutdown(self):
|
||||
"""No-op for sync backend."""
|
||||
self._initialized = False
|
||||
logger.debug("SyncTaskBackend shutdown")
|
||||
|
||||
|
||||
class AsyncIOQueueBackend(TaskBackend):
|
||||
"""
|
||||
Task backend implementation using asyncio queues.
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
"""
|
||||
Hindsight Extensions System.
|
||||
|
||||
Extensions allow customizing and extending Hindsight behavior without modifying core code.
|
||||
Extensions are loaded via environment variables pointing to implementation classes.
|
||||
|
||||
Example:
|
||||
HINDSIGHT_API_OPERATION_VALIDATOR_EXTENSION=mypackage.validators:MyValidator
|
||||
HINDSIGHT_API_OPERATION_VALIDATOR_MAX_RETRIES=3
|
||||
|
||||
HINDSIGHT_API_HTTP_EXTENSION=mypackage.http:MyHttpExtension
|
||||
HINDSIGHT_API_HTTP_SOME_CONFIG=value
|
||||
|
||||
Extensions receive an ExtensionContext that provides a controlled API for interacting
|
||||
with the system (e.g., running migrations for tenant schemas).
|
||||
"""
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
from hindsight_api.extensions.builtin import ApiKeyTenantExtension
|
||||
from hindsight_api.extensions.context import DefaultExtensionContext, ExtensionContext
|
||||
from hindsight_api.extensions.http import HttpExtension
|
||||
from hindsight_api.extensions.loader import load_extension
|
||||
from hindsight_api.extensions.operation_validator import (
|
||||
OperationValidationError,
|
||||
OperationValidatorExtension,
|
||||
RecallContext,
|
||||
RecallResult,
|
||||
ReflectContext,
|
||||
ReflectResultContext,
|
||||
RetainContext,
|
||||
RetainResult,
|
||||
ValidationResult,
|
||||
)
|
||||
from hindsight_api.extensions.tenant import (
|
||||
AuthenticationError,
|
||||
TenantContext,
|
||||
TenantExtension,
|
||||
)
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
__all__ = [
|
||||
# Base
|
||||
"Extension",
|
||||
"load_extension",
|
||||
# Context
|
||||
"ExtensionContext",
|
||||
"DefaultExtensionContext",
|
||||
# HTTP Extension
|
||||
"HttpExtension",
|
||||
# Operation Validator
|
||||
"OperationValidationError",
|
||||
"OperationValidatorExtension",
|
||||
"RecallContext",
|
||||
"RecallResult",
|
||||
"ReflectContext",
|
||||
"ReflectResultContext",
|
||||
"RetainContext",
|
||||
"RetainResult",
|
||||
"ValidationResult",
|
||||
# Tenant/Auth
|
||||
"ApiKeyTenantExtension",
|
||||
"AuthenticationError",
|
||||
"RequestContext",
|
||||
"TenantContext",
|
||||
"TenantExtension",
|
||||
]
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Base Extension class for all Hindsight extensions."""
|
||||
|
||||
from abc import ABC
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.extensions.context import ExtensionContext
|
||||
|
||||
|
||||
class Extension(ABC):
|
||||
"""
|
||||
Base class for all Hindsight extensions.
|
||||
|
||||
Extensions are loaded via environment variables and receive configuration
|
||||
from prefixed environment variables.
|
||||
|
||||
Example:
|
||||
HINDSIGHT_API_MY_EXTENSION=mypackage.ext:MyExtension
|
||||
HINDSIGHT_API_MY_SOME_CONFIG=value
|
||||
|
||||
The extension receives: {"some_config": "value"}
|
||||
|
||||
Extensions also receive an ExtensionContext that provides a controlled API
|
||||
for interacting with the system (e.g., running migrations for tenant schemas).
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict[str, str]):
|
||||
"""
|
||||
Initialize the extension with configuration.
|
||||
|
||||
Args:
|
||||
config: Dictionary of configuration values from environment variables.
|
||||
Keys are lowercased with the prefix stripped.
|
||||
"""
|
||||
self.config = config
|
||||
self._context: "ExtensionContext | None" = None
|
||||
|
||||
def set_context(self, context: "ExtensionContext") -> None:
|
||||
"""
|
||||
Set the extension context.
|
||||
|
||||
Called by the extension loader after instantiation.
|
||||
Extensions should not call this directly.
|
||||
|
||||
Args:
|
||||
context: The ExtensionContext providing system APIs.
|
||||
"""
|
||||
self._context = context
|
||||
|
||||
@property
|
||||
def context(self) -> "ExtensionContext":
|
||||
"""
|
||||
Get the extension context.
|
||||
|
||||
Returns:
|
||||
The ExtensionContext providing system APIs.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If context has not been set yet.
|
||||
"""
|
||||
if self._context is None:
|
||||
raise RuntimeError(
|
||||
"Extension context not set. Context is available after the extension is loaded by the system."
|
||||
)
|
||||
return self._context
|
||||
|
||||
async def on_startup(self) -> None:
|
||||
"""
|
||||
Called when the application starts.
|
||||
|
||||
Override to perform initialization tasks like connecting to external services.
|
||||
"""
|
||||
pass
|
||||
|
||||
async def on_shutdown(self) -> None:
|
||||
"""
|
||||
Called when the application shuts down.
|
||||
|
||||
Override to perform cleanup tasks like closing connections.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,18 @@
|
||||
"""
|
||||
Built-in extension implementations.
|
||||
|
||||
These are ready-to-use implementations of the extension interfaces.
|
||||
They can be used directly or serve as examples for custom implementations.
|
||||
|
||||
Available built-in extensions:
|
||||
- ApiKeyTenantExtension: Simple API key validation with public schema
|
||||
|
||||
Example usage:
|
||||
HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
|
||||
"""
|
||||
|
||||
from hindsight_api.extensions.builtin.tenant import ApiKeyTenantExtension
|
||||
|
||||
__all__ = [
|
||||
"ApiKeyTenantExtension",
|
||||
]
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Built-in tenant extension implementations."""
|
||||
|
||||
from hindsight_api.extensions.tenant import AuthenticationError, TenantContext, TenantExtension
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
class ApiKeyTenantExtension(TenantExtension):
|
||||
"""
|
||||
Built-in tenant extension that validates API key against an environment variable.
|
||||
|
||||
This is a simple implementation that:
|
||||
1. Validates the API key matches HINDSIGHT_API_TENANT_API_KEY
|
||||
2. Returns 'public' as the schema for all authenticated requests
|
||||
|
||||
Configuration:
|
||||
HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
|
||||
HINDSIGHT_API_TENANT_API_KEY=your-secret-key
|
||||
|
||||
For multi-tenant setups with separate schemas per tenant, implement a custom
|
||||
TenantExtension that looks up the schema based on the API key or token claims.
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict[str, str]):
|
||||
super().__init__(config)
|
||||
self.expected_api_key = config.get("api_key")
|
||||
if not self.expected_api_key:
|
||||
raise ValueError("HINDSIGHT_API_TENANT_API_KEY is required when using ApiKeyTenantExtension")
|
||||
|
||||
async def authenticate(self, context: RequestContext) -> TenantContext:
|
||||
"""Validate API key and return public schema context."""
|
||||
if context.api_key != self.expected_api_key:
|
||||
raise AuthenticationError("Invalid API key")
|
||||
return TenantContext(schema_name="public")
|
||||
@@ -0,0 +1,110 @@
|
||||
"""Extension context providing a controlled API for extensions to interact with the system."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.interface import MemoryEngineInterface
|
||||
|
||||
|
||||
class ExtensionContext(ABC):
|
||||
"""
|
||||
Abstract context providing a controlled API for extensions.
|
||||
|
||||
Extensions receive this context instead of direct access to internal
|
||||
components like MemoryEngine or database connections. This provides:
|
||||
- A stable API that won't break when internals change
|
||||
- Security by limiting what extensions can access
|
||||
- Clear documentation of what extensions can do
|
||||
|
||||
Built-in implementation:
|
||||
hindsight_api.extensions.builtin.context.DefaultExtensionContext
|
||||
|
||||
Example usage in an extension:
|
||||
class MyTenantExtension(TenantExtension):
|
||||
async def on_startup(self) -> None:
|
||||
# Run migrations for a new tenant schema
|
||||
await self.context.run_migration("tenant_acme")
|
||||
|
||||
class MyHttpExtension(HttpExtension):
|
||||
def get_router(self, memory):
|
||||
# Use memory engine for custom endpoints
|
||||
engine = self.context.get_memory_engine()
|
||||
...
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def run_migration(self, schema: str) -> None:
|
||||
"""
|
||||
Run database migrations for a specific schema.
|
||||
|
||||
This creates the schema if it doesn't exist and runs all pending
|
||||
migrations. Uses advisory locks to coordinate between distributed workers.
|
||||
|
||||
Args:
|
||||
schema: PostgreSQL schema name (e.g., "tenant_acme").
|
||||
The schema will be created if it doesn't exist.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If migrations fail to complete.
|
||||
|
||||
Example:
|
||||
# Provision a new tenant schema
|
||||
await context.run_migration("tenant_acme")
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def get_memory_engine(self) -> "MemoryEngineInterface":
|
||||
"""
|
||||
Get the memory engine interface.
|
||||
|
||||
Returns the MemoryEngineInterface for performing memory operations
|
||||
like retain, recall, reflect, and entity/document management.
|
||||
|
||||
Returns:
|
||||
MemoryEngineInterface instance.
|
||||
|
||||
Example:
|
||||
engine = context.get_memory_engine()
|
||||
result = await engine.recall_async(bank_id, query)
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class DefaultExtensionContext(ExtensionContext):
|
||||
"""
|
||||
Default implementation of ExtensionContext.
|
||||
|
||||
Uses the system's database URL and migration infrastructure.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
database_url: str,
|
||||
memory_engine: "MemoryEngineInterface | None" = None,
|
||||
):
|
||||
"""
|
||||
Initialize the context.
|
||||
|
||||
Args:
|
||||
database_url: SQLAlchemy database URL for migrations.
|
||||
memory_engine: Optional MemoryEngine instance for memory operations.
|
||||
"""
|
||||
self._database_url = database_url
|
||||
self._memory_engine = memory_engine
|
||||
|
||||
async def run_migration(self, schema: str) -> None:
|
||||
"""Run migrations for a specific schema."""
|
||||
from hindsight_api.migrations import run_migrations
|
||||
|
||||
run_migrations(self._database_url, schema=schema)
|
||||
|
||||
def get_memory_engine(self) -> "MemoryEngineInterface":
|
||||
"""Get the memory engine interface."""
|
||||
if self._memory_engine is None:
|
||||
raise RuntimeError(
|
||||
"Memory engine not configured in ExtensionContext. "
|
||||
"Ensure the context was created with a memory_engine parameter."
|
||||
)
|
||||
return self._memory_engine
|
||||
@@ -0,0 +1,89 @@
|
||||
"""
|
||||
HTTP Extension for adding custom endpoints to the Hindsight API.
|
||||
|
||||
This extension allows adding custom HTTP endpoints under the /ext/ path prefix.
|
||||
The extension provides a FastAPI router that is mounted on the main application.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi import APIRouter
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api import MemoryEngine
|
||||
|
||||
|
||||
class HttpExtension(Extension, ABC):
|
||||
"""
|
||||
Base class for HTTP extensions that add custom API endpoints.
|
||||
|
||||
HTTP extensions provide a FastAPI router that gets mounted under /ext/.
|
||||
The extension has full control over the routes, request/response models, and handlers.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from fastapi import APIRouter
|
||||
from hindsight_api.extensions import HttpExtension
|
||||
|
||||
class MyHttpExtension(HttpExtension):
|
||||
def get_router(self, memory: MemoryEngine) -> APIRouter:
|
||||
router = APIRouter()
|
||||
|
||||
@router.get("/hello")
|
||||
async def hello():
|
||||
return {"message": "Hello from extension!"}
|
||||
|
||||
@router.post("/custom/{bank_id}/action")
|
||||
async def custom_action(bank_id: str):
|
||||
# Access memory engine for database operations
|
||||
pool = await memory._get_pool()
|
||||
# ... custom logic
|
||||
return {"status": "ok"}
|
||||
|
||||
return router
|
||||
```
|
||||
|
||||
The routes will be available at:
|
||||
- GET /ext/hello
|
||||
- POST /ext/custom/{bank_id}/action
|
||||
|
||||
Configuration via environment variables:
|
||||
HINDSIGHT_API_HTTP_EXTENSION=mypackage.ext:MyHttpExtension
|
||||
HINDSIGHT_API_HTTP_SOME_CONFIG=value
|
||||
|
||||
The extension receives config: {"some_config": "value"}
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_router(self, memory: "MemoryEngine") -> APIRouter:
|
||||
"""
|
||||
Return a FastAPI router with custom endpoints.
|
||||
|
||||
The router will be mounted at /ext/ on the main application.
|
||||
All routes defined in the router will be prefixed with /ext/.
|
||||
|
||||
Args:
|
||||
memory: The MemoryEngine instance for database access and core operations.
|
||||
Use this to access the connection pool, run queries, or call
|
||||
memory operations like retain, recall, etc.
|
||||
|
||||
Returns:
|
||||
A FastAPI APIRouter with the custom endpoints defined.
|
||||
|
||||
Example:
|
||||
```python
|
||||
def get_router(self, memory: MemoryEngine) -> APIRouter:
|
||||
router = APIRouter(tags=["My Extension"])
|
||||
|
||||
@router.get("/status")
|
||||
async def status():
|
||||
health = await memory.health_check()
|
||||
return {"extension": "healthy", "memory": health}
|
||||
|
||||
return router
|
||||
```
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,125 @@
|
||||
"""Extension loader utilities."""
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
import os
|
||||
from typing import TYPE_CHECKING, TypeVar
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.extensions.context import ExtensionContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T", bound=Extension)
|
||||
|
||||
|
||||
class ExtensionLoadError(Exception):
|
||||
"""Raised when an extension fails to load."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def load_extension(
|
||||
prefix: str,
|
||||
base_class: type[T],
|
||||
env_prefix: str = "HINDSIGHT_API",
|
||||
context: "ExtensionContext | None" = None,
|
||||
) -> T | None:
|
||||
"""
|
||||
Load an extension from environment variable configuration.
|
||||
|
||||
The extension class is specified via {env_prefix}_{prefix}_EXTENSION environment
|
||||
variable in the format "module.path:ClassName".
|
||||
|
||||
Configuration for the extension is collected from all environment variables
|
||||
matching {env_prefix}_{prefix}_* (excluding the EXTENSION variable itself).
|
||||
|
||||
Args:
|
||||
prefix: The extension prefix (e.g., "OPERATION_VALIDATOR").
|
||||
base_class: The base class that the extension must inherit from.
|
||||
env_prefix: The environment variable prefix (default: "HINDSIGHT_API").
|
||||
context: Optional ExtensionContext to provide system APIs to the extension.
|
||||
|
||||
Returns:
|
||||
An instance of the extension, or None if not configured.
|
||||
|
||||
Raises:
|
||||
ExtensionLoadError: If the extension fails to load or validate.
|
||||
|
||||
Example:
|
||||
HINDSIGHT_API_OPERATION_VALIDATOR_EXTENSION=mypackage.validators:MyValidator
|
||||
HINDSIGHT_API_OPERATION_VALIDATOR_MAX_REQUESTS=100
|
||||
|
||||
ext = load_extension("OPERATION_VALIDATOR", OperationValidatorExtension)
|
||||
# ext.config == {"max_requests": "100"}
|
||||
"""
|
||||
env_var = f"{env_prefix}_{prefix}_EXTENSION"
|
||||
ext_path = os.getenv(env_var)
|
||||
|
||||
if not ext_path:
|
||||
logger.debug(f"No extension configured for {env_var}")
|
||||
return None
|
||||
|
||||
logger.info(f"Loading extension from {env_var}={ext_path}")
|
||||
|
||||
# Parse "module.path:ClassName"
|
||||
if ":" not in ext_path:
|
||||
raise ExtensionLoadError(f"Invalid extension path '{ext_path}'. Expected format: 'module.path:ClassName'")
|
||||
|
||||
module_path, class_name = ext_path.rsplit(":", 1)
|
||||
|
||||
# Import the module
|
||||
try:
|
||||
module = importlib.import_module(module_path)
|
||||
except ImportError as e:
|
||||
raise ExtensionLoadError(f"Failed to import extension module '{module_path}': {e}") from e
|
||||
|
||||
# Get the class
|
||||
try:
|
||||
ext_class = getattr(module, class_name)
|
||||
except AttributeError as e:
|
||||
raise ExtensionLoadError(f"Extension class '{class_name}' not found in module '{module_path}'") from e
|
||||
|
||||
# Validate inheritance
|
||||
if not isinstance(ext_class, type) or not issubclass(ext_class, base_class):
|
||||
raise ExtensionLoadError(f"Extension class '{ext_class.__name__}' must inherit from '{base_class.__name__}'")
|
||||
|
||||
# Collect configuration from environment variables
|
||||
config = _collect_config(env_prefix, prefix)
|
||||
|
||||
logger.info(f"Loaded extension {ext_class.__name__} with config keys: {list(config.keys())}")
|
||||
|
||||
# Instantiate the extension
|
||||
try:
|
||||
extension = ext_class(config)
|
||||
except Exception as e:
|
||||
raise ExtensionLoadError(f"Failed to instantiate extension '{ext_class.__name__}': {e}") from e
|
||||
|
||||
# Set the context if provided
|
||||
if context is not None:
|
||||
extension.set_context(context)
|
||||
logger.debug(f"Set context on extension {ext_class.__name__}")
|
||||
|
||||
return extension
|
||||
|
||||
|
||||
def _collect_config(env_prefix: str, prefix: str) -> dict[str, str]:
|
||||
"""
|
||||
Collect configuration from environment variables.
|
||||
|
||||
Collects all variables matching {env_prefix}_{prefix}_* except for
|
||||
{env_prefix}_{prefix}_EXTENSION, strips the prefix, and lowercases keys.
|
||||
"""
|
||||
config = {}
|
||||
full_prefix = f"{env_prefix}_{prefix}_"
|
||||
extension_var = f"{full_prefix}EXTENSION"
|
||||
|
||||
for key, value in os.environ.items():
|
||||
if key.startswith(full_prefix) and key != extension_var:
|
||||
# Strip prefix and lowercase the key
|
||||
config_key = key[len(full_prefix) :].lower()
|
||||
config[config_key] = value
|
||||
|
||||
return config
|
||||
@@ -0,0 +1,325 @@
|
||||
"""Operation Validator Extension for validating retain/recall/reflect operations."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.engine.response_models import RecallResult as RecallResultModel
|
||||
from hindsight_api.engine.response_models import ReflectResult
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
class OperationValidationError(Exception):
|
||||
"""Raised when an operation fails validation."""
|
||||
|
||||
def __init__(self, reason: str):
|
||||
self.reason = reason
|
||||
super().__init__(f"Operation validation failed: {reason}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ValidationResult:
|
||||
"""Result of an operation validation."""
|
||||
|
||||
allowed: bool
|
||||
reason: str | None = None
|
||||
|
||||
@classmethod
|
||||
def accept(cls) -> "ValidationResult":
|
||||
"""Create an accepted validation result."""
|
||||
return cls(allowed=True)
|
||||
|
||||
@classmethod
|
||||
def reject(cls, reason: str) -> "ValidationResult":
|
||||
"""Create a rejected validation result with a reason."""
|
||||
return cls(allowed=False, reason=reason)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Pre-operation Contexts (all user-provided parameters)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainContext:
|
||||
"""Context for a retain operation validation (pre-operation).
|
||||
|
||||
Contains ALL user-provided parameters for the retain operation.
|
||||
"""
|
||||
|
||||
bank_id: str
|
||||
contents: list[dict] # List of {content, context, event_date, document_id}
|
||||
request_context: "RequestContext"
|
||||
document_id: str | None = None
|
||||
fact_type_override: str | None = None
|
||||
confidence_score: float | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RecallContext:
|
||||
"""Context for a recall operation validation (pre-operation).
|
||||
|
||||
Contains ALL user-provided parameters for the recall operation.
|
||||
"""
|
||||
|
||||
bank_id: str
|
||||
query: str
|
||||
request_context: "RequestContext"
|
||||
budget: "Budget | None" = None
|
||||
max_tokens: int = 4096
|
||||
enable_trace: bool = False
|
||||
fact_types: list[str] = field(default_factory=list)
|
||||
question_date: datetime | None = None
|
||||
include_entities: bool = False
|
||||
max_entity_tokens: int = 500
|
||||
include_chunks: bool = False
|
||||
max_chunk_tokens: int = 8192
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReflectContext:
|
||||
"""Context for a reflect operation validation (pre-operation).
|
||||
|
||||
Contains ALL user-provided parameters for the reflect operation.
|
||||
"""
|
||||
|
||||
bank_id: str
|
||||
query: str
|
||||
request_context: "RequestContext"
|
||||
budget: "Budget | None" = None
|
||||
context: str | None = None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Post-operation Contexts (includes results)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainResult:
|
||||
"""Result context for post-retain hook.
|
||||
|
||||
Contains the operation parameters and the result.
|
||||
"""
|
||||
|
||||
bank_id: str
|
||||
contents: list[dict]
|
||||
request_context: "RequestContext"
|
||||
document_id: str | None
|
||||
fact_type_override: str | None
|
||||
confidence_score: float | None
|
||||
# Result
|
||||
unit_ids: list[list[str]] # List of unit IDs per content item
|
||||
success: bool = True
|
||||
error: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RecallResult:
|
||||
"""Result context for post-recall hook.
|
||||
|
||||
Contains the operation parameters and the result.
|
||||
"""
|
||||
|
||||
bank_id: str
|
||||
query: str
|
||||
request_context: "RequestContext"
|
||||
budget: "Budget | None"
|
||||
max_tokens: int
|
||||
enable_trace: bool
|
||||
fact_types: list[str]
|
||||
question_date: datetime | None
|
||||
include_entities: bool
|
||||
max_entity_tokens: int
|
||||
include_chunks: bool
|
||||
max_chunk_tokens: int
|
||||
# Result
|
||||
result: "RecallResultModel | None" = None
|
||||
success: bool = True
|
||||
error: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReflectResultContext:
|
||||
"""Result context for post-reflect hook.
|
||||
|
||||
Contains the operation parameters and the result.
|
||||
"""
|
||||
|
||||
bank_id: str
|
||||
query: str
|
||||
request_context: "RequestContext"
|
||||
budget: "Budget | None"
|
||||
context: str | None
|
||||
# Result
|
||||
result: "ReflectResult | None" = None
|
||||
success: bool = True
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class OperationValidatorExtension(Extension, ABC):
|
||||
"""
|
||||
Validates and hooks into retain/recall/reflect operations.
|
||||
|
||||
This extension allows implementing custom logic such as:
|
||||
- Rate limiting (pre-operation)
|
||||
- Quota enforcement (pre-operation)
|
||||
- Permission checks (pre-operation)
|
||||
- Content filtering (pre-operation)
|
||||
- Usage tracking (post-operation)
|
||||
- Audit logging (post-operation)
|
||||
- Metrics collection (post-operation)
|
||||
|
||||
Enable via environment variable:
|
||||
HINDSIGHT_API_OPERATION_VALIDATOR_EXTENSION=mypackage.validators:MyValidator
|
||||
|
||||
Configuration is passed from prefixed environment variables:
|
||||
HINDSIGHT_API_OPERATION_VALIDATOR_MAX_REQUESTS=100
|
||||
-> config = {"max_requests": "100"}
|
||||
|
||||
Hook execution order:
|
||||
1. validate_retain/validate_recall/validate_reflect (pre-operation)
|
||||
2. [operation executes]
|
||||
3. on_retain_complete/on_recall_complete/on_reflect_complete (post-operation)
|
||||
"""
|
||||
|
||||
# =========================================================================
|
||||
# Pre-operation validation hooks (abstract - must be implemented)
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
|
||||
"""
|
||||
Validate a retain operation before execution.
|
||||
|
||||
Called before the retain operation is processed. Return ValidationResult.reject()
|
||||
to prevent the operation from executing.
|
||||
|
||||
Args:
|
||||
ctx: Context containing all user-provided parameters:
|
||||
- bank_id: Bank identifier
|
||||
- contents: List of content dicts
|
||||
- request_context: Request context with auth info
|
||||
- document_id: Optional document ID
|
||||
- fact_type_override: Optional fact type override
|
||||
- confidence_score: Optional confidence score
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the operation is allowed.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
|
||||
"""
|
||||
Validate a recall operation before execution.
|
||||
|
||||
Called before the recall operation is processed. Return ValidationResult.reject()
|
||||
to prevent the operation from executing.
|
||||
|
||||
Args:
|
||||
ctx: Context containing all user-provided parameters:
|
||||
- bank_id: Bank identifier
|
||||
- query: Search query
|
||||
- request_context: Request context with auth info
|
||||
- budget: Budget level
|
||||
- max_tokens: Maximum tokens to return
|
||||
- enable_trace: Whether to include trace info
|
||||
- fact_types: List of fact types to search
|
||||
- question_date: Optional date context for query
|
||||
- include_entities: Whether to include entity data
|
||||
- max_entity_tokens: Max tokens for entities
|
||||
- include_chunks: Whether to include chunks
|
||||
- max_chunk_tokens: Max tokens for chunks
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the operation is allowed.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
|
||||
"""
|
||||
Validate a reflect operation before execution.
|
||||
|
||||
Called before the reflect operation is processed. Return ValidationResult.reject()
|
||||
to prevent the operation from executing.
|
||||
|
||||
Args:
|
||||
ctx: Context containing all user-provided parameters:
|
||||
- bank_id: Bank identifier
|
||||
- query: Question to answer
|
||||
- request_context: Request context with auth info
|
||||
- budget: Budget level
|
||||
- context: Optional additional context
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the operation is allowed.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Post-operation hooks (optional - override to implement)
|
||||
# =========================================================================
|
||||
|
||||
async def on_retain_complete(self, result: RetainResult) -> None:
|
||||
"""
|
||||
Called after a retain operation completes (success or failure).
|
||||
|
||||
Override this method to implement post-operation logic such as:
|
||||
- Usage tracking
|
||||
- Audit logging
|
||||
- Metrics collection
|
||||
- Notifications
|
||||
|
||||
Args:
|
||||
result: Result context containing:
|
||||
- All original operation parameters
|
||||
- unit_ids: List of created unit IDs (if success)
|
||||
- success: Whether the operation succeeded
|
||||
- error: Error message (if failed)
|
||||
"""
|
||||
pass
|
||||
|
||||
async def on_recall_complete(self, result: RecallResult) -> None:
|
||||
"""
|
||||
Called after a recall operation completes (success or failure).
|
||||
|
||||
Override this method to implement post-operation logic such as:
|
||||
- Usage tracking
|
||||
- Audit logging
|
||||
- Metrics collection
|
||||
- Query analytics
|
||||
|
||||
Args:
|
||||
result: Result context containing:
|
||||
- All original operation parameters
|
||||
- result: RecallResultModel (if success)
|
||||
- success: Whether the operation succeeded
|
||||
- error: Error message (if failed)
|
||||
"""
|
||||
pass
|
||||
|
||||
async def on_reflect_complete(self, result: ReflectResultContext) -> None:
|
||||
"""
|
||||
Called after a reflect operation completes (success or failure).
|
||||
|
||||
Override this method to implement post-operation logic such as:
|
||||
- Usage tracking
|
||||
- Audit logging
|
||||
- Metrics collection
|
||||
- Response analytics
|
||||
|
||||
Args:
|
||||
result: Result context containing:
|
||||
- All original operation parameters
|
||||
- result: ReflectResult (if success)
|
||||
- success: Whether the operation succeeded
|
||||
- error: Error message (if failed)
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Tenant Extension for multi-tenancy and API key authentication."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
class AuthenticationError(Exception):
|
||||
"""Raised when authentication fails."""
|
||||
|
||||
def __init__(self, reason: str):
|
||||
self.reason = reason
|
||||
super().__init__(f"Authentication failed: {reason}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class TenantContext:
|
||||
"""
|
||||
Tenant context returned by authentication.
|
||||
|
||||
Contains the PostgreSQL schema name for tenant isolation.
|
||||
All database queries will use fully-qualified table names
|
||||
with this schema (e.g., schema_name.memory_units).
|
||||
"""
|
||||
|
||||
schema_name: str
|
||||
|
||||
|
||||
class TenantExtension(Extension, ABC):
|
||||
"""
|
||||
Extension for multi-tenancy and API key authentication.
|
||||
|
||||
This extension validates incoming requests and returns the tenant context
|
||||
including the PostgreSQL schema to use for database operations.
|
||||
|
||||
Built-in implementation:
|
||||
hindsight_api.extensions.builtin.tenant.ApiKeyTenantExtension
|
||||
|
||||
Enable via environment variable:
|
||||
HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
|
||||
HINDSIGHT_API_TENANT_API_KEY=your-secret-key
|
||||
|
||||
The returned schema_name is used for fully-qualified table names in queries,
|
||||
enabling tenant isolation at the database level.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def authenticate(self, context: RequestContext) -> TenantContext:
|
||||
"""
|
||||
Authenticate the action context and return tenant context.
|
||||
|
||||
Args:
|
||||
context: The action context containing API key and other auth data.
|
||||
|
||||
Returns:
|
||||
TenantContext with the schema_name for database operations.
|
||||
|
||||
Raises:
|
||||
AuthenticationError: If authentication fails.
|
||||
"""
|
||||
...
|
||||
@@ -4,6 +4,9 @@ Command-line interface for Hindsight API.
|
||||
Run the server with:
|
||||
hindsight-api
|
||||
|
||||
Run as background daemon:
|
||||
hindsight-api --daemon
|
||||
|
||||
Stop with Ctrl+C.
|
||||
"""
|
||||
|
||||
@@ -21,9 +24,13 @@ from . import MemoryEngine
|
||||
from .api import create_app
|
||||
from .banner import print_banner
|
||||
from .config import HindsightConfig, get_config
|
||||
|
||||
print()
|
||||
print_banner()
|
||||
from .daemon import (
|
||||
DEFAULT_DAEMON_PORT,
|
||||
DEFAULT_IDLE_TIMEOUT,
|
||||
DaemonLock,
|
||||
IdleTimeoutMiddleware,
|
||||
daemonize,
|
||||
)
|
||||
|
||||
# Filter deprecation warnings from third-party libraries
|
||||
warnings.filterwarnings("ignore", message="websockets.legacy is deprecated")
|
||||
@@ -106,8 +113,52 @@ def main():
|
||||
parser.add_argument("--ssl-keyfile", default=None, help="SSL key file")
|
||||
parser.add_argument("--ssl-certfile", default=None, help="SSL certificate file")
|
||||
|
||||
# Daemon mode options
|
||||
parser.add_argument(
|
||||
"--daemon",
|
||||
action="store_true",
|
||||
help=f"Run as background daemon (uses port {DEFAULT_DAEMON_PORT}, auto-exits after idle)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--idle-timeout",
|
||||
type=int,
|
||||
default=DEFAULT_IDLE_TIMEOUT,
|
||||
help=f"Idle timeout in seconds before auto-exit in daemon mode (default: {DEFAULT_IDLE_TIMEOUT})",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Daemon mode handling
|
||||
if args.daemon:
|
||||
# Use fixed daemon port
|
||||
args.port = DEFAULT_DAEMON_PORT
|
||||
args.host = "127.0.0.1" # Only bind to localhost for security
|
||||
|
||||
# Check if another daemon is already running
|
||||
daemon_lock = DaemonLock()
|
||||
if not daemon_lock.acquire():
|
||||
print(f"Daemon already running (PID: {daemon_lock.get_pid()})", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# Fork into background
|
||||
daemonize()
|
||||
|
||||
# Re-acquire lock in child process
|
||||
daemon_lock = DaemonLock()
|
||||
if not daemon_lock.acquire():
|
||||
sys.exit(1)
|
||||
|
||||
# Register cleanup to release lock
|
||||
def release_lock():
|
||||
daemon_lock.release()
|
||||
|
||||
atexit.register(release_lock)
|
||||
|
||||
# Print banner (not in daemon mode)
|
||||
if not args.daemon:
|
||||
print()
|
||||
print_banner()
|
||||
|
||||
# Configure Python logging based on log level
|
||||
# Update config with CLI override if provided
|
||||
if args.log_level != config.log_level:
|
||||
@@ -128,9 +179,12 @@ def main():
|
||||
log_level=args.log_level,
|
||||
mcp_enabled=config.mcp_enabled,
|
||||
graph_retriever=config.graph_retriever,
|
||||
skip_llm_verification=config.skip_llm_verification,
|
||||
lazy_reranker=config.lazy_reranker,
|
||||
)
|
||||
config.configure_logging()
|
||||
config.log_config()
|
||||
if not args.daemon:
|
||||
config.log_config()
|
||||
|
||||
# Register cleanup handlers
|
||||
atexit.register(_cleanup)
|
||||
@@ -149,6 +203,12 @@ def main():
|
||||
initialize_memory=True,
|
||||
)
|
||||
|
||||
# Wrap with idle timeout middleware in daemon mode
|
||||
idle_middleware = None
|
||||
if args.daemon:
|
||||
idle_middleware = IdleTimeoutMiddleware(app, idle_timeout=args.idle_timeout)
|
||||
app = idle_middleware
|
||||
|
||||
# Prepare uvicorn config
|
||||
uvicorn_config = {
|
||||
"app": app,
|
||||
@@ -172,20 +232,40 @@ def main():
|
||||
if args.ssl_certfile:
|
||||
uvicorn_config["ssl_certfile"] = args.ssl_certfile
|
||||
|
||||
from .banner import print_startup_info
|
||||
# Print startup info (not in daemon mode)
|
||||
if not args.daemon:
|
||||
from .banner import print_startup_info
|
||||
|
||||
print_startup_info(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
database_url=config.database_url,
|
||||
llm_provider=config.llm_provider,
|
||||
llm_model=config.llm_model,
|
||||
embeddings_provider=config.embeddings_provider,
|
||||
reranker_provider=config.reranker_provider,
|
||||
mcp_enabled=config.mcp_enabled,
|
||||
)
|
||||
print_startup_info(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
database_url=config.database_url,
|
||||
llm_provider=config.llm_provider,
|
||||
llm_model=config.llm_model,
|
||||
embeddings_provider=config.embeddings_provider,
|
||||
reranker_provider=config.reranker_provider,
|
||||
mcp_enabled=config.mcp_enabled,
|
||||
)
|
||||
|
||||
uvicorn.run(**uvicorn_config)
|
||||
# Start idle checker in daemon mode
|
||||
if idle_middleware is not None:
|
||||
# Start the idle checker in a background thread with its own event loop
|
||||
import threading
|
||||
|
||||
def run_idle_checker():
|
||||
import time
|
||||
|
||||
time.sleep(2) # Wait for uvicorn to start
|
||||
try:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
loop.run_until_complete(idle_middleware._check_idle())
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
threading.Thread(target=run_idle_checker, daemon=True).start()
|
||||
|
||||
uvicorn.run(**uvicorn_config) # type: ignore[invalid-argument-type] - dict kwargs
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -87,6 +87,7 @@ def create_local_mcp_server(bank_id: str, memory=None) -> FastMCP:
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
# Create memory engine with pg0 embedded database if not provided
|
||||
if memory is None:
|
||||
@@ -115,7 +116,11 @@ def create_local_mcp_server(bank_id: str, memory=None) -> FastMCP:
|
||||
|
||||
async def _retain():
|
||||
try:
|
||||
await memory.retain_batch_async(bank_id=bank_id, contents=[{"content": content, "context": context}])
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": content, "context": context}],
|
||||
request_context=RequestContext(),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing memory: {e}", exc_info=True)
|
||||
|
||||
@@ -142,6 +147,7 @@ def create_local_mcp_server(bank_id: str, memory=None) -> FastMCP:
|
||||
fact_type=list(VALID_RECALL_FACT_TYPES),
|
||||
budget=budget_enum,
|
||||
max_tokens=max_tokens,
|
||||
request_context=RequestContext(),
|
||||
)
|
||||
|
||||
return search_result.model_dump()
|
||||
|
||||
@@ -6,12 +6,16 @@ on application startup. It is designed to be safe for concurrent
|
||||
execution using PostgreSQL advisory locks to coordinate between
|
||||
distributed workers.
|
||||
|
||||
Supports multi-tenant schema isolation: migrations can target a specific
|
||||
PostgreSQL schema, allowing each tenant to have isolated tables.
|
||||
|
||||
Important: All migrations must be backward-compatible to allow
|
||||
safe rolling deployments.
|
||||
|
||||
No alembic.ini required - all configuration is done programmatically.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
@@ -26,11 +30,29 @@ logger = logging.getLogger(__name__)
|
||||
MIGRATION_LOCK_ID = 123456789
|
||||
|
||||
|
||||
def _run_migrations_internal(database_url: str, script_location: str) -> None:
|
||||
def _get_schema_lock_id(schema: str) -> int:
|
||||
"""
|
||||
Generate a unique advisory lock ID for a schema.
|
||||
|
||||
Uses hash of schema name to create a deterministic lock ID.
|
||||
"""
|
||||
# Use hash to create a unique lock ID per schema
|
||||
# Keep within PostgreSQL's bigint range
|
||||
hash_bytes = hashlib.sha256(schema.encode()).digest()[:8]
|
||||
return int.from_bytes(hash_bytes, byteorder="big") % (2**31)
|
||||
|
||||
|
||||
def _run_migrations_internal(database_url: str, script_location: str, schema: str | None = None) -> None:
|
||||
"""
|
||||
Internal function to run migrations without locking.
|
||||
|
||||
Args:
|
||||
database_url: SQLAlchemy database URL
|
||||
script_location: Path to alembic scripts
|
||||
schema: Target schema (None for default/public)
|
||||
"""
|
||||
logger.info("Running database migrations to head...")
|
||||
schema_name = schema or "public"
|
||||
logger.info(f"Running database migrations to head for schema '{schema_name}'...")
|
||||
logger.info(f"Database URL: {database_url}")
|
||||
logger.info(f"Script location: {script_location}")
|
||||
|
||||
@@ -50,13 +72,22 @@ def _run_migrations_internal(database_url: str, script_location: str) -> None:
|
||||
# Set path_separator to avoid deprecation warning
|
||||
alembic_cfg.set_main_option("path_separator", "os")
|
||||
|
||||
# Run migrations to head (latest version)
|
||||
# If targeting a specific schema, pass it to env.py via config
|
||||
# env.py will handle setting search_path and version_table_schema
|
||||
if schema:
|
||||
alembic_cfg.set_main_option("target_schema", schema)
|
||||
|
||||
# Run migrations
|
||||
command.upgrade(alembic_cfg, "head")
|
||||
|
||||
logger.info("Database migrations completed successfully")
|
||||
logger.info(f"Database migrations completed successfully for schema '{schema_name}'")
|
||||
|
||||
|
||||
def run_migrations(database_url: str, script_location: str | None = None) -> None:
|
||||
def run_migrations(
|
||||
database_url: str,
|
||||
script_location: str | None = None,
|
||||
schema: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Run database migrations to the latest version using programmatic Alembic configuration.
|
||||
|
||||
@@ -65,19 +96,28 @@ def run_migrations(database_url: str, script_location: str | None = None) -> Non
|
||||
- Other workers wait for the lock, then verify migrations are complete
|
||||
- If schema is already up-to-date, this is a fast no-op
|
||||
|
||||
Supports multi-tenant schema isolation: when a schema is specified, migrations
|
||||
run in that schema instead of public. This allows tenant extensions to provision
|
||||
new tenant schemas with their own isolated tables.
|
||||
|
||||
Args:
|
||||
database_url: SQLAlchemy database URL (e.g., "postgresql://user:pass@host/db")
|
||||
script_location: Path to alembic migrations directory (e.g., "/path/to/alembic").
|
||||
If None, defaults to hindsight-api/alembic directory.
|
||||
schema: Target PostgreSQL schema name. If None, uses default (public).
|
||||
When specified, creates the schema if needed and runs migrations there.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If migrations fail to complete
|
||||
FileNotFoundError: If script_location doesn't exist
|
||||
|
||||
Example:
|
||||
# Using default location (hindsight_api package)
|
||||
# Using default location and public schema
|
||||
run_migrations("postgresql://user:pass@host/db")
|
||||
|
||||
# Run migrations for a specific tenant schema
|
||||
run_migrations("postgresql://user:pass@host/db", schema="tenant_acme")
|
||||
|
||||
# Using custom location (when importing from another project)
|
||||
run_migrations(
|
||||
"postgresql://user:pass@host/db",
|
||||
@@ -99,21 +139,25 @@ def run_migrations(database_url: str, script_location: str | None = None) -> Non
|
||||
f"Alembic script location not found at {script_location}. Database migrations cannot be run."
|
||||
)
|
||||
|
||||
# Use schema-specific lock ID for multi-tenant isolation
|
||||
lock_id = _get_schema_lock_id(schema) if schema else MIGRATION_LOCK_ID
|
||||
schema_name = schema or "public"
|
||||
|
||||
# Use PostgreSQL advisory lock to coordinate between distributed workers
|
||||
engine = create_engine(database_url)
|
||||
with engine.connect() as conn:
|
||||
# pg_advisory_lock blocks until the lock is acquired
|
||||
# The lock is automatically released when the connection closes
|
||||
logger.debug(f"Acquiring migration advisory lock (id={MIGRATION_LOCK_ID})...")
|
||||
conn.execute(text(f"SELECT pg_advisory_lock({MIGRATION_LOCK_ID})"))
|
||||
logger.debug(f"Acquiring migration advisory lock for schema '{schema_name}' (id={lock_id})...")
|
||||
conn.execute(text(f"SELECT pg_advisory_lock({lock_id})"))
|
||||
logger.debug("Migration advisory lock acquired")
|
||||
|
||||
try:
|
||||
# Run migrations while holding the lock
|
||||
_run_migrations_internal(database_url, script_location)
|
||||
_run_migrations_internal(database_url, script_location, schema=schema)
|
||||
finally:
|
||||
# Explicitly release the lock (also released on connection close)
|
||||
conn.execute(text(f"SELECT pg_advisory_unlock({MIGRATION_LOCK_ID})"))
|
||||
conn.execute(text(f"SELECT pg_advisory_unlock({lock_id})"))
|
||||
logger.debug("Migration advisory lock released")
|
||||
|
||||
except FileNotFoundError:
|
||||
|
||||
@@ -2,9 +2,24 @@
|
||||
SQLAlchemy models for the memory system.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from uuid import UUID as PyUUID
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestContext:
|
||||
"""
|
||||
Context for request authentication and authorization.
|
||||
|
||||
This dataclass carries authentication data from HTTP requests to the
|
||||
memory engine operations. It can be extended to include additional
|
||||
context like headers, tokens, user info, etc.
|
||||
"""
|
||||
|
||||
api_key: str | None = None
|
||||
|
||||
|
||||
from pgvector.sqlalchemy import Vector
|
||||
from sqlalchemy import (
|
||||
CheckConstraint,
|
||||
|
||||
@@ -40,7 +40,7 @@ class EmbeddedPostgres:
|
||||
# Only set port if explicitly specified
|
||||
if self.port is not None:
|
||||
kwargs["port"] = self.port
|
||||
self._pg0 = Pg0(**kwargs)
|
||||
self._pg0 = Pg0(**kwargs) # type: ignore[invalid-argument-type] - dict kwargs
|
||||
return self._pg0
|
||||
|
||||
async def start(self, max_retries: int = 5, retry_delay: float = 4.0) -> str:
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api"
|
||||
version = "0.1.11"
|
||||
version = "0.1.12"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -92,6 +92,7 @@ dev = [
|
||||
"python-dotenv>=1.2.1",
|
||||
"filelock>=3.0.0",
|
||||
"ruff>=0.8.0",
|
||||
"ty>=0.0.1",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
@@ -121,3 +122,28 @@ ignore = [
|
||||
[tool.ruff.format]
|
||||
quote-style = "double"
|
||||
indent-style = "space"
|
||||
|
||||
[tool.ty]
|
||||
# Type checking configuration
|
||||
# ty is an extremely fast Python type checker from Astral (same team as ruff/uv)
|
||||
|
||||
[tool.ty.environment]
|
||||
python-version = "3.11"
|
||||
|
||||
[tool.ty.src]
|
||||
exclude = [
|
||||
"tests/",
|
||||
"hindsight_api/alembic/",
|
||||
]
|
||||
|
||||
[tool.ty.rules]
|
||||
# Disable noisy rules while keeping important ones
|
||||
invalid-argument-type = "ignore" # False positives with **kwargs patterns
|
||||
invalid-return-type = "ignore" # Often intentional in async code
|
||||
invalid-parameter-default = "ignore" # Optional params with None default
|
||||
possibly-missing-attribute = "ignore" # Common with Optional types
|
||||
invalid-raise = "ignore" # False positives with exception tracking
|
||||
call-non-callable = "ignore" # False positives with Optional types
|
||||
invalid-key = "ignore" # Pydantic ConfigDict not understood
|
||||
invalid-method-override = "ignore" # Intentional signature differences
|
||||
unresolved-reference = "ignore" # Forward references not always resolved
|
||||
|
||||
@@ -8,7 +8,7 @@ import os
|
||||
import filelock
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv
|
||||
from hindsight_api import MemoryEngine, LLMConfig, LocalSTEmbeddings
|
||||
from hindsight_api import MemoryEngine, LLMConfig, LocalSTEmbeddings, RequestContext
|
||||
|
||||
from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder
|
||||
from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer
|
||||
@@ -99,6 +99,12 @@ def pg0_db_url(db_url, tmp_path_factory, worker_id):
|
||||
return url
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def request_context():
|
||||
"""Provide a default RequestContext for tests."""
|
||||
return RequestContext()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def llm_config():
|
||||
"""
|
||||
|
||||
@@ -3,7 +3,7 @@ Tests for agent management API (profile, disposition, background).
|
||||
"""
|
||||
import pytest
|
||||
import uuid
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api import MemoryEngine, RequestContext
|
||||
from hindsight_api.api import CreateBankRequest, DispositionTraits
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
@@ -17,11 +17,11 @@ class TestAgentProfile:
|
||||
"""Tests for agent profile management."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_agent_profile_creates_default(self, memory: MemoryEngine):
|
||||
async def test_get_agent_profile_creates_default(self, memory: MemoryEngine, request_context):
|
||||
"""Test that getting a profile for a new agent creates default disposition."""
|
||||
bank_id = unique_agent_id("test_profile_default")
|
||||
|
||||
profile = await memory.get_bank_profile(bank_id)
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
assert profile is not None
|
||||
assert "disposition" in profile
|
||||
@@ -35,11 +35,11 @@ class TestAgentProfile:
|
||||
assert profile["background"] == ""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_agent_disposition(self, memory: MemoryEngine):
|
||||
async def test_update_agent_disposition(self, memory: MemoryEngine, request_context):
|
||||
"""Test updating agent disposition traits."""
|
||||
bank_id = unique_agent_id("test_profile_update")
|
||||
|
||||
profile = await memory.get_bank_profile(bank_id)
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
assert profile["disposition"].skepticism == 3
|
||||
|
||||
new_disposition = {
|
||||
@@ -47,26 +47,26 @@ class TestAgentProfile:
|
||||
"literalism": 4,
|
||||
"empathy": 2,
|
||||
}
|
||||
await memory.update_bank_disposition(bank_id, new_disposition)
|
||||
await memory.update_bank_disposition(bank_id, new_disposition, request_context=request_context)
|
||||
|
||||
updated_profile = await memory.get_bank_profile(bank_id)
|
||||
updated_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
disposition = updated_profile["disposition"]
|
||||
assert disposition.skepticism == new_disposition["skepticism"]
|
||||
assert disposition.literalism == new_disposition["literalism"]
|
||||
assert disposition.empathy == new_disposition["empathy"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_agents(self, memory: MemoryEngine):
|
||||
async def test_list_agents(self, memory: MemoryEngine, request_context):
|
||||
"""Test listing all agents."""
|
||||
agent_id_1 = unique_agent_id("test_list")
|
||||
agent_id_2 = unique_agent_id("test_list")
|
||||
agent_id_3 = unique_agent_id("test_list")
|
||||
|
||||
await memory.get_bank_profile(agent_id_1)
|
||||
await memory.get_bank_profile(agent_id_2)
|
||||
await memory.get_bank_profile(agent_id_3)
|
||||
await memory.get_bank_profile(agent_id_1, request_context=request_context)
|
||||
await memory.get_bank_profile(agent_id_2, request_context=request_context)
|
||||
await memory.get_bank_profile(agent_id_3, request_context=request_context)
|
||||
|
||||
agents = await memory.list_banks()
|
||||
agents = await memory.list_banks(request_context=request_context)
|
||||
|
||||
agent_ids = [a["bank_id"] for a in agents]
|
||||
assert agent_id_1 in agent_ids
|
||||
@@ -85,46 +85,50 @@ class TestAgentBackground:
|
||||
"""Tests for agent background management."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_merge_agent_background(self, memory: MemoryEngine):
|
||||
async def test_merge_agent_background(self, memory: MemoryEngine, request_context):
|
||||
"""Test merging agent background information."""
|
||||
bank_id = unique_agent_id("test_profile_merge")
|
||||
|
||||
profile = await memory.get_bank_profile(bank_id)
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
assert profile["background"] == ""
|
||||
|
||||
result1 = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I was born in Texas",
|
||||
update_disposition=False
|
||||
update_disposition=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert "Texas" in result1["background"]
|
||||
|
||||
result2 = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I have 10 years of startup experience",
|
||||
update_disposition=False
|
||||
update_disposition=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert "Texas" in result2["background"] or "startup" in result2["background"]
|
||||
|
||||
final_profile = await memory.get_bank_profile(bank_id)
|
||||
final_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
assert final_profile["background"] != ""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_merge_background_handles_conflicts(self, memory: MemoryEngine):
|
||||
async def test_merge_background_handles_conflicts(self, memory: MemoryEngine, request_context):
|
||||
"""Test that merging background handles conflicts (new overwrites old)."""
|
||||
bank_id = unique_agent_id("test_profile_conflict")
|
||||
|
||||
result1 = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I was born in Colorado",
|
||||
update_disposition=False
|
||||
update_disposition=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert "Colorado" in result1["background"]
|
||||
|
||||
result2 = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"You were born in Texas",
|
||||
update_disposition=False
|
||||
update_disposition=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert "Texas" in result2["background"]
|
||||
|
||||
@@ -133,7 +137,7 @@ class TestAgentEndpoint:
|
||||
"""Tests for agent PUT endpoint logic."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_put_agent_create(self, memory: MemoryEngine):
|
||||
async def test_put_agent_create(self, memory: MemoryEngine, request_context):
|
||||
"""Test creating an agent via PUT endpoint."""
|
||||
bank_id = unique_agent_id("test_put_create")
|
||||
|
||||
@@ -146,12 +150,13 @@ class TestAgentEndpoint:
|
||||
background="I am a creative software engineer"
|
||||
)
|
||||
|
||||
profile = await memory.get_bank_profile(bank_id)
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
if request.disposition is not None:
|
||||
await memory.update_bank_disposition(
|
||||
bank_id,
|
||||
request.disposition.model_dump()
|
||||
request.disposition.model_dump(),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
if request.background is not None:
|
||||
@@ -168,14 +173,14 @@ class TestAgentEndpoint:
|
||||
request.background
|
||||
)
|
||||
|
||||
final_profile = await memory.get_bank_profile(bank_id)
|
||||
final_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
assert final_profile["disposition"].skepticism == 4
|
||||
assert final_profile["disposition"].literalism == 5
|
||||
assert final_profile["background"] == "I am a creative software engineer"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_put_agent_partial_update(self, memory: MemoryEngine):
|
||||
async def test_put_agent_partial_update(self, memory: MemoryEngine, request_context):
|
||||
"""Test updating only background."""
|
||||
bank_id = unique_agent_id("test_put_partial")
|
||||
|
||||
@@ -183,7 +188,7 @@ class TestAgentEndpoint:
|
||||
background="I am a data scientist"
|
||||
)
|
||||
|
||||
profile = await memory.get_bank_profile(bank_id)
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
if request.background is not None:
|
||||
pool = await memory._get_pool()
|
||||
@@ -199,7 +204,7 @@ class TestAgentEndpoint:
|
||||
request.background
|
||||
)
|
||||
|
||||
final_profile = await memory.get_bank_profile(bank_id)
|
||||
final_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
assert final_profile["disposition"].skepticism == 3 # Default
|
||||
assert final_profile["background"] == "I am a data scientist"
|
||||
@@ -209,7 +214,7 @@ class TestAgentDispositionIntegration:
|
||||
"""Tests for disposition integration with other features."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_think_uses_disposition(self, memory: MemoryEngine):
|
||||
async def test_think_uses_disposition(self, memory: MemoryEngine, request_context):
|
||||
"""Test that THINK operation uses agent disposition."""
|
||||
bank_id = unique_agent_id("test_think")
|
||||
|
||||
@@ -218,12 +223,13 @@ class TestAgentDispositionIntegration:
|
||||
"literalism": 4, # High literalism
|
||||
"empathy": 2, # Low empathy
|
||||
}
|
||||
await memory.update_bank_disposition(bank_id, disposition)
|
||||
await memory.update_bank_disposition(bank_id, disposition, request_context=request_context)
|
||||
|
||||
await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am a creative artist who values innovation over tradition",
|
||||
update_disposition=False
|
||||
update_disposition=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.retain_batch_async(
|
||||
@@ -232,13 +238,14 @@ class TestAgentDispositionIntegration:
|
||||
{"content": "Traditional painting techniques have been used for centuries"},
|
||||
{"content": "Modern digital art is changing the art world"}
|
||||
],
|
||||
document_id="art_facts"
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
result = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="What do you think about traditional vs modern art?",
|
||||
budget=Budget.LOW
|
||||
budget=Budget.LOW,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result.text is not None
|
||||
|
||||
@@ -6,7 +6,7 @@ import os
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_large_batch_auto_chunks(memory):
|
||||
async def test_large_batch_auto_chunks(memory, request_context):
|
||||
bank_id = "test_chunking_agent"
|
||||
# Create a large batch that should trigger chunking
|
||||
# Each item is ~2000 chars, so 30 items = 60k chars (exceeds 50k threshold)
|
||||
@@ -24,7 +24,8 @@ async def test_large_batch_auto_chunks(memory):
|
||||
# Ingest the large batch (should auto-chunk)
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify we got results back
|
||||
@@ -33,7 +34,7 @@ async def test_large_batch_auto_chunks(memory):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_small_batch_no_chunking(memory):
|
||||
async def test_small_batch_no_chunking(memory, request_context):
|
||||
bank_id = "test_no_chunking_agent"
|
||||
|
||||
# Create a small batch that should NOT trigger chunking
|
||||
@@ -50,7 +51,8 @@ async def test_small_batch_no_chunking(memory):
|
||||
# Ingest the small batch (should NOT auto-chunk)
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify we got results back
|
||||
|
||||
@@ -10,6 +10,7 @@ import pytest
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api.engine.search.types import RetrievalResult, MergedCandidate, ScoredResult
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
|
||||
class TestRRFNormalization:
|
||||
@@ -125,7 +126,7 @@ class TestCombinedScoringFormula:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_has_normalized_rrf(memory):
|
||||
async def test_trace_has_normalized_rrf(memory, request_context):
|
||||
"""Integration test: verify trace contains normalized RRF values, not raw."""
|
||||
bank_id = f"test_scoring_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -135,21 +136,25 @@ async def test_trace_has_normalized_rrf(memory):
|
||||
bank_id=bank_id,
|
||||
content="Python is a programming language created by Guido van Rossum",
|
||||
context="tech facts",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="JavaScript was created by Brendan Eich at Netscape",
|
||||
context="tech facts",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="The Eiffel Tower is located in Paris, France",
|
||||
context="geography facts",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Mount Everest is the tallest mountain on Earth",
|
||||
context="geography facts",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Search with tracing
|
||||
@@ -160,6 +165,7 @@ async def test_trace_has_normalized_rrf(memory):
|
||||
budget=Budget.LOW,
|
||||
max_tokens=1024,
|
||||
enable_trace=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result.trace is not None, "Trace should be present"
|
||||
@@ -210,11 +216,11 @@ async def test_trace_has_normalized_rrf(memory):
|
||||
print(f" - First result score components: {sc}")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rrf_normalized_not_raw_in_trace(memory):
|
||||
async def test_rrf_normalized_not_raw_in_trace(memory, request_context):
|
||||
"""Verify that raw RRF scores (0.04-0.06 range) don't appear as normalized values."""
|
||||
bank_id = f"test_rrf_raw_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -225,6 +231,7 @@ async def test_rrf_normalized_not_raw_in_trace(memory):
|
||||
bank_id=bank_id,
|
||||
content=f"Test fact number {i} about various topics",
|
||||
context="test context",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
result = await memory.recall_async(
|
||||
@@ -234,6 +241,7 @@ async def test_rrf_normalized_not_raw_in_trace(memory):
|
||||
budget=Budget.LOW,
|
||||
max_tokens=512,
|
||||
enable_trace=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
trace = result.trace
|
||||
@@ -268,11 +276,11 @@ async def test_rrf_normalized_not_raw_in_trace(memory):
|
||||
print("\n✓ RRF raw vs normalized test passed!")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_combined_score_matches_components(memory):
|
||||
async def test_combined_score_matches_components(memory, request_context):
|
||||
"""Verify the final score actually equals the weighted sum of components."""
|
||||
bank_id = f"test_combined_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -281,11 +289,13 @@ async def test_combined_score_matches_components(memory):
|
||||
bank_id=bank_id,
|
||||
content="The quick brown fox jumps over the lazy dog",
|
||||
context="test",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="A quick test of the emergency broadcast system",
|
||||
context="test",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
result = await memory.recall_async(
|
||||
@@ -295,6 +305,7 @@ async def test_combined_score_matches_components(memory):
|
||||
budget=Budget.LOW,
|
||||
max_tokens=512,
|
||||
enable_trace=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
trace = result.trace
|
||||
@@ -320,4 +331,4 @@ async def test_combined_score_matches_components(memory):
|
||||
print("\n✓ Combined score verification test passed!")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -4,10 +4,11 @@ Tests for document tracking and upsert functionality.
|
||||
import logging
|
||||
import pytest
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_creation_and_retrieval(memory):
|
||||
async def test_document_creation_and_retrieval(memory, request_context):
|
||||
"""Test that documents are created and can be retrieved."""
|
||||
bank_id = f"test_doc_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -19,11 +20,12 @@ async def test_document_creation_and_retrieval(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google. Bob works at Microsoft.",
|
||||
context="Team meeting",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Retrieve document
|
||||
doc = await memory.get_document(document_id, bank_id)
|
||||
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
|
||||
assert doc is not None
|
||||
assert doc["id"] == document_id
|
||||
@@ -32,11 +34,11 @@ async def test_document_creation_and_retrieval(memory):
|
||||
assert doc["memory_unit_count"] > 0
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_upsert(memory):
|
||||
async def test_document_upsert(memory, request_context):
|
||||
"""Test that providing the same document_id automatically upserts (deletes old units and creates new ones)."""
|
||||
bank_id = f"test_upsert_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -48,11 +50,12 @@ async def test_document_upsert(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google.",
|
||||
context="Initial",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Get document stats
|
||||
doc_v1 = await memory.get_document(document_id, bank_id)
|
||||
doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
count_v1 = doc_v1["memory_unit_count"]
|
||||
|
||||
# Update with different content (automatic upsert when same document_id is provided)
|
||||
@@ -60,11 +63,12 @@ async def test_document_upsert(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Microsoft. Bob works at Apple.",
|
||||
context="Updated",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Get updated document stats
|
||||
doc_v2 = await memory.get_document(document_id, bank_id)
|
||||
doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
count_v2 = doc_v2["memory_unit_count"]
|
||||
|
||||
# Verify old units were replaced
|
||||
@@ -75,11 +79,11 @@ async def test_document_upsert(memory):
|
||||
assert set(units_v1).isdisjoint(set(units_v2))
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_deletion(memory):
|
||||
async def test_document_deletion(memory, request_context):
|
||||
"""Test that deleting a document cascades to memory units."""
|
||||
bank_id = f"test_delete_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -91,29 +95,30 @@ async def test_document_deletion(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google.",
|
||||
context="Test",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify it exists
|
||||
doc = await memory.get_document(document_id, bank_id)
|
||||
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc is not None
|
||||
assert doc["memory_unit_count"] > 0
|
||||
|
||||
# Delete document
|
||||
result = await memory.delete_document(document_id, bank_id)
|
||||
result = await memory.delete_document(document_id, bank_id, request_context=request_context)
|
||||
assert result["document_deleted"] == 1
|
||||
assert result["memory_units_deleted"] > 0
|
||||
|
||||
# Verify it's gone
|
||||
doc_after = await memory.get_document(document_id, bank_id)
|
||||
doc_after = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc_after is None
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memory_without_document(memory):
|
||||
async def test_memory_without_document(memory, request_context):
|
||||
"""Test that memories can still be created without document tracking."""
|
||||
bank_id = f"test_no_doc_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -122,10 +127,11 @@ async def test_memory_without_document(memory):
|
||||
units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google.",
|
||||
context="Test"
|
||||
context="Test",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(units) > 0
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -0,0 +1,796 @@
|
||||
"""Tests for the Hindsight extensions system."""
|
||||
|
||||
from collections import defaultdict
|
||||
|
||||
import pytest
|
||||
from fastapi import APIRouter
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from hindsight_api.extensions import (
|
||||
ApiKeyTenantExtension,
|
||||
AuthenticationError,
|
||||
Extension,
|
||||
HttpExtension,
|
||||
OperationValidationError,
|
||||
OperationValidatorExtension,
|
||||
RecallContext,
|
||||
RecallResult,
|
||||
ReflectContext,
|
||||
ReflectResultContext,
|
||||
RequestContext,
|
||||
RetainContext,
|
||||
RetainResult,
|
||||
TenantContext,
|
||||
TenantExtension,
|
||||
ValidationResult,
|
||||
load_extension,
|
||||
)
|
||||
|
||||
|
||||
class TestExtensionLoader:
|
||||
"""Tests for extension loading and lifecycle."""
|
||||
|
||||
def test_load_extension_with_config(self, monkeypatch):
|
||||
"""Extension receives config from prefixed env vars and supports lifecycle."""
|
||||
monkeypatch.setenv(
|
||||
"HINDSIGHT_API_TEST_EXTENSION",
|
||||
"tests.test_extensions:LifecycleTestExtension",
|
||||
)
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEST_API_URL", "https://example.com")
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEST_MAX_RETRIES", "5")
|
||||
|
||||
ext = load_extension("TEST", Extension)
|
||||
|
||||
assert ext is not None
|
||||
assert ext.config["api_url"] == "https://example.com"
|
||||
assert ext.config["max_retries"] == "5"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extension_lifecycle(self, monkeypatch):
|
||||
"""Extension on_startup and on_shutdown are called."""
|
||||
monkeypatch.setenv(
|
||||
"HINDSIGHT_API_TEST_EXTENSION",
|
||||
"tests.test_extensions:LifecycleTestExtension",
|
||||
)
|
||||
|
||||
ext = load_extension("TEST", Extension)
|
||||
|
||||
assert not ext.started
|
||||
assert not ext.stopped
|
||||
|
||||
await ext.on_startup()
|
||||
assert ext.started
|
||||
|
||||
await ext.on_shutdown()
|
||||
assert ext.stopped
|
||||
|
||||
|
||||
class LifecycleTestExtension(Extension):
|
||||
"""Test extension for config and lifecycle tests."""
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
self.started = False
|
||||
self.stopped = False
|
||||
|
||||
async def on_startup(self):
|
||||
self.started = True
|
||||
|
||||
async def on_shutdown(self):
|
||||
self.stopped = True
|
||||
|
||||
|
||||
class RateLimitingValidator(OperationValidatorExtension):
|
||||
"""
|
||||
Mock validator that blocks after N attempts per bank_id.
|
||||
|
||||
Used for testing the extension integration with MemoryEngine.
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict):
|
||||
super().__init__(config)
|
||||
self.max_attempts = int(config.get("max_attempts", "2"))
|
||||
self.retain_counts: dict[str, int] = defaultdict(int)
|
||||
self.recall_counts: dict[str, int] = defaultdict(int)
|
||||
self.reflect_counts: dict[str, int] = defaultdict(int)
|
||||
|
||||
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
|
||||
self.retain_counts[ctx.bank_id] += 1
|
||||
if self.retain_counts[ctx.bank_id] > self.max_attempts:
|
||||
return ValidationResult.reject(
|
||||
f"Retain limit exceeded for bank {ctx.bank_id}"
|
||||
)
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
|
||||
self.recall_counts[ctx.bank_id] += 1
|
||||
if self.recall_counts[ctx.bank_id] > self.max_attempts:
|
||||
return ValidationResult.reject(
|
||||
f"Recall limit exceeded for bank {ctx.bank_id}"
|
||||
)
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
|
||||
self.reflect_counts[ctx.bank_id] += 1
|
||||
if self.reflect_counts[ctx.bank_id] > self.max_attempts:
|
||||
return ValidationResult.reject(
|
||||
f"Reflect limit exceeded for bank {ctx.bank_id}"
|
||||
)
|
||||
return ValidationResult.accept()
|
||||
|
||||
|
||||
class TrackingValidator(OperationValidatorExtension):
|
||||
"""
|
||||
Mock validator that tracks all pre and post hook calls with full parameters.
|
||||
|
||||
Used for testing that hooks receive all user-provided parameters.
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict):
|
||||
super().__init__(config)
|
||||
# Pre-hook tracking
|
||||
self.pre_retain_calls: list[RetainContext] = []
|
||||
self.pre_recall_calls: list[RecallContext] = []
|
||||
self.pre_reflect_calls: list[ReflectContext] = []
|
||||
# Post-hook tracking
|
||||
self.post_retain_calls: list[RetainResult] = []
|
||||
self.post_recall_calls: list[RecallResult] = []
|
||||
self.post_reflect_calls: list[ReflectResultContext] = []
|
||||
|
||||
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
|
||||
self.pre_retain_calls.append(ctx)
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
|
||||
self.pre_recall_calls.append(ctx)
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
|
||||
self.pre_reflect_calls.append(ctx)
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def on_retain_complete(self, result: RetainResult) -> None:
|
||||
self.post_retain_calls.append(result)
|
||||
|
||||
async def on_recall_complete(self, result: RecallResult) -> None:
|
||||
self.post_recall_calls.append(result)
|
||||
|
||||
async def on_reflect_complete(self, result: ReflectResultContext) -> None:
|
||||
self.post_reflect_calls.append(result)
|
||||
|
||||
|
||||
class TestMemoryEngineValidation:
|
||||
"""Tests for validation integration with MemoryEngine.
|
||||
|
||||
The OperationValidatorExtension is integrated at the MemoryEngine level,
|
||||
so all interfaces (HTTP API, MCP, SDK) get the same validation behavior.
|
||||
|
||||
For retain, the batch is validated as a whole (all or nothing) using
|
||||
retain_batch_async which is the public method used by the HTTP API.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_batch_validation(self, memory_with_validator):
|
||||
"""Retain batch is validated as a whole - accepts or rejects entire batch."""
|
||||
memory = memory_with_validator
|
||||
bank_id = "test-retain-batch"
|
||||
ctx = RequestContext()
|
||||
|
||||
# First batch should succeed
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{"content": "First item"},
|
||||
{"content": "Second item"},
|
||||
],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
# Second batch should succeed (2nd attempt)
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": "Third item"}],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
# Third batch should be blocked entirely (exceeds limit)
|
||||
with pytest.raises(OperationValidationError) as exc_info:
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{"content": "Should not be stored"},
|
||||
{"content": "Neither should this"},
|
||||
],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert "limit exceeded" in str(exc_info.value).lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_validation(self, memory_with_validator):
|
||||
"""Recall is validated before execution."""
|
||||
memory = memory_with_validator
|
||||
bank_id = "test-recall-validation"
|
||||
ctx = RequestContext()
|
||||
|
||||
# First recall should pass validation
|
||||
await memory.recall_async(bank_id, "test query", fact_type=["world"], request_context=ctx)
|
||||
|
||||
# Second recall should pass validation
|
||||
await memory.recall_async(bank_id, "another query", fact_type=["world"], request_context=ctx)
|
||||
|
||||
# Third recall should be blocked by validator
|
||||
with pytest.raises(OperationValidationError) as exc_info:
|
||||
await memory.recall_async(bank_id, "blocked query", fact_type=["world"], request_context=ctx)
|
||||
|
||||
assert "limit exceeded" in str(exc_info.value).lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_validation(self, memory_with_validator):
|
||||
"""Reflect is validated before execution."""
|
||||
memory = memory_with_validator
|
||||
bank_id = "test-reflect-validation"
|
||||
ctx = RequestContext()
|
||||
|
||||
# First reflect should pass validation (may fail internally but validation passes)
|
||||
try:
|
||||
await memory.reflect_async(bank_id, "test question", request_context=ctx)
|
||||
except OperationValidationError:
|
||||
raise # Re-raise validation errors
|
||||
except Exception:
|
||||
pass # Other errors are fine (e.g., no data)
|
||||
|
||||
# Second reflect should pass validation
|
||||
try:
|
||||
await memory.reflect_async(bank_id, "another question", request_context=ctx)
|
||||
except OperationValidationError:
|
||||
raise
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Third reflect should be blocked by validator
|
||||
with pytest.raises(OperationValidationError) as exc_info:
|
||||
await memory.reflect_async(bank_id, "blocked question", request_context=ctx)
|
||||
|
||||
assert "limit exceeded" in str(exc_info.value).lower()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memory_with_validator(memory):
|
||||
"""Memory engine with a rate-limiting validator (max 2 attempts per bank)."""
|
||||
validator = RateLimitingValidator({"max_attempts": "2"})
|
||||
memory._operation_validator = validator
|
||||
return memory
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memory_with_tracking_validator(memory):
|
||||
"""Memory engine with a tracking validator that records all hook calls."""
|
||||
validator = TrackingValidator({})
|
||||
memory._operation_validator = validator
|
||||
return memory, validator
|
||||
|
||||
|
||||
class TestOperationHooksParameters:
|
||||
"""Tests for pre and post operation hooks receiving all user-provided parameters."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_pre_hook_receives_all_parameters(self, memory_with_tracking_validator):
|
||||
"""Pre-retain hook receives all user-provided parameters."""
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-retain-params"
|
||||
ctx = RequestContext(api_key="test-key")
|
||||
contents = [{"content": "Test content", "context": "test context"}]
|
||||
document_id = "doc-123"
|
||||
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
document_id=document_id,
|
||||
fact_type_override="world",
|
||||
confidence_score=0.9,
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert len(validator.pre_retain_calls) == 1
|
||||
pre_ctx = validator.pre_retain_calls[0]
|
||||
|
||||
# Verify all parameters are present
|
||||
assert pre_ctx.bank_id == bank_id
|
||||
# Note: contents is copied before document_id is applied to individual items
|
||||
assert len(pre_ctx.contents) == len(contents)
|
||||
assert pre_ctx.contents[0]["content"] == contents[0]["content"]
|
||||
assert pre_ctx.document_id == document_id
|
||||
assert pre_ctx.fact_type_override == "world"
|
||||
assert pre_ctx.confidence_score == 0.9
|
||||
assert pre_ctx.request_context == ctx
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_post_hook_receives_all_parameters_and_result(self, memory_with_tracking_validator):
|
||||
"""Post-retain hook receives all parameters plus the result."""
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-retain-post"
|
||||
ctx = RequestContext(api_key="test-key")
|
||||
contents = [{"content": "Test content for post hook"}]
|
||||
document_id = "doc-456"
|
||||
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
document_id=document_id,
|
||||
fact_type_override="experience",
|
||||
confidence_score=0.8,
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert len(validator.post_retain_calls) == 1
|
||||
post_result = validator.post_retain_calls[0]
|
||||
|
||||
# Verify all parameters are present
|
||||
assert post_result.bank_id == bank_id
|
||||
assert post_result.document_id == document_id
|
||||
assert post_result.fact_type_override == "experience"
|
||||
assert post_result.confidence_score == 0.8
|
||||
assert post_result.request_context == ctx
|
||||
|
||||
# Verify result data
|
||||
assert post_result.success is True
|
||||
assert post_result.error is None
|
||||
assert post_result.unit_ids == result # Should match the return value
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_pre_hook_receives_all_parameters(self, memory_with_tracking_validator):
|
||||
"""Pre-recall hook receives all user-provided parameters."""
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-recall-params"
|
||||
ctx = RequestContext(api_key="test-key")
|
||||
query = "test query"
|
||||
question_date = datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
|
||||
await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
budget=Budget.HIGH,
|
||||
max_tokens=2048,
|
||||
enable_trace=True,
|
||||
fact_type=["world", "experience"],
|
||||
question_date=question_date,
|
||||
include_entities=True,
|
||||
max_entity_tokens=300,
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=4096,
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert len(validator.pre_recall_calls) == 1
|
||||
pre_ctx = validator.pre_recall_calls[0]
|
||||
|
||||
# Verify all parameters are present
|
||||
assert pre_ctx.bank_id == bank_id
|
||||
assert pre_ctx.query == query
|
||||
assert pre_ctx.budget == Budget.HIGH
|
||||
assert pre_ctx.max_tokens == 2048
|
||||
assert pre_ctx.enable_trace is True
|
||||
assert pre_ctx.fact_types == ["world", "experience"]
|
||||
assert pre_ctx.question_date == question_date
|
||||
assert pre_ctx.include_entities is True
|
||||
assert pre_ctx.max_entity_tokens == 300
|
||||
assert pre_ctx.include_chunks is True
|
||||
assert pre_ctx.max_chunk_tokens == 4096
|
||||
assert pre_ctx.request_context == ctx
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_post_hook_receives_all_parameters_and_result(self, memory_with_tracking_validator):
|
||||
"""Post-recall hook receives all parameters plus the result."""
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-recall-post"
|
||||
ctx = RequestContext(api_key="test-key")
|
||||
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="test query for post",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=1024,
|
||||
fact_type=["world"],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert len(validator.post_recall_calls) == 1
|
||||
post_result = validator.post_recall_calls[0]
|
||||
|
||||
# Verify all parameters are present
|
||||
assert post_result.bank_id == bank_id
|
||||
assert post_result.query == "test query for post"
|
||||
assert post_result.budget == Budget.LOW
|
||||
assert post_result.max_tokens == 1024
|
||||
assert post_result.fact_types == ["world"]
|
||||
assert post_result.request_context == ctx
|
||||
|
||||
# Verify result data
|
||||
assert post_result.success is True
|
||||
assert post_result.error is None
|
||||
assert post_result.result == result # Should match the return value
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_pre_hook_receives_all_parameters(self, memory_with_tracking_validator):
|
||||
"""Pre-reflect hook receives all user-provided parameters."""
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-reflect-params"
|
||||
ctx = RequestContext(api_key="test-key")
|
||||
|
||||
try:
|
||||
await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="test question",
|
||||
budget=Budget.MID,
|
||||
context="additional context",
|
||||
request_context=ctx,
|
||||
)
|
||||
except Exception:
|
||||
pass # May fail if no data, but pre-hook should still be called
|
||||
|
||||
assert len(validator.pre_reflect_calls) == 1
|
||||
pre_ctx = validator.pre_reflect_calls[0]
|
||||
|
||||
# Verify all parameters are present
|
||||
assert pre_ctx.bank_id == bank_id
|
||||
assert pre_ctx.query == "test question"
|
||||
assert pre_ctx.budget == Budget.MID
|
||||
assert pre_ctx.context == "additional context"
|
||||
assert pre_ctx.request_context == ctx
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_post_hook_receives_all_parameters_and_result(self, memory_with_tracking_validator):
|
||||
"""Post-reflect hook receives all parameters plus the result on success."""
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-reflect-post"
|
||||
ctx = RequestContext(api_key="test-key")
|
||||
|
||||
# Store some content first so reflect has something to work with
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": "Alice is a software engineer at Google."}],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
result = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="What does Alice do?",
|
||||
budget=Budget.LOW,
|
||||
context="work context",
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert len(validator.post_reflect_calls) == 1
|
||||
post_result = validator.post_reflect_calls[0]
|
||||
|
||||
# Verify all parameters are present
|
||||
assert post_result.bank_id == bank_id
|
||||
assert post_result.query == "What does Alice do?"
|
||||
assert post_result.budget == Budget.LOW
|
||||
assert post_result.context == "work context"
|
||||
assert post_result.request_context == ctx
|
||||
|
||||
# Verify result data
|
||||
assert post_result.success is True
|
||||
assert post_result.error is None
|
||||
assert post_result.result == result # Should match the return value
|
||||
assert post_result.result.text is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_hooks_called_in_order_after_pre_hooks(self, memory_with_tracking_validator):
|
||||
"""Post hooks are called after pre hooks and after operation completes."""
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-hook-order"
|
||||
ctx = RequestContext()
|
||||
|
||||
# Retain operation
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": "Test content"}],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
# Pre-hook should be called before post-hook
|
||||
assert len(validator.pre_retain_calls) == 1
|
||||
assert len(validator.post_retain_calls) == 1
|
||||
|
||||
# Recall operation
|
||||
await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="test",
|
||||
fact_type=["world"],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert len(validator.pre_recall_calls) == 1
|
||||
assert len(validator.post_recall_calls) == 1
|
||||
|
||||
|
||||
class TestTenantExtension:
|
||||
"""Tests for TenantExtension and ApiKeyTenantExtension."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_key_tenant_extension_valid_key(self):
|
||||
"""ApiKeyTenantExtension accepts valid API key."""
|
||||
ext = ApiKeyTenantExtension({"api_key": "secret-key-123"})
|
||||
|
||||
result = await ext.authenticate(RequestContext(api_key="secret-key-123"))
|
||||
|
||||
assert result.schema_name == "public"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_key_tenant_extension_invalid_key(self):
|
||||
"""ApiKeyTenantExtension rejects invalid API key."""
|
||||
ext = ApiKeyTenantExtension({"api_key": "secret-key-123"})
|
||||
|
||||
with pytest.raises(AuthenticationError) as exc_info:
|
||||
await ext.authenticate(RequestContext(api_key="wrong-key"))
|
||||
|
||||
assert "Invalid API key" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_key_tenant_extension_missing_key(self):
|
||||
"""ApiKeyTenantExtension rejects missing API key."""
|
||||
ext = ApiKeyTenantExtension({"api_key": "secret-key-123"})
|
||||
|
||||
with pytest.raises(AuthenticationError):
|
||||
await ext.authenticate(RequestContext(api_key=None))
|
||||
|
||||
def test_api_key_tenant_extension_requires_config(self):
|
||||
"""ApiKeyTenantExtension requires api_key in config."""
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
ApiKeyTenantExtension({})
|
||||
|
||||
assert "HINDSIGHT_API_TENANT_API_KEY is required" in str(exc_info.value)
|
||||
|
||||
|
||||
class TestMemoryEngineTenantAuth:
|
||||
"""Tests for tenant authentication in MemoryEngine."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_requires_tenant_request_when_extension_configured(
|
||||
self, memory_with_tenant
|
||||
):
|
||||
"""Retain fails without RequestContext when tenant extension is configured."""
|
||||
memory = memory_with_tenant
|
||||
|
||||
with pytest.raises(AuthenticationError) as exc_info:
|
||||
await memory.retain_batch_async(
|
||||
bank_id="test-bank",
|
||||
contents=[{"content": "test"}],
|
||||
request_context=None, # Missing!
|
||||
)
|
||||
|
||||
assert "RequestContext is required" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_succeeds_with_valid_tenant_request(self, memory_with_tenant):
|
||||
"""Retain succeeds with valid RequestContext."""
|
||||
memory = memory_with_tenant
|
||||
|
||||
# Should not raise
|
||||
await memory.retain_batch_async(
|
||||
bank_id="test-bank-tenant",
|
||||
contents=[{"content": "test content"}],
|
||||
request_context=RequestContext(api_key="test-api-key"),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_fails_with_invalid_api_key(self, memory_with_tenant):
|
||||
"""Retain fails with invalid API key."""
|
||||
memory = memory_with_tenant
|
||||
|
||||
with pytest.raises(AuthenticationError) as exc_info:
|
||||
await memory.retain_batch_async(
|
||||
bank_id="test-bank",
|
||||
contents=[{"content": "test"}],
|
||||
request_context=RequestContext(api_key="wrong-key"),
|
||||
)
|
||||
|
||||
assert "Invalid API key" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_requires_tenant_request_when_extension_configured(
|
||||
self, memory_with_tenant
|
||||
):
|
||||
"""Recall fails without RequestContext when tenant extension is configured."""
|
||||
memory = memory_with_tenant
|
||||
|
||||
with pytest.raises(AuthenticationError):
|
||||
await memory.recall_async(
|
||||
bank_id="test-bank",
|
||||
query="test query",
|
||||
fact_type=["world"],
|
||||
request_context=None,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_tenant_request_needed_without_extension(self, memory):
|
||||
"""Operations work with empty RequestContext when no tenant extension configured."""
|
||||
# Should not raise - no tenant extension configured, just pass empty RequestContext
|
||||
await memory.retain_batch_async(
|
||||
bank_id="test-bank-no-tenant",
|
||||
contents=[{"content": "test content"}],
|
||||
request_context=RequestContext(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memory_with_tenant(memory):
|
||||
"""Memory engine with a tenant extension (API key auth)."""
|
||||
tenant_ext = ApiKeyTenantExtension({"api_key": "test-api-key"})
|
||||
memory._tenant_extension = tenant_ext
|
||||
return memory
|
||||
|
||||
|
||||
class SampleHttpExtension(HttpExtension):
|
||||
"""Sample HTTP extension for testing that provides custom endpoints."""
|
||||
|
||||
def __init__(self, config: dict):
|
||||
super().__init__(config)
|
||||
self.started = False
|
||||
self.stopped = False
|
||||
self.request_count = 0
|
||||
|
||||
async def on_startup(self):
|
||||
self.started = True
|
||||
|
||||
async def on_shutdown(self):
|
||||
self.stopped = True
|
||||
|
||||
def get_router(self, memory) -> APIRouter:
|
||||
router = APIRouter()
|
||||
|
||||
@router.get("/hello")
|
||||
async def hello():
|
||||
self.request_count += 1
|
||||
return {"message": "Hello from extension!"}
|
||||
|
||||
@router.get("/config")
|
||||
async def get_config():
|
||||
return {"config": self.config}
|
||||
|
||||
@router.get("/health-check")
|
||||
async def extension_health():
|
||||
health = await memory.health_check()
|
||||
return {"extension": "healthy", "memory": health}
|
||||
|
||||
@router.post("/echo")
|
||||
async def echo(data: dict):
|
||||
return {"echoed": data}
|
||||
|
||||
return router
|
||||
|
||||
|
||||
class TestHttpExtensionIntegration:
|
||||
"""Tests for HTTP extension integration."""
|
||||
|
||||
def test_load_http_extension(self, monkeypatch):
|
||||
"""HttpExtension can be loaded from environment variable."""
|
||||
monkeypatch.setenv(
|
||||
"HINDSIGHT_API_HTTP_EXTENSION",
|
||||
"tests.test_extensions:SampleHttpExtension",
|
||||
)
|
||||
monkeypatch.setenv("HINDSIGHT_API_HTTP_CUSTOM_PARAM", "custom_value")
|
||||
|
||||
ext = load_extension("HTTP", HttpExtension)
|
||||
|
||||
assert ext is not None
|
||||
assert isinstance(ext, SampleHttpExtension)
|
||||
assert ext.config["custom_param"] == "custom_value"
|
||||
|
||||
def test_http_extension_router_mounted_at_ext(self, memory):
|
||||
"""HTTP extension router is mounted at /ext/."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
ext = SampleHttpExtension({"test_key": "test_value"})
|
||||
app = create_app(memory, initialize_memory=False, http_extension=ext)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Extension endpoint should be accessible at /ext/
|
||||
response = client.get("/ext/hello")
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"message": "Hello from extension!"}
|
||||
|
||||
# Should track request count
|
||||
assert ext.request_count == 1
|
||||
|
||||
# Old path should NOT work
|
||||
response = client.get("/extension/hello")
|
||||
assert response.status_code == 404
|
||||
|
||||
def test_http_extension_config_endpoint(self, memory):
|
||||
"""Extension can expose its config via custom endpoint."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
ext = SampleHttpExtension({"api_key": "secret", "limit": "100"})
|
||||
app = create_app(memory, initialize_memory=False, http_extension=ext)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.get("/ext/config")
|
||||
assert response.status_code == 200
|
||||
assert response.json()["config"]["api_key"] == "secret"
|
||||
assert response.json()["config"]["limit"] == "100"
|
||||
|
||||
def test_http_extension_can_access_memory(self, memory):
|
||||
"""Extension endpoints can access memory engine."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
ext = SampleHttpExtension({})
|
||||
app = create_app(memory, initialize_memory=False, http_extension=ext)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.get("/ext/health-check")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["extension"] == "healthy"
|
||||
assert "memory" in data
|
||||
|
||||
def test_http_extension_post_endpoint(self, memory):
|
||||
"""Extension can handle POST requests with JSON body."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
ext = SampleHttpExtension({})
|
||||
app = create_app(memory, initialize_memory=False, http_extension=ext)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/ext/echo", json={"key": "value", "number": 42})
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"echoed": {"key": "value", "number": 42}}
|
||||
|
||||
def test_http_extension_not_mounted_when_none(self, memory):
|
||||
"""No extension routes when http_extension is None."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory, initialize_memory=False, http_extension=None)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Extension endpoint should not exist
|
||||
response = client.get("/ext/hello")
|
||||
assert response.status_code == 404
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_extension_lifecycle(self):
|
||||
"""HTTP extension on_startup and on_shutdown are called."""
|
||||
ext = SampleHttpExtension({})
|
||||
|
||||
assert not ext.started
|
||||
assert not ext.stopped
|
||||
|
||||
await ext.on_startup()
|
||||
assert ext.started
|
||||
|
||||
await ext.on_shutdown()
|
||||
assert ext.stopped
|
||||
|
||||
def test_core_routes_still_work_with_extension(self, memory):
|
||||
"""Core API routes still work when extension is mounted."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
ext = SampleHttpExtension({})
|
||||
app = create_app(memory, initialize_memory=False, http_extension=ext)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Health endpoint should work
|
||||
response = client.get("/health")
|
||||
assert response.status_code in (200, 503) # May be unhealthy if DB not connected
|
||||
|
||||
# Banks list endpoint should work
|
||||
response = client.get("/v1/default/banks")
|
||||
assert response.status_code in (200, 500) # May fail if DB not ready
|
||||
@@ -897,7 +897,7 @@ class TestDispositionInference:
|
||||
"""Tests for LLM-based disposition trait inference from background."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_merge_with_disposition_inference(self, memory):
|
||||
async def test_background_merge_with_disposition_inference(self, memory, request_context):
|
||||
"""Test that background merge infers disposition traits by default."""
|
||||
import uuid
|
||||
bank_id = f"test_infer_{uuid.uuid4().hex[:8]}"
|
||||
@@ -905,7 +905,8 @@ class TestDispositionInference:
|
||||
result = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am a creative software engineer who loves innovation and trying new technologies",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert "background" in result
|
||||
@@ -923,30 +924,31 @@ class TestDispositionInference:
|
||||
assert 1 <= disposition[trait] <= 5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_merge_without_disposition_inference(self, memory):
|
||||
async def test_background_merge_without_disposition_inference(self, memory, request_context):
|
||||
"""Test that background merge skips disposition inference when disabled."""
|
||||
import uuid
|
||||
bank_id = f"test_no_infer_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
initial_profile = await memory.get_bank_profile(bank_id)
|
||||
initial_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
initial_disposition = initial_profile["disposition"]
|
||||
|
||||
result = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am a data scientist",
|
||||
update_disposition=False
|
||||
update_disposition=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert "background" in result
|
||||
assert "disposition" not in result
|
||||
|
||||
final_profile = await memory.get_bank_profile(bank_id)
|
||||
final_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
final_disposition = final_profile["disposition"]
|
||||
|
||||
assert initial_disposition == final_disposition
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disposition_inference_for_lawyer(self, memory):
|
||||
async def test_disposition_inference_for_lawyer(self, memory, request_context):
|
||||
"""Test disposition inference for lawyer profile (high skepticism, high literalism)."""
|
||||
import uuid
|
||||
bank_id = f"test_lawyer_{uuid.uuid4().hex[:8]}"
|
||||
@@ -954,7 +956,8 @@ class TestDispositionInference:
|
||||
result = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am a lawyer who focuses on contract details and never takes claims at face value",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
disposition = result["disposition"]
|
||||
@@ -964,7 +967,7 @@ class TestDispositionInference:
|
||||
assert disposition["literalism"] >= 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disposition_inference_for_therapist(self, memory):
|
||||
async def test_disposition_inference_for_therapist(self, memory, request_context):
|
||||
"""Test disposition inference for therapist profile (high empathy)."""
|
||||
import uuid
|
||||
bank_id = f"test_therapist_{uuid.uuid4().hex[:8]}"
|
||||
@@ -972,7 +975,8 @@ class TestDispositionInference:
|
||||
result = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am a therapist who deeply understands and connects with people's emotional struggles",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
disposition = result["disposition"]
|
||||
@@ -981,7 +985,7 @@ class TestDispositionInference:
|
||||
assert disposition["empathy"] >= 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disposition_updates_in_database(self, memory):
|
||||
async def test_disposition_updates_in_database(self, memory, request_context):
|
||||
"""Test that inferred disposition is actually stored in database."""
|
||||
import uuid
|
||||
bank_id = f"test_db_update_{uuid.uuid4().hex[:8]}"
|
||||
@@ -989,12 +993,13 @@ class TestDispositionInference:
|
||||
result = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am an innovative designer",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
inferred_disposition = result["disposition"]
|
||||
|
||||
profile = await memory.get_bank_profile(bank_id)
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
db_disposition = profile["disposition"]
|
||||
|
||||
# Compare values (db_disposition is a Pydantic model)
|
||||
@@ -1003,7 +1008,7 @@ class TestDispositionInference:
|
||||
assert db_disposition.empathy == inferred_disposition["empathy"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_background_merges_update_disposition(self, memory):
|
||||
async def test_multiple_background_merges_update_disposition(self, memory, request_context):
|
||||
"""Test that each background merge can update disposition."""
|
||||
import uuid
|
||||
bank_id = f"test_multi_merge_{uuid.uuid4().hex[:8]}"
|
||||
@@ -1011,14 +1016,16 @@ class TestDispositionInference:
|
||||
result1 = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am a software engineer",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
disposition1 = result1["disposition"]
|
||||
|
||||
result2 = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I love creative problem solving and innovation",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
disposition2 = result2["disposition"]
|
||||
|
||||
@@ -1026,7 +1033,7 @@ class TestDispositionInference:
|
||||
assert "creative" in result2["background"].lower() or "innovation" in result2["background"].lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_merge_conflict_resolution_with_disposition(self, memory):
|
||||
async def test_background_merge_conflict_resolution_with_disposition(self, memory, request_context):
|
||||
"""Test that conflicts are resolved and disposition reflects final background."""
|
||||
import uuid
|
||||
bank_id = f"test_conflict_{uuid.uuid4().hex[:8]}"
|
||||
@@ -1034,13 +1041,15 @@ class TestDispositionInference:
|
||||
await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I was born in Colorado and prefer stability",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
result = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"You were born in Texas and are very skeptical of people",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
background = result["background"]
|
||||
|
||||
@@ -7,24 +7,24 @@ distinguish between things said earlier vs later.
|
||||
"""
|
||||
import pytest
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api import MemoryEngine, RequestContext
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
import os
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fact_ordering_within_conversation(memory):
|
||||
async def test_fact_ordering_within_conversation(memory, request_context):
|
||||
bank_id = "test_ordering_agent"
|
||||
|
||||
# Get/create agent (auto-creates with defaults)
|
||||
await memory.get_bank_profile(bank_id)
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Update disposition to match Marcus
|
||||
await memory.update_bank_disposition(bank_id, {
|
||||
"skepticism": 3,
|
||||
"literalism": 3,
|
||||
"empathy": 3
|
||||
})
|
||||
}, request_context=request_context)
|
||||
|
||||
# A conversation where Marcus changes his position
|
||||
conversation = """
|
||||
@@ -43,7 +43,8 @@ Marcus: Yeah, I realized I was being too optimistic about their defense.
|
||||
content=conversation,
|
||||
context="podcast discussion about NFL game",
|
||||
event_date=base_event_date,
|
||||
document_id="test_conv_1"
|
||||
document_id="test_conv_1",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Search for all facts about Marcus's predictions
|
||||
@@ -52,7 +53,8 @@ Marcus: Yeah, I realized I was being too optimistic about their defense.
|
||||
query="Marcus prediction Rams",
|
||||
fact_type=['opinion', 'experience', 'world'],
|
||||
budget=Budget.LOW,
|
||||
max_tokens=8192
|
||||
max_tokens=8192,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Retrieved {len(results.results)} facts ===")
|
||||
@@ -113,17 +115,17 @@ Marcus: Yeah, I realized I was being too optimistic about their defense.
|
||||
print(f"\n✅ Temporal ordering preserved: First prediction came before changed prediction")
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
print(f"\n✅ Test passed: Fact ordering within conversation is preserved")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_documents_ordering(memory):
|
||||
async def test_multiple_documents_ordering(memory, request_context):
|
||||
|
||||
bank_id = "test_multi_doc_agent"
|
||||
|
||||
await memory.get_bank_profile(bank_id) # Auto-creates with defaults
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context) # Auto-creates with defaults
|
||||
|
||||
# Two separate conversations with same base time
|
||||
base_time = datetime(2024, 11, 14, 10, 0, 0, tzinfo=timezone.utc)
|
||||
@@ -146,7 +148,8 @@ Alice: I reconsidered the team's experience level.
|
||||
contents=[
|
||||
{"content": conv1, "context": "project discussion 1", "event_date": base_time},
|
||||
{"content": conv2, "context": "project discussion 2", "event_date": base_time}
|
||||
]
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Search for Alice's preferences
|
||||
@@ -155,7 +158,8 @@ Alice: I reconsidered the team's experience level.
|
||||
query="Alice preference React Vue",
|
||||
fact_type=['opinion', 'experience'],
|
||||
budget=Budget.LOW,
|
||||
max_tokens=8192
|
||||
max_tokens=8192,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Retrieved {len(results.results)} agent facts ===")
|
||||
@@ -175,6 +179,6 @@ Alice: I reconsidered the team's experience level.
|
||||
print(f"\n✅ Facts from {len(agent_facts)} statements have {len(unique_timestamps)} unique timestamps")
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
print(f"\n✅ Test passed: Multiple documents maintain separate ordering")
|
||||
|
||||
@@ -26,6 +26,9 @@ MODEL_MATRIX = [
|
||||
("gemini", "gemini-2.5-flash"),
|
||||
("gemini", "gemini-2.5-flash-lite"),
|
||||
("gemini", "gemini-3-pro-preview"),
|
||||
# Ollama models (local)
|
||||
("ollama", "gemma3:12b"),
|
||||
("ollama", "gemma3:1b"),
|
||||
]
|
||||
|
||||
|
||||
@@ -48,12 +51,18 @@ async def test_llm_provider_memory_operations(provider: str, model: str):
|
||||
All models must pass this test.
|
||||
"""
|
||||
api_key = get_api_key_for_provider(provider)
|
||||
if not api_key:
|
||||
|
||||
# Skip Ollama tests in CI (no models available)
|
||||
if provider == "ollama" and os.getenv("CI"):
|
||||
pytest.skip(f"Skipping {provider}/{model}: Ollama not available in CI")
|
||||
|
||||
# Other providers need an API key
|
||||
if provider != "ollama" and not api_key:
|
||||
pytest.skip(f"Skipping {provider}/{model}: no API key available")
|
||||
|
||||
llm = LLMProvider(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
api_key=api_key or "",
|
||||
base_url="",
|
||||
model=model,
|
||||
)
|
||||
|
||||
@@ -3,11 +3,12 @@ Test observation generation and entity state functionality.
|
||||
"""
|
||||
import pytest
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import RequestContext
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_observation_generation_on_put(memory):
|
||||
async def test_observation_generation_on_put(memory, request_context):
|
||||
"""
|
||||
Test that observations are generated SYNCHRONOUSLY when new facts are added.
|
||||
|
||||
@@ -36,7 +37,8 @@ async def test_observation_generation_on_put(memory):
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="work info",
|
||||
event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Observations are generated SYNCHRONOUSLY during retain,
|
||||
@@ -75,7 +77,7 @@ async def test_observation_generation_on_put(memory):
|
||||
print(f"Entity: {entity_name} (id: {entity_id})")
|
||||
|
||||
# Get observations for the entity - should be available immediately
|
||||
observations = await memory.get_entity_observations(bank_id, entity_id, limit=10)
|
||||
observations = await memory.get_entity_observations(bank_id, entity_id, limit=10, request_context=request_context)
|
||||
|
||||
print(f"\n=== Observations for {entity_name} ===")
|
||||
print(f"Total observations: {len(observations)}")
|
||||
@@ -102,7 +104,7 @@ async def test_observation_generation_on_put(memory):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regenerate_entity_observations(memory):
|
||||
async def test_regenerate_entity_observations(memory, request_context):
|
||||
"""
|
||||
Test explicit regeneration of observations for an entity.
|
||||
"""
|
||||
@@ -114,7 +116,8 @@ async def test_regenerate_entity_observations(memory):
|
||||
bank_id=bank_id,
|
||||
content="Sarah is a product manager who loves user research and data analysis.",
|
||||
context="work info",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.wait_for_background_tasks()
|
||||
@@ -140,14 +143,15 @@ async def test_regenerate_entity_observations(memory):
|
||||
created_ids = await memory.regenerate_entity_observations(
|
||||
bank_id=bank_id,
|
||||
entity_id=entity_id,
|
||||
entity_name=entity_name
|
||||
entity_name=entity_name,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Regenerated Observations ===")
|
||||
print(f"Created {len(created_ids)} observations for {entity_name}")
|
||||
|
||||
# Get the observations
|
||||
observations = await memory.get_entity_observations(bank_id, entity_id, limit=10)
|
||||
observations = await memory.get_entity_observations(bank_id, entity_id, limit=10, request_context=request_context)
|
||||
for obs in observations:
|
||||
print(f" - {obs.text}")
|
||||
|
||||
@@ -170,7 +174,7 @@ async def test_regenerate_entity_observations(memory):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_with_include_entities(memory):
|
||||
async def test_search_with_include_entities(memory, request_context):
|
||||
"""
|
||||
Test that search with include_entities=True returns entity observations.
|
||||
|
||||
@@ -196,7 +200,8 @@ async def test_search_with_include_entities(memory):
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="work info",
|
||||
event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Observations are generated synchronously during retain, no need to wait
|
||||
@@ -209,7 +214,8 @@ async def test_search_with_include_entities(memory):
|
||||
budget=Budget.LOW,
|
||||
max_tokens=2000,
|
||||
include_entities=True,
|
||||
max_entity_tokens=500
|
||||
max_entity_tokens=500,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Search Results ===")
|
||||
@@ -263,7 +269,7 @@ async def test_search_with_include_entities(memory):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_entity_state(memory):
|
||||
async def test_get_entity_state(memory, request_context):
|
||||
"""
|
||||
Test getting the full state of an entity.
|
||||
"""
|
||||
@@ -275,7 +281,8 @@ async def test_get_entity_state(memory):
|
||||
bank_id=bank_id,
|
||||
content="Bob is a frontend developer who specializes in React and TypeScript.",
|
||||
context="work info",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.wait_for_background_tasks()
|
||||
@@ -302,7 +309,8 @@ async def test_get_entity_state(memory):
|
||||
bank_id=bank_id,
|
||||
entity_id=entity_id,
|
||||
entity_name=entity_name,
|
||||
limit=10
|
||||
limit=10,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Entity State for {entity_name} ===")
|
||||
@@ -324,7 +332,7 @@ async def test_get_entity_state(memory):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_observation_fact_type_in_database(memory):
|
||||
async def test_observation_fact_type_in_database(memory, request_context):
|
||||
"""
|
||||
Test that observations are stored with correct fact_type in database.
|
||||
"""
|
||||
@@ -336,7 +344,8 @@ async def test_observation_fact_type_in_database(memory):
|
||||
bank_id=bank_id,
|
||||
content="Charlie is a DevOps engineer who manages the Kubernetes infrastructure.",
|
||||
context="work info",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.wait_for_background_tasks()
|
||||
@@ -374,7 +383,7 @@ async def test_observation_fact_type_in_database(memory):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_entity_prioritized_for_observations(memory):
|
||||
async def test_user_entity_prioritized_for_observations(memory, request_context):
|
||||
"""
|
||||
Test that the 'user' entity gets observations even when many other entities exist.
|
||||
|
||||
@@ -410,7 +419,8 @@ async def test_user_entity_prioritized_for_observations(memory):
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="personal info",
|
||||
event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Observations are generated synchronously during retain
|
||||
@@ -466,7 +476,7 @@ async def test_user_entity_prioritized_for_observations(memory):
|
||||
f"User entity should have at least 5 facts, but has {user_fact_count}"
|
||||
|
||||
# Get observations for user entity
|
||||
observations = await memory.get_entity_observations(bank_id, user_entity_id, limit=10)
|
||||
observations = await memory.get_entity_observations(bank_id, user_entity_id, limit=10, request_context=request_context)
|
||||
|
||||
print(f"\n=== User Entity Observations ===")
|
||||
print(f"Total observations: {len(observations)}")
|
||||
|
||||
+157
-104
@@ -5,12 +5,13 @@ import pytest
|
||||
import logging
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_with_chunks(memory):
|
||||
async def test_retain_with_chunks(memory, request_context):
|
||||
"""
|
||||
Test that retain function:
|
||||
1. Stores facts with associated chunks
|
||||
@@ -41,7 +42,8 @@ async def test_retain_with_chunks(memory):
|
||||
content=long_content,
|
||||
context="team overview",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Retained {len(unit_ids)} facts ===")
|
||||
@@ -56,7 +58,8 @@ async def test_retain_with_chunks(memory):
|
||||
fact_type=["world"], # Search for world facts
|
||||
include_entities=False, # Disable entities for simpler test
|
||||
include_chunks=True, # Enable chunks
|
||||
max_chunk_tokens=8192
|
||||
max_chunk_tokens=8192,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Recall Results (with chunks) ===")
|
||||
@@ -88,12 +91,12 @@ async def test_retain_with_chunks(memory):
|
||||
|
||||
finally:
|
||||
# Cleanup - delete the test bank
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
print(f"\n=== Cleaned up bank: {bank_id} ===")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunks_and_entities_follow_fact_order(memory):
|
||||
async def test_chunks_and_entities_follow_fact_order(memory, request_context):
|
||||
"""
|
||||
Test that chunks and entities in recall results follow the same order as facts.
|
||||
This is critical because token limits may truncate later items.
|
||||
@@ -130,7 +133,8 @@ async def test_chunks_and_entities_follow_fact_order(memory):
|
||||
content=item["content"],
|
||||
context=item["context"],
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
document_id=item["document_id"]
|
||||
document_id=item["document_id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print("\n=== Stored 3 separate documents ===")
|
||||
@@ -144,7 +148,8 @@ async def test_chunks_and_entities_follow_fact_order(memory):
|
||||
fact_type=["world"],
|
||||
include_entities=True,
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=8192
|
||||
max_chunk_tokens=8192,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Recall Results ===")
|
||||
@@ -214,12 +219,12 @@ async def test_chunks_and_entities_follow_fact_order(memory):
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
print(f"\n=== Cleaned up bank: {bank_id} ===")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_event_date_storage(memory):
|
||||
async def test_event_date_storage(memory, request_context):
|
||||
"""
|
||||
Test that event_date is correctly stored as occurred_start.
|
||||
Verifies that we can track when events actually happened vs when they were stored.
|
||||
@@ -235,7 +240,8 @@ async def test_event_date_storage(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice completed the Q2 product launch on June 15th, 2023.",
|
||||
context="project history",
|
||||
event_date=past_event_date
|
||||
event_date=past_event_date,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should have created at least one memory unit"
|
||||
@@ -246,7 +252,8 @@ async def test_event_date_storage(memory):
|
||||
query="When did Alice complete the product launch?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall the stored fact"
|
||||
@@ -268,11 +275,11 @@ async def test_event_date_storage(memory):
|
||||
print(f"\n✓ Event date correctly stored: {occurred_dt}")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporal_ordering(memory):
|
||||
async def test_temporal_ordering(memory, request_context):
|
||||
"""
|
||||
Test that facts can be stored and retrieved with correct temporal ordering.
|
||||
Stores facts with different event_dates and verifies temporal relationships.
|
||||
@@ -305,7 +312,8 @@ async def test_temporal_ordering(memory):
|
||||
bank_id=bank_id,
|
||||
content=event["content"],
|
||||
context=event["context"],
|
||||
event_date=event["event_date"]
|
||||
event_date=event["event_date"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print("\n=== Stored 3 events with different temporal dates ===")
|
||||
@@ -316,7 +324,8 @@ async def test_temporal_ordering(memory):
|
||||
query="Tell me about Alice's career progression",
|
||||
budget=Budget.MID,
|
||||
max_tokens=1000,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) >= 3, f"Should recall all 3 events, got {len(result.results)}"
|
||||
@@ -345,11 +354,11 @@ async def test_temporal_ordering(memory):
|
||||
print(f"\n✓ Temporal ordering preserved: {min_date.date()} to {max_date.date()}")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mentioned_at_vs_occurred(memory):
|
||||
async def test_mentioned_at_vs_occurred(memory, request_context):
|
||||
"""
|
||||
Test distinction between when fact occurred vs when it was mentioned.
|
||||
|
||||
@@ -369,7 +378,8 @@ async def test_mentioned_at_vs_occurred(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice graduated from MIT in March 2020.",
|
||||
context="education history",
|
||||
event_date=conversation_date # When this conversation happened
|
||||
event_date=conversation_date, # When this conversation happened
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create memory unit"
|
||||
@@ -380,7 +390,8 @@ async def test_mentioned_at_vs_occurred(memory):
|
||||
query="Where did Alice go to school?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall the fact"
|
||||
@@ -415,11 +426,11 @@ async def test_mentioned_at_vs_occurred(memory):
|
||||
print(f"✓ Test passed: Historical conversation correctly ingested with event_date=2020")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_occurred_dates_not_defaulted(memory):
|
||||
async def test_occurred_dates_not_defaulted(memory, request_context):
|
||||
"""
|
||||
Test that occurred_start and occurred_end are NOT defaulted to mentioned_at.
|
||||
|
||||
@@ -441,7 +452,8 @@ async def test_occurred_dates_not_defaulted(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice likes coffee. The weather is sunny today.",
|
||||
context="current observations",
|
||||
event_date=event_date
|
||||
event_date=event_date,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create memory unit"
|
||||
@@ -452,7 +464,8 @@ async def test_occurred_dates_not_defaulted(memory):
|
||||
query="What does Alice like?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world", "opinion"]
|
||||
fact_type=["world", "opinion"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall the fact"
|
||||
@@ -504,11 +517,11 @@ async def test_occurred_dates_not_defaulted(memory):
|
||||
print(f"✓ Test passed: occurred dates are not incorrectly defaulted to mentioned_at")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mentioned_at_from_context_string(memory):
|
||||
async def test_mentioned_at_from_context_string(memory, request_context):
|
||||
"""
|
||||
Test that mentioned_at is extracted from context string by LLM.
|
||||
|
||||
@@ -527,7 +540,8 @@ async def test_mentioned_at_from_context_string(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice mentioned she loves hiking in the mountains.",
|
||||
context=f"Session ABC123 - you are the assistant in this conversation - happened on {session_date.strftime('%Y-%m-%d %H:%M:%S')} UTC.",
|
||||
event_date=None # Not providing event_date - should default to now() if LLM doesn't extract
|
||||
event_date=None, # Not providing event_date - should default to now() if LLM doesn't extract
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create memory unit"
|
||||
@@ -538,7 +552,8 @@ async def test_mentioned_at_from_context_string(memory):
|
||||
query="What does Alice like?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall the fact"
|
||||
@@ -574,7 +589,7 @@ async def test_mentioned_at_from_context_string(memory):
|
||||
print(f"✓ mentioned_at is always set (never None)")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -582,7 +597,7 @@ async def test_mentioned_at_from_context_string(memory):
|
||||
# ============================================================
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_preservation(memory):
|
||||
async def test_context_preservation(memory, request_context):
|
||||
"""
|
||||
Test that context is preserved and retrievable.
|
||||
Context helps understand why/how memory was formed.
|
||||
@@ -597,7 +612,8 @@ async def test_context_preservation(memory):
|
||||
bank_id=bank_id,
|
||||
content="The team decided to prioritize mobile development for next quarter.",
|
||||
context=specific_context,
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create at least one memory unit"
|
||||
@@ -608,7 +624,8 @@ async def test_context_preservation(memory):
|
||||
query="What did the team decide?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall the stored fact"
|
||||
@@ -620,11 +637,11 @@ async def test_context_preservation(memory):
|
||||
print(f" Retrieved {len(result.results)} facts")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_with_batch(memory):
|
||||
async def test_context_with_batch(memory, request_context):
|
||||
"""
|
||||
Test that each item in a batch can have different contexts.
|
||||
"""
|
||||
@@ -650,7 +667,8 @@ async def test_context_with_batch(memory):
|
||||
"context": "incident response",
|
||||
"event_date": datetime(2024, 1, 12, tzinfo=timezone.utc)
|
||||
}
|
||||
]
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should have created facts from all items
|
||||
@@ -661,7 +679,7 @@ async def test_context_with_batch(memory):
|
||||
print(f" Created {total_units} total memory units")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -669,7 +687,7 @@ async def test_context_with_batch(memory):
|
||||
# ============================================================
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metadata_storage_and_retrieval(memory):
|
||||
async def test_metadata_storage_and_retrieval(memory, request_context):
|
||||
"""
|
||||
Test that user-defined metadata is preserved.
|
||||
Metadata allows arbitrary key-value data to be stored with facts.
|
||||
@@ -692,7 +710,8 @@ async def test_metadata_storage_and_retrieval(memory):
|
||||
bank_id=bank_id,
|
||||
content="The product launch is scheduled for March 1st.",
|
||||
context="planning meeting",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create memory units"
|
||||
@@ -703,7 +722,8 @@ async def test_metadata_storage_and_retrieval(memory):
|
||||
query="When is the product launch?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall stored facts"
|
||||
@@ -712,7 +732,7 @@ async def test_metadata_storage_and_retrieval(memory):
|
||||
print(f" (Note: Metadata support depends on API implementation)")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -720,7 +740,7 @@ async def test_metadata_storage_and_retrieval(memory):
|
||||
# ============================================================
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_batch(memory):
|
||||
async def test_empty_batch(memory, request_context):
|
||||
"""
|
||||
Test that empty batch is handled gracefully without errors.
|
||||
"""
|
||||
@@ -730,7 +750,8 @@ async def test_empty_batch(memory):
|
||||
# Attempt to store empty batch
|
||||
unit_ids = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[]
|
||||
contents=[],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should return empty list or handle gracefully
|
||||
@@ -741,11 +762,11 @@ async def test_empty_batch(memory):
|
||||
|
||||
finally:
|
||||
# Clean up (though nothing should be stored)
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_item_batch(memory):
|
||||
async def test_single_item_batch(memory, request_context):
|
||||
"""
|
||||
Test that batch with one item works correctly.
|
||||
"""
|
||||
@@ -761,7 +782,8 @@ async def test_single_item_batch(memory):
|
||||
"context": "deployment log",
|
||||
"event_date": datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
}
|
||||
]
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) == 1, "Should return one list of unit IDs"
|
||||
@@ -770,11 +792,11 @@ async def test_single_item_batch(memory):
|
||||
print(f"✓ Single-item batch created {len(unit_ids[0])} units")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mixed_content_batch(memory):
|
||||
async def test_mixed_content_batch(memory, request_context):
|
||||
"""
|
||||
Test batch with varying content sizes (short and long).
|
||||
"""
|
||||
@@ -798,7 +820,8 @@ async def test_mixed_content_batch(memory):
|
||||
{"content": short_content, "context": "onboarding"},
|
||||
{"content": long_content, "context": "performance review"},
|
||||
{"content": "Charlie is on vacation this week.", "context": "team status"}
|
||||
]
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# All items should be processed
|
||||
@@ -813,11 +836,11 @@ async def test_mixed_content_batch(memory):
|
||||
print(f" Long content: {long_units} units")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_with_missing_optional_fields(memory):
|
||||
async def test_batch_with_missing_optional_fields(memory, request_context):
|
||||
"""
|
||||
Test that batch handles items with missing optional fields.
|
||||
"""
|
||||
@@ -842,7 +865,8 @@ async def test_batch_with_missing_optional_fields(memory):
|
||||
"context": "code review",
|
||||
# No event_date
|
||||
}
|
||||
]
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# All items should be processed successfully
|
||||
@@ -852,7 +876,7 @@ async def test_batch_with_missing_optional_fields(memory):
|
||||
print(f"✓ Batch with mixed optional fields created {total_units} total units")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -860,7 +884,7 @@ async def test_batch_with_missing_optional_fields(memory):
|
||||
# ============================================================
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_batch_multiple_documents(memory):
|
||||
async def test_single_batch_multiple_documents(memory, request_context):
|
||||
"""
|
||||
Test storing multiple distinct documents in a single batch call.
|
||||
Each should be tracked separately.
|
||||
@@ -876,21 +900,24 @@ async def test_single_batch_multiple_documents(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice's resume: 10 years Python experience, worked at Google.",
|
||||
context="resume review",
|
||||
document_id="resume_alice"
|
||||
document_id="resume_alice",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
doc2_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob's resume: 5 years JavaScript experience, worked at Meta.",
|
||||
context="resume review",
|
||||
document_id="resume_bob"
|
||||
document_id="resume_bob",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
doc3_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Charlie's resume: 8 years Go experience, worked at Amazon.",
|
||||
context="resume review",
|
||||
document_id="resume_charlie"
|
||||
document_id="resume_charlie",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# All documents should be stored
|
||||
@@ -907,17 +934,18 @@ async def test_single_batch_multiple_documents(memory):
|
||||
query="Who worked at Google?",
|
||||
budget=Budget.MID,
|
||||
max_tokens=1000,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should find facts about Alice"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_upsert_behavior(memory):
|
||||
async def test_document_upsert_behavior(memory, request_context):
|
||||
"""
|
||||
Test that upserting a document replaces the old content.
|
||||
"""
|
||||
@@ -930,7 +958,8 @@ async def test_document_upsert_behavior(memory):
|
||||
bank_id=bank_id,
|
||||
content="Project is in planning phase. Alice is the lead.",
|
||||
context="status update v1",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(v1_units) > 0, "Should create units for v1"
|
||||
@@ -940,7 +969,8 @@ async def test_document_upsert_behavior(memory):
|
||||
bank_id=bank_id,
|
||||
content="Project is in development phase. Bob has joined as co-lead.",
|
||||
context="status update v2",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(v2_units) > 0, "Should create units for v2"
|
||||
@@ -951,7 +981,8 @@ async def test_document_upsert_behavior(memory):
|
||||
query="What is the project status?",
|
||||
budget=Budget.MID,
|
||||
max_tokens=1000,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall facts"
|
||||
@@ -959,7 +990,7 @@ async def test_document_upsert_behavior(memory):
|
||||
print(f"✓ Document upsert created v1: {len(v1_units)} units, v2: {len(v2_units)} units")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -967,7 +998,7 @@ async def test_document_upsert_behavior(memory):
|
||||
# ============================================================
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_fact_mapping(memory):
|
||||
async def test_chunk_fact_mapping(memory, request_context):
|
||||
"""
|
||||
Test that facts correctly reference their source chunks via chunk_id.
|
||||
"""
|
||||
@@ -990,7 +1021,8 @@ async def test_chunk_fact_mapping(memory):
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="technical documentation",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create memory units"
|
||||
@@ -1003,7 +1035,8 @@ async def test_chunk_fact_mapping(memory):
|
||||
max_tokens=1000,
|
||||
fact_type=["world"],
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=8192
|
||||
max_chunk_tokens=8192,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall facts"
|
||||
@@ -1026,11 +1059,11 @@ async def test_chunk_fact_mapping(memory):
|
||||
print(f" Returned {len(result.chunks)} chunks matching fact references")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_ordering_preservation(memory):
|
||||
async def test_chunk_ordering_preservation(memory, request_context):
|
||||
"""
|
||||
Test that chunk_index reflects the correct order within a document.
|
||||
"""
|
||||
@@ -1070,7 +1103,8 @@ async def test_chunk_ordering_preservation(memory):
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="multi-section document",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create units"
|
||||
@@ -1083,7 +1117,8 @@ async def test_chunk_ordering_preservation(memory):
|
||||
max_tokens=2000,
|
||||
fact_type=["world"],
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=8192
|
||||
max_chunk_tokens=8192,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
if result.chunks:
|
||||
@@ -1103,11 +1138,11 @@ async def test_chunk_ordering_preservation(memory):
|
||||
print("✓ Content stored (may have created single chunk or no chunks returned)")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunks_truncation_behavior(memory):
|
||||
async def test_chunks_truncation_behavior(memory, request_context):
|
||||
"""
|
||||
Test that when chunks exceed max_chunk_tokens, truncation is indicated.
|
||||
"""
|
||||
@@ -1165,7 +1200,8 @@ async def test_chunks_truncation_behavior(memory):
|
||||
bank_id=bank_id,
|
||||
content=large_content,
|
||||
context="large document test",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create units"
|
||||
@@ -1178,7 +1214,8 @@ async def test_chunks_truncation_behavior(memory):
|
||||
max_tokens=1000,
|
||||
fact_type=["world"],
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=500 # Small limit to test truncation
|
||||
max_chunk_tokens=500, # Small limit to test truncation
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
if result.chunks:
|
||||
@@ -1198,7 +1235,7 @@ async def test_chunks_truncation_behavior(memory):
|
||||
print("✓ No chunks returned (may be under token limit)")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -1206,7 +1243,7 @@ async def test_chunks_truncation_behavior(memory):
|
||||
# ============================================================
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporal_links_creation(memory):
|
||||
async def test_temporal_links_creation(memory, request_context):
|
||||
"""
|
||||
Test that temporal links are created between facts with nearby event dates.
|
||||
|
||||
@@ -1223,7 +1260,8 @@ async def test_temporal_links_creation(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice started working on the authentication module.",
|
||||
context="daily standup",
|
||||
event_date=base_date
|
||||
event_date=base_date,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Fact 2 at 2:00 PM same day (4 hours later)
|
||||
@@ -1231,7 +1269,8 @@ async def test_temporal_links_creation(memory):
|
||||
bank_id=bank_id,
|
||||
content="Bob reviewed the API design document.",
|
||||
context="daily standup",
|
||||
event_date=base_date.replace(hour=14)
|
||||
event_date=base_date.replace(hour=14),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Fact 3 at 9:00 AM next day (23 hours later)
|
||||
@@ -1239,7 +1278,8 @@ async def test_temporal_links_creation(memory):
|
||||
bank_id=bank_id,
|
||||
content="Charlie deployed the new database schema.",
|
||||
context="daily standup",
|
||||
event_date=base_date.replace(day=16, hour=9)
|
||||
event_date=base_date.replace(day=16, hour=9),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0
|
||||
@@ -1278,11 +1318,11 @@ async def test_temporal_links_creation(memory):
|
||||
logger.info("Temporal links created successfully with proper weights")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semantic_links_creation(memory):
|
||||
async def test_semantic_links_creation(memory, request_context):
|
||||
"""
|
||||
Test that semantic links are created between facts with similar content.
|
||||
|
||||
@@ -1295,21 +1335,24 @@ async def test_semantic_links_creation(memory):
|
||||
unit_ids_1 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice is an expert in Python programming and has built many web applications.",
|
||||
context="team skills"
|
||||
context="team skills",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Similar content - should create semantic link
|
||||
unit_ids_2 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob is proficient in Python development and specializes in building APIs.",
|
||||
context="team skills"
|
||||
context="team skills",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Different content - less likely to create strong semantic link
|
||||
unit_ids_3 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="The quarterly sales meeting is scheduled for next Tuesday at 3 PM.",
|
||||
context="calendar events"
|
||||
context="calendar events",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0
|
||||
@@ -1349,11 +1392,11 @@ async def test_semantic_links_creation(memory):
|
||||
logger.info("Semantic links created successfully between similar content")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_entity_links_creation(memory):
|
||||
async def test_entity_links_creation(memory, request_context):
|
||||
"""
|
||||
Test that entity links are created between facts that mention the same entities.
|
||||
|
||||
@@ -1367,28 +1410,32 @@ async def test_entity_links_creation(memory):
|
||||
unit_ids_1 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice joined Google as a software engineer in 2020.",
|
||||
context="career history"
|
||||
context="career history",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Mentions same entity (Alice) - should create entity link
|
||||
unit_ids_2 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice led the development of the new authentication system.",
|
||||
context="project updates"
|
||||
context="project updates",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Mentions same entity (Google) - should create entity link
|
||||
unit_ids_3 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Google announced new cloud services at their annual conference.",
|
||||
context="tech news"
|
||||
context="tech news",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Different entities - no entity link expected
|
||||
unit_ids_4 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob works at Meta on machine learning infrastructure.",
|
||||
context="career history"
|
||||
context="career history",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0 and len(unit_ids_4) > 0
|
||||
@@ -1445,11 +1492,11 @@ async def test_entity_links_creation(memory):
|
||||
logger.info("Entity links are properly bidirectional")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_causal_links_creation(memory):
|
||||
async def test_causal_links_creation(memory, request_context):
|
||||
"""
|
||||
Test that causal links are created between facts with causal relationships.
|
||||
|
||||
@@ -1471,7 +1518,8 @@ async def test_causal_links_creation(memory):
|
||||
unit_ids = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="project timeline"
|
||||
context="project timeline",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should have created facts"
|
||||
@@ -1517,11 +1565,11 @@ async def test_causal_links_creation(memory):
|
||||
logger.info("Test completed (causal link extraction is LLM-dependent)")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_link_types_together(memory):
|
||||
async def test_all_link_types_together(memory, request_context):
|
||||
"""
|
||||
Integration test: Verify all link types can be created in a single retain operation.
|
||||
|
||||
@@ -1539,7 +1587,8 @@ async def test_all_link_types_together(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice completed the Python backend service for the authentication system.",
|
||||
context="sprint review",
|
||||
event_date=base_date
|
||||
event_date=base_date,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Fact 2: Related to Alice, similar topic (Python), close in time
|
||||
@@ -1547,7 +1596,8 @@ async def test_all_link_types_together(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice optimized the Python code and improved the authentication performance by 40%.",
|
||||
context="sprint review",
|
||||
event_date=base_date.replace(hour=14) # Same day, 4 hours later
|
||||
event_date=base_date.replace(hour=14), # Same day, 4 hours later
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Fact 3: Related to Alice, different topic but same entity
|
||||
@@ -1555,7 +1605,8 @@ async def test_all_link_types_together(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice presented the security architecture at the team meeting.",
|
||||
context="team meeting",
|
||||
event_date=base_date.replace(day=16) # Next day
|
||||
event_date=base_date.replace(day=16), # Next day
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0
|
||||
@@ -1594,11 +1645,11 @@ async def test_all_link_types_together(memory):
|
||||
logger.info("All major link types (temporal, semantic, entity) are working correctly")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semantic_links_within_same_batch(memory):
|
||||
async def test_semantic_links_within_same_batch(memory, request_context):
|
||||
"""
|
||||
Test that semantic links are created between facts retained in the SAME batch.
|
||||
|
||||
@@ -1617,7 +1668,8 @@ async def test_semantic_links_within_same_batch(memory):
|
||||
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Flatten the list of lists
|
||||
@@ -1652,11 +1704,11 @@ async def test_semantic_links_within_same_batch(memory):
|
||||
logger.info(f" Semantic link: {str(link['from_unit_id'])[:8]}... -> {str(link['to_unit_id'])[:8]}... (weight: {link['weight']:.3f})")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporal_links_within_same_batch(memory):
|
||||
async def test_temporal_links_within_same_batch(memory, request_context):
|
||||
"""
|
||||
Test that temporal links are created between facts retained in the SAME batch.
|
||||
|
||||
@@ -1689,7 +1741,8 @@ async def test_temporal_links_within_same_batch(memory):
|
||||
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Flatten the list of lists
|
||||
@@ -1724,4 +1777,4 @@ async def test_temporal_links_within_same_batch(memory):
|
||||
logger.info(f" Temporal link: {str(link['from_unit_id'])[:8]}... -> {str(link['to_unit_id'])[:8]}... (weight: {link['weight']:.3f})")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -0,0 +1,406 @@
|
||||
"""
|
||||
Tests for multi-tenant schema isolation.
|
||||
|
||||
Verifies that concurrent retain operations from different tenants
|
||||
are properly isolated in their respective PostgreSQL schemas.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.extensions import RequestContext, TenantContext, TenantExtension
|
||||
from hindsight_api.engine.memory_engine import _current_schema, fq_table
|
||||
from hindsight_api.migrations import run_migrations
|
||||
|
||||
|
||||
class MultiSchemaTestTenantExtension(TenantExtension):
|
||||
"""
|
||||
Test tenant extension that maps API keys to schema names.
|
||||
|
||||
API keys are in format: "key-{schema_name}"
|
||||
Provisions schemas on first access using run_migrations(schema=name).
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict):
|
||||
super().__init__(config)
|
||||
self.db_url = config.get("db_url")
|
||||
# Pre-configured valid schemas for test
|
||||
self.valid_schemas = config.get("valid_schemas", set())
|
||||
# Track provisioned schemas
|
||||
self._provisioned: set[str] = set()
|
||||
|
||||
async def authenticate(self, context: RequestContext) -> TenantContext:
|
||||
if not context.api_key:
|
||||
from hindsight_api.extensions import AuthenticationError
|
||||
|
||||
raise AuthenticationError("API key required")
|
||||
|
||||
# Parse schema from API key (format: "key-{schema}")
|
||||
if context.api_key.startswith("key-"):
|
||||
schema = context.api_key[4:] # Remove "key-" prefix
|
||||
if schema in self.valid_schemas:
|
||||
# Provision schema on first access
|
||||
if schema not in self._provisioned and self.db_url:
|
||||
run_migrations(self.db_url, schema=schema)
|
||||
self._provisioned.add(schema)
|
||||
return TenantContext(schema_name=schema)
|
||||
|
||||
from hindsight_api.extensions import AuthenticationError
|
||||
|
||||
raise AuthenticationError(f"Unknown API key: {context.api_key}")
|
||||
|
||||
|
||||
async def drop_schema(conn, schema_name: str) -> None:
|
||||
"""Drop a schema and all its contents."""
|
||||
await conn.execute(f'DROP SCHEMA IF EXISTS "{schema_name}" CASCADE')
|
||||
|
||||
|
||||
async def count_memories_in_schema(conn, schema_name: str, bank_id: str) -> int:
|
||||
"""Count memory units in a specific schema for a bank."""
|
||||
result = await conn.fetchval(
|
||||
f'SELECT COUNT(*) FROM "{schema_name}".memory_units WHERE bank_id = $1',
|
||||
bank_id,
|
||||
)
|
||||
return result or 0
|
||||
|
||||
|
||||
async def get_memory_texts_in_schema(conn, schema_name: str, bank_id: str) -> list[str]:
|
||||
"""Get all memory texts in a specific schema for a bank."""
|
||||
rows = await conn.fetch(
|
||||
f'SELECT text FROM "{schema_name}".memory_units WHERE bank_id = $1 ORDER BY text',
|
||||
bank_id,
|
||||
)
|
||||
return [row["text"] for row in rows]
|
||||
|
||||
|
||||
class TestSchemaIsolation:
|
||||
"""Tests for multi-tenant schema isolation."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_inserts_isolated_by_schema(self, memory, pg0_db_url):
|
||||
"""
|
||||
Multiple concurrent database operations from different tenants
|
||||
should store data in their respective schemas without cross-contamination.
|
||||
|
||||
Uses run_migrations(schema=x) to provision schemas like a real extension.
|
||||
"""
|
||||
import asyncpg
|
||||
|
||||
# Test schemas
|
||||
schemas = ["tenant_alpha", "tenant_beta", "tenant_gamma"]
|
||||
bank_id = f"test-isolation-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Clean up any existing schemas
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
await drop_schema(conn, schema)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
# Configure tenant extension that provisions schemas via run_migrations
|
||||
tenant_ext = MultiSchemaTestTenantExtension({
|
||||
"db_url": pg0_db_url,
|
||||
"valid_schemas": set(schemas),
|
||||
})
|
||||
memory._tenant_extension = tenant_ext
|
||||
|
||||
# Define concurrent insert tasks for each tenant
|
||||
async def insert_for_tenant(schema_name: str, content_prefix: str):
|
||||
"""Insert memories for a specific tenant using schema context."""
|
||||
# Authenticate to set the schema context
|
||||
tenant_request = RequestContext(api_key=f"key-{schema_name}")
|
||||
await memory._authenticate_tenant(tenant_request)
|
||||
|
||||
# Now fq_table will use the correct schema
|
||||
pool = await memory._get_pool()
|
||||
from hindsight_api.engine.db_utils import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Insert 3 memories for this tenant
|
||||
for i in range(3):
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table('memory_units')} (bank_id, text, event_date, fact_type)
|
||||
VALUES ($1, $2, now(), 'world')
|
||||
""",
|
||||
bank_id,
|
||||
f"MARKER_{content_prefix}_DOC{i}: Memory for {schema_name}",
|
||||
)
|
||||
|
||||
# Run concurrent inserts for all tenants
|
||||
await asyncio.gather(
|
||||
insert_for_tenant("tenant_alpha", "ALPHA"),
|
||||
insert_for_tenant("tenant_beta", "BETA"),
|
||||
insert_for_tenant("tenant_gamma", "GAMMA"),
|
||||
)
|
||||
|
||||
# Verify isolation - each schema should only have its own data
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
texts = await get_memory_texts_in_schema(conn, schema, bank_id)
|
||||
prefix = schema.replace("tenant_", "").upper()
|
||||
|
||||
# Should have exactly 3 memories
|
||||
assert len(texts) == 3, f"Schema {schema} should have 3 memories, got {len(texts)}"
|
||||
|
||||
# All texts should contain the schema's marker
|
||||
for text in texts:
|
||||
assert f"MARKER_{prefix}" in text, (
|
||||
f"Memory in {schema} missing its marker: {text}"
|
||||
)
|
||||
|
||||
# Should NOT contain other tenants' markers
|
||||
other_prefixes = ["ALPHA", "BETA", "GAMMA"]
|
||||
other_prefixes.remove(prefix)
|
||||
for other in other_prefixes:
|
||||
for text in texts:
|
||||
assert f"MARKER_{other}" not in text, (
|
||||
f"Cross-contamination! Schema {schema} has {other}'s marker: {text}"
|
||||
)
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
for schema in schemas:
|
||||
await drop_schema(conn, schema)
|
||||
await conn.close()
|
||||
|
||||
# Reset tenant extension
|
||||
memory._tenant_extension = None
|
||||
_current_schema.set("public")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_schema_context_isolation_in_concurrent_tasks(self, pg0_db_url):
|
||||
"""
|
||||
Verify that _current_schema contextvar is properly isolated
|
||||
between concurrent async tasks.
|
||||
"""
|
||||
results = {}
|
||||
errors = []
|
||||
|
||||
async def check_schema_context(schema_name: str, delay: float):
|
||||
"""Set schema context, wait, then verify it's still correct."""
|
||||
try:
|
||||
# Set the schema
|
||||
_current_schema.set(schema_name)
|
||||
|
||||
# Small delay to allow interleaving
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
# Verify schema is still correct
|
||||
current = _current_schema.get()
|
||||
if current != schema_name:
|
||||
errors.append(f"Expected {schema_name}, got {current}")
|
||||
|
||||
# Verify fq_table uses correct schema
|
||||
table = fq_table("memory_units")
|
||||
expected = f"{schema_name}.memory_units"
|
||||
if table != expected:
|
||||
errors.append(f"Expected {expected}, got {table}")
|
||||
|
||||
results[schema_name] = current
|
||||
|
||||
except Exception as e:
|
||||
errors.append(f"Error in {schema_name}: {e}")
|
||||
|
||||
# Run many concurrent tasks with different schemas
|
||||
tasks = []
|
||||
for i in range(10):
|
||||
for schema in ["schema_a", "schema_b", "schema_c"]:
|
||||
# Vary delays to create interleaving
|
||||
delay = 0.01 * (i % 3)
|
||||
tasks.append(check_schema_context(f"{schema}_{i}", delay))
|
||||
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# No errors should have occurred
|
||||
assert not errors, f"Schema context isolation errors: {errors}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_memories_respects_schema(self, memory, pg0_db_url):
|
||||
"""
|
||||
list_memory_units should only return memories from the current schema.
|
||||
|
||||
Uses run_migrations(schema=x) to provision schemas.
|
||||
"""
|
||||
import asyncpg
|
||||
|
||||
schemas = ["tenant_list_a", "tenant_list_b"]
|
||||
bank_id = f"test-list-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Clean up any existing schemas and provision via migrations
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
await drop_schema(conn, schema)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
# Provision schemas using run_migrations
|
||||
for schema in schemas:
|
||||
run_migrations(pg0_db_url, schema=schema)
|
||||
|
||||
# Insert test data directly into each schema
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO "{schema}".memory_units (bank_id, text, event_date, fact_type)
|
||||
VALUES ($1, $2, now(), 'world')
|
||||
""",
|
||||
bank_id,
|
||||
f"Direct insert for {schema}",
|
||||
)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
# Configure tenant extension
|
||||
tenant_ext = MultiSchemaTestTenantExtension({
|
||||
"db_url": pg0_db_url,
|
||||
"valid_schemas": set(schemas),
|
||||
})
|
||||
memory._tenant_extension = tenant_ext
|
||||
|
||||
try:
|
||||
# Query as tenant_list_a - should only see tenant_list_a's data
|
||||
tenant_a_request = RequestContext(api_key="key-tenant_list_a")
|
||||
await memory._authenticate_tenant(tenant_a_request)
|
||||
|
||||
result_a = await memory.list_memory_units(bank_id=bank_id, request_context=tenant_a_request)
|
||||
texts_a = [item["text"] for item in result_a.get("items", [])]
|
||||
|
||||
assert len(texts_a) == 1, f"Expected 1 memory for tenant_list_a, got {len(texts_a)}"
|
||||
assert "tenant_list_a" in texts_a[0], f"Wrong content: {texts_a[0]}"
|
||||
|
||||
# Query as tenant_list_b - should only see tenant_list_b's data
|
||||
tenant_b_request = RequestContext(api_key="key-tenant_list_b")
|
||||
await memory._authenticate_tenant(tenant_b_request)
|
||||
|
||||
result_b = await memory.list_memory_units(bank_id=bank_id, request_context=tenant_b_request)
|
||||
texts_b = [item["text"] for item in result_b.get("items", [])]
|
||||
|
||||
assert len(texts_b) == 1, f"Expected 1 memory for tenant_list_b, got {len(texts_b)}"
|
||||
assert "tenant_list_b" in texts_b[0], f"Wrong content: {texts_b[0]}"
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
await drop_schema(conn, schema)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
memory._tenant_extension = None
|
||||
_current_schema.set("public")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_high_concurrency_schema_isolation(self, memory, pg0_db_url):
|
||||
"""
|
||||
Stress test: Many concurrent operations across multiple schemas
|
||||
should maintain perfect isolation.
|
||||
|
||||
Uses run_migrations(schema=x) to provision schemas like a real extension.
|
||||
"""
|
||||
import asyncpg
|
||||
|
||||
# Create more schemas for stress test
|
||||
num_schemas = 5
|
||||
ops_per_schema = 10
|
||||
schemas = [f"stress_tenant_{i}" for i in range(num_schemas)]
|
||||
bank_id = f"test-stress-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Clean up any existing schemas first
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
await drop_schema(conn, schema)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
# Provision schemas using run_migrations
|
||||
for schema in schemas:
|
||||
run_migrations(pg0_db_url, schema=schema)
|
||||
|
||||
# Configure tenant extension (schemas already provisioned)
|
||||
tenant_ext = MultiSchemaTestTenantExtension({
|
||||
"db_url": pg0_db_url,
|
||||
"valid_schemas": set(schemas),
|
||||
})
|
||||
# Mark schemas as already provisioned so extension doesn't re-run migrations
|
||||
tenant_ext._provisioned = set(schemas)
|
||||
memory._tenant_extension = tenant_ext
|
||||
|
||||
errors = []
|
||||
|
||||
async def insert_one(schema: str, item_id: int):
|
||||
"""Single insert operation for tracking."""
|
||||
try:
|
||||
# Authenticate to set the schema context
|
||||
tenant_request = RequestContext(api_key=f"key-{schema}")
|
||||
await memory._authenticate_tenant(tenant_request)
|
||||
|
||||
# Insert using fq_table
|
||||
pool = await memory._get_pool()
|
||||
from hindsight_api.engine.db_utils import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table('memory_units')} (bank_id, text, event_date, fact_type)
|
||||
VALUES ($1, $2, now(), 'world')
|
||||
""",
|
||||
bank_id,
|
||||
f"STRESS_MARKER_{schema}_ITEM{item_id}: Memory for {schema}",
|
||||
)
|
||||
except Exception as e:
|
||||
errors.append(f"Insert error for {schema}: {e}")
|
||||
|
||||
# Run many concurrent operations
|
||||
tasks = []
|
||||
for i in range(ops_per_schema):
|
||||
for schema in schemas:
|
||||
tasks.append(insert_one(schema, i))
|
||||
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# Check for errors during insert
|
||||
assert not errors, f"Errors during insert: {errors}"
|
||||
|
||||
# Verify no cross-contamination
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
texts = await get_memory_texts_in_schema(conn, schema, bank_id)
|
||||
|
||||
# Should have exactly ops_per_schema memories
|
||||
assert len(texts) == ops_per_schema, (
|
||||
f"Schema {schema} should have {ops_per_schema} memories, got {len(texts)}"
|
||||
)
|
||||
|
||||
# All memories should reference this schema only
|
||||
for text in texts:
|
||||
# Check it contains our schema marker
|
||||
assert f"STRESS_MARKER_{schema}" in text, (
|
||||
f"Memory in {schema} doesn't contain schema marker: {text}"
|
||||
)
|
||||
|
||||
# Check it doesn't contain other schema markers
|
||||
for other_schema in schemas:
|
||||
if other_schema != schema:
|
||||
assert f"STRESS_MARKER_{other_schema}" not in text, (
|
||||
f"Cross-contamination! {schema} has {other_schema}'s data: {text}"
|
||||
)
|
||||
finally:
|
||||
# Cleanup
|
||||
for schema in schemas:
|
||||
await drop_schema(conn, schema)
|
||||
await conn.close()
|
||||
|
||||
memory._tenant_extension = None
|
||||
_current_schema.set("public")
|
||||
@@ -3,12 +3,12 @@ Test search tracing functionality.
|
||||
"""
|
||||
import pytest
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import SearchTrace
|
||||
from hindsight_api import SearchTrace, RequestContext
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_with_trace(memory):
|
||||
async def test_search_with_trace(memory, request_context):
|
||||
"""Test that search with enable_trace=True returns a valid SearchTrace."""
|
||||
# Generate a unique agent ID for this test
|
||||
bank_id = f"test_trace_{datetime.now(timezone.utc).timestamp()}"
|
||||
@@ -20,16 +20,19 @@ async def test_search_with_trace(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google in Mountain View",
|
||||
context="test context",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob also works at Google but in New York",
|
||||
context="test context",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Charlie founded a startup called TechCorp",
|
||||
context="test context",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Search with tracing enabled
|
||||
@@ -40,6 +43,7 @@ async def test_search_with_trace(memory):
|
||||
budget=Budget.LOW, # 20,
|
||||
max_tokens=512,
|
||||
enable_trace=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify results
|
||||
@@ -102,11 +106,11 @@ async def test_search_with_trace(memory):
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_without_trace(memory):
|
||||
async def test_search_without_trace(memory, request_context):
|
||||
"""Test that search with enable_trace=False returns None for trace."""
|
||||
bank_id = f"test_no_trace_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -117,6 +121,7 @@ async def test_search_without_trace(memory):
|
||||
bank_id=bank_id,
|
||||
content="Test memory without trace",
|
||||
context="test",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Search without tracing
|
||||
@@ -127,6 +132,7 @@ async def test_search_without_trace(memory):
|
||||
budget=Budget.LOW, # 10,
|
||||
max_tokens=512,
|
||||
enable_trace=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify trace is None
|
||||
@@ -137,4 +143,4 @@ async def test_search_without_trace(memory):
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
"""
|
||||
Safety tests to ensure all SQL queries use fully-qualified table names.
|
||||
|
||||
This prevents cross-tenant data access by ensuring every table reference
|
||||
includes the schema prefix (e.g., public.memory_units instead of just memory_units).
|
||||
"""
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
# All tables that MUST be schema-qualified in SQL queries
|
||||
TABLES = [
|
||||
"memory_units",
|
||||
"memory_links",
|
||||
"unit_entities",
|
||||
"entities",
|
||||
"entity_cooccurrences",
|
||||
"banks",
|
||||
"documents",
|
||||
"chunks",
|
||||
"async_operations",
|
||||
]
|
||||
|
||||
# Files to scan for SQL queries
|
||||
SCAN_PATHS = [
|
||||
"hindsight_api/engine",
|
||||
"hindsight_api/api",
|
||||
]
|
||||
|
||||
# Files to exclude (e.g., migrations, tests)
|
||||
EXCLUDE_PATTERNS = [
|
||||
"alembic",
|
||||
"__pycache__",
|
||||
"test_",
|
||||
]
|
||||
|
||||
|
||||
def get_python_files() -> list[Path]:
|
||||
"""Get all Python files to scan."""
|
||||
root = Path(__file__).parent.parent
|
||||
files = []
|
||||
for scan_path in SCAN_PATHS:
|
||||
path = root / scan_path
|
||||
if path.exists():
|
||||
for py_file in path.rglob("*.py"):
|
||||
# Check exclusions
|
||||
if any(excl in str(py_file) for excl in EXCLUDE_PATTERNS):
|
||||
continue
|
||||
files.append(py_file)
|
||||
return files
|
||||
|
||||
|
||||
def find_unqualified_table_refs(content: str, filename: str) -> list[tuple[int, str, str]]:
|
||||
"""
|
||||
Find SQL statements with unqualified table references.
|
||||
|
||||
Returns list of (line_number, table_name, line_content).
|
||||
"""
|
||||
violations = []
|
||||
|
||||
# Patterns that indicate SQL context
|
||||
sql_keywords = r"(?:FROM|JOIN|INTO|UPDATE|DELETE\s+FROM)\s+"
|
||||
|
||||
# Additional SQL indicators to confirm this is actually SQL, not prose
|
||||
sql_indicators = re.compile(
|
||||
r"(SELECT|INSERT|DELETE|UPDATE|CREATE|ALTER|DROP|WHERE|SET|VALUES|"
|
||||
r'f"""|f\'\'\'|""".*SELECT|\'\'\'.*SELECT)',
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
lines = content.split("\n")
|
||||
for line_num, line in enumerate(lines, 1):
|
||||
# Skip comments and strings that are clearly not SQL
|
||||
stripped = line.strip()
|
||||
if stripped.startswith("#"):
|
||||
continue
|
||||
|
||||
for table in TABLES:
|
||||
# Pattern: SQL keyword followed by unqualified table name
|
||||
# Should match: FROM memory_units, JOIN memory_units, INTO memory_units
|
||||
# Should NOT match: FROM public.memory_units, FROM {schema}.memory_units
|
||||
# Should NOT match: fq_table("memory_units")
|
||||
|
||||
# Check for unqualified table after SQL keyword
|
||||
pattern = rf"{sql_keywords}{table}(?:\s|$|,|\))"
|
||||
|
||||
if re.search(pattern, line, re.IGNORECASE):
|
||||
# Check if it's actually qualified (has schema prefix)
|
||||
qualified_pattern = rf"\.\s*{table}(?:\s|$|,|\))"
|
||||
fq_table_pattern = rf'fq_table\s*\(\s*["\']?{table}'
|
||||
|
||||
if not re.search(qualified_pattern, line) and not re.search(
|
||||
fq_table_pattern, line
|
||||
):
|
||||
# Additional check: line must have SQL indicators
|
||||
# This avoids false positives in docstrings like "split into chunks"
|
||||
if sql_indicators.search(line):
|
||||
violations.append((line_num, table, stripped))
|
||||
|
||||
return violations
|
||||
|
||||
|
||||
class TestSQLSchemaSafety:
|
||||
"""Ensure all SQL uses schema-qualified table names."""
|
||||
|
||||
def test_no_unqualified_table_references(self):
|
||||
"""All SQL queries must use fq_table() or schema.table format."""
|
||||
all_violations = []
|
||||
|
||||
for py_file in get_python_files():
|
||||
content = py_file.read_text()
|
||||
violations = find_unqualified_table_refs(content, py_file.name)
|
||||
|
||||
for line_num, table, line in violations:
|
||||
all_violations.append(
|
||||
f"{py_file.relative_to(py_file.parent.parent)}:{line_num} - "
|
||||
f"unqualified '{table}': {line[:80]}..."
|
||||
)
|
||||
|
||||
if all_violations:
|
||||
msg = (
|
||||
f"Found {len(all_violations)} unqualified table references!\n"
|
||||
"These could cause cross-tenant data access.\n"
|
||||
"Use fq_table('table_name') for all table references.\n\n"
|
||||
+ "\n".join(all_violations[:20]) # Show first 20
|
||||
)
|
||||
if len(all_violations) > 20:
|
||||
msg += f"\n... and {len(all_violations) - 20} more"
|
||||
pytest.fail(msg)
|
||||
|
||||
def test_tables_list_is_complete(self):
|
||||
"""Verify we're checking for all tables (sanity check)."""
|
||||
# This is a sanity check - if you add a new table, add it to TABLES
|
||||
assert len(TABLES) >= 9, "Update TABLES list if you added new tables"
|
||||
@@ -3,16 +3,17 @@ import asyncio
|
||||
from datetime import datetime, timezone, timedelta
|
||||
import pytest
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporal_ranges_are_written(memory):
|
||||
async def test_temporal_ranges_are_written(memory, request_context):
|
||||
"""Test that occurred_start, occurred_end, and mentioned_at are actually written to database."""
|
||||
bank_id = "test_temporal_ranges"
|
||||
|
||||
# Clean up any existing data
|
||||
try:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -23,7 +24,8 @@ async def test_temporal_ranges_are_written(memory):
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=text1,
|
||||
event_date=conversation_date
|
||||
event_date=conversation_date,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Test 2: Period event (month range)
|
||||
@@ -32,7 +34,8 @@ async def test_temporal_ranges_are_written(memory):
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=text2,
|
||||
event_date=conversation_date
|
||||
event_date=conversation_date,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Give it a moment for async processing
|
||||
@@ -114,7 +117,8 @@ async def test_temporal_ranges_are_written(memory):
|
||||
query="pottery workshop",
|
||||
fact_type=["world", "experience"],
|
||||
budget=Budget.LOW,
|
||||
max_tokens=4096
|
||||
max_tokens=4096,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"Found {len(search_result.results)} search results")
|
||||
@@ -132,4 +136,4 @@ async def test_temporal_ranges_are_written(memory):
|
||||
print("⚠ Temporal fields not yet populated in search results (known issue)")
|
||||
|
||||
# Clean up
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -4,10 +4,11 @@ Test think function for opinion generation and consistency.
|
||||
import pytest
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_think_opinion_consistency(memory):
|
||||
async def test_think_opinion_consistency(memory, request_context):
|
||||
"""
|
||||
Test that think function:
|
||||
1. Generates an opinion
|
||||
@@ -23,14 +24,16 @@ async def test_think_opinion_consistency(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice is a software engineer who has worked on 5 major projects. She always delivers on time and writes clean, well-documented code.",
|
||||
context="performance review",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob recently joined the team. He missed his first deadline and his code had many bugs.",
|
||||
context="performance review",
|
||||
event_date=datetime(2024, 2, 1, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 2, 1, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# First think call - should generate opinions
|
||||
@@ -39,6 +42,7 @@ async def test_think_opinion_consistency(memory):
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
budget=Budget.LOW,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== First Think Call ===")
|
||||
@@ -82,6 +86,7 @@ async def test_think_opinion_consistency(memory):
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
budget=Budget.LOW,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Second Think Call ===")
|
||||
@@ -122,13 +127,13 @@ async def test_think_opinion_consistency(memory):
|
||||
finally:
|
||||
# Clean up agent data
|
||||
try:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception as e:
|
||||
print(f"Warning: Error during cleanup: {e}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_think_without_prior_context(memory):
|
||||
async def test_think_without_prior_context(memory, request_context):
|
||||
"""
|
||||
Test that think function handles queries when there's no relevant context.
|
||||
"""
|
||||
@@ -139,6 +144,7 @@ async def test_think_without_prior_context(memory):
|
||||
bank_id=bank_id,
|
||||
query="What is the capital of France?",
|
||||
budget=Budget.LOW,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Think Without Context ===")
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "hindsight-cli"
|
||||
version = "0.1.11"
|
||||
version = "0.1.12"
|
||||
edition = "2021"
|
||||
authors = ["Hindsight Team"]
|
||||
description = "A beautiful CLI for Hindsight - semantic memory system"
|
||||
|
||||
@@ -64,13 +64,24 @@ pub struct ApiClient {
|
||||
}
|
||||
|
||||
impl ApiClient {
|
||||
pub fn new(base_url: String) -> Result<Self> {
|
||||
pub fn new(base_url: String, api_key: Option<String>) -> Result<Self> {
|
||||
let runtime = std::sync::Arc::new(tokio::runtime::Runtime::new()?);
|
||||
|
||||
// Create HTTP client with 2-minute timeout
|
||||
let http_client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.build()?;
|
||||
// Create HTTP client with 2-minute timeout and optional auth header
|
||||
let mut client_builder = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120));
|
||||
|
||||
if let Some(key) = api_key {
|
||||
let mut headers = reqwest::header::HeaderMap::new();
|
||||
let auth_value = format!("Bearer {}", key);
|
||||
headers.insert(
|
||||
reqwest::header::AUTHORIZATION,
|
||||
reqwest::header::HeaderValue::from_str(&auth_value)?,
|
||||
);
|
||||
client_builder = client_builder.default_headers(headers);
|
||||
}
|
||||
|
||||
let http_client = client_builder.build()?;
|
||||
|
||||
let client = AsyncClient::new_with_client(&base_url, http_client);
|
||||
Ok(ApiClient { client, runtime })
|
||||
|
||||
+38
-12
@@ -10,6 +10,7 @@ const CONFIG_DIR_NAME: &str = ".hindsight";
|
||||
|
||||
pub struct Config {
|
||||
pub api_url: String,
|
||||
pub api_key: Option<String>,
|
||||
pub source: ConfigSource,
|
||||
}
|
||||
|
||||
@@ -32,22 +33,27 @@ impl std::fmt::Display for ConfigSource {
|
||||
|
||||
impl Config {
|
||||
/// Load configuration with the following priority:
|
||||
/// 1. Environment variable (HINDSIGHT_API_URL) - highest priority, for overrides
|
||||
/// 1. Environment variable (HINDSIGHT_API_URL, HINDSIGHT_API_KEY) - highest priority, for overrides
|
||||
/// 2. Local config file (~/.hindsight/config.toml)
|
||||
/// 3. Default (http://localhost:8888)
|
||||
pub fn load() -> Result<Self> {
|
||||
// Load API key from environment (highest priority)
|
||||
let env_api_key = env::var("HINDSIGHT_API_KEY").ok();
|
||||
|
||||
// 1. Environment variable takes highest priority (for overrides)
|
||||
if let Ok(api_url) = env::var("HINDSIGHT_API_URL") {
|
||||
return Self::validate_and_create(api_url, ConfigSource::Environment);
|
||||
return Self::validate_and_create(api_url, env_api_key, ConfigSource::Environment);
|
||||
}
|
||||
|
||||
// 2. Try local config file
|
||||
if let Some(api_url) = Self::load_from_file()? {
|
||||
return Self::validate_and_create(api_url, ConfigSource::LocalFile);
|
||||
if let Some((api_url, file_api_key)) = Self::load_from_file()? {
|
||||
// Environment api_key takes precedence over file api_key
|
||||
let api_key = env_api_key.or(file_api_key);
|
||||
return Self::validate_and_create(api_url, api_key, ConfigSource::LocalFile);
|
||||
}
|
||||
|
||||
// 3. Fall back to default
|
||||
Self::validate_and_create(DEFAULT_API_URL.to_string(), ConfigSource::Default)
|
||||
Self::validate_and_create(DEFAULT_API_URL.to_string(), env_api_key, ConfigSource::Default)
|
||||
}
|
||||
|
||||
/// Legacy method for backwards compatibility
|
||||
@@ -55,14 +61,14 @@ impl Config {
|
||||
Self::load()
|
||||
}
|
||||
|
||||
fn validate_and_create(api_url: String, source: ConfigSource) -> Result<Self> {
|
||||
fn validate_and_create(api_url: String, api_key: Option<String>, source: ConfigSource) -> Result<Self> {
|
||||
if !api_url.starts_with("http://") && !api_url.starts_with("https://") {
|
||||
anyhow::bail!(
|
||||
"Invalid API URL: {}. Must start with http:// or https://",
|
||||
api_url
|
||||
);
|
||||
}
|
||||
Ok(Config { api_url, source })
|
||||
Ok(Config { api_url, api_key, source })
|
||||
}
|
||||
|
||||
fn config_dir() -> Option<PathBuf> {
|
||||
@@ -73,7 +79,7 @@ impl Config {
|
||||
Self::config_dir().map(|dir| dir.join(CONFIG_FILE_NAME))
|
||||
}
|
||||
|
||||
fn load_from_file() -> Result<Option<String>> {
|
||||
fn load_from_file() -> Result<Option<(String, Option<String>)>> {
|
||||
let config_path = match Self::config_file_path() {
|
||||
Some(path) => path,
|
||||
None => return Ok(None),
|
||||
@@ -86,23 +92,40 @@ impl Config {
|
||||
let content = fs::read_to_string(&config_path)
|
||||
.with_context(|| format!("Failed to read config file: {}", config_path.display()))?;
|
||||
|
||||
// Simple TOML parsing for api_url
|
||||
let mut api_url: Option<String> = None;
|
||||
let mut api_key: Option<String> = None;
|
||||
|
||||
// Simple TOML parsing for api_url and api_key
|
||||
for line in content.lines() {
|
||||
let line = line.trim();
|
||||
if line.starts_with("api_url") {
|
||||
if let Some(value) = line.split('=').nth(1) {
|
||||
let value = value.trim().trim_matches('"').trim_matches('\'');
|
||||
if !value.is_empty() {
|
||||
return Ok(Some(value.to_string()));
|
||||
api_url = Some(value.to_string());
|
||||
}
|
||||
}
|
||||
} else if line.starts_with("api_key") {
|
||||
if let Some(value) = line.split('=').nth(1) {
|
||||
let value = value.trim().trim_matches('"').trim_matches('\'');
|
||||
if !value.is_empty() {
|
||||
api_key = Some(value.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
match api_url {
|
||||
Some(url) => Ok(Some((url, api_key))),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn save_api_url(api_url: &str) -> Result<PathBuf> {
|
||||
Self::save_config(api_url, None)
|
||||
}
|
||||
|
||||
pub fn save_config(api_url: &str, api_key: Option<&str>) -> Result<PathBuf> {
|
||||
let config_dir = Self::config_dir()
|
||||
.ok_or_else(|| anyhow::anyhow!("Could not determine home directory"))?;
|
||||
|
||||
@@ -113,7 +136,10 @@ impl Config {
|
||||
}
|
||||
|
||||
let config_path = config_dir.join(CONFIG_FILE_NAME);
|
||||
let content = format!("api_url = \"{}\"\n", api_url);
|
||||
let mut content = format!("api_url = \"{}\"\n", api_url);
|
||||
if let Some(key) = api_key {
|
||||
content.push_str(&format!("api_key = \"{}\"\n", key));
|
||||
}
|
||||
|
||||
fs::write(&config_path, content)
|
||||
.with_context(|| format!("Failed to write config file: {}", config_path.display()))?;
|
||||
|
||||
@@ -94,12 +94,15 @@ enum Commands {
|
||||
/// Launch the web-based control plane UI
|
||||
Ui,
|
||||
|
||||
/// Configure the CLI (API URL, etc.)
|
||||
#[command(after_help = "Configuration priority:\n 1. Environment variable (HINDSIGHT_API_URL) - highest priority\n 2. Config file (~/.hindsight/config)\n 3. Default (http://localhost:8888)")]
|
||||
/// Configure the CLI (API URL, API key, etc.)
|
||||
#[command(after_help = "Configuration priority:\n 1. Environment variables (HINDSIGHT_API_URL, HINDSIGHT_API_KEY) - highest priority\n 2. Config file (~/.hindsight/config)\n 3. Default (http://localhost:8888)")]
|
||||
Configure {
|
||||
/// API URL to connect to (interactive prompt if not provided)
|
||||
#[arg(long)]
|
||||
api_url: Option<String>,
|
||||
/// API key for authentication (sent as Bearer token)
|
||||
#[arg(long)]
|
||||
api_key: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -372,8 +375,8 @@ fn run() -> Result<()> {
|
||||
let verbose = cli.verbose;
|
||||
|
||||
// Handle configure command before loading full config (it doesn't need API client)
|
||||
if let Commands::Configure { api_url } = cli.command {
|
||||
return handle_configure(api_url, output_format);
|
||||
if let Commands::Configure { api_url, api_key } = cli.command {
|
||||
return handle_configure(api_url, api_key, output_format);
|
||||
}
|
||||
|
||||
// Handle ui command - needs config but not API client
|
||||
@@ -389,9 +392,10 @@ fn run() -> Result<()> {
|
||||
});
|
||||
|
||||
let api_url = config.api_url().to_string();
|
||||
let api_key = config.api_key.clone();
|
||||
|
||||
// Create API client
|
||||
let client = ApiClient::new(api_url.clone()).unwrap_or_else(|e| {
|
||||
let client = ApiClient::new(api_url.clone(), api_key).unwrap_or_else(|e| {
|
||||
errors::handle_api_error(e, &api_url);
|
||||
});
|
||||
|
||||
@@ -476,7 +480,7 @@ fn run() -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle_configure(api_url: Option<String>, output_format: OutputFormat) -> Result<()> {
|
||||
fn handle_configure(api_url: Option<String>, api_key: Option<String>, output_format: OutputFormat) -> Result<()> {
|
||||
// Load current config to show current state
|
||||
let current_config = Config::load().ok();
|
||||
|
||||
@@ -487,6 +491,15 @@ fn handle_configure(api_url: Option<String>, output_format: OutputFormat) -> Res
|
||||
// Show current configuration
|
||||
if let Some(ref config) = current_config {
|
||||
println!(" Current API URL: {}", config.api_url);
|
||||
if let Some(ref key) = config.api_key {
|
||||
// Mask the API key for display
|
||||
let masked = if key.len() > 8 {
|
||||
format!("{}...{}", &key[..4], &key[key.len()-4..])
|
||||
} else {
|
||||
"****".to_string()
|
||||
};
|
||||
println!(" Current API Key: {}", masked);
|
||||
}
|
||||
println!(" Source: {}", config.source);
|
||||
println!();
|
||||
}
|
||||
@@ -511,18 +524,30 @@ fn handle_configure(api_url: Option<String>, output_format: OutputFormat) -> Res
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Use provided api_key, or keep existing one if not provided
|
||||
let new_api_key = api_key.or_else(|| current_config.as_ref().and_then(|c| c.api_key.clone()));
|
||||
|
||||
// Save to config file
|
||||
let config_path = Config::save_api_url(&new_api_url)?;
|
||||
let config_path = Config::save_config(&new_api_url, new_api_key.as_deref())?;
|
||||
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_success(&format!("Configuration saved to {}", config_path.display()));
|
||||
println!();
|
||||
println!(" API URL: {}", new_api_url);
|
||||
if let Some(ref key) = new_api_key {
|
||||
let masked = if key.len() > 8 {
|
||||
format!("{}...{}", &key[..4], &key[key.len()-4..])
|
||||
} else {
|
||||
"****".to_string()
|
||||
};
|
||||
println!(" API Key: {}", masked);
|
||||
}
|
||||
println!();
|
||||
println!("Note: Environment variable HINDSIGHT_API_URL will override this setting.");
|
||||
println!("Note: Environment variables HINDSIGHT_API_URL and HINDSIGHT_API_KEY will override these settings.");
|
||||
} else {
|
||||
let result = serde_json::json!({
|
||||
"api_url": new_api_url,
|
||||
"api_key_set": new_api_key.is_some(),
|
||||
"config_path": config_path.display().to_string(),
|
||||
});
|
||||
output::print_output(&result, output_format)?;
|
||||
|
||||
@@ -5,7 +5,7 @@ use crate::output::OutputFormat;
|
||||
|
||||
/// Get API client from config
|
||||
pub fn get_client(config: &Config) -> Result<ApiClient> {
|
||||
ApiClient::new(config.api_url.clone())
|
||||
ApiClient::new(config.api_url.clone(), config.api_key.clone())
|
||||
.context("Failed to create API client")
|
||||
}
|
||||
|
||||
|
||||
@@ -44,8 +44,12 @@ class Hindsight:
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
# Without authentication
|
||||
client = Hindsight(base_url="http://localhost:8888")
|
||||
|
||||
# With API key authentication
|
||||
client = Hindsight(base_url="http://localhost:8888", api_key="your-api-key")
|
||||
|
||||
# Store a memory
|
||||
client.retain(bank_id="alice", content="Alice loves AI")
|
||||
|
||||
@@ -59,15 +63,16 @@ class Hindsight:
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, base_url: str, timeout: float = 30.0):
|
||||
def __init__(self, base_url: str, api_key: Optional[str] = None, timeout: float = 30.0):
|
||||
"""
|
||||
Initialize the Hindsight client.
|
||||
|
||||
Args:
|
||||
base_url: The base URL of the Hindsight API server
|
||||
api_key: Optional API key for authentication (sent as Bearer token)
|
||||
timeout: Request timeout in seconds (default: 30.0)
|
||||
"""
|
||||
config = hindsight_client_api.Configuration(host=base_url)
|
||||
config = hindsight_client_api.Configuration(host=base_url, access_token=api_key)
|
||||
self._api_client = hindsight_client_api.ApiClient(config)
|
||||
self._api = default_api.DefaultApi(self._api_client)
|
||||
|
||||
@@ -80,9 +85,21 @@ class Hindsight:
|
||||
self.close()
|
||||
|
||||
def close(self):
|
||||
"""Close the API client."""
|
||||
"""Close the API client (sync version - use aclose() in async code)."""
|
||||
if self._api_client:
|
||||
_run_async(self._api_client.close())
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
# We're in an async context - schedule but don't wait
|
||||
# The caller should use aclose() instead
|
||||
loop.create_task(self._api_client.close())
|
||||
except RuntimeError:
|
||||
# No running loop - safe to run synchronously
|
||||
_run_async(self._api_client.close())
|
||||
|
||||
async def aclose(self):
|
||||
"""Close the API client (async version)."""
|
||||
if self._api_client:
|
||||
await self._api_client.close()
|
||||
|
||||
# Simplified methods for main operations
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "hindsight-client"
|
||||
version = "0.1.11"
|
||||
version = "0.1.12"
|
||||
description = "Python client for Hindsight - Semantic memory system with personality-driven thinking"
|
||||
authors = [
|
||||
{name = "Hindsight Team"}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-client",
|
||||
"version": "0.1.11",
|
||||
"version": "0.1.12",
|
||||
"description": "TypeScript client for Hindsight - Semantic memory system with personality-driven thinking",
|
||||
"main": "./dist/src/index.js",
|
||||
"types": "./dist/src/index.d.ts",
|
||||
|
||||
@@ -5,8 +5,15 @@
|
||||
* ```typescript
|
||||
* import { HindsightClient } from '@vectorize-io/hindsight-client';
|
||||
*
|
||||
* // Without authentication
|
||||
* const client = new HindsightClient({ baseUrl: 'http://localhost:8888' });
|
||||
*
|
||||
* // With API key authentication
|
||||
* const client = new HindsightClient({
|
||||
* baseUrl: 'http://localhost:8888',
|
||||
* apiKey: 'your-api-key'
|
||||
* });
|
||||
*
|
||||
* // Retain a memory
|
||||
* await client.retain('alice', 'Alice loves AI');
|
||||
*
|
||||
@@ -37,6 +44,10 @@ import type {
|
||||
|
||||
export interface HindsightClientOptions {
|
||||
baseUrl: string;
|
||||
/**
|
||||
* Optional API key for authentication (sent as Bearer token in Authorization header)
|
||||
*/
|
||||
apiKey?: string;
|
||||
}
|
||||
|
||||
export interface MemoryItemInput {
|
||||
@@ -54,6 +65,9 @@ export class HindsightClient {
|
||||
this.client = createClient(
|
||||
createConfig({
|
||||
baseUrl: options.baseUrl,
|
||||
headers: options.apiKey
|
||||
? { Authorization: `Bearer ${options.apiKey}` }
|
||||
: undefined,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-control-plane",
|
||||
"version": "0.1.11",
|
||||
"version": "0.1.12",
|
||||
"description": "Control plane for Hindsight - Semantic memory system",
|
||||
"bin": {
|
||||
"hindsight-control-plane": "./bin/cli.js"
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -3,28 +3,28 @@ LoComo-specific benchmark implementations.
|
||||
|
||||
Provides dataset, answer generator, and evaluator for the LoComo benchmark.
|
||||
"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from benchmarks.common.benchmark_runner import BenchmarkRunner
|
||||
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from typing import List, Dict, Any, Tuple, Optional
|
||||
import asyncio
|
||||
import pydantic
|
||||
from openai import AsyncOpenAI
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from benchmarks.common.benchmark_runner import BenchmarkDataset, LLMAnswerGenerator, LLMAnswerEvaluator
|
||||
import pydantic
|
||||
from hindsight_api.engine.llm_wrapper import LLMConfig
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from benchmarks.common.benchmark_runner import BenchmarkDataset, BenchmarkRunner, LLMAnswerEvaluator, LLMAnswerGenerator
|
||||
|
||||
|
||||
class LoComoDataset(BenchmarkDataset):
|
||||
"""LoComo dataset implementation."""
|
||||
|
||||
def load(self, path: Path, max_items: Optional[int] = None) -> List[Dict[str, Any]]:
|
||||
"""Load LoComo dataset from JSON file."""
|
||||
with open(path, 'r') as f:
|
||||
with open(path, "r") as f:
|
||||
dataset = json.load(f)
|
||||
|
||||
if max_items:
|
||||
@@ -34,7 +34,7 @@ class LoComoDataset(BenchmarkDataset):
|
||||
|
||||
def get_item_id(self, item: Dict) -> str:
|
||||
"""Get sample ID from LoComo item."""
|
||||
return item['sample_id']
|
||||
return item["sample_id"]
|
||||
|
||||
def prepare_sessions_for_ingestion(self, item: Dict) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
@@ -45,12 +45,12 @@ class LoComoDataset(BenchmarkDataset):
|
||||
Returns:
|
||||
List of session dicts, each containing 'content', 'context', 'event_date', 'document_id'
|
||||
"""
|
||||
conv = item['conversation']
|
||||
speaker_a = conv['speaker_a']
|
||||
speaker_b = conv['speaker_b']
|
||||
conv = item["conversation"]
|
||||
speaker_a = conv["speaker_a"]
|
||||
speaker_b = conv["speaker_b"]
|
||||
|
||||
# Get all session keys sorted
|
||||
session_keys = sorted([k for k in conv.keys() if k.startswith('session_') and not k.endswith('_date_time')])
|
||||
session_keys = sorted([k for k in conv.keys() if k.startswith("session_") and not k.endswith("_date_time")])
|
||||
|
||||
session_items = []
|
||||
|
||||
@@ -65,12 +65,14 @@ class LoComoDataset(BenchmarkDataset):
|
||||
session_date = self._parse_date(conv.get(date_key))
|
||||
session_content = json.dumps(session_data)
|
||||
document_id = f"{item['sample_id']}_{session_key}"
|
||||
session_items.append({
|
||||
"content": session_content,
|
||||
"context": f"Conversation between {speaker_a} and {speaker_b} ({session_key} of {item['sample_id']})",
|
||||
"event_date": session_date,
|
||||
"document_id": document_id
|
||||
})
|
||||
session_items.append(
|
||||
{
|
||||
"content": session_content,
|
||||
"context": f"Conversation between {speaker_a} and {speaker_b} ({session_key} of {item['sample_id']})",
|
||||
"event_date": session_date,
|
||||
"document_id": document_id,
|
||||
}
|
||||
)
|
||||
|
||||
return session_items
|
||||
|
||||
@@ -81,7 +83,7 @@ class LoComoDataset(BenchmarkDataset):
|
||||
Returns:
|
||||
List of QA dicts with 'question', 'answer', 'category'
|
||||
"""
|
||||
return item['qa']
|
||||
return item["qa"]
|
||||
|
||||
def _parse_date(self, date_string: str) -> datetime:
|
||||
"""Parse LoComo date format to datetime."""
|
||||
@@ -95,6 +97,7 @@ class LoComoDataset(BenchmarkDataset):
|
||||
|
||||
class QuestionAnswer(pydantic.BaseModel):
|
||||
"""Answer format for LoComo questions."""
|
||||
|
||||
answer: str
|
||||
reasoning: str
|
||||
|
||||
@@ -113,7 +116,7 @@ class LoComoAnswerGenerator(LLMAnswerGenerator):
|
||||
question: str,
|
||||
recall_result: Dict[str, Any],
|
||||
question_date: Optional[datetime] = None,
|
||||
question_type: Optional[str] = None
|
||||
question_type: Optional[str] = None,
|
||||
) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
|
||||
"""
|
||||
Generate answer from retrieved memories using Groq.
|
||||
@@ -141,7 +144,7 @@ class LoComoAnswerGenerator(LLMAnswerGenerator):
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful expert assistant answering questions from lme_experiment users based on the provided context."
|
||||
"content": "You are a helpful expert assistant answering questions from lme_experiment users based on the provided context.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
@@ -165,11 +168,11 @@ Context:
|
||||
Question: {question}
|
||||
Answer:
|
||||
|
||||
"""
|
||||
}
|
||||
""",
|
||||
},
|
||||
],
|
||||
response_format=QuestionAnswer,
|
||||
scope="memory"
|
||||
scope="memory",
|
||||
)
|
||||
return answer_obj.answer, answer_obj.reasoning, None
|
||||
except Exception as e:
|
||||
@@ -183,7 +186,7 @@ class LoComoThinkAnswerGenerator(LLMAnswerGenerator):
|
||||
so it doesn't need external search to be performed by the benchmark runner.
|
||||
"""
|
||||
|
||||
def __init__(self, memory: 'MemoryEngine', agent_id: str, thinking_budget: int = 500):
|
||||
def __init__(self, memory: "MemoryEngine", agent_id: str, thinking_budget: int = 500):
|
||||
"""Initialize with memory instance and agent_id.
|
||||
|
||||
Args:
|
||||
@@ -204,7 +207,7 @@ class LoComoThinkAnswerGenerator(LLMAnswerGenerator):
|
||||
question: str,
|
||||
recall_result: Dict[str, Any],
|
||||
question_date: Optional[datetime] = None,
|
||||
question_type: Optional[str] = None
|
||||
question_type: Optional[str] = None,
|
||||
) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
|
||||
"""
|
||||
Generate answer using the integrated think API.
|
||||
@@ -235,9 +238,9 @@ class LoComoThinkAnswerGenerator(LLMAnswerGenerator):
|
||||
|
||||
# Extract memories from based_on
|
||||
based_on = result.based_on
|
||||
world_facts = based_on.get('world', [])
|
||||
agent_facts = based_on.get('agent', [])
|
||||
opinion_facts = based_on.get('opinion', [])
|
||||
world_facts = based_on.get("world", [])
|
||||
agent_facts = based_on.get("agent", [])
|
||||
opinion_facts = based_on.get("opinion", [])
|
||||
|
||||
# Combine all facts into retrieved_memories
|
||||
retrieved_memories = []
|
||||
@@ -271,7 +274,7 @@ async def run_benchmark(
|
||||
api_url: str = None,
|
||||
max_concurrent_questions_override: int = None,
|
||||
only_failed: bool = False,
|
||||
only_invalid: bool = False
|
||||
only_invalid: bool = False,
|
||||
):
|
||||
"""
|
||||
Run the LoComo benchmark.
|
||||
@@ -287,6 +290,7 @@ async def run_benchmark(
|
||||
only_invalid: If True, only run conversations that have invalid questions (is_invalid=True)
|
||||
"""
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
# Load previous results if filtering for failed/invalid conversations
|
||||
@@ -294,35 +298,41 @@ async def run_benchmark(
|
||||
invalid_conversation_ids = set()
|
||||
if only_failed or only_invalid:
|
||||
suffix = "_think" if use_think else ""
|
||||
results_filename = f'benchmark_results{suffix}.json'
|
||||
results_path = Path(__file__).parent / 'results' / results_filename
|
||||
results_filename = f"benchmark_results{suffix}.json"
|
||||
results_path = Path(__file__).parent / "results" / results_filename
|
||||
|
||||
if not results_path.exists():
|
||||
console.print(f"[red]Error: Cannot use --only-failed or --only-invalid without existing results file[/red]")
|
||||
console.print("[red]Error: Cannot use --only-failed or --only-invalid without existing results file[/red]")
|
||||
console.print(f"[yellow]Results file not found: {results_path}[/yellow]")
|
||||
return
|
||||
|
||||
with open(results_path, 'r') as f:
|
||||
with open(results_path, "r") as f:
|
||||
previous_results = json.load(f)
|
||||
|
||||
# Extract conversation IDs that have failed or invalid questions
|
||||
for item_result in previous_results.get('item_results', []):
|
||||
item_id = item_result['item_id']
|
||||
for detail in item_result['metrics'].get('detailed_results', []):
|
||||
if only_failed and detail.get('is_correct') == False and not detail.get('is_invalid', False):
|
||||
for item_result in previous_results.get("item_results", []):
|
||||
item_id = item_result["item_id"]
|
||||
for detail in item_result["metrics"].get("detailed_results", []):
|
||||
if only_failed and detail.get("is_correct") == False and not detail.get("is_invalid", False):
|
||||
failed_conversation_ids.add(item_id)
|
||||
if only_invalid and detail.get('is_invalid', False):
|
||||
if only_invalid and detail.get("is_invalid", False):
|
||||
invalid_conversation_ids.add(item_id)
|
||||
|
||||
if only_failed:
|
||||
console.print(f"[cyan]Filtering to {len(failed_conversation_ids)} conversations with failed questions (is_correct=False)[/cyan]")
|
||||
console.print(
|
||||
f"[cyan]Filtering to {len(failed_conversation_ids)} conversations with failed questions (is_correct=False)[/cyan]"
|
||||
)
|
||||
if only_invalid:
|
||||
console.print(f"[cyan]Filtering to {len(invalid_conversation_ids)} conversations with invalid questions (is_invalid=True)[/cyan]")
|
||||
console.print(
|
||||
f"[cyan]Filtering to {len(invalid_conversation_ids)} conversations with invalid questions (is_invalid=True)[/cyan]"
|
||||
)
|
||||
|
||||
target_ids = failed_conversation_ids if only_failed else invalid_conversation_ids
|
||||
if not target_ids:
|
||||
filter_type = "failed" if only_failed else "invalid"
|
||||
console.print(f"[yellow]No conversations with {filter_type} questions found in previous results. Nothing to run.[/yellow]")
|
||||
console.print(
|
||||
f"[yellow]No conversations with {filter_type} questions found in previous results. Nothing to run.[/yellow]"
|
||||
)
|
||||
return
|
||||
|
||||
# Initialize components
|
||||
@@ -331,18 +341,16 @@ async def run_benchmark(
|
||||
# Use remote API client if api_url is provided, otherwise use local memory
|
||||
if api_url:
|
||||
from benchmarks.common.benchmark_runner import HindsightClientAdapter
|
||||
|
||||
memory = HindsightClientAdapter(base_url=api_url)
|
||||
await memory.initialize()
|
||||
else:
|
||||
from benchmarks.common.benchmark_runner import create_memory_engine
|
||||
|
||||
memory = await create_memory_engine()
|
||||
|
||||
if use_think:
|
||||
answer_generator = LoComoThinkAnswerGenerator(
|
||||
memory=memory,
|
||||
agent_id="locomo",
|
||||
thinking_budget=500
|
||||
)
|
||||
answer_generator = LoComoThinkAnswerGenerator(memory=memory, agent_id="locomo", thinking_budget=500)
|
||||
max_concurrent_questions = max_concurrent_questions_override or 4
|
||||
eval_semaphore_size = 4
|
||||
else:
|
||||
@@ -356,14 +364,11 @@ async def run_benchmark(
|
||||
|
||||
# Create benchmark runner
|
||||
runner = BenchmarkRunner(
|
||||
dataset=dataset,
|
||||
answer_generator=answer_generator,
|
||||
answer_evaluator=answer_evaluator,
|
||||
memory=memory
|
||||
dataset=dataset, answer_generator=answer_generator, answer_evaluator=answer_evaluator, memory=memory
|
||||
)
|
||||
|
||||
# Filter dataset if using --only-failed or --only-invalid
|
||||
dataset_path = Path(__file__).parent / 'datasets' / 'locomo10.json'
|
||||
dataset_path = Path(__file__).parent / "datasets" / "locomo10.json"
|
||||
|
||||
if only_failed or only_invalid:
|
||||
# Load and filter dataset
|
||||
@@ -374,14 +379,16 @@ async def run_benchmark(
|
||||
|
||||
# Temporarily replace dataset's load method
|
||||
original_load = dataset.load
|
||||
|
||||
def filtered_load(path: Path, max_items: Optional[int] = None):
|
||||
return filtered_items[:max_items] if max_items else filtered_items
|
||||
|
||||
dataset.load = filtered_load
|
||||
|
||||
# Determine output filename based on mode
|
||||
suffix = "_think" if use_think else ""
|
||||
results_filename = f'benchmark_results{suffix}.json'
|
||||
output_path = Path(__file__).parent / 'results' / results_filename
|
||||
results_filename = f"benchmark_results{suffix}.json"
|
||||
output_path = Path(__file__).parent / "results" / results_filename
|
||||
|
||||
# Create results directory if it doesn't exist
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
@@ -406,7 +413,7 @@ async def run_benchmark(
|
||||
clear_agent_per_item=True, # Use unique agent ID per conversation
|
||||
max_concurrent_items=3, # Process up to 3 conversations in parallel
|
||||
output_path=output_path, # Save results incrementally
|
||||
merge_with_existing=merge_with_existing
|
||||
merge_with_existing=merge_with_existing,
|
||||
)
|
||||
|
||||
# Display results (final save already happened incrementally)
|
||||
@@ -430,14 +437,10 @@ def generate_markdown_table(results: dict, use_think: bool = False):
|
||||
4 = Open-domain
|
||||
"""
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
category_names = {
|
||||
'1': 'Multi-hop',
|
||||
'2': 'Single-hop',
|
||||
'3': 'Temporal',
|
||||
'4': 'Open-domain'
|
||||
}
|
||||
category_names = {"1": "Multi-hop", "2": "Single-hop", "3": "Temporal", "4": "Open-domain"}
|
||||
|
||||
# Build markdown content
|
||||
lines = []
|
||||
@@ -446,33 +449,41 @@ def generate_markdown_table(results: dict, use_think: bool = False):
|
||||
lines.append("")
|
||||
|
||||
# Add model configuration
|
||||
if 'model_config' in results:
|
||||
config = results['model_config']
|
||||
if "model_config" in results:
|
||||
config = results["model_config"]
|
||||
lines.append("## Model Configuration")
|
||||
lines.append("")
|
||||
lines.append(f"- **Hindsight**: {config['hindsight']['provider']}/{config['hindsight']['model']}")
|
||||
lines.append(f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}")
|
||||
lines.append(
|
||||
f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}"
|
||||
)
|
||||
lines.append(f"- **LLM Judge**: {config['judge']['provider']}/{config['judge']['model']}")
|
||||
lines.append("")
|
||||
|
||||
lines.append(f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})")
|
||||
lines.append(
|
||||
f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})"
|
||||
)
|
||||
lines.append("")
|
||||
lines.append("| Sample ID | Sessions | Questions | Correct | Accuracy | Multi-hop | Single-hop | Temporal | Open-domain |")
|
||||
lines.append("|-----------|----------|-----------|---------|----------|-----------|------------|----------|-------------|")
|
||||
lines.append(
|
||||
"| Sample ID | Sessions | Questions | Correct | Accuracy | Multi-hop | Single-hop | Temporal | Open-domain |"
|
||||
)
|
||||
lines.append(
|
||||
"|-----------|----------|-----------|---------|----------|-----------|------------|----------|-------------|"
|
||||
)
|
||||
|
||||
for item_result in results['item_results']:
|
||||
item_id = item_result['item_id']
|
||||
num_sessions = item_result['num_sessions']
|
||||
metrics = item_result['metrics']
|
||||
for item_result in results["item_results"]:
|
||||
item_id = item_result["item_id"]
|
||||
num_sessions = item_result["num_sessions"]
|
||||
metrics = item_result["metrics"]
|
||||
|
||||
# Calculate category accuracies
|
||||
cat_stats = metrics.get('category_stats', {})
|
||||
cat_stats = metrics.get("category_stats", {})
|
||||
cat_accuracies = {}
|
||||
|
||||
for cat_id in ['1', '2', '3', '4']:
|
||||
for cat_id in ["1", "2", "3", "4"]:
|
||||
if cat_id in cat_stats:
|
||||
stats = cat_stats[cat_id]
|
||||
acc = (stats['correct'] / stats['total'] * 100) if stats['total'] > 0 else 0
|
||||
acc = (stats["correct"] / stats["total"] * 100) if stats["total"] > 0 else 0
|
||||
cat_accuracies[cat_id] = f"{acc:.1f}% ({stats['correct']}/{stats['total']})"
|
||||
else:
|
||||
cat_accuracies[cat_id] = "N/A"
|
||||
@@ -485,28 +496,48 @@ def generate_markdown_table(results: dict, use_think: bool = False):
|
||||
|
||||
# Write to file with suffix
|
||||
suffix = "_think" if use_think else ""
|
||||
output_file = Path(__file__).parent / 'results' / f'results_table{suffix}.md'
|
||||
output_file = Path(__file__).parent / "results" / f"results_table{suffix}.md"
|
||||
output_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_file.write_text('\n'.join(lines))
|
||||
output_file.write_text("\n".join(lines))
|
||||
console.print(f"\n[green]✓[/green] Results table saved to {output_file}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import logging
|
||||
import argparse
|
||||
import logging
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
|
||||
parser = argparse.ArgumentParser(description='Run LoComo benchmark')
|
||||
parser.add_argument('--max-conversations', type=int, default=None, help='Maximum conversations to evaluate')
|
||||
parser.add_argument('--max-questions', type=int, default=None, help='Maximum questions per conversation')
|
||||
parser.add_argument('--skip-ingestion', action='store_true', help='Skip ingestion and use existing data')
|
||||
parser.add_argument('--use-think', action='store_true', help='Use think API instead of search + LLM')
|
||||
parser.add_argument('--conversation', type=str, default=None, help='Run only specific conversation (e.g., "conv-26")')
|
||||
parser.add_argument('--api-url', type=str, default=None, help='Hindsight API URL (default: use local memory, example: http://localhost:8888)')
|
||||
parser.add_argument('--max-concurrent-questions', type=int, default=None, help='Max concurrent questions per conversation (default: 4 for think, 10 for search)')
|
||||
parser.add_argument('--only-failed', action='store_true', help='Only run conversations that have failed questions (is_correct=False). Requires existing results file.')
|
||||
parser.add_argument('--only-invalid', action='store_true', help='Only run conversations that have invalid questions (is_invalid=True). Requires existing results file.')
|
||||
parser = argparse.ArgumentParser(description="Run LoComo benchmark")
|
||||
parser.add_argument("--max-conversations", type=int, default=None, help="Maximum conversations to evaluate")
|
||||
parser.add_argument("--max-questions", type=int, default=None, help="Maximum questions per conversation")
|
||||
parser.add_argument("--skip-ingestion", action="store_true", help="Skip ingestion and use existing data")
|
||||
parser.add_argument("--use-think", action="store_true", help="Use think API instead of search + LLM")
|
||||
parser.add_argument(
|
||||
"--conversation", type=str, default=None, help='Run only specific conversation (e.g., "conv-26")'
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api-url",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Hindsight API URL (default: use local memory, example: http://localhost:8888)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-concurrent-questions",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Max concurrent questions per conversation (default: 4 for think, 10 for search)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-failed",
|
||||
action="store_true",
|
||||
help="Only run conversations that have failed questions (is_correct=False). Requires existing results file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-invalid",
|
||||
action="store_true",
|
||||
help="Only run conversations that have invalid questions (is_invalid=True). Requires existing results file.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -514,14 +545,16 @@ if __name__ == "__main__":
|
||||
if args.only_failed and args.only_invalid:
|
||||
parser.error("Cannot use both --only-failed and --only-invalid at the same time")
|
||||
|
||||
results = asyncio.run(run_benchmark(
|
||||
max_conversations=args.max_conversations,
|
||||
max_questions_per_conv=args.max_questions,
|
||||
skip_ingestion=args.skip_ingestion,
|
||||
use_think=args.use_think,
|
||||
conversation=args.conversation,
|
||||
api_url=args.api_url,
|
||||
max_concurrent_questions_override=args.max_concurrent_questions,
|
||||
only_failed=args.only_failed,
|
||||
only_invalid=args.only_invalid
|
||||
))
|
||||
results = asyncio.run(
|
||||
run_benchmark(
|
||||
max_conversations=args.max_conversations,
|
||||
max_questions_per_conv=args.max_questions,
|
||||
skip_ingestion=args.skip_ingestion,
|
||||
use_think=args.use_think,
|
||||
conversation=args.conversation,
|
||||
api_url=args.api_url,
|
||||
max_concurrent_questions_override=args.max_concurrent_questions,
|
||||
only_failed=args.only_failed,
|
||||
only_invalid=args.only_invalid,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -3,21 +3,20 @@ LongMemEval-specific benchmark implementations.
|
||||
|
||||
Provides dataset, answer generator, and evaluator for the LongMemEval benchmark.
|
||||
"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from benchmarks.common.benchmark_runner import BenchmarkRunner
|
||||
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from typing import List, Dict, Any, Tuple, Optional
|
||||
import asyncio
|
||||
import pydantic
|
||||
from openai import AsyncOpenAI
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from benchmarks.common.benchmark_runner import BenchmarkDataset, LLMAnswerGenerator, LLMAnswerEvaluator
|
||||
import pydantic
|
||||
from hindsight_api.engine.llm_wrapper import LLMConfig
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from benchmarks.common.benchmark_runner import BenchmarkDataset, BenchmarkRunner, LLMAnswerEvaluator, LLMAnswerGenerator
|
||||
|
||||
|
||||
class LongMemEvalDataset(BenchmarkDataset):
|
||||
@@ -25,7 +24,7 @@ class LongMemEvalDataset(BenchmarkDataset):
|
||||
|
||||
def load(self, path: Path, max_items: Optional[int] = None) -> List[Dict[str, Any]]:
|
||||
"""Load LongMemEval dataset from JSON file."""
|
||||
with open(path, 'r') as f:
|
||||
with open(path, "r") as f:
|
||||
dataset = json.load(f)
|
||||
|
||||
if max_items:
|
||||
@@ -67,7 +66,7 @@ class LongMemEvalDataset(BenchmarkDataset):
|
||||
for turn in session_turns:
|
||||
if isinstance(turn, dict):
|
||||
# Create a copy without has_answer
|
||||
cleaned_turn = {k: v for k, v in turn.items() if k != 'has_answer'}
|
||||
cleaned_turn = {k: v for k, v in turn.items() if k != "has_answer"}
|
||||
cleaned_turns.append(cleaned_turn)
|
||||
else:
|
||||
cleaned_turns.append(turn)
|
||||
@@ -75,12 +74,14 @@ class LongMemEvalDataset(BenchmarkDataset):
|
||||
session_content = json.dumps(cleaned_turns)
|
||||
question_id = item.get("question_id", "unknown")
|
||||
document_id = f"{question_id}_{session_id}"
|
||||
batch_contents.append({
|
||||
"content": session_content,
|
||||
"context": f"Session {document_id} - you are the assistant in this conversation - happened on {session_date.strftime('%Y-%m-%d %H:%M:%S')} UTC.",
|
||||
"event_date": session_date,
|
||||
"document_id": document_id
|
||||
})
|
||||
batch_contents.append(
|
||||
{
|
||||
"content": session_content,
|
||||
"context": f"Session {document_id} - you are the assistant in this conversation - happened on {session_date.strftime('%Y-%m-%d %H:%M:%S')} UTC.",
|
||||
"event_date": session_date,
|
||||
"document_id": document_id,
|
||||
}
|
||||
)
|
||||
|
||||
return batch_contents
|
||||
|
||||
@@ -95,22 +96,24 @@ class LongMemEvalDataset(BenchmarkDataset):
|
||||
"""
|
||||
# Parse question_date if available
|
||||
question_date = None
|
||||
if 'question_date' in item:
|
||||
question_date = self._parse_date(item['question_date'])
|
||||
if "question_date" in item:
|
||||
question_date = self._parse_date(item["question_date"])
|
||||
|
||||
return [{
|
||||
'question': item.get("question", ""),
|
||||
'answer': item.get("answer", ""),
|
||||
'category': item.get("question_type", "unknown"),
|
||||
'question_date': question_date
|
||||
}]
|
||||
return [
|
||||
{
|
||||
"question": item.get("question", ""),
|
||||
"answer": item.get("answer", ""),
|
||||
"category": item.get("question_type", "unknown"),
|
||||
"question_date": question_date,
|
||||
}
|
||||
]
|
||||
|
||||
def _parse_date(self, date_str: str) -> datetime:
|
||||
"""Parse date string to datetime object."""
|
||||
try:
|
||||
# LongMemEval format: "2023/05/20 (Sat) 02:21"
|
||||
# Try to parse the main part before the day name
|
||||
date_str_cleaned = date_str.split('(')[0].strip() if '(' in date_str else date_str
|
||||
date_str_cleaned = date_str.split("(")[0].strip() if "(" in date_str else date_str
|
||||
|
||||
# Try multiple formats
|
||||
for fmt in ["%Y/%m/%d %H:%M", "%Y-%m-%d %H:%M:%S", "%Y-%m-%d", "%Y/%m/%d"]:
|
||||
@@ -121,7 +124,7 @@ class LongMemEvalDataset(BenchmarkDataset):
|
||||
continue
|
||||
|
||||
# Fallback: try ISO format
|
||||
return datetime.fromisoformat(date_str.replace('Z', '+00:00'))
|
||||
return datetime.fromisoformat(date_str.replace("Z", "+00:00"))
|
||||
except Exception:
|
||||
raise ValueError(f"Failed to parse date string: {date_str}")
|
||||
|
||||
@@ -130,6 +133,7 @@ class QuestionAnswer(pydantic.BaseModel):
|
||||
answer: str
|
||||
reasoning: Optional[str] = None
|
||||
|
||||
|
||||
class LongMemEvalAnswerGenerator(LLMAnswerGenerator):
|
||||
"""LongMemEval-specific answer generator using configurable LLM provider."""
|
||||
|
||||
@@ -202,10 +206,7 @@ class LongMemEvalAnswerGenerator(LLMAnswerGenerator):
|
||||
chunk_text = chunk_info.get("chunk_text", "")
|
||||
|
||||
# Build the formatted fact entry
|
||||
entry_parts = [
|
||||
f"Fact {i} ({fact_type}): {fact_text}",
|
||||
f"When: {when_str}"
|
||||
]
|
||||
entry_parts = [f"Fact {i} ({fact_type}): {fact_text}", f"When: {when_str}"]
|
||||
|
||||
# Add context field if present
|
||||
context = fact.get("context")
|
||||
@@ -217,7 +218,7 @@ class LongMemEvalAnswerGenerator(LLMAnswerGenerator):
|
||||
# Truncate very long chunks
|
||||
if len(chunk_text) > 1000:
|
||||
chunk_text = chunk_text[:1000] + "..."
|
||||
entry_parts.append(f"Source chunk:\n \"{chunk_text}\"")
|
||||
entry_parts.append(f'Source chunk:\n "{chunk_text}"')
|
||||
|
||||
formatted_parts.append("\n".join(entry_parts))
|
||||
|
||||
@@ -326,43 +327,43 @@ The context contains memory facts extracted from previous conversations, each wi
|
||||
return ""
|
||||
|
||||
async def generate_answer(
|
||||
self,
|
||||
question: str,
|
||||
recall_result: Dict[str, Any],
|
||||
question_date: Optional[datetime] = None,
|
||||
question_type: Optional[str] = None
|
||||
) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
|
||||
"""
|
||||
Generate answer from retrieved memories using Groq.
|
||||
self,
|
||||
question: str,
|
||||
recall_result: Dict[str, Any],
|
||||
question_date: Optional[datetime] = None,
|
||||
question_type: Optional[str] = None,
|
||||
) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
|
||||
"""
|
||||
Generate answer from retrieved memories using Groq.
|
||||
|
||||
Args:
|
||||
question: The question text
|
||||
recall_result: Full RecallResult dict containing results, entities, chunks, and trace
|
||||
question_date: Date when the question was asked (for temporal context)
|
||||
question_type: Question category (e.g., 'single-session-user', 'multi-session-assistant')
|
||||
Args:
|
||||
question: The question text
|
||||
recall_result: Full RecallResult dict containing results, entities, chunks, and trace
|
||||
question_date: Date when the question was asked (for temporal context)
|
||||
question_type: Question category (e.g., 'single-session-user', 'multi-session-assistant')
|
||||
|
||||
Returns:
|
||||
Tuple of (answer, reasoning, None)
|
||||
- None indicates to use the memories from recall_result
|
||||
"""
|
||||
# Format context based on selected mode
|
||||
if self.context_format == "structured":
|
||||
context = self._format_context_structured(recall_result)
|
||||
else:
|
||||
context = self._format_context_json(recall_result)
|
||||
Returns:
|
||||
Tuple of (answer, reasoning, None)
|
||||
- None indicates to use the memories from recall_result
|
||||
"""
|
||||
# Format context based on selected mode
|
||||
if self.context_format == "structured":
|
||||
context = self._format_context_structured(recall_result)
|
||||
else:
|
||||
context = self._format_context_json(recall_result)
|
||||
|
||||
context_instructions = self._get_context_instructions()
|
||||
context_instructions = self._get_context_instructions()
|
||||
|
||||
# Format question date if provided
|
||||
formatted_question_date = question_date.strftime('%Y-%m-%d %H:%M:%S UTC') if question_date else "Not specified"
|
||||
# Format question date if provided
|
||||
formatted_question_date = question_date.strftime("%Y-%m-%d %H:%M:%S UTC") if question_date else "Not specified"
|
||||
|
||||
# Use LLM to generate answer
|
||||
try:
|
||||
answer_obj = await self.llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"""You are a helpful assistant that must answer user questions based on the previous conversations.
|
||||
# Use LLM to generate answer
|
||||
try:
|
||||
answer_obj = await self.llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"""You are a helpful assistant that must answer user questions based on the previous conversations.
|
||||
|
||||
{context_instructions}**Answer Guidelines:**
|
||||
1. Start by scanning retrieved context to understand the facts and events that happened and the timeline.
|
||||
@@ -393,20 +394,20 @@ Retrieved Context:
|
||||
|
||||
|
||||
Answer:
|
||||
"""
|
||||
}
|
||||
],
|
||||
response_format=QuestionAnswer,
|
||||
scope="memory",
|
||||
max_completion_tokens=32768,
|
||||
)
|
||||
reasoning_text = answer_obj.reasoning or ""
|
||||
if reasoning_text:
|
||||
reasoning_text = reasoning_text + " "
|
||||
reasoning_text += f"(question date: {formatted_question_date})"
|
||||
return answer_obj.answer, reasoning_text, None
|
||||
except Exception as e:
|
||||
return f"Error generating answer: {str(e)}", "Error occurred during answer generation.", None
|
||||
""",
|
||||
}
|
||||
],
|
||||
response_format=QuestionAnswer,
|
||||
scope="memory",
|
||||
max_completion_tokens=32768,
|
||||
)
|
||||
reasoning_text = answer_obj.reasoning or ""
|
||||
if reasoning_text:
|
||||
reasoning_text = reasoning_text + " "
|
||||
reasoning_text += f"(question date: {formatted_question_date})"
|
||||
return answer_obj.answer, reasoning_text, None
|
||||
except Exception as e:
|
||||
return f"Error generating answer: {str(e)}", "Error occurred during answer generation.", None
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
@@ -425,7 +426,7 @@ async def run_benchmark(
|
||||
max_concurrent_items: int = 1,
|
||||
results_filename: str = "benchmark_results.json",
|
||||
context_format: str = "json",
|
||||
source_results: str = None
|
||||
source_results: str = None,
|
||||
):
|
||||
"""
|
||||
Run the LongMemEval benchmark.
|
||||
@@ -449,13 +450,16 @@ async def run_benchmark(
|
||||
source_results: Source results file to read failed/invalid questions from (for --only-failed/--only-invalid). Defaults to benchmark_results.json.
|
||||
"""
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
# Validate mutually exclusive arguments
|
||||
# --max-instances-per-category can't be combined with --max-instances or --category
|
||||
# But --category CAN be combined with --max-instances (to limit questions within a category)
|
||||
if max_instances_per_category is not None and (max_instances is not None or category is not None):
|
||||
console.print("[red]Error: --max-questions-per-category cannot be combined with --max-instances or --category[/red]")
|
||||
console.print(
|
||||
"[red]Error: --max-questions-per-category cannot be combined with --max-instances or --category[/red]"
|
||||
)
|
||||
return
|
||||
|
||||
# Validate --only-ingested can't be combined with other dataset filters
|
||||
@@ -480,8 +484,10 @@ async def run_benchmark(
|
||||
dataset_path = Path(__file__).parent / "datasets" / "longmemeval_s_cleaned.json"
|
||||
if not dataset_path.exists():
|
||||
if not download_dataset(dataset_path):
|
||||
console.print(f"[red]Failed to download dataset. Please download manually:[/red]")
|
||||
console.print("[yellow]curl -L 'https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json' -o benchmarks/longmemeval/datasets/longmemeval_s_cleaned.json[/yellow]")
|
||||
console.print("[red]Failed to download dataset. Please download manually:[/red]")
|
||||
console.print(
|
||||
"[yellow]curl -L 'https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json' -o benchmarks/longmemeval/datasets/longmemeval_s_cleaned.json[/yellow]"
|
||||
)
|
||||
return
|
||||
|
||||
# Initialize components
|
||||
@@ -499,9 +505,10 @@ async def run_benchmark(
|
||||
|
||||
# Group by category and take max_instances_per_category from each
|
||||
from collections import defaultdict
|
||||
|
||||
category_items = defaultdict(list)
|
||||
for item in original_dataset_items:
|
||||
cat = item.get('question_type', 'unknown')
|
||||
cat = item.get("question_type", "unknown")
|
||||
category_items[cat].append(item)
|
||||
|
||||
# Take up to max_instances_per_category from each category
|
||||
@@ -518,30 +525,34 @@ async def run_benchmark(
|
||||
invalid_question_ids = set()
|
||||
if only_failed or only_invalid:
|
||||
# Use source_results if specified, otherwise default to benchmark_results.json
|
||||
source_file = source_results if source_results else 'benchmark_results.json'
|
||||
results_path = Path(__file__).parent / 'results' / source_file
|
||||
source_file = source_results if source_results else "benchmark_results.json"
|
||||
results_path = Path(__file__).parent / "results" / source_file
|
||||
if not results_path.exists():
|
||||
console.print(f"[red]Error: Cannot use --only-failed or --only-invalid without existing results file[/red]")
|
||||
console.print("[red]Error: Cannot use --only-failed or --only-invalid without existing results file[/red]")
|
||||
console.print(f"[yellow]Results file not found: {results_path}[/yellow]")
|
||||
return
|
||||
|
||||
console.print(f"[cyan]Reading failed/invalid questions from: {source_file}[/cyan]")
|
||||
with open(results_path, 'r') as f:
|
||||
with open(results_path, "r") as f:
|
||||
previous_results = json.load(f)
|
||||
|
||||
# Extract question IDs that failed or are invalid
|
||||
for item_result in previous_results.get('item_results', []):
|
||||
item_id = item_result['item_id']
|
||||
for detail in item_result['metrics'].get('detailed_results', []):
|
||||
if only_failed and detail.get('is_correct') == False and not detail.get('is_invalid', False):
|
||||
for item_result in previous_results.get("item_results", []):
|
||||
item_id = item_result["item_id"]
|
||||
for detail in item_result["metrics"].get("detailed_results", []):
|
||||
if only_failed and detail.get("is_correct") == False and not detail.get("is_invalid", False):
|
||||
failed_question_ids.add(item_id)
|
||||
if only_invalid and detail.get('is_invalid', False):
|
||||
if only_invalid and detail.get("is_invalid", False):
|
||||
invalid_question_ids.add(item_id)
|
||||
|
||||
if only_failed:
|
||||
console.print(f"[cyan]Filtering to {len(failed_question_ids)} questions that failed (is_correct=False)[/cyan]")
|
||||
console.print(
|
||||
f"[cyan]Filtering to {len(failed_question_ids)} questions that failed (is_correct=False)[/cyan]"
|
||||
)
|
||||
if only_invalid:
|
||||
console.print(f"[cyan]Filtering to {len(invalid_question_ids)} questions that were invalid (is_invalid=True)[/cyan]")
|
||||
console.print(
|
||||
f"[cyan]Filtering to {len(invalid_question_ids)} questions that were invalid (is_invalid=True)[/cyan]"
|
||||
)
|
||||
|
||||
# Filter dataset by category if specified
|
||||
if category:
|
||||
@@ -550,11 +561,11 @@ async def run_benchmark(
|
||||
# Load full dataset without max_instances limit for filtering
|
||||
original_dataset_items = dataset.load(dataset_path, max_items=None)
|
||||
|
||||
filtered_items = [item for item in original_dataset_items if item.get('question_type') == category]
|
||||
filtered_items = [item for item in original_dataset_items if item.get("question_type") == category]
|
||||
|
||||
if not filtered_items:
|
||||
console.print(f"[yellow]No questions found for category '{category}'. Available categories:[/yellow]")
|
||||
available_categories = set(item.get('question_type', 'unknown') for item in original_dataset_items)
|
||||
available_categories = set(item.get("question_type", "unknown") for item in original_dataset_items)
|
||||
for cat in sorted(available_categories):
|
||||
console.print(f" - {cat}")
|
||||
return
|
||||
@@ -562,7 +573,9 @@ async def run_benchmark(
|
||||
total_found = len(filtered_items)
|
||||
will_run = min(total_found, max_instances) if max_instances else total_found
|
||||
if max_instances and total_found > max_instances:
|
||||
console.print(f"[green]Found {total_found} questions for category '{category}' (will run {will_run} due to --max-instances)[/green]")
|
||||
console.print(
|
||||
f"[green]Found {total_found} questions for category '{category}' (will run {will_run} due to --max-instances)[/green]"
|
||||
)
|
||||
else:
|
||||
console.print(f"[green]Found {total_found} questions for category '{category}'[/green]")
|
||||
|
||||
@@ -588,7 +601,9 @@ async def run_benchmark(
|
||||
total_found = len(filtered_items)
|
||||
will_run = min(total_found, max_instances) if max_instances else total_found
|
||||
if max_instances and total_found > max_instances:
|
||||
console.print(f"[green]Found {total_found} {filter_type} items to re-evaluate (will run {will_run} due to --max-instances)[/green]")
|
||||
console.print(
|
||||
f"[green]Found {total_found} {filter_type} items to re-evaluate (will run {will_run} due to --max-instances)[/green]"
|
||||
)
|
||||
else:
|
||||
console.print(f"[green]Found {total_found} {filter_type} items to re-evaluate[/green]")
|
||||
|
||||
@@ -600,6 +615,7 @@ async def run_benchmark(
|
||||
|
||||
# Create local memory engine
|
||||
from benchmarks.common.benchmark_runner import create_memory_engine
|
||||
|
||||
memory = await create_memory_engine()
|
||||
|
||||
# Filter by only_ingested: only run items whose memory bank already exists
|
||||
@@ -623,10 +639,9 @@ async def run_benchmark(
|
||||
# Check if bank has any memory units
|
||||
async with pool.acquire() as conn:
|
||||
result = await conn.fetchrow(
|
||||
"SELECT COUNT(*) as count FROM memory_units WHERE bank_id = $1 LIMIT 1",
|
||||
agent_id
|
||||
"SELECT COUNT(*) as count FROM memory_units WHERE bank_id = $1 LIMIT 1", agent_id
|
||||
)
|
||||
if result['count'] > 0:
|
||||
if result["count"] > 0:
|
||||
ingested_items.append(item)
|
||||
|
||||
filtered_items = ingested_items
|
||||
@@ -638,34 +653,43 @@ async def run_benchmark(
|
||||
|
||||
# Create benchmark runner
|
||||
runner = BenchmarkRunner(
|
||||
dataset=dataset,
|
||||
answer_generator=answer_generator,
|
||||
answer_evaluator=answer_evaluator,
|
||||
memory=memory
|
||||
dataset=dataset, answer_generator=answer_generator, answer_evaluator=answer_evaluator, memory=memory
|
||||
)
|
||||
|
||||
# If filtering by category, failed, invalid, only_ingested, or max_instances_per_category, we need to use a custom dataset that only returns those items
|
||||
# We'll temporarily replace the dataset's load method
|
||||
if filtered_items is not None:
|
||||
original_load = dataset.load
|
||||
|
||||
def filtered_load(path: Path, max_items: Optional[int] = None):
|
||||
return filtered_items[:max_items] if max_items else filtered_items
|
||||
|
||||
dataset.load = filtered_load
|
||||
|
||||
# Run benchmark
|
||||
# Single-phase approach: each question gets its own isolated agent_id
|
||||
# This ensures each question only has access to its own context
|
||||
output_path = Path(__file__).parent / 'results' / results_filename
|
||||
output_path = Path(__file__).parent / "results" / results_filename
|
||||
|
||||
# Create results directory if it doesn't exist
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
merge_with_existing = (filln or question_id is not None or only_failed or only_invalid or only_ingested or category is not None or max_instances_per_category is not None)
|
||||
merge_with_existing = (
|
||||
filln
|
||||
or question_id is not None
|
||||
or only_failed
|
||||
or only_invalid
|
||||
or only_ingested
|
||||
or category is not None
|
||||
or max_instances_per_category is not None
|
||||
)
|
||||
|
||||
results = await runner.run(
|
||||
dataset_path=dataset_path,
|
||||
agent_id="longmemeval", # Will be suffixed with question_id per item
|
||||
max_items=max_instances if not max_instances_per_category else None, # Don't apply max_items when using per-category limit
|
||||
max_items=max_instances
|
||||
if not max_instances_per_category
|
||||
else None, # Don't apply max_items when using per-category limit
|
||||
max_questions_per_item=max_questions_per_instance,
|
||||
thinking_budget=thinking_budget,
|
||||
max_tokens=max_tokens,
|
||||
@@ -678,7 +702,7 @@ async def run_benchmark(
|
||||
specific_item=question_id, # Optional filter for specific question ID
|
||||
max_concurrent_items=max_concurrent_items, # Parallel instance processing
|
||||
output_path=output_path, # Save results incrementally
|
||||
merge_with_existing=merge_with_existing # Merge when using --fill, --category, --only-failed, --only-invalid flags or specific question
|
||||
merge_with_existing=merge_with_existing, # Merge when using --fill, --category, --only-failed, --only-invalid flags or specific question
|
||||
)
|
||||
|
||||
# Display results (final save already happened incrementally)
|
||||
@@ -702,12 +726,14 @@ def download_dataset(dataset_path: Path) -> bool:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
import subprocess
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
url = "https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json"
|
||||
|
||||
console.print(f"[yellow]Dataset not found. Downloading from HuggingFace...[/yellow]")
|
||||
console.print("[yellow]Dataset not found. Downloading from HuggingFace...[/yellow]")
|
||||
console.print(f"[dim]URL: {url}[/dim]")
|
||||
console.print(f"[dim]Destination: {dataset_path}[/dim]")
|
||||
|
||||
@@ -720,18 +746,18 @@ def download_dataset(dataset_path: Path) -> bool:
|
||||
["curl", "-L", "-o", str(dataset_path), url],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=300 # 5 minute timeout
|
||||
timeout=300, # 5 minute timeout
|
||||
)
|
||||
|
||||
if result.returncode == 0 and dataset_path.exists():
|
||||
console.print(f"[green]✓ Dataset downloaded successfully[/green]")
|
||||
console.print("[green]✓ Dataset downloaded successfully[/green]")
|
||||
return True
|
||||
else:
|
||||
console.print(f"[red]✗ Download failed: {result.stderr}[/red]")
|
||||
return False
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
console.print(f"[red]✗ Download timed out after 5 minutes[/red]")
|
||||
console.print("[red]✗ Download timed out after 5 minutes[/red]")
|
||||
return False
|
||||
except Exception as e:
|
||||
console.print(f"[red]✗ Download error: {e}[/red]")
|
||||
@@ -740,22 +766,23 @@ def download_dataset(dataset_path: Path) -> bool:
|
||||
|
||||
def generate_type_report(results: dict):
|
||||
"""Generate a detailed report by question type."""
|
||||
from rich.table import Table
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
console = Console()
|
||||
|
||||
# Aggregate stats by question type
|
||||
type_stats = {}
|
||||
|
||||
for item_result in results['item_results']:
|
||||
metrics = item_result['metrics']
|
||||
by_category = metrics.get('category_stats', {})
|
||||
for item_result in results["item_results"]:
|
||||
metrics = item_result["metrics"]
|
||||
by_category = metrics.get("category_stats", {})
|
||||
|
||||
for qtype, stats in by_category.items():
|
||||
if qtype not in type_stats:
|
||||
type_stats[qtype] = {'total': 0, 'correct': 0}
|
||||
type_stats[qtype]['total'] += stats['total']
|
||||
type_stats[qtype]['correct'] += stats['correct']
|
||||
type_stats[qtype] = {"total": 0, "correct": 0}
|
||||
type_stats[qtype]["total"] += stats["total"]
|
||||
type_stats[qtype]["correct"] += stats["correct"]
|
||||
|
||||
# Display table
|
||||
table = Table(title="Performance by Question Type")
|
||||
@@ -765,13 +792,8 @@ def generate_type_report(results: dict):
|
||||
table.add_column("Accuracy", justify="right", style="magenta")
|
||||
|
||||
for qtype, stats in sorted(type_stats.items()):
|
||||
acc = (stats['correct'] / stats['total'] * 100) if stats['total'] > 0 else 0
|
||||
table.add_row(
|
||||
qtype,
|
||||
str(stats['total']),
|
||||
str(stats['correct']),
|
||||
f"{acc:.1f}%"
|
||||
)
|
||||
acc = (stats["correct"] / stats["total"] * 100) if stats["total"] > 0 else 0
|
||||
table.add_row(qtype, str(stats["total"]), str(stats["correct"]), f"{acc:.1f}%")
|
||||
|
||||
console.print("\n")
|
||||
console.print(table)
|
||||
@@ -780,21 +802,22 @@ def generate_type_report(results: dict):
|
||||
def generate_markdown_table(results: dict, json_output_path: Path):
|
||||
"""Generate a markdown results table with model configuration."""
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
# Aggregate stats by question type
|
||||
type_stats = {}
|
||||
|
||||
for item_result in results['item_results']:
|
||||
metrics = item_result['metrics']
|
||||
by_category = metrics.get('category_stats', {})
|
||||
for item_result in results["item_results"]:
|
||||
metrics = item_result["metrics"]
|
||||
by_category = metrics.get("category_stats", {})
|
||||
|
||||
for qtype, stats in by_category.items():
|
||||
if qtype not in type_stats:
|
||||
type_stats[qtype] = {'total': 0, 'correct': 0, 'invalid': 0}
|
||||
type_stats[qtype]['total'] += stats['total']
|
||||
type_stats[qtype]['correct'] += stats['correct']
|
||||
type_stats[qtype]['invalid'] += stats.get('invalid', 0)
|
||||
type_stats[qtype] = {"total": 0, "correct": 0, "invalid": 0}
|
||||
type_stats[qtype]["total"] += stats["total"]
|
||||
type_stats[qtype]["correct"] += stats["correct"]
|
||||
type_stats[qtype]["invalid"] += stats.get("invalid", 0)
|
||||
|
||||
# Build markdown content
|
||||
lines = []
|
||||
@@ -802,16 +825,20 @@ def generate_markdown_table(results: dict, json_output_path: Path):
|
||||
lines.append("")
|
||||
|
||||
# Add model configuration
|
||||
if 'model_config' in results:
|
||||
config = results['model_config']
|
||||
if "model_config" in results:
|
||||
config = results["model_config"]
|
||||
lines.append("## Model Configuration")
|
||||
lines.append("")
|
||||
lines.append(f"- **Hindsight**: {config['hindsight']['provider']}/{config['hindsight']['model']}")
|
||||
lines.append(f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}")
|
||||
lines.append(
|
||||
f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}"
|
||||
)
|
||||
lines.append(f"- **LLM Judge**: {config['judge']['provider']}/{config['judge']['model']}")
|
||||
lines.append("")
|
||||
|
||||
lines.append(f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})")
|
||||
lines.append(
|
||||
f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})"
|
||||
)
|
||||
lines.append("")
|
||||
|
||||
# Results by question type
|
||||
@@ -822,34 +849,36 @@ def generate_markdown_table(results: dict, json_output_path: Path):
|
||||
|
||||
for qtype in sorted(type_stats.keys()):
|
||||
stats = type_stats[qtype]
|
||||
valid_total = stats['total'] - stats['invalid']
|
||||
acc = (stats['correct'] / valid_total * 100) if valid_total > 0 else 0
|
||||
invalid_str = str(stats['invalid']) if stats['invalid'] > 0 else "-"
|
||||
valid_total = stats["total"] - stats["invalid"]
|
||||
acc = (stats["correct"] / valid_total * 100) if valid_total > 0 else 0
|
||||
invalid_str = str(stats["invalid"]) if stats["invalid"] > 0 else "-"
|
||||
lines.append(f"| {qtype} | {stats['total']} | {stats['correct']} | {invalid_str} | {acc:.1f}% |")
|
||||
|
||||
# Add overall row
|
||||
total_invalid = results.get('total_invalid', 0)
|
||||
total_invalid = results.get("total_invalid", 0)
|
||||
invalid_str = str(total_invalid) if total_invalid > 0 else "-"
|
||||
lines.append(f"| **OVERALL** | **{results['total_questions']}** | **{results['total_correct']}** | **{invalid_str}** | **{results['overall_accuracy']:.1f}%** |")
|
||||
lines.append(
|
||||
f"| **OVERALL** | **{results['total_questions']}** | **{results['total_correct']}** | **{invalid_str}** | **{results['overall_accuracy']:.1f}%** |"
|
||||
)
|
||||
|
||||
# Write to file (same directory as JSON, but .md extension)
|
||||
md_output_path = json_output_path.with_suffix('.md')
|
||||
md_output_path.write_text('\n'.join(lines))
|
||||
md_output_path = json_output_path.with_suffix(".md")
|
||||
md_output_path.write_text("\n".join(lines))
|
||||
console.print(f"\n[green]✓[/green] Results table saved to {md_output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import logging
|
||||
import argparse
|
||||
import logging
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
|
||||
parser = argparse.ArgumentParser(description="Run LongMemEval benchmark")
|
||||
parser.add_argument(
|
||||
"--max-instances",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Limit TOTAL number of questions to evaluate (default: all 500). For per-category limits, use --max-questions-per-category instead."
|
||||
help="Limit TOTAL number of questions to evaluate (default: all 500). For per-category limits, use --max-questions-per-category instead.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-instances-per-category",
|
||||
@@ -857,87 +886,72 @@ if __name__ == "__main__":
|
||||
type=int,
|
||||
default=None,
|
||||
dest="max_instances_per_category",
|
||||
help="Limit number of questions per category (e.g., 20 = 20 questions from each of the 6 categories = 120 total). Cannot be combined with --max-instances or --category."
|
||||
help="Limit number of questions per category (e.g., 20 = 20 questions from each of the 6 categories = 120 total). Cannot be combined with --max-instances or --category.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-questions",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Limit number of questions per instance (for quick testing)"
|
||||
"--max-questions", type=int, default=None, help="Limit number of questions per instance (for quick testing)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--thinking-budget",
|
||||
type=int,
|
||||
default=500,
|
||||
help="Thinking budget for spreading activation search"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
default=8192,
|
||||
help="Maximum tokens to retrieve from memories"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-ingestion",
|
||||
action="store_true",
|
||||
help="Skip ingestion and use existing data"
|
||||
"--thinking-budget", type=int, default=500, help="Thinking budget for spreading activation search"
|
||||
)
|
||||
parser.add_argument("--max-tokens", type=int, default=8192, help="Maximum tokens to retrieve from memories")
|
||||
parser.add_argument("--skip-ingestion", action="store_true", help="Skip ingestion and use existing data")
|
||||
parser.add_argument(
|
||||
"--fill",
|
||||
action="store_true",
|
||||
help="Only process questions not already in results file (for resuming interrupted runs)"
|
||||
help="Only process questions not already in results file (for resuming interrupted runs)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--question-id",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Filter to specific question ID (e.g., 'e47becba'). Useful with --skip-ingestion to test a single question."
|
||||
help="Filter to specific question ID (e.g., 'e47becba'). Useful with --skip-ingestion to test a single question.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-failed",
|
||||
action="store_true",
|
||||
help="Only run questions that were previously marked as incorrect (is_correct=False). Requires existing results file."
|
||||
help="Only run questions that were previously marked as incorrect (is_correct=False). Requires existing results file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-invalid",
|
||||
action="store_true",
|
||||
help="Only run questions that were previously marked as invalid (is_invalid=True). Requires existing results file."
|
||||
help="Only run questions that were previously marked as invalid (is_invalid=True). Requires existing results file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-ingested",
|
||||
action="store_true",
|
||||
help="Only run questions whose memory bank already exists (has been ingested). Automatically skips ingestion. Cannot be combined with --only-failed, --only-invalid, --category, --question-id, or --max-instances-per-category."
|
||||
help="Only run questions whose memory bank already exists (has been ingested). Automatically skips ingestion. Cannot be combined with --only-failed, --only-invalid, --category, --question-id, or --max-instances-per-category.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--category",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Filter questions by category/question_type. Available categories: 'single-session-user', 'multi-session', 'single-session-preference', 'temporal-reasoning', 'knowledge-update', 'single-session-assistant'. Can be combined with --max-instances to limit questions within the category."
|
||||
help="Filter questions by category/question_type. Available categories: 'single-session-user', 'multi-session', 'single-session-preference', 'temporal-reasoning', 'knowledge-update', 'single-session-assistant'. Can be combined with --max-instances to limit questions within the category.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--parallel",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of instances to process in parallel (default: 1 for sequential). Higher values speed up evaluation but use more memory."
|
||||
help="Number of instances to process in parallel (default: 1 for sequential). Higher values speed up evaluation but use more memory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--results-filename",
|
||||
type=str,
|
||||
default="benchmark_results.json",
|
||||
help="Filename for results output (default: benchmark_results.json). Saved in results/ directory."
|
||||
help="Filename for results output (default: benchmark_results.json). Saved in results/ directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--context-format",
|
||||
type=str,
|
||||
choices=["json", "structured"],
|
||||
default="json",
|
||||
help="How to format context for answer generation. 'json' (raw JSON dump, original behavior) or 'structured' (human-readable format with facts grouped with source chunks). Default: json."
|
||||
help="How to format context for answer generation. 'json' (raw JSON dump, original behavior) or 'structured' (human-readable format with facts grouped with source chunks). Default: json.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source-results",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Source results file to read failed/invalid questions from (for --only-failed/--only-invalid). Defaults to benchmark_results.json if not specified."
|
||||
help="Source results file to read failed/invalid questions from (for --only-failed/--only-invalid). Defaults to benchmark_results.json if not specified.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
@@ -951,21 +965,23 @@ if __name__ == "__main__":
|
||||
if args.max_instances_per_category is not None and (args.max_instances is not None or args.category is not None):
|
||||
parser.error("--max-questions-per-category cannot be combined with --max-instances or --category")
|
||||
|
||||
results = asyncio.run(run_benchmark(
|
||||
max_instances=args.max_instances,
|
||||
max_instances_per_category=args.max_instances_per_category,
|
||||
max_questions_per_instance=args.max_questions,
|
||||
thinking_budget=args.thinking_budget,
|
||||
max_tokens=args.max_tokens,
|
||||
skip_ingestion=args.skip_ingestion,
|
||||
filln=args.fill,
|
||||
question_id=args.question_id,
|
||||
only_failed=args.only_failed,
|
||||
only_invalid=args.only_invalid,
|
||||
only_ingested=args.only_ingested,
|
||||
category=args.category,
|
||||
max_concurrent_items=args.parallel,
|
||||
results_filename=args.results_filename,
|
||||
context_format=args.context_format,
|
||||
source_results=args.source_results
|
||||
))
|
||||
results = asyncio.run(
|
||||
run_benchmark(
|
||||
max_instances=args.max_instances,
|
||||
max_instances_per_category=args.max_instances_per_category,
|
||||
max_questions_per_instance=args.max_questions,
|
||||
thinking_budget=args.thinking_budget,
|
||||
max_tokens=args.max_tokens,
|
||||
skip_ingestion=args.skip_ingestion,
|
||||
filln=args.fill,
|
||||
question_id=args.question_id,
|
||||
only_failed=args.only_failed,
|
||||
only_invalid=args.only_invalid,
|
||||
only_ingested=args.only_ingested,
|
||||
category=args.category,
|
||||
max_concurrent_items=args.parallel,
|
||||
results_filename=args.results_filename,
|
||||
context_format=args.context_format,
|
||||
source_results=args.source_results,
|
||||
)
|
||||
)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -5,6 +5,7 @@ Generate changelog entry for a new release.
|
||||
This script fetches the commit diff between releases, uses an LLM to summarize,
|
||||
and prepends the entry to the changelog page.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
@@ -29,6 +30,7 @@ CHANGELOG_PATH = REPO_PATH / "hindsight-docs" / "docs" / "changelog" / "index.md
|
||||
|
||||
class ChangelogEntry(BaseModel):
|
||||
"""A single changelog entry."""
|
||||
|
||||
category: str # "feature", "improvement", "bugfix", "breaking", "other"
|
||||
summary: str # Brief description of the change
|
||||
commit_id: str # Short commit hash
|
||||
@@ -36,12 +38,14 @@ class ChangelogEntry(BaseModel):
|
||||
|
||||
class ChangelogResponse(BaseModel):
|
||||
"""Structured response from LLM."""
|
||||
|
||||
entries: list[ChangelogEntry]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Commit:
|
||||
"""Parsed commit from git log."""
|
||||
|
||||
hash: str
|
||||
message: str
|
||||
|
||||
@@ -151,10 +155,7 @@ def analyze_commits_with_llm(
|
||||
file_diff: str,
|
||||
) -> list[ChangelogEntry]:
|
||||
"""Use LLM to analyze commits and return structured changelog entries."""
|
||||
commits_json = json.dumps(
|
||||
[{"commit_id": c.hash, "message": c.message} for c in commits],
|
||||
indent=2
|
||||
)
|
||||
commits_json = json.dumps([{"commit_id": c.hash, "message": c.message} for c in commits], indent=2)
|
||||
|
||||
prompt = f"""Analyze the following git commits for release {version} of Hindsight (an AI memory system).
|
||||
|
||||
@@ -252,8 +253,8 @@ For full release details, see [GitHub Releases](https://github.com/vectorize-io/
|
||||
|
||||
match = re.search(r"^## ", content, re.MULTILINE)
|
||||
if match:
|
||||
header = content[:match.start()].rstrip() + "\n\n"
|
||||
releases = content[match.start():]
|
||||
header = content[: match.start()].rstrip() + "\n\n"
|
||||
releases = content[match.start() :]
|
||||
else:
|
||||
header = content.rstrip() + "\n\n"
|
||||
releases = ""
|
||||
@@ -283,7 +284,7 @@ def generate_changelog_entry(
|
||||
tag = version if version.startswith("v") else f"v{version}"
|
||||
display_version = version.lstrip("v")
|
||||
|
||||
console.print(f"[blue]Fetching tags from repository...[/blue]")
|
||||
console.print("[blue]Fetching tags from repository...[/blue]")
|
||||
existing_tags = get_git_tags()
|
||||
|
||||
if tag not in existing_tags and display_version not in existing_tags:
|
||||
@@ -300,7 +301,7 @@ def generate_changelog_entry(
|
||||
else:
|
||||
console.print("[yellow]No previous version found, will include all commits[/yellow]")
|
||||
|
||||
console.print(f"[blue]Getting commits...[/blue]")
|
||||
console.print("[blue]Getting commits...[/blue]")
|
||||
commits = get_commits(previous_tag, actual_tag)
|
||||
file_diff = get_detailed_diff(previous_tag, actual_tag)
|
||||
|
||||
|
||||
@@ -4,13 +4,15 @@ Generate OpenAPI specification from FastAPI app.
|
||||
|
||||
This script imports the FastAPI app and exports its OpenAPI schema to a JSON file.
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
|
||||
def generate_openapi_spec(output_path: str = None):
|
||||
"""Generate OpenAPI spec and save to file."""
|
||||
@@ -34,7 +36,7 @@ def generate_openapi_spec(output_path: str = None):
|
||||
|
||||
# Write to file
|
||||
output_file = Path(output_path)
|
||||
with open(output_file, 'w') as f:
|
||||
with open(output_file, "w") as f:
|
||||
json.dump(openapi_schema, f, indent=2)
|
||||
|
||||
print(f"✓ OpenAPI specification generated: {output_file.absolute()}")
|
||||
@@ -44,14 +46,15 @@ def generate_openapi_spec(output_path: str = None):
|
||||
|
||||
# List endpoints
|
||||
print("\n Endpoints:")
|
||||
for path, methods in openapi_schema['paths'].items():
|
||||
for path, methods in openapi_schema["paths"].items():
|
||||
for method in methods.keys():
|
||||
if method.upper() in ['GET', 'POST', 'PUT', 'DELETE', 'PATCH']:
|
||||
if method.upper() in ["GET", "POST", "PUT", "DELETE", "PATCH"]:
|
||||
endpoint_info = methods[method]
|
||||
summary = endpoint_info.get('summary', 'No summary')
|
||||
tags = ', '.join(endpoint_info.get('tags', ['untagged']))
|
||||
summary = endpoint_info.get("summary", "No summary")
|
||||
tags = ", ".join(endpoint_info.get("tags", ["untagged"]))
|
||||
print(f" {method.upper():6} {path:30} [{tags}] - {summary}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
output = sys.argv[1] if len(sys.argv) > 1 else "openapi.json"
|
||||
generate_openapi_spec(output)
|
||||
|
||||
@@ -212,9 +212,7 @@ This recipe is available as an interactive Jupyter notebook.
|
||||
# Insert callout after first heading
|
||||
first_heading_match = re.search(r"^(#\s+.+\n)", md_content, re.MULTILINE)
|
||||
if first_heading_match:
|
||||
idx = md_content.index(first_heading_match.group(0)) + len(
|
||||
first_heading_match.group(0)
|
||||
)
|
||||
idx = md_content.index(first_heading_match.group(0)) + len(first_heading_match.group(0))
|
||||
final_content = md_content[:idx] + "\n" + callout + "\n" + md_content[idx:]
|
||||
else:
|
||||
final_content = callout + "\n" + md_content
|
||||
@@ -247,9 +245,7 @@ def process_applications(cookbook_dir: Path, apps_dir: Path) -> list[dict]:
|
||||
continue
|
||||
|
||||
slug = entry.name
|
||||
title = extract_title_from_readme(readme_path) or " ".join(
|
||||
word.capitalize() for word in slug.split("-")
|
||||
)
|
||||
title = extract_title_from_readme(readme_path) or " ".join(word.capitalize() for word in slug.split("-"))
|
||||
|
||||
print(f" Processing app: {entry.name} → {slug}.md")
|
||||
|
||||
@@ -275,12 +271,8 @@ This is a complete, runnable application demonstrating Hindsight integration.
|
||||
# Insert callout after first heading
|
||||
first_heading_match = re.search(r"^(#\s+.+\n)", readme_content, re.MULTILINE)
|
||||
if first_heading_match:
|
||||
idx = readme_content.index(first_heading_match.group(0)) + len(
|
||||
first_heading_match.group(0)
|
||||
)
|
||||
final_content = (
|
||||
readme_content[:idx] + "\n" + callout + "\n" + readme_content[idx:]
|
||||
)
|
||||
idx = readme_content.index(first_heading_match.group(0)) + len(first_heading_match.group(0))
|
||||
final_content = readme_content[:idx] + "\n" + callout + "\n" + readme_content[idx:]
|
||||
else:
|
||||
final_content = callout + "\n" + readme_content
|
||||
|
||||
@@ -407,26 +399,20 @@ def clean_description(desc: str) -> str:
|
||||
return desc
|
||||
|
||||
|
||||
def update_cookbook_index(
|
||||
recipes: list[dict], apps: list[dict], docs_dir: Path
|
||||
):
|
||||
def update_cookbook_index(recipes: list[dict], apps: list[dict], docs_dir: Path):
|
||||
"""Update cookbook/index.mdx with recipe and app carousels."""
|
||||
# Build recipe items for the carousel
|
||||
recipe_items = []
|
||||
for r in recipes:
|
||||
title = r["title"].replace('"', '\\"')
|
||||
recipe_items.append(
|
||||
f' {{ title: "{title}", href: "/cookbook/recipes/{r["slug"]}" }}'
|
||||
)
|
||||
recipe_items.append(f' {{ title: "{title}", href: "/cookbook/recipes/{r["slug"]}" }}')
|
||||
recipes_json = ",\n".join(recipe_items)
|
||||
|
||||
# Build app items for the carousel
|
||||
app_items = []
|
||||
for a in apps:
|
||||
title = a["title"].replace('"', '\\"')
|
||||
app_items.append(
|
||||
f' {{ title: "{title}", href: "/cookbook/applications/{a["slug"]}" }}'
|
||||
)
|
||||
app_items.append(f' {{ title: "{title}", href: "/cookbook/applications/{a["slug"]}" }}')
|
||||
apps_json = ",\n".join(app_items)
|
||||
|
||||
content = f"""---
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-dev"
|
||||
version = "0.1.11"
|
||||
version = "0.1.12"
|
||||
description = "Development utilities for Hindsight"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
@@ -27,3 +27,58 @@ generate-openapi = "hindsight_dev.generate_openapi:generate_openapi_spec"
|
||||
generate-changelog = "hindsight_dev.generate_changelog:main"
|
||||
sync-cookbook = "hindsight_dev.sync_cookbook:main"
|
||||
generate-llms-full = "hindsight_dev.generate_llms_full:main"
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"ruff>=0.8.0",
|
||||
"ty>=0.0.1",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 120
|
||||
target-version = "py311"
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = [
|
||||
"E", # pycodestyle errors
|
||||
"W", # pycodestyle warnings
|
||||
"F", # Pyflakes
|
||||
"I", # isort
|
||||
]
|
||||
ignore = [
|
||||
"E501", # line too long (handled by formatter)
|
||||
"E402", # module import not at top of file
|
||||
"E712", # avoid equality comparisons to False (intentional for optional bools)
|
||||
"F401", # unused import (too noisy during development)
|
||||
"F403", # star imports (fasthtml uses this pattern)
|
||||
"F405", # may be undefined from star imports (fasthtml uses this pattern)
|
||||
"F841", # unused variable (too noisy during development)
|
||||
"F811", # redefined while unused
|
||||
"F821", # undefined name (forward references in type hints)
|
||||
"W291", # trailing whitespace (in multiline strings)
|
||||
"W293", # blank line contains whitespace (in multiline strings)
|
||||
]
|
||||
|
||||
[tool.ruff.format]
|
||||
quote-style = "double"
|
||||
indent-style = "space"
|
||||
|
||||
[tool.ty]
|
||||
# Type checking configuration
|
||||
# ty is an extremely fast Python type checker from Astral (same team as ruff/uv)
|
||||
|
||||
[tool.ty.environment]
|
||||
python-version = "3.11"
|
||||
|
||||
[tool.ty.rules]
|
||||
# Disable noisy rules
|
||||
invalid-argument-type = "ignore" # Too many false positives
|
||||
invalid-return-type = "ignore" # Often intentional
|
||||
# missing-argument and unknown-argument are enabled (default)
|
||||
invalid-parameter-default = "ignore" # Optional params with None default
|
||||
possibly-missing-attribute = "ignore" # Common with Optional types
|
||||
unsupported-operator = "ignore" # Pandas DataFrame operations
|
||||
unresolved-reference = "ignore" # Forward references
|
||||
unresolved-import = "ignore" # Dynamic imports
|
||||
invalid-assignment = "ignore" # Intentional monkey-patching
|
||||
no-matching-overload = "ignore" # Complex generic issues
|
||||
|
||||
@@ -0,0 +1,226 @@
|
||||
# Extensions
|
||||
|
||||
Extensions allow you to customize and extend Hindsight behavior without modifying core code. They enable multi-tenancy, custom authentication, additional HTTP endpoints, and operation hooks.
|
||||
|
||||
---
|
||||
|
||||
## Available Extensions
|
||||
|
||||
### TenantExtension
|
||||
|
||||
Handles multi-tenancy and API key authentication. Validates incoming requests and determines which PostgreSQL schema to use for database operations, enabling tenant isolation at the database level.
|
||||
|
||||
**Built-in: ApiKeyTenantExtension**
|
||||
|
||||
A simple implementation that validates API keys against an environment variable and uses the `public` schema for all authenticated requests.
|
||||
|
||||
```bash
|
||||
HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
|
||||
HINDSIGHT_API_TENANT_API_KEY=your-secret-key
|
||||
```
|
||||
|
||||
For multi-tenant setups with separate schemas per tenant (e.g., JWT-based auth with per-tenant schemas), implement a custom `TenantExtension`.
|
||||
|
||||
---
|
||||
|
||||
### HttpExtension
|
||||
|
||||
Adds custom HTTP endpoints under the `/ext/` path prefix. Useful for adding domain-specific APIs that integrate with Hindsight's memory engine.
|
||||
|
||||
**No built-in implementation** - implement your own to add custom endpoints.
|
||||
|
||||
```bash
|
||||
HINDSIGHT_API_HTTP_EXTENSION=mypackage.ext:MyHttpExtension
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### OperationValidatorExtension
|
||||
|
||||
Hooks into retain/recall/reflect operations for validation and monitoring. Use cases include:
|
||||
- Rate limiting and quota enforcement
|
||||
- Permission checks and content filtering
|
||||
- Audit logging and usage tracking
|
||||
- Custom metrics collection
|
||||
|
||||
**No built-in implementation** - implement your own based on your requirements.
|
||||
|
||||
```bash
|
||||
HINDSIGHT_API_OPERATION_VALIDATOR_EXTENSION=mypackage.validators:MyValidator
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Writing Custom Extensions
|
||||
|
||||
### Extension Basics
|
||||
|
||||
Extensions are Python classes loaded via environment variables:
|
||||
|
||||
```bash
|
||||
HINDSIGHT_API_<TYPE>_EXTENSION=mypackage.module:MyExtensionClass
|
||||
```
|
||||
|
||||
Configuration is passed via prefixed environment variables:
|
||||
|
||||
```bash
|
||||
HINDSIGHT_API_<TYPE>_SOME_CONFIG=value
|
||||
# Extension receives: {"some_config": "value"}
|
||||
```
|
||||
|
||||
All extensions support lifecycle hooks:
|
||||
- `on_startup()` - Called when the application starts
|
||||
- `on_shutdown()` - Called when the application shuts down
|
||||
|
||||
Extensions have access to an `ExtensionContext` that provides:
|
||||
- `run_migration(schema)` - Run database migrations for a schema
|
||||
- `get_memory_engine()` - Get the MemoryEngine interface
|
||||
|
||||
### Example: Custom TenantExtension with JWT
|
||||
|
||||
```python
|
||||
import jwt
|
||||
from hindsight_api.extensions import TenantExtension, TenantContext, AuthenticationError
|
||||
|
||||
class JwtTenantExtension(TenantExtension):
|
||||
def __init__(self, config: dict[str, str]):
|
||||
super().__init__(config)
|
||||
self.jwt_secret = config.get("jwt_secret")
|
||||
if not self.jwt_secret:
|
||||
raise ValueError("HINDSIGHT_API_TENANT_JWT_SECRET is required")
|
||||
|
||||
async def authenticate(self, context: RequestContext) -> TenantContext:
|
||||
token = context.api_key
|
||||
if not token:
|
||||
raise AuthenticationError("Bearer token required")
|
||||
|
||||
try:
|
||||
payload = jwt.decode(token, self.jwt_secret, algorithms=["HS256"])
|
||||
tenant_id = payload.get("tenant_id")
|
||||
if not tenant_id:
|
||||
raise AuthenticationError("Missing tenant_id in token")
|
||||
return TenantContext(schema_name=f"tenant_{tenant_id}")
|
||||
except jwt.InvalidTokenError as e:
|
||||
raise AuthenticationError(str(e))
|
||||
```
|
||||
|
||||
### Example: Custom HttpExtension
|
||||
|
||||
```python
|
||||
from fastapi import APIRouter
|
||||
from hindsight_api.extensions import HttpExtension
|
||||
|
||||
class MyHttpExtension(HttpExtension):
|
||||
def get_router(self, memory: MemoryEngine) -> APIRouter:
|
||||
router = APIRouter()
|
||||
|
||||
@router.get("/hello")
|
||||
async def hello():
|
||||
return {"message": "Hello from extension!"}
|
||||
|
||||
@router.post("/custom/{bank_id}/action")
|
||||
async def custom_action(bank_id: str):
|
||||
# Access memory engine for database operations
|
||||
pool = await memory._get_pool()
|
||||
# ... custom logic
|
||||
return {"status": "ok"}
|
||||
|
||||
return router
|
||||
```
|
||||
|
||||
Routes are available at `/ext/hello`, `/ext/custom/{bank_id}/action`, etc.
|
||||
|
||||
### Example: Custom OperationValidatorExtension
|
||||
|
||||
```python
|
||||
from hindsight_api.extensions import (
|
||||
OperationValidatorExtension,
|
||||
ValidationResult,
|
||||
RetainContext,
|
||||
RecallContext,
|
||||
ReflectContext,
|
||||
RetainResult,
|
||||
)
|
||||
|
||||
class MyValidator(OperationValidatorExtension):
|
||||
# Pre-operation validation (required)
|
||||
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
|
||||
# Implement your validation logic
|
||||
return ValidationResult.accept()
|
||||
# Or reject: return ValidationResult.reject("Reason")
|
||||
|
||||
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
|
||||
return ValidationResult.accept()
|
||||
|
||||
# Post-operation hooks (optional)
|
||||
async def on_retain_complete(self, result: RetainResult) -> None:
|
||||
# Log usage, update metrics, send notifications, etc.
|
||||
pass
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Deploying Custom Extensions
|
||||
|
||||
### With Docker
|
||||
|
||||
Mount your extension package as a volume and set the environment variable:
|
||||
|
||||
```yaml
|
||||
# docker-compose.yml
|
||||
services:
|
||||
hindsight-api:
|
||||
image: vectorize/hindsight-api:latest
|
||||
volumes:
|
||||
- ./my_extensions:/app/my_extensions
|
||||
environment:
|
||||
- HINDSIGHT_API_TENANT_EXTENSION=my_extensions.auth:JwtTenantExtension
|
||||
- HINDSIGHT_API_TENANT_JWT_SECRET=${JWT_SECRET}
|
||||
- PYTHONPATH=/app
|
||||
```
|
||||
|
||||
Or build a custom image with your extensions:
|
||||
|
||||
```dockerfile
|
||||
FROM vectorize/hindsight-api:latest
|
||||
COPY my_extensions /app/my_extensions
|
||||
ENV PYTHONPATH=/app
|
||||
```
|
||||
|
||||
### Bare Metal
|
||||
|
||||
Install your extension package in the same Python environment as Hindsight:
|
||||
|
||||
```bash
|
||||
# Install Hindsight
|
||||
pip install hindsight-api
|
||||
|
||||
# Install your extension package
|
||||
pip install ./my-extensions
|
||||
# or
|
||||
pip install my-extensions-package
|
||||
|
||||
# Configure
|
||||
export HINDSIGHT_API_TENANT_EXTENSION=my_extensions.auth:JwtTenantExtension
|
||||
export HINDSIGHT_API_TENANT_JWT_SECRET=your-secret
|
||||
|
||||
# Run
|
||||
hindsight-api
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Contributing Extensions
|
||||
|
||||
Custom extensions that solve common use cases are welcome contributions to the Hindsight project. If you've built an extension for:
|
||||
|
||||
- Authentication providers (OAuth, SAML, API gateways)
|
||||
- Rate limiting or quota management
|
||||
- Audit logging integrations
|
||||
- Metrics exporters (Datadog, New Relic, etc.)
|
||||
- Custom HTTP endpoints for specific platforms
|
||||
|
||||
Consider contributing it to the `hindsight_api.extensions.builtin` package. Open an issue or pull request on [GitHub](https://github.com/vectorize-io/hindsight) to discuss your extension.
|
||||
@@ -116,7 +116,18 @@ const config: Config = {
|
||||
],
|
||||
],
|
||||
|
||||
themes: ['@docusaurus/theme-mermaid'],
|
||||
themes: [
|
||||
'@docusaurus/theme-mermaid',
|
||||
[
|
||||
'@easyops-cn/docusaurus-search-local',
|
||||
{
|
||||
hashed: true,
|
||||
docsRouteBasePath: '/',
|
||||
indexBlog: false,
|
||||
highlightSearchTermsOnTargetPage: false,
|
||||
},
|
||||
],
|
||||
],
|
||||
|
||||
themeConfig: {
|
||||
...(ANNOUNCEMENT_BAR && {
|
||||
|
||||
@@ -19,6 +19,7 @@
|
||||
"@docusaurus/preset-classic": "3.9.2",
|
||||
"@docusaurus/theme-common": "^3.9.2",
|
||||
"@docusaurus/theme-mermaid": "^3.9.2",
|
||||
"@easyops-cn/docusaurus-search-local": "^0.52.2",
|
||||
"@mdx-js/react": "^3.0.0",
|
||||
"clsx": "^2.0.0",
|
||||
"prism-react-renderer": "^2.3.0",
|
||||
|
||||
@@ -111,6 +111,11 @@ const sidebars: SidebarsConfig = {
|
||||
id: 'developer/configuration',
|
||||
label: 'Configuration',
|
||||
},
|
||||
{
|
||||
type: 'doc',
|
||||
id: 'developer/extensions',
|
||||
label: 'Extensions',
|
||||
},
|
||||
{
|
||||
type: 'doc',
|
||||
id: 'developer/models',
|
||||
|
||||
@@ -890,3 +890,309 @@ div[class*="announcementBar"] a:hover {
|
||||
background: rgba(255, 255, 255, 0.3);
|
||||
}
|
||||
|
||||
/* ===== Local Search Plugin Styling ===== */
|
||||
|
||||
/* Search bar in navbar */
|
||||
.navbar__search {
|
||||
margin-left: auto;
|
||||
}
|
||||
|
||||
|
||||
/* Hide keyboard shortcut hints */
|
||||
[class*="searchHint"],
|
||||
[class*="searchHintContainer"] {
|
||||
display: none !important;
|
||||
}
|
||||
|
||||
/* Minimal search button */
|
||||
[class*="searchBarContainer"] button,
|
||||
[class*="searchBar_"] {
|
||||
font-family: 'Space Grotesk', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
|
||||
font-size: 0.75rem !important;
|
||||
font-weight: 500 !important;
|
||||
background: transparent !important;
|
||||
border: none !important;
|
||||
border-bottom: 1px solid transparent !important;
|
||||
border-image: linear-gradient(90deg, #0074d9, #009296) 1 !important;
|
||||
border-radius: 0 !important;
|
||||
padding: 0.25rem 0 !important;
|
||||
color: var(--ifm-color-emphasis-600) !important;
|
||||
transition: all 0.2s ease !important;
|
||||
height: auto !important;
|
||||
min-height: unset !important;
|
||||
}
|
||||
|
||||
[class*="searchBarContainer"] button:hover,
|
||||
[class*="searchBar_"]:hover {
|
||||
color: var(--ifm-font-color-base) !important;
|
||||
}
|
||||
|
||||
[data-theme='dark'] [class*="searchBarContainer"] button,
|
||||
[data-theme='dark'] [class*="searchBar_"] {
|
||||
color: var(--ifm-color-emphasis-500) !important;
|
||||
}
|
||||
|
||||
[data-theme='dark'] [class*="searchBarContainer"] button:hover,
|
||||
[data-theme='dark'] [class*="searchBar_"]:hover {
|
||||
color: #e2e8f0 !important;
|
||||
}
|
||||
|
||||
/* Search input field - remove ALL borders */
|
||||
[class*="searchInput"],
|
||||
[class*="searchQueryInput"],
|
||||
[class*="searchInput"] *,
|
||||
[class*="searchQueryInput"] *,
|
||||
[class*="searchBarContainer"] input,
|
||||
[class*="searchBarContainer"] input * {
|
||||
font-family: 'Space Grotesk', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
|
||||
font-size: 0.875rem !important;
|
||||
font-weight: 500 !important;
|
||||
background: transparent !important;
|
||||
border: 0 !important;
|
||||
border-top: 0 !important;
|
||||
border-right: 0 !important;
|
||||
border-bottom: 0 !important;
|
||||
border-left: 0 !important;
|
||||
border-radius: 0 !important;
|
||||
color: var(--ifm-font-color-base) !important;
|
||||
padding: 0.5rem 0 !important;
|
||||
box-shadow: none !important;
|
||||
outline: none !important;
|
||||
}
|
||||
|
||||
/* Add single underline via pseudo or direct */
|
||||
[class*="searchQueryInput"],
|
||||
[class*="searchBarContainer"] input[type="search"] {
|
||||
border-bottom: 1px solid var(--ifm-toc-border-color) !important;
|
||||
}
|
||||
|
||||
[data-theme='dark'] [class*="searchQueryInput"],
|
||||
[data-theme='dark'] [class*="searchBarContainer"] input[type="search"] {
|
||||
border-bottom: 1px solid #27272a !important;
|
||||
}
|
||||
|
||||
[class*="searchQueryInput"]:focus,
|
||||
[class*="searchBarContainer"] input[type="search"]:focus {
|
||||
border-bottom: 1px solid #0074d9 !important;
|
||||
}
|
||||
|
||||
[class*="searchInput"]::placeholder,
|
||||
[class*="searchQueryInput"]::placeholder,
|
||||
[class*="searchBarContainer"] input::placeholder {
|
||||
font-family: 'Space Grotesk', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
|
||||
font-weight: 500 !important;
|
||||
color: var(--ifm-color-emphasis-400) !important;
|
||||
}
|
||||
|
||||
/* Minimal dropdown */
|
||||
[class*="dropdownMenu"],
|
||||
[class*="suggestionsContainer"],
|
||||
[class*="searchResultsContainer"],
|
||||
[class*="suggestions"],
|
||||
ul[class*="suggestion"] {
|
||||
font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
|
||||
background: #ffffff !important;
|
||||
background-color: #ffffff !important;
|
||||
border: none !important;
|
||||
border-radius: 0.25rem !important;
|
||||
box-shadow: 0 4px 20px -4px rgba(0, 0, 0, 0.15) !important;
|
||||
padding: 0.25rem !important;
|
||||
}
|
||||
|
||||
[data-theme='dark'] [class*="dropdownMenu"],
|
||||
[data-theme='dark'] [class*="suggestionsContainer"],
|
||||
[data-theme='dark'] [class*="searchResultsContainer"],
|
||||
[data-theme='dark'] [class*="suggestions"],
|
||||
[data-theme='dark'] ul[class*="suggestion"] {
|
||||
background: #0f0f11 !important;
|
||||
background-color: #0f0f11 !important;
|
||||
box-shadow: 0 4px 20px -4px rgba(0, 0, 0, 0.5) !important;
|
||||
}
|
||||
|
||||
/* Minimal result items */
|
||||
[class*="searchResultItem"],
|
||||
[class*="suggestion"] {
|
||||
font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
|
||||
background: transparent !important;
|
||||
border-radius: 0.25rem !important;
|
||||
padding: 0.5rem 0.75rem !important;
|
||||
margin: 0 !important;
|
||||
transition: background-color 0.1s ease !important;
|
||||
cursor: pointer !important;
|
||||
border-left: 2px solid transparent !important;
|
||||
}
|
||||
|
||||
/* No background hover - keep transparent */
|
||||
[class*="searchResultItem"]:hover,
|
||||
[class*="searchResultItem"][class*="cursor"],
|
||||
[class*="suggestion"]:hover,
|
||||
[class*="suggestion"][class*="cursor"] {
|
||||
background: transparent !important;
|
||||
border-left-color: transparent !important;
|
||||
}
|
||||
|
||||
/* Icons - default color */
|
||||
[class*="hitIcon"],
|
||||
[class*="hitTree"] {
|
||||
color: #71717a !important;
|
||||
transition: all 0.15s ease !important;
|
||||
}
|
||||
|
||||
[class*="hitIcon"] svg path,
|
||||
[class*="hitIcon"] svg,
|
||||
[class*="hitTree"] svg path,
|
||||
[class*="hitTree"] svg {
|
||||
stroke: #71717a !important;
|
||||
color: #71717a !important;
|
||||
transition: all 0.15s ease !important;
|
||||
}
|
||||
|
||||
/* Icons - hover gradient effect */
|
||||
[class*="suggestion"]:hover [class*="hitIcon"] svg path,
|
||||
[class*="suggestion"]:hover [class*="hitIcon"] svg,
|
||||
[class*="suggestion"]:hover [class*="hitTree"] svg path,
|
||||
[class*="suggestion"]:hover [class*="hitTree"] svg {
|
||||
stroke: #0074d9 !important;
|
||||
color: #0074d9 !important;
|
||||
}
|
||||
|
||||
[data-theme='dark'] [class*="hitIcon"] svg path,
|
||||
[data-theme='dark'] [class*="hitIcon"] svg,
|
||||
[data-theme='dark'] [class*="hitTree"] svg path,
|
||||
[data-theme='dark'] [class*="hitTree"] svg {
|
||||
stroke: #71717a !important;
|
||||
color: #71717a !important;
|
||||
}
|
||||
|
||||
[data-theme='dark'] [class*="suggestion"]:hover [class*="hitIcon"] svg path,
|
||||
[data-theme='dark'] [class*="suggestion"]:hover [class*="hitIcon"] svg,
|
||||
[data-theme='dark'] [class*="suggestion"]:hover [class*="hitTree"] svg path,
|
||||
[data-theme='dark'] [class*="suggestion"]:hover [class*="hitTree"] svg {
|
||||
stroke: #3396e8 !important;
|
||||
color: #3396e8 !important;
|
||||
}
|
||||
|
||||
/* === PAGE-LEVEL RESULTS (doc icon, no tree) === */
|
||||
/* Title is the page name - GRADIENT */
|
||||
[class*="suggestion"]:not(:has([class*="hitTree"])) [class*="hitTitle"] {
|
||||
font-family: 'Space Grotesk', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
|
||||
font-weight: 600 !important;
|
||||
font-size: 0.9375rem !important;
|
||||
background: linear-gradient(90deg, #0074d9, #009296) !important;
|
||||
-webkit-background-clip: text !important;
|
||||
background-clip: text !important;
|
||||
-webkit-text-fill-color: transparent !important;
|
||||
margin: 0 !important;
|
||||
}
|
||||
|
||||
/* === SECTION-LEVEL RESULTS (hashtag icon, has tree) === */
|
||||
/* Title is the section name - regular gray text */
|
||||
[class*="suggestion"]:has([class*="hitTree"]) [class*="hitTitle"] {
|
||||
font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
|
||||
font-weight: 500 !important;
|
||||
font-size: 0.8125rem !important;
|
||||
background: none !important;
|
||||
-webkit-background-clip: unset !important;
|
||||
background-clip: unset !important;
|
||||
color: #3f3f46 !important;
|
||||
-webkit-text-fill-color: #3f3f46 !important;
|
||||
margin: 0 !important;
|
||||
}
|
||||
|
||||
[data-theme='dark'] [class*="suggestion"]:has([class*="hitTree"]) [class*="hitTitle"] {
|
||||
color: #d4d4d8 !important;
|
||||
-webkit-text-fill-color: #d4d4d8 !important;
|
||||
}
|
||||
|
||||
/* Path is the page name - GRADIENT */
|
||||
[class*="suggestion"]:has([class*="hitTree"]) [class*="hitPath"] {
|
||||
font-family: 'Space Grotesk', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
|
||||
font-weight: 600 !important;
|
||||
font-size: 0.75rem !important;
|
||||
background: linear-gradient(90deg, #0074d9, #009296) !important;
|
||||
-webkit-background-clip: text !important;
|
||||
background-clip: text !important;
|
||||
-webkit-text-fill-color: transparent !important;
|
||||
text-transform: none !important;
|
||||
letter-spacing: 0 !important;
|
||||
margin-top: 0.125rem !important;
|
||||
display: block !important;
|
||||
}
|
||||
|
||||
/* === FALLBACK for browsers without :has() === */
|
||||
[class*="hitPath"] {
|
||||
font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
|
||||
font-size: 0.625rem !important;
|
||||
font-weight: 400 !important;
|
||||
color: #71717a !important;
|
||||
-webkit-text-fill-color: #71717a !important;
|
||||
text-transform: none !important;
|
||||
letter-spacing: 0 !important;
|
||||
}
|
||||
|
||||
[class*="hitTitle"] {
|
||||
font-family: 'Space Grotesk', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
|
||||
font-weight: 600 !important;
|
||||
font-size: 0.9375rem !important;
|
||||
background: linear-gradient(90deg, #0074d9, #009296) !important;
|
||||
-webkit-background-clip: text !important;
|
||||
background-clip: text !important;
|
||||
-webkit-text-fill-color: transparent !important;
|
||||
margin: 0 !important;
|
||||
}
|
||||
|
||||
/* Highlight matched text - must override color rules */
|
||||
[class*="searchResultItem"] mark,
|
||||
[class*="suggestion"] mark,
|
||||
.search-result-match mark,
|
||||
[class*="searchResultItem"]:hover mark,
|
||||
[class*="suggestion"]:hover mark,
|
||||
[class*="searchResultItem"][class*="cursor"] mark,
|
||||
[class*="suggestion"][class*="cursor"] mark {
|
||||
background: transparent !important;
|
||||
color: #0074d9 !important;
|
||||
padding: 0 !important;
|
||||
border-radius: 0 !important;
|
||||
font-weight: 600 !important;
|
||||
text-decoration: none !important;
|
||||
}
|
||||
|
||||
[data-theme='dark'] [class*="searchResultItem"] mark,
|
||||
[data-theme='dark'] [class*="suggestion"] mark,
|
||||
[data-theme='dark'] .search-result-match mark,
|
||||
[data-theme='dark'] [class*="searchResultItem"]:hover mark,
|
||||
[data-theme='dark'] [class*="suggestion"]:hover mark,
|
||||
[data-theme='dark'] [class*="searchResultItem"][class*="cursor"] mark,
|
||||
[data-theme='dark'] [class*="suggestion"][class*="cursor"] mark {
|
||||
color: #66b3f0 !important;
|
||||
}
|
||||
|
||||
/* No results */
|
||||
[class*="noResults"],
|
||||
[class*="searchNoResult"] {
|
||||
font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
|
||||
font-size: 0.75rem !important;
|
||||
color: var(--ifm-color-emphasis-500) !important;
|
||||
padding: 1rem !important;
|
||||
text-align: center !important;
|
||||
}
|
||||
|
||||
/* Loading */
|
||||
[class*="searchIndexLoading"],
|
||||
[class*="loadingRing"] {
|
||||
font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
|
||||
font-size: 0.75rem !important;
|
||||
color: var(--ifm-color-emphasis-400) !important;
|
||||
}
|
||||
|
||||
/* Hide footer */
|
||||
[class*="hitFooter"],
|
||||
[class*="searchFooter"] {
|
||||
display: none !important;
|
||||
}
|
||||
|
||||
/* Hide action icons */
|
||||
[class*="hitAction"] {
|
||||
display: none !important;
|
||||
}
|
||||
|
||||
|
||||
Executable
+266
@@ -0,0 +1,266 @@
|
||||
#!/bin/bash
|
||||
#
|
||||
# Install Hindsight Agent Skill
|
||||
#
|
||||
# Usage:
|
||||
# curl -fsSL https://hindsight.vectorize.io/get-skill | bash
|
||||
#
|
||||
# Options:
|
||||
# --app <app> Target app: claude, opencode, codex
|
||||
#
|
||||
# Examples:
|
||||
# curl -fsSL https://hindsight.vectorize.io/get-skill | bash -s -- --app claude
|
||||
# curl -fsSL https://hindsight.vectorize.io/get-skill | bash -s -- --app opencode
|
||||
#
|
||||
|
||||
set -e
|
||||
|
||||
# Colors for output
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
YELLOW='\033[1;33m'
|
||||
BLUE='\033[0;34m'
|
||||
CYAN='\033[0;36m'
|
||||
BOLD='\033[1m'
|
||||
DIM='\033[2m'
|
||||
NC='\033[0m' # No Color
|
||||
|
||||
print_info() {
|
||||
echo -e "${BLUE}ℹ${NC} $1"
|
||||
}
|
||||
|
||||
print_success() {
|
||||
echo -e "${GREEN}✓${NC} $1"
|
||||
}
|
||||
|
||||
print_error() {
|
||||
echo -e "${RED}✗${NC} $1"
|
||||
exit 1
|
||||
}
|
||||
|
||||
print_warning() {
|
||||
echo -e "${YELLOW}⚠${NC} $1"
|
||||
}
|
||||
|
||||
print_step() {
|
||||
echo ""
|
||||
echo -e "${BOLD}${CYAN}▸ $1${NC}"
|
||||
echo ""
|
||||
}
|
||||
|
||||
print_banner() {
|
||||
echo ""
|
||||
# ANSI logo
|
||||
echo -e " \033[38;2;9;127;184m▄\033[0m\033[48;2;8;130;178m\033[38;2;5;133;186m▄\033[0m \033[48;2;10;143;160m\033[38;2;10;143;165m▄\033[0m\033[38;2;7;140;156m▄\033[0m "
|
||||
echo -e " \033[38;2;8;125;192m▄\033[0m \033[38;2;3;132;191m▀\033[0m\033[38;2;2;133;192m▄\033[0m \033[38;2;3;132;180m▄\033[0m\033[38;2;1;137;184m▄\033[0m\033[38;2;3;133;174m▄\033[0m \033[38;2;3;142;176m▄\033[0m\033[38;2;4;142;169m▀\033[0m \033[38;2;10;144;164m▄\033[0m "
|
||||
echo -e "\033[38;2;6;121;195m▀\033[0m\033[38;2;5;128;203m▀\033[0m\033[48;2;5;124;195m\033[38;2;3;125;200m▄\033[0m\033[38;2;2;126;196m▄\033[0m\033[48;2;3;128;188m\033[38;2;1;131;196m▄\033[0m\033[48;2;0;152;219m\033[38;2;2;131;191m▄\033[0m\033[38;2;1;141;196m▀\033[0m\033[38;2;1;135;183m▀\033[0m\033[38;2;1;148;198m▀\033[0m\033[48;2;1;156;202m\033[38;2;2;135;180m▄\033[0m\033[48;2;4;134;169m\033[38;2;1;137;177m▄\033[0m\033[38;2;3;138;173m▄\033[0m\033[48;2;6;137;165m\033[38;2;2;140;170m▄\033[0m\033[38;2;7;144;169m▀\033[0m\033[38;2;7;139;158m▀\033[0m"
|
||||
echo -e " \033[48;2;2;128;202m\033[38;2;2;124;201m▄\033[0m\033[48;2;1;130;201m\033[38;2;0;135;212m▄\033[0m\033[38;2;2;128;196m▄\033[0m \033[48;2;2;142;204m\033[38;2;7;138;199m▄\033[0m \033[38;2;1;135;186m▄\033[0m\033[48;2;1;142;186m\033[38;2;2;144;194m▄\033[0m\033[48;2;3;138;176m\033[38;2;2;134;176m▄\033[0m "
|
||||
echo -e " \033[48;2;8;118;200m\033[38;2;8;121;209m▄\033[0m\033[38;2;3;121;203m▀\033[0m \033[38;2;3;122;192m▀\033[0m\033[38;2;1;138;216m▀\033[0m\033[48;2;0;138;210m\033[38;2;3;128;198m▄\033[0m\033[48;2;0;126;188m\033[38;2;2;131;198m▄\033[0m\033[48;2;0;142;205m\033[38;2;3;132;193m▄\033[0m\033[38;2;1;140;196m▀\033[0m \033[38;2;4;134;175m▀\033[0m\033[48;2;13;135;167m\033[38;2;8;136;174m▄\033[0m "
|
||||
echo ""
|
||||
echo -e " ${BOLD}HINDSIGHT SKILL INSTALLER${NC}"
|
||||
echo -e " ${DIM}Give your AI agent persistent memory${NC}"
|
||||
echo ""
|
||||
}
|
||||
|
||||
# Embedded SKILL.md content
|
||||
SKILL_CONTENT='---
|
||||
name: hindsight
|
||||
description: Give your agent persistent memory that works like human memory. Store facts, preferences, and context that persist across sessions.
|
||||
---
|
||||
|
||||
# Hindsight Memory Skill
|
||||
|
||||
You have access to persistent memory via the `hindsight-embed` CLI. Use it to remember important information about the user and recall it when relevant.
|
||||
|
||||
## Setup (first time only)
|
||||
|
||||
Run: `uvx hindsight-embed configure`
|
||||
|
||||
This will configure your LLM provider and start a local daemon that manages your memory bank.
|
||||
|
||||
## Commands
|
||||
|
||||
The CLI uses a bank ID to organize memories. Use `default` for general memories or create project-specific banks.
|
||||
|
||||
### Store a memory
|
||||
|
||||
Use `memory retain` to store important facts, preferences, decisions, or context:
|
||||
|
||||
```bash
|
||||
uvx hindsight-embed memory retain default "User prefers dark mode for all UIs"
|
||||
uvx hindsight-embed memory retain default "Project uses Python 3.11 with FastAPI" --context work
|
||||
uvx hindsight-embed memory retain myproject "API uses JWT authentication"
|
||||
```
|
||||
|
||||
### Recall memories
|
||||
|
||||
Use `memory recall` to search for relevant memories before starting tasks:
|
||||
|
||||
```bash
|
||||
uvx hindsight-embed memory recall default "What are the user'"'"'s UI preferences?"
|
||||
uvx hindsight-embed memory recall default "What tech stack does this project use?"
|
||||
```
|
||||
|
||||
### Reflect on memories
|
||||
|
||||
Use `memory reflect` for contextual answers that synthesize multiple memories:
|
||||
|
||||
```bash
|
||||
uvx hindsight-embed memory reflect default "How should I set up the dev environment?"
|
||||
```
|
||||
|
||||
### Other commands
|
||||
|
||||
```bash
|
||||
uvx hindsight-embed bank list # List all memory banks
|
||||
uvx hindsight-embed daemon status # Check daemon status
|
||||
uvx hindsight-embed --help # Full CLI help
|
||||
```
|
||||
|
||||
## When to Use
|
||||
|
||||
### Store memories when you learn:
|
||||
- User preferences (coding style, tools, UI preferences)
|
||||
- Project context (tech stack, architecture decisions)
|
||||
- Personal information the user shares (name, role, company)
|
||||
- Important decisions or outcomes
|
||||
|
||||
### Recall memories when:
|
||||
- Starting a new task (get relevant context first)
|
||||
- Making decisions that should consider user preferences
|
||||
- Working on a project where past context would help
|
||||
|
||||
## Best Practices
|
||||
|
||||
1. **Be specific**: Store "User prefers 2-space indentation" not "User has preferences"
|
||||
2. **Recall first**: Before starting tasks, recall relevant context
|
||||
3. **Use context tags**: Organize with `--context` (work, personal, preferences)
|
||||
4. **Use project banks**: Create separate banks for different projects
|
||||
'
|
||||
|
||||
# Get skills directory for app (bash 3.x compatible)
|
||||
get_skills_dir() {
|
||||
case "$1" in
|
||||
claude) echo "$HOME/.claude/skills" ;;
|
||||
opencode) echo "$HOME/.opencode/skills" ;;
|
||||
codex) echo "$HOME/.codex/skills" ;;
|
||||
*) echo "" ;;
|
||||
esac
|
||||
}
|
||||
|
||||
# Get app display name (bash 3.x compatible)
|
||||
get_app_name() {
|
||||
case "$1" in
|
||||
claude) echo "Claude Code" ;;
|
||||
opencode) echo "OpenCode" ;;
|
||||
codex) echo "Codex CLI" ;;
|
||||
*) echo "$1" ;;
|
||||
esac
|
||||
}
|
||||
|
||||
# Parse arguments
|
||||
APP=""
|
||||
|
||||
show_usage() {
|
||||
echo "Usage: $0 [--app <app>]"
|
||||
echo ""
|
||||
echo "Options:"
|
||||
echo " --app <app> Target app: claude, opencode, codex"
|
||||
echo ""
|
||||
echo "Examples:"
|
||||
echo " $0 --app claude"
|
||||
echo " $0 --app opencode"
|
||||
exit 1
|
||||
}
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
--app)
|
||||
APP="$2"
|
||||
shift 2
|
||||
;;
|
||||
--help|-h)
|
||||
show_usage
|
||||
;;
|
||||
*)
|
||||
print_error "Unknown option: $1"
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# Show banner
|
||||
print_banner
|
||||
|
||||
# Validate app parameter
|
||||
if [ -z "$APP" ]; then
|
||||
echo -e "${DIM}Select your AI coding assistant:${NC}"
|
||||
echo ""
|
||||
echo " ${BOLD}1)${NC} Claude Code"
|
||||
echo " ${BOLD}2)${NC} OpenCode"
|
||||
echo " ${BOLD}3)${NC} Codex CLI"
|
||||
echo ""
|
||||
read -p "Enter choice [1]: " app_choice
|
||||
app_choice=${app_choice:-1}
|
||||
|
||||
case $app_choice in
|
||||
1) APP="claude" ;;
|
||||
2) APP="opencode" ;;
|
||||
3) APP="codex" ;;
|
||||
*) APP="claude" ;;
|
||||
esac
|
||||
echo ""
|
||||
fi
|
||||
|
||||
# Get skills directory for selected app
|
||||
SKILLS_DIR=$(get_skills_dir "$APP")
|
||||
APP_NAME=$(get_app_name "$APP")
|
||||
if [ -z "$SKILLS_DIR" ]; then
|
||||
print_error "Unknown app '$APP'. Supported: claude, opencode, codex"
|
||||
fi
|
||||
|
||||
print_info "Installing for ${BOLD}$APP_NAME${NC}"
|
||||
|
||||
# Step 1: Check for Python/uvx
|
||||
print_step "Checking prerequisites"
|
||||
|
||||
if ! command -v python3 &> /dev/null && ! command -v uvx &> /dev/null; then
|
||||
print_error "Python 3 or uvx is required.\nInstall from https://python.org or https://docs.astral.sh/uv/"
|
||||
fi
|
||||
print_success "Python/uvx available"
|
||||
|
||||
# Step 2: Configure LLM provider using the CLI
|
||||
print_step "Configuring LLM provider"
|
||||
|
||||
# Install/run hindsight-embed configure
|
||||
if command -v uvx &> /dev/null; then
|
||||
uvx hindsight-embed configure
|
||||
else
|
||||
pip install -q hindsight-embed
|
||||
hindsight-embed configure
|
||||
fi
|
||||
|
||||
# Step 3: Install skill to app's skills directory
|
||||
print_step "Installing skill to $APP_NAME"
|
||||
|
||||
mkdir -p "$SKILLS_DIR/hindsight"
|
||||
|
||||
# Write embedded SKILL.md content
|
||||
echo "$SKILL_CONTENT" > "$SKILLS_DIR/hindsight/SKILL.md"
|
||||
print_success "Installed to $SKILLS_DIR/hindsight/"
|
||||
|
||||
# Done!
|
||||
echo ""
|
||||
echo -e "${GREEN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
|
||||
echo -e "${GREEN} ✓ Installation Complete!${NC}"
|
||||
echo -e "${GREEN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
|
||||
echo ""
|
||||
echo -e " The Hindsight skill is now available in ${BOLD}$APP_NAME${NC}."
|
||||
echo ""
|
||||
echo -e " ${DIM}Test the CLI:${NC}"
|
||||
echo -e " ${CYAN}uvx hindsight-embed memory retain default \"Test memory\"${NC}"
|
||||
echo -e " ${CYAN}uvx hindsight-embed memory recall default \"test\"${NC}"
|
||||
echo ""
|
||||
echo -e " ${DIM}$APP_NAME will automatically use the skill when relevant.${NC}"
|
||||
echo ""
|
||||
echo -e " ${DIM}Documentation:${NC} ${BLUE}https://hindsight.vectorize.io${NC}"
|
||||
echo ""
|
||||
@@ -0,0 +1,167 @@
|
||||
# hindsight-embed
|
||||
|
||||
Hindsight embedded CLI - local memory operations with automatic daemon management.
|
||||
|
||||
This package provides a simple CLI for storing and recalling memories using Hindsight's memory engine. It automatically manages a background daemon for fast operations - no manual server setup required.
|
||||
|
||||
## How It Works
|
||||
|
||||
`hindsight-embed` uses a background daemon architecture for optimal performance:
|
||||
|
||||
1. **First command**: Automatically starts a local daemon (first run downloads dependencies and loads ML models - can take 1-3 minutes)
|
||||
2. **Subsequent commands**: Near-instant responses (~1-2s) since daemon is already running
|
||||
3. **Auto-shutdown**: Daemon automatically exits after 5 minutes of inactivity
|
||||
|
||||
The daemon runs on `localhost:8889` and uses an embedded PostgreSQL database (pg0) - everything stays local on your machine.
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install hindsight-embed
|
||||
# or with uvx (no install needed)
|
||||
uvx hindsight-embed --help
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
# Interactive setup (recommended)
|
||||
hindsight-embed configure
|
||||
|
||||
# Or set your LLM API key manually
|
||||
export OPENAI_API_KEY=sk-...
|
||||
|
||||
# Store a memory (bank_id = "default")
|
||||
hindsight-embed memory retain default "User prefers dark mode"
|
||||
|
||||
# Recall memories
|
||||
hindsight-embed memory recall default "What are user preferences?"
|
||||
```
|
||||
|
||||
## Commands
|
||||
|
||||
### configure
|
||||
|
||||
Interactive setup wizard:
|
||||
|
||||
```bash
|
||||
hindsight-embed configure
|
||||
```
|
||||
|
||||
This will:
|
||||
- Let you choose an LLM provider (OpenAI, Groq, Google, Ollama)
|
||||
- Configure your API key
|
||||
- Set the model and memory bank ID
|
||||
- Start the daemon with your configuration
|
||||
|
||||
### memory retain
|
||||
|
||||
Store a memory:
|
||||
|
||||
```bash
|
||||
hindsight-embed memory retain default "User prefers dark mode"
|
||||
hindsight-embed memory retain default "Meeting on Monday" --context work
|
||||
hindsight-embed memory retain myproject "API uses JWT authentication"
|
||||
```
|
||||
|
||||
### memory recall
|
||||
|
||||
Search memories:
|
||||
|
||||
```bash
|
||||
hindsight-embed memory recall default "user preferences"
|
||||
hindsight-embed memory recall default "upcoming events"
|
||||
```
|
||||
|
||||
Use `-o json` for JSON output:
|
||||
```bash
|
||||
hindsight-embed memory recall default "user preferences" -o json
|
||||
```
|
||||
|
||||
### memory reflect
|
||||
|
||||
Get contextual answers that synthesize multiple memories:
|
||||
|
||||
```bash
|
||||
hindsight-embed memory reflect default "How should I set up the dev environment?"
|
||||
```
|
||||
|
||||
### bank list
|
||||
|
||||
List all memory banks:
|
||||
|
||||
```bash
|
||||
hindsight-embed bank list
|
||||
```
|
||||
|
||||
### daemon
|
||||
|
||||
Manage the background daemon:
|
||||
|
||||
```bash
|
||||
hindsight-embed daemon status # Check if daemon is running
|
||||
hindsight-embed daemon start # Start the daemon
|
||||
hindsight-embed daemon stop # Stop the daemon
|
||||
hindsight-embed daemon logs # View last 50 lines of logs
|
||||
hindsight-embed daemon logs -f # Follow logs in real-time
|
||||
hindsight-embed daemon logs -n 100 # View last 100 lines
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
### Interactive Setup
|
||||
|
||||
Run `hindsight-embed configure` for a guided setup that saves to `~/.hindsight/embed`.
|
||||
|
||||
### Environment Variables
|
||||
|
||||
| Variable | Description | Default |
|
||||
|----------|-------------|---------|
|
||||
| `HINDSIGHT_EMBED_LLM_API_KEY` | LLM API key (or use `OPENAI_API_KEY`) | Required |
|
||||
| `HINDSIGHT_EMBED_LLM_PROVIDER` | LLM provider (`openai`, `groq`, `google`, `ollama`) | `openai` |
|
||||
| `HINDSIGHT_EMBED_LLM_MODEL` | LLM model | `gpt-4o-mini` |
|
||||
| `HINDSIGHT_EMBED_BANK_ID` | Memory bank ID | `default` |
|
||||
|
||||
### Files
|
||||
|
||||
| Path | Description |
|
||||
|------|-------------|
|
||||
| `~/.hindsight/embed` | Configuration file |
|
||||
| `~/.hindsight/config.env` | Alternative config file location |
|
||||
| `~/.hindsight/daemon.log` | Daemon logs |
|
||||
| `~/.hindsight/daemon.lock` | Daemon lock file (PID) |
|
||||
|
||||
## Use with AI Coding Assistants
|
||||
|
||||
This CLI is designed to work with AI coding assistants like Claude Code, Cursor, and Windsurf. Install the Hindsight skill:
|
||||
|
||||
```bash
|
||||
curl -fsSL https://hindsight.vectorize.io/get-skill | bash
|
||||
```
|
||||
|
||||
This will configure the LLM provider and install the skill to your assistant's skills directory.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
**Daemon won't start:**
|
||||
```bash
|
||||
# Check logs for errors
|
||||
hindsight-embed daemon logs
|
||||
|
||||
# Stop any stuck daemon and restart
|
||||
hindsight-embed daemon stop
|
||||
hindsight-embed daemon start
|
||||
```
|
||||
|
||||
**Slow first command:**
|
||||
This is expected - the first command needs to download dependencies, start the daemon, and load ML models. First run can take 1-3 minutes depending on network speed. Subsequent commands will be fast (~1-2s).
|
||||
|
||||
**Change configuration:**
|
||||
```bash
|
||||
# Re-run configure (automatically restarts daemon)
|
||||
hindsight-embed configure
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
Apache 2.0
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Hindsight embedded CLI - local memory operations without a server."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
@@ -0,0 +1,417 @@
|
||||
"""
|
||||
Hindsight Embedded CLI.
|
||||
|
||||
A wrapper CLI that manages a local daemon and forwards commands to hindsight-cli.
|
||||
No external server required - runs everything locally with automatic daemon management.
|
||||
|
||||
Usage:
|
||||
hindsight-embed configure # Interactive setup
|
||||
hindsight-embed retain "User prefers dark mode"
|
||||
hindsight-embed recall "What are user preferences?"
|
||||
hindsight-embed daemon status # Check daemon status
|
||||
|
||||
Environment variables:
|
||||
HINDSIGHT_EMBED_LLM_API_KEY: Required. API key for LLM provider.
|
||||
HINDSIGHT_EMBED_LLM_PROVIDER: Optional. LLM provider (default: "openai").
|
||||
HINDSIGHT_EMBED_LLM_MODEL: Optional. LLM model (default: "gpt-4o-mini").
|
||||
HINDSIGHT_EMBED_BANK_ID: Optional. Memory bank ID (default: "default").
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
CONFIG_DIR = Path.home() / ".hindsight"
|
||||
CONFIG_FILE = CONFIG_DIR / "embed"
|
||||
CONFIG_FILE_ALT = CONFIG_DIR / "config.env" # Alternative config file location
|
||||
|
||||
|
||||
def setup_logging(verbose: bool = False):
|
||||
"""Configure logging."""
|
||||
level_str = os.environ.get("HINDSIGHT_EMBED_LOG_LEVEL", "info").lower()
|
||||
if verbose:
|
||||
level_str = "debug"
|
||||
|
||||
level_map = {
|
||||
"debug": logging.DEBUG,
|
||||
"info": logging.INFO,
|
||||
"warning": logging.WARNING,
|
||||
"error": logging.ERROR,
|
||||
}
|
||||
level = level_map.get(level_str, logging.INFO)
|
||||
|
||||
logging.basicConfig(
|
||||
level=level,
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
stream=sys.stderr,
|
||||
)
|
||||
return logging.getLogger(__name__)
|
||||
|
||||
|
||||
def load_config_file():
|
||||
"""Load configuration from file if it exists."""
|
||||
# Check both config file locations
|
||||
config_files = [CONFIG_FILE, CONFIG_FILE_ALT]
|
||||
for config_path in config_files:
|
||||
if config_path.exists():
|
||||
with open(config_path) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line and not line.startswith("#") and "=" in line:
|
||||
# Handle 'export VAR=value' format
|
||||
if line.startswith("export "):
|
||||
line = line[7:]
|
||||
key, value = line.split("=", 1)
|
||||
if key not in os.environ: # Don't override env vars
|
||||
os.environ[key] = value
|
||||
|
||||
|
||||
def get_config():
|
||||
"""Get configuration from environment variables."""
|
||||
load_config_file()
|
||||
return {
|
||||
"llm_api_key": os.environ.get("HINDSIGHT_EMBED_LLM_API_KEY")
|
||||
or os.environ.get("HINDSIGHT_API_LLM_API_KEY")
|
||||
or os.environ.get("OPENAI_API_KEY"),
|
||||
"llm_provider": os.environ.get("HINDSIGHT_EMBED_LLM_PROVIDER")
|
||||
or os.environ.get("HINDSIGHT_API_LLM_PROVIDER", "openai"),
|
||||
"llm_model": os.environ.get("HINDSIGHT_EMBED_LLM_MODEL")
|
||||
or os.environ.get("HINDSIGHT_API_LLM_MODEL", "gpt-4o-mini"),
|
||||
"bank_id": os.environ.get("HINDSIGHT_EMBED_BANK_ID", "default"),
|
||||
}
|
||||
|
||||
|
||||
def do_configure(args):
|
||||
"""Interactive configuration setup with beautiful TUI."""
|
||||
import questionary
|
||||
from questionary import Style
|
||||
|
||||
# If stdin is not a terminal (e.g., running via curl | bash),
|
||||
# reopen stdin from /dev/tty for interactive prompts
|
||||
if not sys.stdin.isatty():
|
||||
try:
|
||||
sys.stdin = open('/dev/tty', 'r')
|
||||
except OSError:
|
||||
print("Error: Cannot run interactive configuration without a terminal.", file=sys.stderr)
|
||||
print("Run directly: uvx hindsight-embed configure", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
# Custom style for the prompts
|
||||
custom_style = Style([
|
||||
('qmark', 'fg:cyan bold'),
|
||||
('question', 'fg:white bold'),
|
||||
('answer', 'fg:cyan'),
|
||||
('pointer', 'fg:cyan bold'),
|
||||
('highlighted', 'fg:cyan bold'),
|
||||
('selected', 'fg:green'),
|
||||
('text', 'fg:white'),
|
||||
])
|
||||
|
||||
print()
|
||||
print("\033[1m\033[36m ╭─────────────────────────────────────╮\033[0m")
|
||||
print("\033[1m\033[36m │ Hindsight Embed Configuration │\033[0m")
|
||||
print("\033[1m\033[36m ╰─────────────────────────────────────╯\033[0m")
|
||||
print()
|
||||
|
||||
# Check existing config
|
||||
if CONFIG_FILE.exists():
|
||||
if not questionary.confirm(
|
||||
"Existing configuration found. Reconfigure?",
|
||||
default=False,
|
||||
style=custom_style,
|
||||
).ask():
|
||||
print("\n\033[32m✓\033[0m Keeping existing configuration.")
|
||||
return 0
|
||||
print()
|
||||
|
||||
# Provider selection with descriptions
|
||||
providers = [
|
||||
questionary.Choice("OpenAI (recommended)", value=("openai", "o3-mini", "OpenAI")),
|
||||
questionary.Choice("Groq (fast & free tier)", value=("groq", "openai/gpt-oss-20b", "Groq")),
|
||||
questionary.Choice("Google Gemini", value=("google", "gemini-2.0-flash", "Google")),
|
||||
questionary.Choice("Ollama (local, no API key)", value=("ollama", "llama3.2", None)),
|
||||
]
|
||||
|
||||
result = questionary.select(
|
||||
"Select your LLM provider:",
|
||||
choices=providers,
|
||||
style=custom_style,
|
||||
).ask()
|
||||
|
||||
if result is None: # User cancelled
|
||||
print("\n\033[33m⚠\033[0m Configuration cancelled.")
|
||||
return 1
|
||||
|
||||
provider, default_model, key_name = result
|
||||
|
||||
# API key
|
||||
api_key = ""
|
||||
if key_name:
|
||||
env_keys = {
|
||||
"OpenAI": "OPENAI_API_KEY",
|
||||
"Groq": "GROQ_API_KEY",
|
||||
"Google": "GOOGLE_API_KEY",
|
||||
}
|
||||
env_key = env_keys.get(key_name, "")
|
||||
existing = os.environ.get(env_key, "")
|
||||
|
||||
if existing:
|
||||
masked = existing[:8] + "..." + existing[-4:] if len(existing) > 12 else "***"
|
||||
if questionary.confirm(
|
||||
f"Found {key_name} key in ${env_key} ({masked}). Use it?",
|
||||
default=True,
|
||||
style=custom_style,
|
||||
).ask():
|
||||
api_key = existing
|
||||
|
||||
if not api_key:
|
||||
api_key = questionary.password(
|
||||
f"Enter your {key_name} API key:",
|
||||
style=custom_style,
|
||||
).ask()
|
||||
|
||||
if not api_key:
|
||||
print("\n\033[31m✗\033[0m API key is required.", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
# Model selection
|
||||
model = questionary.text(
|
||||
"Model name:",
|
||||
default=default_model,
|
||||
style=custom_style,
|
||||
).ask()
|
||||
|
||||
if model is None:
|
||||
return 1
|
||||
|
||||
# Bank ID
|
||||
bank_id = questionary.text(
|
||||
"Memory bank ID:",
|
||||
default="default",
|
||||
style=custom_style,
|
||||
).ask()
|
||||
|
||||
if bank_id is None:
|
||||
return 1
|
||||
|
||||
# Save configuration
|
||||
CONFIG_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(CONFIG_FILE, "w") as f:
|
||||
f.write("# Hindsight Embed Configuration\n")
|
||||
f.write(f"# Generated by hindsight-embed configure\n\n")
|
||||
f.write(f"HINDSIGHT_EMBED_LLM_PROVIDER={provider}\n")
|
||||
f.write(f"HINDSIGHT_EMBED_LLM_MODEL={model}\n")
|
||||
f.write(f"HINDSIGHT_EMBED_BANK_ID={bank_id}\n")
|
||||
if api_key:
|
||||
f.write(f"HINDSIGHT_EMBED_LLM_API_KEY={api_key}\n")
|
||||
|
||||
CONFIG_FILE.chmod(0o600)
|
||||
|
||||
# Stop existing daemon if running (it needs to pick up new config)
|
||||
from . import daemon_client
|
||||
|
||||
if daemon_client._is_daemon_running():
|
||||
print("\n \033[2mRestarting daemon with new configuration...\033[0m")
|
||||
daemon_client.stop_daemon()
|
||||
|
||||
# Start daemon with new config
|
||||
new_config = {
|
||||
"llm_api_key": api_key,
|
||||
"llm_provider": provider,
|
||||
"llm_model": model,
|
||||
"bank_id": bank_id,
|
||||
}
|
||||
if daemon_client.ensure_daemon_running(new_config):
|
||||
print(" \033[32m✓ Daemon started\033[0m")
|
||||
else:
|
||||
print(" \033[33m⚠ Failed to start daemon (will start on first command)\033[0m")
|
||||
|
||||
print()
|
||||
print("\033[32m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\033[0m")
|
||||
print("\033[32m ✓ Configuration saved!\033[0m")
|
||||
print("\033[32m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\033[0m")
|
||||
print()
|
||||
print(f" \033[2mConfig:\033[0m {CONFIG_FILE}")
|
||||
print()
|
||||
print(" \033[2mTest with:\033[0m")
|
||||
print(' \033[36mhindsight-embed retain "Test memory"\033[0m')
|
||||
print(' \033[36mhindsight-embed recall "test"\033[0m')
|
||||
print()
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
def do_daemon(args, config: dict, logger):
|
||||
"""Handle daemon subcommands."""
|
||||
from pathlib import Path
|
||||
from . import daemon_client
|
||||
|
||||
daemon_log_path = Path.home() / ".hindsight" / "daemon.log"
|
||||
|
||||
if args.daemon_command == "start":
|
||||
if daemon_client._is_daemon_running():
|
||||
print("Daemon is already running")
|
||||
return 0
|
||||
|
||||
print("Starting daemon...")
|
||||
if daemon_client.ensure_daemon_running(config):
|
||||
print("Daemon started successfully")
|
||||
print(f" Port: {daemon_client.DAEMON_PORT}")
|
||||
print(f" Logs: {daemon_log_path}")
|
||||
return 0
|
||||
else:
|
||||
print("Failed to start daemon", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
elif args.daemon_command == "stop":
|
||||
if not daemon_client._is_daemon_running():
|
||||
print("Daemon is not running")
|
||||
return 0
|
||||
|
||||
print("Stopping daemon...")
|
||||
if daemon_client.stop_daemon():
|
||||
print("Daemon stopped")
|
||||
return 0
|
||||
else:
|
||||
print("Failed to stop daemon", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
elif args.daemon_command == "status":
|
||||
if daemon_client._is_daemon_running():
|
||||
# Get PID from lockfile
|
||||
lockfile = Path.home() / ".hindsight" / "daemon.lock"
|
||||
pid = "unknown"
|
||||
if lockfile.exists():
|
||||
try:
|
||||
pid = lockfile.read_text().strip()
|
||||
except Exception:
|
||||
pass
|
||||
print(f"Daemon is running (PID: {pid})")
|
||||
print(f" URL: http://127.0.0.1:{daemon_client.DAEMON_PORT}")
|
||||
print(f" Logs: {daemon_log_path}")
|
||||
return 0
|
||||
else:
|
||||
print("Daemon is not running")
|
||||
return 1
|
||||
|
||||
elif args.daemon_command == "logs":
|
||||
if not daemon_log_path.exists():
|
||||
print("No daemon logs found", file=sys.stderr)
|
||||
print(f" Expected at: {daemon_log_path}")
|
||||
return 1
|
||||
|
||||
if args.follow:
|
||||
# Follow mode - like tail -f
|
||||
import subprocess
|
||||
try:
|
||||
subprocess.run(["tail", "-f", str(daemon_log_path)])
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
return 0
|
||||
else:
|
||||
# Show last N lines
|
||||
try:
|
||||
with open(daemon_log_path) as f:
|
||||
lines = f.readlines()
|
||||
for line in lines[-args.lines:]:
|
||||
print(line, end="")
|
||||
return 0
|
||||
except Exception as e:
|
||||
print(f"Error reading logs: {e}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
else:
|
||||
print("Usage: hindsight-embed daemon {start|stop|status|logs}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point."""
|
||||
# Check for built-in commands first (before argparse)
|
||||
# This allows us to forward unknown commands to hindsight-cli
|
||||
if len(sys.argv) > 1:
|
||||
command = sys.argv[1]
|
||||
|
||||
# Handle configure
|
||||
if command == "configure":
|
||||
logger = setup_logging(False)
|
||||
exit_code = do_configure(None)
|
||||
sys.exit(exit_code)
|
||||
|
||||
# Handle daemon subcommands
|
||||
if command == "daemon":
|
||||
# Parse daemon subcommand
|
||||
parser = argparse.ArgumentParser(prog="hindsight-embed daemon")
|
||||
subparsers = parser.add_subparsers(dest="daemon_command")
|
||||
subparsers.add_parser("start", help="Start the daemon")
|
||||
subparsers.add_parser("stop", help="Stop the daemon")
|
||||
subparsers.add_parser("status", help="Check daemon status")
|
||||
logs_parser = subparsers.add_parser("logs", help="View daemon logs")
|
||||
logs_parser.add_argument("--follow", "-f", action="store_true")
|
||||
logs_parser.add_argument("--lines", "-n", type=int, default=50)
|
||||
|
||||
args = parser.parse_args(sys.argv[2:])
|
||||
logger = setup_logging(False)
|
||||
config = get_config()
|
||||
exit_code = do_daemon(args, config, logger)
|
||||
sys.exit(exit_code)
|
||||
|
||||
# Handle --help / -h
|
||||
if command in ("--help", "-h"):
|
||||
print_help()
|
||||
sys.exit(0)
|
||||
|
||||
# Forward all other commands to hindsight-cli
|
||||
config = get_config()
|
||||
|
||||
# Check for LLM API key
|
||||
if not config["llm_api_key"]:
|
||||
print("Error: LLM API key is required.", file=sys.stderr)
|
||||
print("Run 'hindsight-embed configure' to set up.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
from . import daemon_client
|
||||
|
||||
# Forward to hindsight-cli (handles daemon startup and CLI installation)
|
||||
exit_code = daemon_client.run_cli(sys.argv[1:], config)
|
||||
sys.exit(exit_code)
|
||||
|
||||
# No command - show help
|
||||
print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def print_help():
|
||||
"""Print help message."""
|
||||
print("""Hindsight Embedded CLI - local memory operations with automatic daemon management.
|
||||
|
||||
Usage: hindsight-embed <command> [options]
|
||||
|
||||
Built-in commands:
|
||||
configure Interactive configuration setup
|
||||
daemon start Start the background daemon
|
||||
daemon stop Stop the daemon
|
||||
daemon status Check daemon status
|
||||
daemon logs [-f] [-n] View daemon logs
|
||||
|
||||
CLI commands (forwarded to hindsight-cli):
|
||||
memory retain <bank> <content> Store a memory
|
||||
memory recall <bank> <query> Search memories
|
||||
memory reflect <bank> <query> Generate contextual answer
|
||||
bank list List memory banks
|
||||
... Run 'hindsight --help' for all commands
|
||||
|
||||
Examples:
|
||||
hindsight-embed configure
|
||||
hindsight-embed memory retain default "User prefers dark mode"
|
||||
hindsight-embed memory recall default "user preferences"
|
||||
hindsight-embed daemon status
|
||||
hindsight-embed daemon logs -f
|
||||
""")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,298 @@
|
||||
"""
|
||||
Client for communicating with the Hindsight daemon.
|
||||
|
||||
Handles daemon lifecycle (start if needed) and API requests via the Python client.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import httpx # Used only for health check
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DAEMON_PORT = 8889
|
||||
DAEMON_URL = f"http://127.0.0.1:{DAEMON_PORT}"
|
||||
DAEMON_STARTUP_TIMEOUT = 180 # seconds - needs to be long for first run (downloads dependencies)
|
||||
DAEMON_IDLE_TIMEOUT = 300 # 5 minutes - auto-exit after idle
|
||||
|
||||
# CLI paths - check multiple locations
|
||||
CLI_INSTALL_DIRS = [
|
||||
Path.home() / ".local" / "bin", # Standard location from get-cli installer
|
||||
Path.home() / ".hindsight" / "bin", # Alternative location
|
||||
]
|
||||
CLI_INSTALLER_URL = "https://hindsight.vectorize.io/get-cli"
|
||||
|
||||
|
||||
def _find_hindsight_api_command() -> list[str]:
|
||||
"""Find the command to run hindsight-api."""
|
||||
# Check if we're in development mode (local hindsight-api available)
|
||||
# Path: daemon_client.py -> hindsight_embed/ -> hindsight-embed/ -> memory-poc/
|
||||
dev_api_path = Path(__file__).parent.parent.parent / "hindsight-api"
|
||||
if dev_api_path.exists() and (dev_api_path / "pyproject.toml").exists():
|
||||
# Use uv run with the local project
|
||||
return ["uv", "run", "--project", str(dev_api_path), "hindsight-api"]
|
||||
|
||||
# Fall back to uvx for installed version
|
||||
return ["uvx", "hindsight-api"]
|
||||
|
||||
|
||||
def _is_daemon_running() -> bool:
|
||||
"""Check if daemon is running and responsive."""
|
||||
try:
|
||||
with httpx.Client(timeout=2) as client:
|
||||
response = client.get(f"{DAEMON_URL}/health")
|
||||
return response.status_code == 200
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _start_daemon(config: dict) -> bool:
|
||||
"""
|
||||
Start the daemon in background.
|
||||
|
||||
Returns True if daemon started successfully.
|
||||
"""
|
||||
import sys
|
||||
|
||||
logger.info("Starting daemon...")
|
||||
|
||||
# Build environment with LLM config
|
||||
env = os.environ.copy()
|
||||
if config.get("llm_api_key"):
|
||||
env["HINDSIGHT_API_LLM_API_KEY"] = config["llm_api_key"]
|
||||
if config.get("llm_provider"):
|
||||
env["HINDSIGHT_API_LLM_PROVIDER"] = config["llm_provider"]
|
||||
if config.get("llm_model"):
|
||||
env["HINDSIGHT_API_LLM_MODEL"] = config["llm_model"]
|
||||
|
||||
# Use pg0 database specific to bank
|
||||
bank_id = config.get("bank_id", "default")
|
||||
env["HINDSIGHT_API_DATABASE_URL"] = f"pg0://hindsight-embed-{bank_id}"
|
||||
|
||||
# Optimization flags for faster startup
|
||||
env["HINDSIGHT_API_SKIP_LLM_VERIFICATION"] = "true"
|
||||
env["HINDSIGHT_API_LOG_LEVEL"] = "warning"
|
||||
|
||||
cmd = _find_hindsight_api_command() + ["--daemon", "--idle-timeout", str(DAEMON_IDLE_TIMEOUT)]
|
||||
|
||||
# Create log directory
|
||||
log_dir = Path.home() / ".hindsight"
|
||||
log_dir.mkdir(parents=True, exist_ok=True)
|
||||
daemon_log = log_dir / "daemon.log"
|
||||
daemon_stderr = log_dir / "daemon.stderr"
|
||||
|
||||
print(f"Starting daemon with command: {' '.join(cmd)}", file=sys.stderr)
|
||||
print(f" Log file: {daemon_log}", file=sys.stderr)
|
||||
|
||||
try:
|
||||
# Start daemon in background, but capture initial stderr for debugging
|
||||
with open(daemon_stderr, "w") as stderr_file:
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
env=env,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=stderr_file,
|
||||
start_new_session=True,
|
||||
)
|
||||
|
||||
# Wait for daemon to be ready
|
||||
# Note: With --daemon flag, the parent process forks and exits immediately (code 0).
|
||||
# The child process (actual daemon) continues running. So we can't rely on process.poll()
|
||||
# to detect failures - we must use the health check.
|
||||
start_time = time.time()
|
||||
last_check_time = start_time
|
||||
while time.time() - start_time < DAEMON_STARTUP_TIMEOUT:
|
||||
if _is_daemon_running():
|
||||
logger.info("Daemon started successfully")
|
||||
return True
|
||||
|
||||
# Periodically log progress
|
||||
if time.time() - last_check_time > 5:
|
||||
elapsed = int(time.time() - start_time)
|
||||
print(f" Still waiting for daemon... ({elapsed}s elapsed)", file=sys.stderr)
|
||||
last_check_time = time.time()
|
||||
|
||||
time.sleep(0.5)
|
||||
|
||||
logger.error("Daemon failed to start within timeout")
|
||||
# Show logs on timeout
|
||||
if daemon_log.exists():
|
||||
log_content = daemon_log.read_text()
|
||||
if log_content:
|
||||
print(f"Daemon log:\n{log_content[-2000:]}", file=sys.stderr) # Last 2000 chars
|
||||
if daemon_stderr.exists():
|
||||
stderr_content = daemon_stderr.read_text()
|
||||
if stderr_content:
|
||||
print(f"Daemon stderr:\n{stderr_content}", file=sys.stderr)
|
||||
return False
|
||||
|
||||
except FileNotFoundError as e:
|
||||
print(f"Command not found: {cmd[0]}", file=sys.stderr)
|
||||
print(f" Full command: {' '.join(cmd)}", file=sys.stderr)
|
||||
logger.error("hindsight-api command not found. Install with: pip install hindsight-api")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"Failed to start daemon: {e}", file=sys.stderr)
|
||||
logger.error(f"Failed to start daemon: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def ensure_daemon_running(config: dict) -> bool:
|
||||
"""
|
||||
Ensure daemon is running, starting it if needed.
|
||||
|
||||
Returns True if daemon is running.
|
||||
"""
|
||||
if _is_daemon_running():
|
||||
logger.debug("Daemon already running")
|
||||
return True
|
||||
|
||||
return _start_daemon(config)
|
||||
|
||||
|
||||
def stop_daemon() -> bool:
|
||||
"""Stop the running daemon."""
|
||||
# Try to kill by PID from lockfile
|
||||
lockfile = Path.home() / ".hindsight" / "daemon.lock"
|
||||
if lockfile.exists():
|
||||
try:
|
||||
pid = int(lockfile.read_text().strip())
|
||||
os.kill(pid, 15) # SIGTERM
|
||||
# Wait for process to exit
|
||||
for _ in range(50):
|
||||
time.sleep(0.1)
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
except OSError:
|
||||
return True
|
||||
except (ValueError, OSError):
|
||||
pass
|
||||
|
||||
return not _is_daemon_running()
|
||||
|
||||
|
||||
def find_cli_binary() -> Path | None:
|
||||
"""Find the hindsight CLI binary in known locations or PATH."""
|
||||
import shutil
|
||||
|
||||
# Check standard install locations
|
||||
for install_dir in CLI_INSTALL_DIRS:
|
||||
binary = install_dir / "hindsight"
|
||||
if binary.exists() and os.access(binary, os.X_OK):
|
||||
return binary
|
||||
|
||||
# Check PATH
|
||||
path_binary = shutil.which("hindsight")
|
||||
if path_binary:
|
||||
return Path(path_binary)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def is_cli_installed() -> bool:
|
||||
"""Check if the hindsight CLI is installed."""
|
||||
return find_cli_binary() is not None
|
||||
|
||||
|
||||
def install_cli() -> bool:
|
||||
"""
|
||||
Install the hindsight CLI using the official installer.
|
||||
|
||||
Returns True if installation succeeded.
|
||||
"""
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
print("Installing hindsight CLI...")
|
||||
print(f" Installer URL: {CLI_INSTALLER_URL}")
|
||||
|
||||
try:
|
||||
# Download and run installer
|
||||
result = subprocess.run(
|
||||
["bash", "-c", f"curl -fsSL {CLI_INSTALLER_URL} | bash"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
print(f"CLI installation failed (exit code {result.returncode}):", file=sys.stderr)
|
||||
if result.stdout:
|
||||
print(f" stdout: {result.stdout}", file=sys.stderr)
|
||||
if result.stderr:
|
||||
print(f" stderr: {result.stderr}", file=sys.stderr)
|
||||
return False
|
||||
|
||||
cli_binary = find_cli_binary()
|
||||
if cli_binary:
|
||||
print(f"CLI installed to {cli_binary}")
|
||||
return True
|
||||
else:
|
||||
print("CLI installation completed but binary not found", file=sys.stderr)
|
||||
print(f" stdout: {result.stdout}", file=sys.stderr)
|
||||
print(f" stderr: {result.stderr}", file=sys.stderr)
|
||||
# Check known locations
|
||||
for install_dir in CLI_INSTALL_DIRS:
|
||||
binary = install_dir / "hindsight"
|
||||
print(f" Checking {binary}: exists={binary.exists()}", file=sys.stderr)
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"CLI installation failed: {e}", file=sys.stderr)
|
||||
return False
|
||||
|
||||
|
||||
def ensure_cli_installed() -> bool:
|
||||
"""Ensure CLI is installed, installing if needed."""
|
||||
if is_cli_installed():
|
||||
return True
|
||||
return install_cli()
|
||||
|
||||
|
||||
def run_cli(args: list[str], config: dict) -> int:
|
||||
"""
|
||||
Run the hindsight CLI with the given arguments.
|
||||
|
||||
Ensures daemon is running and passes the API URL.
|
||||
|
||||
Args:
|
||||
args: CLI arguments (e.g., ["memory", "retain", "bank", "content"])
|
||||
config: Configuration dict with llm settings
|
||||
|
||||
Returns:
|
||||
Exit code from CLI
|
||||
"""
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
# Ensure CLI is installed
|
||||
if not ensure_cli_installed():
|
||||
return 1
|
||||
|
||||
cli_binary = find_cli_binary()
|
||||
if not cli_binary:
|
||||
print("Error: hindsight CLI not found", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
# Ensure daemon is running
|
||||
if not ensure_daemon_running(config):
|
||||
print("Error: Failed to start daemon", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
# Build environment with API URL pointing to daemon
|
||||
env = os.environ.copy()
|
||||
env["HINDSIGHT_API_URL"] = DAEMON_URL
|
||||
|
||||
# Run CLI
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[str(cli_binary)] + args,
|
||||
env=env,
|
||||
)
|
||||
return result.returncode
|
||||
except Exception as e:
|
||||
print(f"Error running CLI: {e}", file=sys.stderr)
|
||||
return 1
|
||||
@@ -0,0 +1,20 @@
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-embed"
|
||||
version = "0.1.0"
|
||||
description = "Hindsight embedded CLI - local memory operations without a server"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"questionary>=2.0.0",
|
||||
"httpx>=0.27.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
hindsight-embed = "hindsight_embed.cli:main"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["hindsight_embed"]
|
||||
Executable
+236
@@ -0,0 +1,236 @@
|
||||
#!/bin/bash
|
||||
#
|
||||
# Smoke test for hindsight-embed CLI with daemon mode
|
||||
# Tests retain and recall operations via the background daemon
|
||||
#
|
||||
|
||||
set -e
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
API_DIR="$(cd "$SCRIPT_DIR/../hindsight-api" && pwd)"
|
||||
|
||||
echo "=== Hindsight Embed Smoke Test (Daemon Mode) ==="
|
||||
|
||||
# Check required environment (load from config if not set)
|
||||
if [ -f ~/.hindsight/config.env ]; then
|
||||
source ~/.hindsight/config.env
|
||||
fi
|
||||
|
||||
if [ -z "$HINDSIGHT_EMBED_LLM_API_KEY" ] && [ -z "$OPENAI_API_KEY" ]; then
|
||||
echo "Error: HINDSIGHT_EMBED_LLM_API_KEY or OPENAI_API_KEY is required"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Use a unique bank ID for this test run
|
||||
BANK_ID="test-$$-$(date +%s)"
|
||||
echo "Using bank ID: $BANK_ID"
|
||||
echo "Script dir: $SCRIPT_DIR"
|
||||
echo "API dir: $API_DIR"
|
||||
|
||||
# Debug: Check if hindsight CLI is available
|
||||
echo ""
|
||||
echo "Checking hindsight CLI availability..."
|
||||
if command -v hindsight &> /dev/null; then
|
||||
echo " hindsight CLI found at: $(which hindsight)"
|
||||
else
|
||||
echo " hindsight CLI not in PATH"
|
||||
if [ -f ~/.local/bin/hindsight ]; then
|
||||
echo " Found at ~/.local/bin/hindsight"
|
||||
else
|
||||
echo " Not found at ~/.local/bin/hindsight - attempting installation..."
|
||||
# Try to install the CLI
|
||||
if curl -fsSL https://hindsight.vectorize.io/get-cli | bash; then
|
||||
echo " CLI installation completed"
|
||||
if [ -f ~/.local/bin/hindsight ]; then
|
||||
echo " CLI now available at ~/.local/bin/hindsight"
|
||||
export PATH="$HOME/.local/bin:$PATH"
|
||||
else
|
||||
echo " WARNING: CLI still not found after installation"
|
||||
ls -la ~/.local/bin/ 2>/dev/null || echo " ~/.local/bin does not exist"
|
||||
fi
|
||||
else
|
||||
echo " WARNING: CLI installation failed"
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
||||
# Show environment info for debugging
|
||||
echo ""
|
||||
echo "Environment:"
|
||||
echo " HINDSIGHT_EMBED_LLM_PROVIDER: ${HINDSIGHT_EMBED_LLM_PROVIDER:-not set}"
|
||||
echo " HINDSIGHT_EMBED_LLM_MODEL: ${HINDSIGHT_EMBED_LLM_MODEL:-not set}"
|
||||
echo " HINDSIGHT_EMBED_LLM_API_KEY: ${HINDSIGHT_EMBED_LLM_API_KEY:+set (hidden)}"
|
||||
echo " PATH includes ~/.local/bin: $(echo $PATH | grep -q "$HOME/.local/bin" && echo yes || echo no)"
|
||||
|
||||
# Final check that CLI is available before proceeding
|
||||
echo ""
|
||||
echo "Final CLI check before tests..."
|
||||
CLI_PATH=""
|
||||
if [ -f ~/.local/bin/hindsight ]; then
|
||||
CLI_PATH="$HOME/.local/bin/hindsight"
|
||||
elif command -v hindsight &> /dev/null; then
|
||||
CLI_PATH="$(which hindsight)"
|
||||
fi
|
||||
|
||||
if [ -n "$CLI_PATH" ]; then
|
||||
echo " CLI found at: $CLI_PATH"
|
||||
echo " CLI version: $($CLI_PATH --version 2>&1 || echo 'unknown')"
|
||||
else
|
||||
echo " ERROR: hindsight CLI not found. Tests cannot proceed."
|
||||
echo " The hindsight-embed package forwards commands to the hindsight CLI."
|
||||
echo " Please ensure the CLI is installed or check the get-cli installer output above."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Stop any existing daemon
|
||||
echo ""
|
||||
echo "Stopping any existing daemon..."
|
||||
uv run --project "$SCRIPT_DIR" hindsight-embed daemon stop 2>/dev/null || true
|
||||
sleep 1
|
||||
|
||||
# Test 1: Retain (this should start the daemon)
|
||||
echo ""
|
||||
echo "Test 1: Retaining a memory (first call - daemon will start)..."
|
||||
START_TIME=$(python3 -c "import time; print(time.time())")
|
||||
set +e # Temporarily disable exit on error to capture output
|
||||
OUTPUT=$(uv run --project "$SCRIPT_DIR" hindsight-embed memory retain "$BANK_ID" "The user's favorite color is blue" 2>&1)
|
||||
EXIT_CODE=$?
|
||||
set -e
|
||||
END_TIME=$(python3 -c "import time; print(time.time())")
|
||||
DURATION=$(python3 -c "print(f'{$END_TIME - $START_TIME:.2f}')")
|
||||
echo "$OUTPUT"
|
||||
echo "Duration: ${DURATION}s"
|
||||
echo "Exit code: $EXIT_CODE"
|
||||
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
echo "FAIL: Command exited with code $EXIT_CODE"
|
||||
echo ""
|
||||
echo "Checking daemon logs..."
|
||||
if [ -f ~/.hindsight/daemon.log ]; then
|
||||
echo "=== daemon.log ==="
|
||||
tail -50 ~/.hindsight/daemon.log
|
||||
else
|
||||
echo "No daemon.log found"
|
||||
fi
|
||||
if [ -f ~/.hindsight/daemon.stderr ]; then
|
||||
echo ""
|
||||
echo "=== daemon.stderr ==="
|
||||
cat ~/.hindsight/daemon.stderr
|
||||
else
|
||||
echo "No daemon.stderr found"
|
||||
fi
|
||||
echo ""
|
||||
echo "Checking for hindsight-api..."
|
||||
which hindsight-api 2>/dev/null || echo "hindsight-api not in PATH"
|
||||
which uvx 2>/dev/null || echo "uvx not in PATH"
|
||||
which uv 2>/dev/null || echo "uv not in PATH"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ! echo "$OUTPUT" | grep -qi "retained"; then
|
||||
echo "FAIL: Expected 'retained' in output"
|
||||
exit 1
|
||||
fi
|
||||
echo "PASS: Memory retained successfully"
|
||||
|
||||
# Test 2: Recall (daemon already running - should be faster)
|
||||
echo ""
|
||||
echo "Test 2: Recalling memories (daemon already running)..."
|
||||
START_TIME=$(python3 -c "import time; print(time.time())")
|
||||
set +e
|
||||
OUTPUT=$(uv run --project "$SCRIPT_DIR" hindsight-embed memory recall "$BANK_ID" "What is the user's favorite color?" 2>&1)
|
||||
EXIT_CODE=$?
|
||||
set -e
|
||||
END_TIME=$(python3 -c "import time; print(time.time())")
|
||||
DURATION=$(python3 -c "print(f'{$END_TIME - $START_TIME:.2f}')")
|
||||
echo "$OUTPUT"
|
||||
echo "Duration: ${DURATION}s"
|
||||
echo "Exit code: $EXIT_CODE"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
echo "FAIL: Command exited with code $EXIT_CODE"
|
||||
exit 1
|
||||
fi
|
||||
if ! echo "$OUTPUT" | grep -qi "blue"; then
|
||||
echo "FAIL: Expected 'blue' in recall output"
|
||||
exit 1
|
||||
fi
|
||||
echo "PASS: Memory recalled successfully"
|
||||
|
||||
# Test 3: Retain with context (daemon should still be running)
|
||||
echo ""
|
||||
echo "Test 3: Retaining memory with context..."
|
||||
START_TIME=$(python3 -c "import time; print(time.time())")
|
||||
set +e
|
||||
OUTPUT=$(uv run --project "$SCRIPT_DIR" hindsight-embed memory retain "$BANK_ID" "User prefers Python over JavaScript" --context work 2>&1)
|
||||
EXIT_CODE=$?
|
||||
set -e
|
||||
END_TIME=$(python3 -c "import time; print(time.time())")
|
||||
DURATION=$(python3 -c "print(f'{$END_TIME - $START_TIME:.2f}')")
|
||||
echo "$OUTPUT"
|
||||
echo "Duration: ${DURATION}s"
|
||||
echo "Exit code: $EXIT_CODE"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
echo "FAIL: Command exited with code $EXIT_CODE"
|
||||
exit 1
|
||||
fi
|
||||
if ! echo "$OUTPUT" | grep -qi "retained"; then
|
||||
echo "FAIL: Expected 'retained' in output"
|
||||
exit 1
|
||||
fi
|
||||
echo "PASS: Memory with context retained successfully"
|
||||
|
||||
# Test 4: Recall with JSON output
|
||||
echo ""
|
||||
echo "Test 4: Recalling with JSON output..."
|
||||
START_TIME=$(python3 -c "import time; print(time.time())")
|
||||
set +e
|
||||
JSON_OUTPUT=$(uv run --project "$SCRIPT_DIR" hindsight-embed memory recall "$BANK_ID" "programming preferences" -o json 2>&1)
|
||||
EXIT_CODE=$?
|
||||
set -e
|
||||
END_TIME=$(python3 -c "import time; print(time.time())")
|
||||
DURATION=$(python3 -c "print(f'{$END_TIME - $START_TIME:.2f}')")
|
||||
echo "$JSON_OUTPUT"
|
||||
echo "Duration: ${DURATION}s"
|
||||
echo "Exit code: $EXIT_CODE"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
echo "FAIL: Command exited with code $EXIT_CODE"
|
||||
exit 1
|
||||
fi
|
||||
if ! echo "$JSON_OUTPUT" | grep -qi "python"; then
|
||||
echo "FAIL: Expected 'Python' in recall output"
|
||||
exit 1
|
||||
fi
|
||||
if ! echo "$JSON_OUTPUT" | python3 -c "import sys, json; json.load(sys.stdin)" 2>/dev/null; then
|
||||
echo "FAIL: Expected valid JSON output"
|
||||
exit 1
|
||||
fi
|
||||
echo "PASS: Memory recalled with JSON format successfully"
|
||||
|
||||
# Test 5: Check daemon is running
|
||||
echo ""
|
||||
echo "Test 5: Verifying daemon is running..."
|
||||
if curl -s http://127.0.0.1:8889/health | grep -q "healthy"; then
|
||||
echo "PASS: Daemon is running and healthy"
|
||||
else
|
||||
echo "FAIL: Daemon is not running"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Test 6: Daemon status command
|
||||
echo ""
|
||||
echo "Test 6: Testing daemon status command..."
|
||||
STATUS_OUTPUT=$(uv run --project "$SCRIPT_DIR" hindsight-embed daemon status 2>&1)
|
||||
echo "$STATUS_OUTPUT"
|
||||
if ! echo "$STATUS_OUTPUT" | grep -qi "running"; then
|
||||
echo "FAIL: Expected 'running' in daemon status output"
|
||||
exit 1
|
||||
fi
|
||||
echo "PASS: Daemon status command works"
|
||||
|
||||
# Cleanup: Stop daemon
|
||||
echo ""
|
||||
echo "Stopping daemon..."
|
||||
uv run --project "$SCRIPT_DIR" hindsight-embed daemon stop 2>/dev/null || true
|
||||
|
||||
echo ""
|
||||
echo "=== All tests passed! ==="
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "hindsight-litellm"
|
||||
version = "0.1.11"
|
||||
version = "0.1.12"
|
||||
description = "Universal LLM memory integration via LiteLLM - works with 100+ providers"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.1.11"
|
||||
version = "0.1.12"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
|
||||
Generated
+865
-5467
File diff suppressed because it is too large
Load Diff
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.uv.workspace]
|
||||
members = ["hindsight", "hindsight-api", "hindsight-dev", "hindsight-mcp-server", "hindsight-openai", "hindsight-langmem", "hindsight-clients/python"]
|
||||
members = ["hindsight", "hindsight-api", "hindsight-dev", "hindsight-mcp-server", "hindsight-openai", "hindsight-langmem", "hindsight-clients/python", "hindsight-embed"]
|
||||
|
||||
[tool.uv]
|
||||
dev-dependencies = []
|
||||
|
||||
@@ -1,22 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Lint Node/TypeScript code with ESLint and Prettier
|
||||
|
||||
REPO_ROOT="$(git rev-parse --show-toplevel)"
|
||||
cd "$REPO_ROOT/hindsight-control-plane"
|
||||
|
||||
echo " Linting Node/TS with ESLint and Prettier..."
|
||||
|
||||
# Capture output and only show on failure
|
||||
OUTPUT=$(npx eslint --fix "src/**/*.{ts,tsx}" 2>&1)
|
||||
if [ $? -ne 0 ]; then
|
||||
echo "$OUTPUT"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
OUTPUT=$(npx prettier --write "src/**/*.{ts,tsx}" 2>&1)
|
||||
if [ $? -ne 0 ]; then
|
||||
echo "$OUTPUT"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo " Node lint complete"
|
||||
@@ -1,22 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Lint Python code with Ruff
|
||||
|
||||
REPO_ROOT="$(git rev-parse --show-toplevel)"
|
||||
cd "$REPO_ROOT/hindsight-api"
|
||||
|
||||
echo " Linting Python with Ruff..."
|
||||
|
||||
# Capture output and only show on failure
|
||||
OUTPUT=$(uv run ruff check --fix . 2>&1)
|
||||
if [ $? -ne 0 ]; then
|
||||
echo "$OUTPUT"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
OUTPUT=$(uv run ruff format . 2>&1)
|
||||
if [ $? -ne 0 ]; then
|
||||
echo "$OUTPUT"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo " Python lint complete"
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user