Compare commits

...
Author SHA1 Message Date
Nicolò Boschi 3f978b446d fix: remove unused PluginState import from tools.ts 2026-04-07 10:05:36 +02:00
Nicolò Boschi a4a9d32480 fix: review fixes for opencode integration
- Rename CI job from build-opencode-integration to test-opencode-integration
  to match naming convention for integrations that run tests
- Fix tsconfig module resolution to Node16 (consistent with other integrations)
- Extract shared makeConfig test helper to avoid duplication across 3 test files
2026-04-07 09:49:19 +02:00
DK09876andClaude Opus 4.6 8b4c7eeb8e fix: recall retry semantics and README bank scoping clarity
1. recallForContext now returns { context, ok } to distinguish
   "no results" (ok=true) from "API error" (ok=false). System
   transform consumes the session on ok=true even with 0 results,
   so empty banks don't cause repeated queries. Only transient API
   failures preserve retry.

2. README clarifies that channel/user bank dimensions are process-
   scoped (set via env vars before launch), not per-session dynamic
   within a running OpenCode process.

89 tests pass.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-04-07 09:45:38 +02:00
DK09876andClaude Opus 4.6 311b96a192 fix: docs/tools findings from second review round
1. Remove "session" from supported dynamic bank fields in docs —
   the implementation can't vary bank ID per session since it's
   derived once at plugin startup.

2. Explicit tools (retain, reflect) now call ensureBankMission()
   before API calls, so bankMission/retainMission are applied even
   when the agent uses tools exclusively without triggering hooks.

3. Added tests for mission setup via tools path.

88 tests pass.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-04-07 09:45:38 +02:00
DK09876andClaude Opus 4.6 fdc48b1544 fix: address review findings for opencode integration
1. Pre-compaction retain now uses shared retainSession() helper,
   respecting retainMode, documentId, and session_id metadata
   consistently with idle-retain (was bypassing retention policy).

2. System transform recall is only consumed after successful injection.
   If Hindsight is briefly unavailable, the plugin retries on the next
   LLM call instead of permanently skipping recall for the session.

3. Config validation for retainMode and recallBudget — typos like
   "full_session" or "maximum" now log a warning and fall back to
   the default instead of silently changing retention semantics.

85 tests (6 new covering compaction documentId, recall retry, and
config validation).

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-04-07 09:45:38 +02:00
DK09876andClaude Opus 4.6 ebb05cf9ca feat: add OpenCode persistent memory plugin
Add hindsight-opencode integration with:
- Three custom tools: hindsight_retain, hindsight_recall, hindsight_reflect
- Auto-retain on session.idle with document_id deduplication
- Memory injection on session start via system transform hook
- Memory preservation during context window compaction
- Sliding window retain with retainOverlapTurns support
- 4-level config hierarchy (defaults, user file, plugin options, env vars)
- Dynamic bank ID derivation (agent, project, channel, user dimensions)
- CI job, release script entry, docs page

79 tests across 6 test files.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-04-07 09:45:38 +02:00
Nicolò Boschi 66cbdda3cb test: add regression tests for #874 and #894 (#901)
Add tests for None event_date in fact extraction (AttributeError fix)
and for _register_profile skipping .env overwrite with short config keys.
2026-04-07 09:43:25 +02:00
Nicolò Boschi cf4bd598b4 fix: make bank_id metric label opt-in to prevent OTel memory leak (#898)
* fix: make bank_id metric label opt-in to prevent OTel memory leak

bank_id as an OTel metric attribute creates unbounded histogram growth
since each unique bank_id produces never-evicted time series. Default
to excluding it; opt in with HINDSIGHT_API_METRICS_INCLUDE_BANK_ID=true
for deployments with few banks.

Closes #850

* refactor: use config.py for metrics_include_bank_id setting

Move HINDSIGHT_API_METRICS_INCLUDE_BANK_ID from direct os.getenv in
metrics.py to the standard HindsightConfig path. Add configuration
documentation.
2026-04-07 09:42:59 +02:00
Nicolò Boschi 443c94c827 fix(mcp): auto-coerce string-encoded JSON in tool arguments (#849) (#899)
LLM agents frequently serialize list/dict tool arguments as JSON strings
instead of native types (e.g., tags='["a","b"]' instead of tags=["a","b"]),
causing Pydantic validation failures. This extends _make_tools_tolerant to
detect array/object parameters from the JSON Schema and auto-coerce string
values via json.loads before validation.

Also fixes _make_tools_tolerant compatibility with FastMCP 3.x by adding
a _get_mcp_tools helper that supports both 2.x and 3.x internal APIs.
2026-04-07 09:33:12 +02:00
Abdulkadirklc 26794aab09 feat(recall): add proof_count boost to combined scoring (#821)
* feat(recall): add proof_count boost to combined scoring

Observations with more supporting evidence now rank slightly higher
in recall results. proof_count is threaded through the retrieval
pipeline and applied as a multiplicative boost in reranking:

- types.py: add proof_count field to RetrievalResult
- retrieval.py: include proof_count in SELECT columns
- reranking.py: add log1p-normalized proof_count boost (alpha=0.1)

The boost uses the same multiplicative pattern as recency and temporal
signals. proof_count=1 is neutral, proof_count=50 gives ~+5% boost.
Non-observation fact types are unaffected (neutral 0.5).

* fix(retrieval): Apply proof_count boost to graph and temporal retrieval, normalize scaling

* fix(retrieval): correct proof_norm math to zero-center at count 1

* fix(retrieval): Apply proof_count boost to link_expansion retrieval

* fix: remove BFS zombie, clamp proof_norm to [0,1], fix test comment (log1p->math.log)
2026-04-07 09:32:44 +02:00
Nicolò Boschi 7863ffeb49 fix(paperclip): address review fixes for paperclip integration (#900)
- Add CI job for paperclip integration tests with change detection
- Add paperclip to valid release integrations
- Validate hindsightApiUrl is set in loadConfig()
- Log warnings on recall/retain failures instead of silently swallowing
- Remove hardcoded timeout from reflect call
- Fix tsconfig module resolution to Node16
- Update tests to pass required hindsightApiUrl
2026-04-07 09:32:24 +02:00
Octopus 9e2890ba81 fix(embed): skip profile .env overwrite when config has no HINDSIGHT_API_* keys (#896)
When the daemon is already running, ensure_running() calls _register_profile()
with a config dict using short keys (llm_api_key, llm_provider, etc.) that do
not match the HINDSIGHT_API_* prefix filter. This caused api_config to always
be empty, and create_profile() would overwrite the existing .env with an empty
file on every CLI command.

Add an early return guard so _register_profile() skips the create_profile()
call when api_config is empty, preserving any existing profile configuration.

Fixes #894
2026-04-07 09:28:53 +02:00
Chris Bartholomew e0e65c44f6 fix(query_analyzer): handle dateparser internal crashes gracefully (#893)
DateparserQueryAnalyzer.analyze() called dateparser.search.search_dates()
without any error handling, so internal bugs in the third-party library
propagated all the way up the search/consolidation pipeline and failed
the calling task.

Observed traceback:

  File ".../engine/query_analyzer.py", line 140, in analyze
    results = self._search_dates(query, settings=settings)
  File ".../dateparser/search/search.py", line 294, in search_dates
    "Dates": self.search.search_parse(...)
  File ".../dateparser/search/search.py", line 168, in search_parse
    translated, original = self.search(shortname, text, settings)
  File ".../dateparser/languages/locale.py", line 224, in translate_search
    [original_tokens[i], original_tokens[i + 1]],
  IndexError: list index out of range

Wrap the call in a try/except so any parser failure is treated as
"no temporal constraint found" — the caller can then fall back to
non-temporal retrieval instead of erroring out the whole task. The
failure is logged at WARNING level so we still notice it.

Add a regression test that monkey-patches _search_dates to raise an
IndexError and asserts the analyzer returns an empty constraint and
emits a warning log.
2026-04-07 09:26:27 +02:00
dependabot[bot]anddependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> 6881f63781 chore(deps): bump actions/github-script from 7 to 8 (#879)
Bumps [actions/github-script](https://github.com/actions/github-script) from 7 to 8.
- [Release notes](https://github.com/actions/github-script/releases)
- [Commits](https://github.com/actions/github-script/compare/v7...v8)

---
updated-dependencies:
- dependency-name: actions/github-script
  dependency-version: '8'
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-04-07 09:25:55 +02:00
Daniyar 6cb309f72b Fix AttributeError when event_date is None in fact_extraction (#875)
* Fix AttributeError when event_date is None in fact_extraction

`_extract_facts_from_chunk` crashes with `'NoneType' object has no
attribute 'isoformat'` when retaining documents without a timestamp.

Two locations fixed:
- Line 1058: debug log called `event_date.isoformat()` without a None
  check
- Line 921: `parse_datetime_flexible()` can return None, so re-check
  before calling `.strftime()` / `.isoformat()`

Fixes #874

* Revert unnecessary None guard on line 921

The original `if event_date is not None:` already guards that block.
Only line 1058 needed the fix.
2026-04-07 09:22:07 +02:00
shun yiandyishun.eason f9fe6953a3 fix: Windows compatibility for hindsight-embed (#867)
- Add cross-platform file locking support
- Use fcntl on Unix-like systems, msvcrt on Windows
- Add detailed documentation explaining why we don't use external libraries
- Fixes issue where module couldn't be imported on Windows due to missing fcntl

Co-authored-by: yishun.eason <[email protected]>
2026-04-07 09:15:59 +02:00
Volodymyr Prypeshniuk 07de798c3b feat(google): add support for google embeddings and reranker (#863)
* Add support for google embeddings gemini/vertex and google reranker via vertex search api

* Add reference docs
2026-04-07 09:15:31 +02:00
Byeonghoon YooandClaude Opus 4.6 cefa75545a feat(helm): add persistent volume for local model cache (#861)
* feat(helm): add persistent volume for local model cache

When using local reranker (e.g., BAAI/bge-reranker-v2-m3) or local
embedding models, the models are downloaded to /home/hindsight/.cache
on every pod restart, causing slow startup and unnecessary bandwidth.

Add optional persistent volume support:
- api: PVC mounted at /home/hindsight/.cache
- worker: volumeClaimTemplate (StatefulSet) at same path

Disabled by default. Enable via:
  api.persistence.modelCache.enabled: true
  worker.persistence.modelCache.enabled: true

Closes #860

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* feat(helm): add extraVolumes and extraVolumeMounts for api and worker

Allow users to mount arbitrary volumes (configMaps, secrets, emptyDir,
etc.) into api and worker pods via values, following common helm chart
library conventions.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
2026-04-07 09:14:09 +02:00
Octopus cd99eef4c5 fix: use max_tokens for OpenAI-compatible endpoints with custom base URL (#858)
Mistral (and several other providers) reject 'max_completion_tokens' with a 422
because they haven't adopted the newer OpenAI parameter name. When the openai
provider is configured with a custom base_url (e.g. Mistral, Together AI),
fall back to the widely-supported 'max_tokens' parameter.

Native OpenAI (no custom base_url) and Groq still use 'max_completion_tokens'.

Fixes #852
2026-04-07 09:13:08 +02:00
Ben cd4b3e96e2 blog: Persistent Memory for AutoGen Agents with Hindsight (#883)
* Add AutoGen persistent memory blog post
2026-04-06 14:59:43 -04:00
Ben e02e7ad3d4 blog: Hindsight is now a native memory provider in Hermes Agent (#882)
* Add Hermes native memory provider blog post
2026-04-06 10:55:48 -04:00
Ben 98fee1e380 docs(hermes): update integration docs for plugin overhaul (hermes-agent#5094) (#881)
* docs(hermes): update integration docs for hermes-agent plugin overhaul
2026-04-06 10:54:59 -04:00
Nicolò Boschi 906b740dd7 fix(docs): add missing SEO frontmatter to paperclip integration 2026-04-02 17:37:02 +02:00
Nicolò Boschi 7990381f6a fix(ci): resolve all CI failures (#847)
* fix(ci): resolve all CI failures — unversioned integrations, test retries

- Move integration docs to separate unversioned docs plugin (docs-integrations/)
  so new integrations don't need to be duplicated across versioned_docs
- Remove integration pages from versioned_docs (v0.3, v0.4) — sidebar
  entries now use links instead of doc refs
- Add missing title/description SEO frontmatter to autogen.md
- Add retry logic (2 attempts) to test-doc-examples.sh for transient
  LLM timeouts
- Add pytest-rerunfailures to test-api with --reruns 2 for flaky
  Gemini-dependent integration tests

* ci: retrigger

* fix: graph entity inheritance, SyncTaskBackend error propagation, fact_type test regressions

- Fix observation entity inheritance in get_graph_data: the unit_entities
  query only fetched entities for visible observation IDs, not their source
  memory IDs, so the inheritance loop always found an empty entity_map
- Remove error swallowing in SyncTaskBackend._execute_task so test failures
  surface instead of being silently logged
- Wrap remaining consolidation submission call sites with try/except since
  consolidation is non-critical for those operations
- Fix test_sync_backend test to expect errors to propagate
- Remove fact_type=["world"] filter from test_document_upsert_behavior and
  test_mentioned_at_from_context_string (same PR #848 regression)
- Remove flaky marker from consolidation test (now deterministic)
2026-04-02 17:17:42 +02:00
Ben 045e8910d1 Blog: Hindsight Is #1 on BEAM — the Benchmark That Tests Memory at 10M Tokens (#851)
* Add BEAM SOTA blog post
2026-04-02 11:02:16 -04:00
Ben 81441ee9af feat(paperclip): add hindsight-paperclip TypeScript integration (#773)
* feat(paperclip): add hindsight-paperclip TypeScript integration

Adds long-term memory for Paperclip AI agents via a lightweight
TypeScript/Node.js npm package with no runtime dependencies.

- recall() / retain() functions for heartbeat lifecycle hooks
- createMemoryMiddleware() for Express HTTP adapter agents
- Bank ID strategy: paperclip::{companyId}::{agentId} (configurable)
- Skill file for agents to call Hindsight REST API directly
- 27 unit tests covering bank derivation, recall, and retain
- Docs page at sdks/integrations/paperclip

* Remove skills file from paperclip integration

* Rename package to @vectorize-io/hindsight-paperclip
2026-04-02 14:26:45 +02:00
Nicolò Boschi 30a319a6ab feat: bank template import/export with Template Hub (#819)
* feat(api): add bank template import/export endpoints

Add POST /banks/{bank_id}/import and GET /banks/{bank_id}/export
endpoints for declarative bank setup via JSON manifests.

A template manifest (version 1) can include bank config overrides
and mental model definitions. Import creates or updates mental
models matched by id, applies config as per-bank overrides, and
returns async operation IDs for content generation.

Export dumps a bank's explicit overrides and mental models as a
manifest that can be re-imported into another bank.

Includes control plane UI: bank creation dialog now accepts an
optional template JSON to pre-configure the bank on creation.

* docs: add Template Gallery page and bank templates reference

- Template Gallery (/templates) with search, category filter, manifest
  preview modal with copy-to-clipboard
- 5 starter templates: Customer Support, Research Assistant, Personal
  Journal, Code Review Buddy, Meeting Notes
- Bank Templates API reference doc (developer/api/bank-templates)
- Sidebar entry under API section

* docs: add Template Gallery links to navbar and sidebar

- Top navbar: "Templates" link between Integrations and Changelog
- Sidebar: "Template Gallery" in Resources section

* fix(docs): remove emoji icons, autofocus search, fix placeholder in template gallery

* docs: rename to Bank Templates, move to Resources sidebar only

* docs: add Bank Templates to Resources navbar dropdown

* feat(api): add directives to bank template import/export

- Add BankTemplateDirective model with name, content, priority, is_active, tags
- Import creates/updates directives matched by name
- Export includes all directives (active and inactive)
- Validation: duplicate names rejected, empty name/content caught
- Tests: 24 tests covering directives create/update, existing vs new
  bank import, validation, export with directives, full round-trip

* docs: add directives to bank templates docs and sample templates

* feat(api): add JSON Schema endpoint for bank template validation

- GET /v1/default/bank-template-schema returns the JSON Schema
  auto-generated from the Pydantic BankTemplateManifest model
- Static schema file at docs/static/bank-template-schema.json
- Docs updated with schema endpoint, static file link, and
  validation examples (Python jsonschema, Node ajv-cli)

* feat(api): live schema validation on import, fix schema endpoint path

- Move schema endpoint to /v1/bank-template-schema (system-level, not per-bank)
- Import endpoint now accepts raw JSON and validates with Pydantic manually,
  returning clean 400 errors instead of raw 422s for all validation failures
- All validation (schema + semantic) returns consistent 400 with detailed messages

* docs: add interactive JSON Schema viewer to Bank Templates page

Renders the Pydantic-generated schema as a collapsible property tree
with types, required badges, defaults, and descriptions. The schema
is imported from the static bank-template-schema.json file.

* ui: add template toggle switch and browse link to bank creation dialog

- Replace always-visible textarea with a switch toggle ("Import from template")
- Textarea only shows when switch is on, keeping the dialog clean by default
- Add "Browse templates" link pointing to hindsight.vectorize.io/templates
- Reset template state when switch is toggled off or dialog is cancelled

* ui: add empty state with Add Document CTA to data view

When a bank has 0 memories, the data view (all tabs: constellation,
graph, table, timeline) shows a centered empty state with a CTA
button that opens the Add Document dialog.

* docs: replace templates with Conversation and Coding Agent

Remove generic placeholder templates. Add two practical templates
based on actual integration patterns:

- Conversation: for chat agents (LiteLLM, LangGraph, Pydantic AI,
  Vercel AI SDK). Tracks user preferences, open threads.
- Coding Agent: for Claude Code/Codex. Tracks technical decisions,
  project context, developer preferences. High literalism.

* docs: rename gallery to Bank Templates Hub, keep API doc as Bank Templates

* docs: register layout-template and file-json icons in navbar and sidebar

* docs: register layout-template icon in DefaultNavbarItem for dropdown items

* docs: show integration icons on template cards

Templates now have an optional `integrations` field referencing
integration IDs from integrations.json. Icons are resolved at render
time and shown in the card header next to the category badge.

* docs: add Personal Assistant template for OpenClaw, Hermes, NemoClaw

* feat: add Export Template to bank actions + map all integrations to templates

- Add "Export Template" to the bank Actions dropdown — exports config,
  mental models, and directives as JSON, copies to clipboard
- Add export API route and client method
- Map remaining integrations to templates: CrewAI, AG2, Agno, Strands,
  LlamaIndex, local-mcp, skills → Conversation; hindclaw → Personal Assistant

* feat: add --template flag to LoCoMo benchmark + remove schema from Hub

- LoCoMo benchmark accepts --template <path> to apply a bank template
  manifest (config, mental models, directives) before ingestion
- Template is applied per-bank in both single-phase and two-phase modes
- BenchmarkRunner.apply_template() reuses the same engine methods as
  the /import API endpoint
- Remove Manifest Schema section from Bank Templates Hub page
  (schema stays in the API reference doc)

* refactor: remove description field from bank template manifest

* docs: remove tags, fact_types, and directives from starter templates

* docs: remove reflect_mission and disposition fields from starter templates

* build: validate template manifests against JSON Schema during docs build

* cleanup: remove unused JsonSchemaViewer component

* docs: remove retain_extraction_mode from starter templates

* ui: enable word wrap in template manifest preview

* docs: add link to Bank Templates reference doc from Hub page

* docs: convert bank templates doc to mdx with multi-language code snippets

- Convert bank-templates.md to .mdx with Tabs/CodeSnippet components
- Add example files: bank-templates.py, .mjs, .sh, .go with doc markers
- Examples cover import, dry-run, export, round-trip, and schema
- Regenerate OpenAPI spec and all client SDKs (Python, TS, Rust, Go)

* fix: migration revision collision + use typed models in benchmark template

- Rename merge migration d6e7f8a9b0c1 -> d6e7f8a9b0c2 to resolve
  revision ID collision with case_insensitive_entities_trgm_index
- Update a4b5c6d7e8f9 down_revision to point to the renamed migration
- Fix f-string lint in case_insensitive migration
- BenchmarkRunner.apply_template() now validates manifest through
  BankTemplateManifest Pydantic model instead of raw dict access
- Remove redundant inline imports (json, Path already at module top)

* fix(docs): add missing Go tab to dry-run code snippet

* ci: retrigger

* fix: sync skills openapi.json + fix bankId null type error in export

- Copy updated openapi.json to skills/hindsight-docs/references/
- Add null guard for bankId in Export Template onClick handler

* fix: sync generated files (memory_engine formatting, docs skill references)

* cleanup: remove obsolete migration collision workaround
2026-04-02 12:21:53 +02:00
Nicolò Boschi 9cfdd464a9 fix(retain): preserve normalized experience fact types (#848)
* fix(retain): preserve normalized experience fact types and remove deprecated opinion type

The ExtractedFactType conversion was re-checking for raw "assistant" fact_type
after the parsing layer had already normalized it to "experience". Since
fact_from_llm.fact_type was always "experience" (never "assistant"), the ternary
always fell through to "world", silently losing experience classification.

Also removes the deprecated "opinion" fact type from internal extraction models,
database constraints/indexes (via migration), and dead code paths. The public API
surface (descriptions, response models, backwards-compat filter) is unchanged.

* refactor(retain): drop unused confidence_score column

The confidence_score column was only ever non-null for opinion facts
(which are now removed). It was always written as NULL and never read
back from the database. Remove it from:
- DB model and migration (DROP COLUMN)
- INSERT queries in fact_storage.py
- retain_async/retain_batch_async parameters
- RetainContext/RetainResult extension models
- RetainBatch dataclass
2026-04-02 12:20:37 +02:00
Nicolò Boschi 8d1bfbbd2b feat: add detail parameter to list/get mental models (#846)
* feat: add detail parameter to list/get mental models (#825)

Add a `detail` query parameter (metadata|content|full) to both list and get
mental model endpoints (HTTP + MCP) to control response size. This reduces
payload for agent boot flows and MCP clients where context budget is limited.

Closes #825

* fix: update Rust CLI for optional mental model fields

The generated Rust client now has content/source_query as Option<String>
after the detail parameter was added. Update CLI code to handle optionals.
2026-04-02 11:52:45 +02:00
Nicolò Boschi 7d6c570a3a fix(embed): clear stale daemon on port before starting (#843)
* fix(embed): clear stale daemon on port before starting new one (#843)

When `uvx hindsight-embed@latest` resolves to a new version, the old
daemon may still be bound to the port, causing EADDRINUSE. Before
starting a daemon, check if the port is occupied, verify it's a
hindsight process via /health, and SIGTERM it if so.

* chore: remove unused signal import from test

* refactor: use cross-platform port check instead of lsof-only

Use socket for port check (works on all platforms), extract PID lookup
into a helper with Windows (netstat) and Unix (lsof) paths, and
extract kill logic into a testable static method.

* refactor: reuse cross-platform helpers in stop() and stop_ui()
2026-04-02 10:57:28 +02:00
Nicolò Boschi 26a64cc00e fix(api): clear memories endpoint no longer deletes the bank profile (#837)
DELETE /v1/default/banks/{id}/memories and the MCP clear_memories tool
were calling delete_bank() without distinguishing from the actual delete-bank
endpoint. When no fact_type filter was provided, the bank row itself was
deleted along with its memories.

Add a delete_bank_profile parameter to delete_bank() (default True) and
pass False from all clear-memories callers so the bank profile, disposition,
and background are preserved.
2026-04-01 18:34:39 +02:00
087545cc1b feat(openclaw): JSONL-backed retain queue for external API resilience (#740)
When the external Hindsight API is unreachable, retain requests are
buffered as JSON lines in a local file and automatically flushed once
connectivity is restored. Queue survives process restarts.

- Only active in external API mode (local daemon handles its own persistence)
- Zero dependencies — uses only Node built-ins (fs, crypto)
- Bulk removal via removeMany() for O(1) file rewrites during flush
- Cached item count so size() is O(1)
- Configurable: retainQueuePath, retainQueueMaxAgeMs (-1 = forever),
  retainQueueFlushIntervalMs (default 60s)
- Flushes on successful retain and on a periodic timer
- All logging routed through structured logger (api.logger)

Co-authored-by: billy <[email protected]>
Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
Co-authored-by: Antoine Khater <[email protected]>
2026-04-01 18:06:57 +02:00
Nicolò Boschi 7415ebff7c fix: resolve 25 test regressions from streaming retain pipeline (#722) (#836)
The 3-phase retain pipeline (914ba796) introduced several regressions:

1. **Per-content tags lost** — streaming pipeline used `contents[0].tags`
   for ALL chunks, breaking tag-based visibility. Fixed by tracking
   chunk-to-content mapping so each chunk uses its source content's tags.

2. **Multi-document batches broken** — batches with per-content
   `document_id` values were merged into a single document. Fixed by
   grouping by document_id and processing each group independently.

3. **Migration ID collision** — `d6e7f8a9b0c1` was used by both
   `drop_documents_metadata` and `case_insensitive_entities_trgm_index`.
   Renamed trgm migration to `e8f9a0b1c2d3`, fixed chain, added missing
   schema prefix on DROP INDEX.

4. **Graph entity inheritance** — `get_graph_data` queried entities for
   observation IDs only, but observations inherit entities from source
   memories. Fixed by querying `all_relevant_ids`.

5. **Docstring false positives** — link_utils.py docstrings triggered
   the SQL schema safety test's unqualified table reference check.

6. **Config test count** — `retain_chunk_batch_size` added to
   `_CONFIGURABLE_FIELDS` without updating the test assertion.
2026-04-01 17:59:10 +02:00
Nicolò Boschi 0c97b555ab release(autogen): v0.1.1 2026-04-01 17:51:46 +02:00
Nicolò Boschi 4d117cc274 chore: add autogen to changelog valid integrations list 2026-04-01 17:51:05 +02:00
DK09876andClaude Opus 4.6 a757765ab2 feat: add AutoGen integration for Hindsight (#719)
* feat: add AutoGen integration for Hindsight

Adds hindsight-autogen package providing FunctionTool instances that give
AutoGen agents persistent long-term memory via retain/recall/reflect APIs.

- Package: hindsight_autogen with create_hindsight_tools() factory
- 31 unit tests covering tool creation, invocation, config fallback, errors
- Docs page and integrations.json entry
- README with quickstart and configuration reference

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address PR review feedback for autogen integration

- Fix install instructions to include autogen-agentchat and autogen-ext[openai]
- Add autogen.svg icon to prevent broken image in integrations grid
- Change icon reference from .png to .svg in integrations.json

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: add sleep between retain/recall and close clients in examples

- Add time.sleep(3) between retain and recall to wait for async processing
- Close Hindsight client and model client to avoid unclosed session warnings

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: use asyncio.sleep instead of time.sleep in async examples

time.sleep blocks the event loop; asyncio.sleep yields control.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address PR review feedback - validation, defaults, release script

- Add autogen to VALID_INTEGRATIONS in release-integration.sh
- Remove unused verbose config field
- Extract DEFAULT_BUDGET/MAX_TOKENS/RECALL_TAGS_MATCH constants in config.py,
  import from tools.py to eliminate default duplication
- Add Literal types for budget and recall_tags_match validation
- Modernize type hints to X | None with from __future__ import annotations
- Add [tool.ruff] line-length = 120 to match monorepo convention
- Add py.typed PEP 561 marker
- Re-raise HindsightError before broad Exception catch
- Expand asyncio.sleep(3) comment explaining when/why it's needed
- Remove verbose from docs configure() reference table

Co-Authored-By: Claude Opus 4.6 <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-04-01 17:48:11 +02:00
Derek Bouius 300d089b6a fix: resolve remaining Dependabot security alerts (#833)
* fix: resolve remaining Dependabot security alerts

- Regenerate package-lock.json so npm overrides take effect
  (serialize-javascript, handlebars, path-to-regexp, brace-expansion)
- Upgrade Pygments 2.19.2 -> 2.20.0 in crewai and integration-tests
  lockfiles (fixes ReDoS via GUID matching)

* fix: resolve duplicate alembic revision ID d6e7f8a9b0c1

Two migrations shared the same revision ID: the merge migration
(drop_documents_metadata_column) and the trigram index migration
(case_insensitive_entities_trgm_index). Assign a new unique ID
to the trigram migration and update the downstream dependency.

* chore: fix lint formatting for generated and existing files
2026-04-01 17:22:38 +02:00
Ben 1a1fb35cb0 Add OpenClaw shared memory team setup guide (#788)
* Add blog post: Shared Memory for OpenClaw
2026-04-01 09:33:15 -04:00
Nicolò Boschi 914ba7962c perf: 3-phase retain pipeline — fix deadlocks, cap temporal links, query-time entity expansion (#722)
* perf: 3-phase retain pipeline — fix deadlocks, cap temporal links, query-time entity expansion

Major retain pipeline overhaul addressing deadlocks, write amplification,
and TimeoutErrors. Restructures retain into three phases:

Phase 1: Entity resolution on separate connection (read-heavy)
Phase 2: Core write transaction (atomic) — facts, unit_entities, links
Phase 3: Best-effort display data (error-isolated) — entity viz links, stats

Key changes:
- Sorted bulk INSERT FROM unnest() prevents deadlocks
- Temporal links capped to top-20 per unit (95% reduction)
- Batched semantic ANN via temp table + LATERAL
- Query-time entity expansion via unit_entities self-join
- Entity viz links moved to Phase 3 (post-transaction)
- HINDSIGHT_API_RETAIN_MAX_CONCURRENT config (default: 32)

* fix: increase semantic link top_k from 5 to 20

The hardcoded top_k=5 was artificially limiting semantic link creation.
Link expansion retrieval can consume up to budget (50-200) semantic
neighbors per seed set, but each fact only had 5 outgoing edges — making
the bidirectional graph very sparse.

Increasing to 20 gives retrieval 4x more edges to work with. The ANN
probe cost is unchanged (same HNSW traversal per fact, just returning
more rows). INSERT cost is negligible (~14k rows via bulk INSERT).

Also: all 18 TimeoutErrors in the latest benchmark (beam-1m-u20) were
from Gemini LLM calls, zero from the database — confirming the entity
resolution split eliminated DB timeouts entirely.

* perf: move semantic ANN search to Phase 1 to avoid transaction timeouts

The batched LATERAL ANN query (700 HNSW probes) was the last remaining
source of DB TimeoutErrors — all 29 in the latest benchmark were from
create_semantic_links_batch inside the Phase 2 write transaction.

Split semantic link creation into three phases:
- Phase 1 (separate conn, autocommit): ANN search via temp table + LATERAL.
  No transaction locks, no contention with concurrent writers.
- Phase 2 (write transaction): within-batch numpy similarities (instant) +
  INSERT of both within-batch and Phase 1 ANN results. No DB reads.
- Phase 3 (flush_pending_stats): future hook point for re-checking ANN
  results after commit to catch links missed by concurrent batches.

Also adds 7 unit tests for compute_semantic_links_within_batch covering
empty input, identical/orthogonal embeddings, threshold filtering, top_k
cap, and tuple structure validation.

* fix: handle placeholder unit_ids in Phase 1 ANN search (not valid UUIDs)

* test: add Phase 1 ANN cross-batch test + configurable test PG port

- New test_semantic_links_phase1_ann_cross_batch verifies that the Phase 1
  ANN search with placeholder unit IDs correctly creates cross-batch
  semantic links after remapping to real IDs.
- Test PG port now configurable via HINDSIGHT_TEST_PG_PORT env var
  (default: 5556) to avoid conflicts with running benchmark daemons.

* perf: remove retry_with_backoff from retain, set semaphore default to 4

Remove retry_with_backoff from _run_db_work and _run_delta_db_work:
- Deadlocks are prevented by sorted bulk INSERT (no need for retry)
- Transient timeouts are handled by the worker poller's task-level retry
  (3 attempts, 60s spacing) which is better than rapid internal retries
  that amplify I/O pressure during contention storms

Set HINDSIGHT_API_RETAIN_MAX_CONCURRENT default from 32 to 4:
- The semaphore gates Phase 1 (ANN + entity resolution) + Phase 2 (writes)
- At 4 concurrent, HNSW index I/O is manageable; at 10+ concurrent the
  probes saturate disk and cause cascading timeouts
- LLM extraction still runs at full parallelism (semaphore acquired after)

* fix: add fact_type filter to Phase 1 ANN query to use per-bank HNSW indexes

The LATERAL ANN query was falling back to sequential scan + sort (90ms/probe)
because the per-bank HNSW indexes are partial indexes filtered on fact_type.
Without fact_type in the WHERE clause, PostgreSQL couldn't use them.

Fix: iterate over ('world', 'experience') and run one HNSW-indexed ANN per
type. EXPLAIN shows 8ms/probe (was 90ms) — 11x faster.

700 probes × 8ms × 2 types = ~11s total (was ~63s via seq scan).

* fix: scope temporal links by fact_type + add integration tests

Temporal links now filter by fact_type in the LATERAL query — world facts
only link to world facts, experience to experience. This matches how
retrieval filters results and avoids wasted cross-type link rows.

New integration tests:
- test_semantic_ann_uses_hnsw_index: verifies Phase 1 ANN creates
  cross-batch semantic links (tests fact_type filter + placeholder remap)
- test_temporal_links_scoped_by_fact_type: verifies world facts get
  temporal links to other world facts but NOT to experience facts

* fix: tolerate individual chunk LLM failures instead of failing entire batch

Changed asyncio.gather(*tasks) to asyncio.gather(*tasks, return_exceptions=True)
in both chunk-level and content-level fact extraction. A single chunk timeout
(e.g., Gemini >90s) no longer discards all other successfully extracted facts.

For a 50MB document with 17k chunks, even a 2% chunk failure rate previously
caused 0 completions (entire batch discarded). Now 16,700 facts are extracted
and only the 300 failed chunks are skipped with a warning log.

* fix: batch temporal LATERAL query for large documents (16k+ chunks)

The LATERAL query for temporal links passed all unit_ids at once into
unnest(), causing PostgreSQL timeouts on documents with 16k+ chunks.
Split into batches of 500 units per query to keep each under the
command_timeout.

Also identified: HNSW index creation on shared pg0 instances with
50k+ existing units exceeds the 60s command_timeout. This is a
test infrastructure issue (shared pg0 accumulates data) but also
affects production when creating new banks on large instances.

* feat: streaming chunk batching for large documents (RETAIN_CHUNK_BATCH_SIZE)

Process chunks in mini-batches of N (default 500), committing each batch
to the DB before starting the next. This prevents OOM kills on large
documents (50MB / 17k+ chunks) by keeping only ~500 facts + embeddings
in memory at a time instead of 50k+.

Each mini-batch goes through the full Phase 1 → 2 → 3 pipeline
independently, sharing the same document_id. On recovery (process dies
mid-way), delta retain detects already-committed chunks via content_hash
and skips them — only remaining chunks get re-extracted.

Config: HINDSIGHT_API_RETAIN_CHUNK_BATCH_SIZE (default: 500, 0 to disable)
Per-bank configurable via the hierarchical config system.

Tests:
- test_streaming_chunk_batching_produces_same_facts
- test_streaming_chunk_batching_recovery (delta retain skips committed chunks)
- test_streaming_disabled_for_small_docs

* perf(retain): producer-consumer pipeline + deferred semantic ANN

Replace the sequential streaming loop with a producer-consumer pipeline:
- LLM producer fires concurrent chunk extractions (semaphore-bounded)
- DB consumer drains queue in batches, runs Phase 1+2+3 per batch
- LLM and DB work overlap instead of running sequentially

Defer semantic links to a single final ANN pass after all batches commit:
- Remove within-batch semantic links from Phase 2 (was 2.6s/batch)
- Run parallel ANN (4 connections) after all facts committed
- top_k reduced from 50 to 20 (recall uses at most 20 neighbors)
- Recovery via operation result_metadata checkpoint

Additional optimizations:
- skip_exists_check on temporal/causal link INSERT (saves ~0.5s/batch)
- WHERE EXISTS guard on semantic link INSERT (handles document upsert)
- timeout=300s on ANN queries and bulk INSERT for large banks
- Demote [ANN] debug logs to logger.debug()
- Fix docstring typos (agent_id → bank_id)
- Fix content_index remapping in producer-consumer batches
- Fix delta retain passing contents vs delta_contents

50MB benchmark (mock LLM): 9.2 min (was 23 min) — 2.5x faster.
BEAM 10m benchmark: zero deadlocks, zero DB errors.

* refactor(retain): remove legacy fallback code paths

- Remove process_entities_batch (legacy single-connection entity processing)
- Remove extract_entities_batch_optimized (only caller was the above)
- Remove fallback entity processing inside Phase 2 transaction
- Remove legacy ANN inline fallback in create_semantic_links_batch
- Remove fallback entity_links direct-insert path in Phase 3
- Make resolved_entity_ids/entity_to_unit/unit_to_entity_ids required params

* refactor(retain): replace tuple returns with dataclasses, remove dead code

- Add EntityResolutionResult and Phase1Result dataclasses in types.py
- Replace 4-tuple return from _pre_resolve_phase1 with Phase1Result
- Remove dead `entity_links = []` variables in retain_batch and _try_delta_retain
- Remove unused `confidence_score` parameter from orchestrator.retain_batch
  and _retain_batch_async_internal (was accepted but never used)

* fix(entity-resolver): remove LIKE full-scan fallbacks, use index-only trigram matching

The entity resolution query had LIKE '%...' substring conditions that bypassed
the GIN trigram index, causing full sequential scans of the entities table.
On banks with 10k+ entities, this caused TimeoutErrors (observed in BEAM 10m).

Changes:
- Remove LIKE fallbacks, use trigram % operator only (GIN index-based)
- Lower similarity threshold from 0.3 to 0.15 to catch substring relationships
- Use LOWER() on both sides for case-insensitive matching
- Migration: recreate GIN trigram index on LOWER(canonical_name)

* fix: remove schema prefix from index names in trigram migration

* fix(delta-retain): use same chunk_size as streaming path (3000 vs 120000)

_chunk_contents_for_delta defaulted to chunk_size=120000 while the streaming
path used 3000. On retry, delta re-chunked the document with different
boundaries, found 0 matching chunks, and fell through to full re-extraction.
This wasted all LLM calls on already-committed chunks.

Fix: use the same default (3000) so chunk hashes match on recovery.

* fix(retain): persist generated document_id in operation metadata for retry recovery

When no document_id is provided, retain generates a UUID. On retry, a new UUID
was generated, making delta retain and streaming chunk-hash recovery unable to
find previously committed chunks. All LLM extraction was wasted on retry.

Fix: resolve document_id early in retain_batch (before delta), persist it to
operation result_metadata, and recover it on retry. Both delta and streaming
paths now see the same document_id across attempts.

* refactor(retain): unify into single streaming pipeline, remove non-streaming path

All retains now go through the producer-consumer streaming pipeline,
regardless of document size. Small documents are processed as a single batch.
This eliminates the maintenance burden of two separate code paths.

Also fix document upsert: compare content hash to distinguish recovery
(same content, partially committed) from update (different content, needs
cascade-delete). Previously, existing chunks always triggered recovery mode.

* refactor(retain): remove dead code, replace raw dicts with Phase3Context dataclass

- Remove dead _handle_zero_facts_documents (no callers after path unification)
- Remove unused imports: defaultdict, EntityLink
- Replace raw dict phase3_context with typed Phase3Context dataclass
- Update _build_and_insert_entity_links_phase3 to use typed parameter
2026-04-01 12:52:49 +02:00
Nicolò Boschi 6f173b10a7 fix(consolidation): improve observation quality with structured processing rules (#814)
Rewrite consolidation prompt rules to produce clean, single-facet observations:
- One observation per distinct facet (count, named entity, relationship)
- Match updates by entity/facet, not topic similarity
- No computation — never infer/calculate values not explicitly stated
- Cascade state changes to all affected observations
- Preserve event history (sold, died, moved) — conservative deletes
- Include dates on state changes when available
- Keep observations concise — no cross-facet narrative bloat

Add test_horse_observations.py exercising a realistic sequence of retain
operations (farm with horses being named, sold, dying) and verifying that
observations track history correctly and mental models can synthesize them.
2026-04-01 12:44:52 +02:00
Nicolò Boschi ea834bc7dc breaking: remove BFS and MPFP graph retrieval strategies (#767)
Remove the BFS spreading activation and MPFP (Multi-Path Fact Propagation)
graph retrieval strategies, leaving link_expansion as the sole graph
retrieval algorithm. Rename MPFPTimings to GraphRetrievalTimings and
mpfp_timings field to graph_timings since the timing struct is used by
LinkExpansionRetriever.

Deleted:
- hindsight-api-slim/hindsight_api/engine/search/mpfp_retrieval.py
- hindsight-api-slim/tests/test_mpfp_retrieval.py

Removed config: HINDSIGHT_API_MPFP_TOP_K_NEIGHBORS
2026-04-01 12:44:40 +02:00
Nicolò Boschi 4fd7c5d1f8 fix(db): respect vector extension config in per-bank index migration (#832)
* fix(db): respect vector extension config in per-bank index migration

Migration d5e6f7a8b9c0 hardcoded HNSW when creating per-bank partial
vector indexes, ignoring HINDSIGHT_API_VECTOR_EXTENSION. This caused
banks migrated from pre-v0.4.18 to get HNSW indexes even when
pgvectorscale (DiskANN) or vchord was configured.

- Fix the original migration to read the vector extension config
- Add migration a4b5c6d7e8f9 to detect and recreate mismatched indexes
  (skipped entirely when extension is pgvector, since those are correct)

* chore: regenerate openapi.json for v0.4.22 version bump
2026-04-01 12:22:08 +02:00
Nicolò Boschi 36783df320 feat(control-plane): add Constellation view with Pretext canvas rendering (#831)
Add a new "Constellation" memory visualization as the default view in the
control plane, powered by @chenglou/pretext for DOM-free text layout on canvas.

- Canvas-rendered zoomable/pannable memory map with spatial label deconfliction
- Nodes colored by link-count heat gradient (Hindsight brand teal→cyan→blue)
- Star-like rendering with varied size/opacity based on connectivity
- Hover shows rich tooltip with full memory metadata (text, entities, tags, dates)
- Hover highlights connected nodes and their links, dims the rest
- Click to select and view memory details in the side panel
- Fullscreen mode toggle
- Link type legend and heat gradient legend on the HUD

Also optimizes the graph API endpoint:
- Entity query now filters by visible unit IDs (was doing full table scan)
- Links query caps at 10k edges sorted by weight (was returning 500k+ uncapped)
- Replaced expensive DISTINCT ON with LEAST/GREATEST sort with simple ORDER BY
2026-04-01 11:15:14 +02:00
Derek Bouius ee4510a762 fix(deps): address critical and high severity security vulnerabilities (#827)
* fix(deps): address critical and high severity security vulnerabilities

Bump vulnerable dependencies to patched versions across the monorepo:

Python (critical/high):
- fastmcp >=2.14.0 → >=3.2.0 (SSRF, path traversal, OAuth confused deputy, command injection)
- langchain-core >=1.2.11 → >=1.2.22 (path traversal in legacy load_prompt)

Python (low):
- cryptography >=46.0.5 → >=46.0.6 (incomplete DNS name constraint enforcement)
- pygments: add >=2.20.0 pin (ReDoS via GUID regex)

Node.js:
- serialize-javascript ^7.0.3 → ^7.0.5 (CPU exhaustion DoS)
- handlebars: add >=4.7.9 override (JS injection via AST type confusion)
- path-to-regexp: add >=0.1.13 override (ReDoS via route params)
- brace-expansion: add version range override (process hang/memory exhaustion)

Also adds type: ignore comments for FastMCP 2.x private attribute access that
ty now flags since FastMCP 3.x removed _tool_manager (guarded by try/except
and hasattr at runtime).

Regenerated all lock files across API, integrations, and tests.

* fix(deps): add ajv v8 scoped overrides for schema-utils and ajv-keywords

The global ajv ^6.14.0 override caused schema-utils and ajv-keywords to
receive ajv v6, but they require ajv v8 (for dist/compile/codegen). Add
scoped overrides to ensure these packages get ajv v8 while the global
override remains for packages that need v6.

* fix(tests): remove stateless_http from FastMCP() constructor calls

FastMCP 3.x no longer accepts stateless_http in the constructor. The
tests call tools directly without HTTP transport, so the parameter is
not needed.

* fix: update MCP tests for FastMCP 3.x _tool_manager removal

FastMCP 3.x removed _tool_manager. Tests now use
_local_provider._components for sync tool dict access and
mcp.list_tools() for async filtered tool listing.

* fix: resolve docusaurus build failures (ajv overrides + missing blog date)

- Remove global ajv ^6.14.0 override and scoped ajv-keywords/schema-utils
  overrides that caused webpack compilation errors manifesting as
  "Cannot read properties of undefined (reading 'date')" during SSR
  and "these parameters are deprecated" warnings. Natural version
  resolution (v6.12.6+ for v6 consumers, v8+ for v8 consumers) already
  satisfies the security fix (>= 6.12.3).
- Add missing date frontmatter to learning-capabilities blog post.

* chore: regenerate openapi spec and docs skill
2026-04-01 09:20:34 +02:00
f3f2c6b023 Fix timeline group sort: localeCompare → numeric Date comparison (#820)
* Initial plan

* Fix timeline sort to use numeric datetime comparison instead of string localeCompare

* chore: remove accidentally committed root package-lock.json

Agent-Logs-Url: https://github.com/ThePlenkov/hindsight/sessions/d02f10c5-cc48-4977-84a9-48870f9460ec

Co-authored-by: ThePlenkov <[email protected]>

* chore: restore package-lock.json to its original state from main

Agent-Logs-Url: https://github.com/ThePlenkov/hindsight/sessions/2ede4783-55ef-4f36-8ea7-7d65c5362a0a

Co-authored-by: ThePlenkov <[email protected]>

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: ThePlenkov <[email protected]>
2026-03-31 21:53:08 +02:00
Nicolò Boschi 6f7437be21 blog: What's New in Hindsight 0.4.22 release notes and changelog (#818) 2026-03-31 18:47:30 +02:00
Nicolò Boschi d7f6723546 Release v0.4.22
- Update version to 0.4.22 in all components
- Regenerate OpenAPI spec and client SDKs
- Python packages: hindsight-api, hindsight-dev, hindsight-all, hindsight-embed
- Python client: hindsight-clients/python
- TypeScript client: hindsight-clients/typescript
- Rust CLI: hindsight-cli
- Control Plane: hindsight-control-plane
- Helm chart
- Sync documentation to version-0.4
2026-03-31 18:14:59 +02:00
Nicolò Boschi 2c32ffadc9 fix(mental-models): add tags_match and tag_groups to trigger config (#786) (#804)
When a mental model has tags, refresh_mental_model hardcoded
tags_match="all_strict", causing empty results when most memories
are untagged. Add configurable tags_match and tag_groups fields
to MentalModelTrigger so users can control refresh filtering.

- Add tags_match (any/all/any_strict/all_strict) to override default
- Add tag_groups for compound boolean tag expressions during refresh
- Default behavior unchanged (all_strict when tags present)
- Update both refresh paths (task-based and direct)
- Add UI controls in Create/Update mental model dialogs
- Regenerate OpenAPI spec and client SDKs
2026-03-31 18:09:01 +02:00
Nicolò Boschi baf5447de2 refactor: replace LLMProvider classmethods with from_env() and document missing config fields (#816)
CI failures are unrelated to this PR:
- test_mental_models_dimension_change_empty_table: database OID error (infrastructure flake)
- test_reflect_searches_mental_models_when_available: LLM-dependent assertion (flaky)
2026-03-31 18:00:41 +02:00
KaguraandClaude Opus 4.6 84985ee9bc fix(reranker): use httpx for Cohere Azure endpoints to avoid 404 errors (#790)
When using Azure AI Foundry Cohere rerank endpoints, the Cohere SDK
incorrectly appends /v1/rerank to the base_url, but Azure endpoints
already include the full path (e.g., /models/.../invoke). This causes
double-pathing and 404 errors.

This commit modifies CohereCrossEncoder to detect when base_url is
provided and use httpx directly for custom endpoints, while keeping
the native Cohere SDK for standard API usage. The Azure Cohere API
response format is compatible with the native format.

Fixes #783

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-31 17:44:42 +02:00
emirhan-gaziandEMIRHAN GAZI ecaa1ad1e0 feat(api): add HINDSIGHT_API_LLM_EXTRA_BODY config for custom model params (#781)
Enable passing arbitrary extra_body parameters to OpenAI-compatible API
calls via a JSON-encoded env var. This supports custom model servers
(e.g. vLLM) that need parameters like chat_template_kwargs to control
thinking mode.

Co-authored-by: EMIRHAN GAZI <[email protected]>
2026-03-31 17:03:33 +02:00
dependabot[bot]anddependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> ea0c616240 chore(deps): bump dorny/paths-filter from 3 to 4 (#762)
Bumps [dorny/paths-filter](https://github.com/dorny/paths-filter) from 3 to 4.
- [Release notes](https://github.com/dorny/paths-filter/releases)
- [Changelog](https://github.com/dorny/paths-filter/blob/master/CHANGELOG.md)
- [Commits](https://github.com/dorny/paths-filter/compare/v3...v4)

---
updated-dependencies:
- dependency-name: dorny/paths-filter
  dependency-version: '4'
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-03-31 17:02:01 +02:00
dependabot[bot]anddependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> 0b29378eb8 chore(deps): bump azure/setup-helm from 4 to 5 (#761)
Bumps [azure/setup-helm](https://github.com/azure/setup-helm) from 4 to 5.
- [Release notes](https://github.com/azure/setup-helm/releases)
- [Changelog](https://github.com/Azure/setup-helm/blob/main/CHANGELOG.md)
- [Commits](https://github.com/azure/setup-helm/compare/v4...v5)

---
updated-dependencies:
- dependency-name: azure/setup-helm
  dependency-version: '5'
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-03-31 17:01:52 +02:00
dependabot[bot]anddependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> c2e801ccb0 chore(deps): bump actions/deploy-pages from 4 to 5 (#763)
Bumps [actions/deploy-pages](https://github.com/actions/deploy-pages) from 4 to 5.
- [Release notes](https://github.com/actions/deploy-pages/releases)
- [Commits](https://github.com/actions/deploy-pages/compare/v4...v5)

---
updated-dependencies:
- dependency-name: actions/deploy-pages
  dependency-version: '5'
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-03-31 17:01:44 +02:00
Ben 410c208746 What's New: multi-org support and credit transfers (March 29) (#815)
* Add What's New post: multi-org support and credit transfers
2026-03-31 10:19:08 -04:00
Nicolò Boschi c475c6bb56 ci: trigger full CI on PR approval instead of safe-to-test label (#813)
Replace the `pull_request_target` + `safe-to-test` label mechanism with
`pull_request_review` (submitted, approved). External contributor PRs now
get basic builds/lints on open, and full secret-dependent CI only after
a maintainer approves — no manual labeling needed.
2026-03-31 15:15:14 +02:00
Amin Bolakhrif f841bcb92d feat: add optional LiteLLM SDK embedding output dimensions (#809)
* feat: add optional LiteLLM SDK embedding output dimensions

Allow configuring an optional output dimension for litellm-sdk embeddings and pass it through only when set, while preserving default behavior.

Made-with: Cursor

* test: assert wrapped init error for invalid dimensions

Add a LiteLLM SDK embeddings test that verifies invalid OpenAI dimensions fail during initialize() and preserve provider error details in the wrapped RuntimeError.

Made-with: Cursor
2026-03-31 14:58:02 +02:00
Maxim Kremmnev fa82efc886 fix(claude-code): disable built-in tools to prevent MCP tool deferral (#784) 2026-03-31 14:20:17 +02:00
Nicolò Boschi 627ec5d524 feat: expose document_metadata in API and control plane (#798)
* feat: expose document_metadata in API and control plane

Add document_metadata (sourced from retain_params.metadata) to both
list and get document endpoints. Display it in the control plane
documents table and detail panel. Drop the unused metadata column
from the documents table (was always stored as empty {}).

* fix: code review fixes for document_metadata feature

- Remove unnecessary `import json as _json` (json already imported at module level)
- Simplify redundant truthiness checks in retain_params parsing
- Regenerate OpenAPI spec and client SDKs (Python, TypeScript, Go)
- Add tests for document_metadata in get_document and list_documents

* feat(ui): improve documents table and detail panel

- Relative timestamps with full date on hover
- Remove context column from table
- Metadata shown as k=v badges (blue, like tags)
- Size in bytes instead of chars
- Document IDs wrap instead of truncating
- Detail panel wider (560px)
- Retain params: context, event_date, metadata badges
2026-03-31 11:42:02 +02:00
Nicolò Boschi bdb33c58d1 feat: add /code-review skill with project standards (#806)
* feat: add /code-review skill for automated code quality checks

Adds a Claude Code skill that reviews changes against project standards:
missing tests, dead code, type safety, lint, and CLAUDE.md conventions.
CLAUDE.md now instructs contributors to run /code-review after implementation.

* refactor: move code standards from CLAUDE.md into /code-review skill

Single source of truth for coding conventions (Python style, type safety,
TypeScript style) is now .claude/skills/code-review.md. CLAUDE.md points
to the skill for reading before coding and running after implementation.

* feat: add code comments convention to /code-review skill

Require comments explaining non-trivial technical decisions, with history
of previous approaches. Review step checks for missing reasoning comments,
stale comments, and undocumented approach changes.

* fix: move skill to directory structure for Claude Code discovery

Claude Code requires .claude/skills/<name>/SKILL.md, not loose .md files.

* feat: add branch hygiene checks to /code-review skill

Review step 1 now verifies branch is based on recent origin/main and
all commits are relevant to the feature. Unrelated commits flagged as
must-fix.

* feat: strengthen code review rules and fix stale CLAUDE.md references

- Enforce no multi-item tuple returns and no raw dicts even for internal code
- Add mandatory /code-review gate before push/PR
- Add integration completeness checklist (tests, CI job, release-integration.sh)
- Fix stale references: remove hindsight/ dir, update integrations list,
  update LLM providers, remove hardcoded file sizes, fix _HIERARCHICAL_FIELDS
  -> _CONFIGURABLE_FIELDS

* docs: add ./scripts/dev/start.sh for local dev in CLAUDE.md
2026-03-31 11:09:41 +02:00
Nicolò Boschi 1dbbe39ea1 ci: report safe-to-test CI results on PR (#807)
* feat(api): warn on unknown request parameters via X-Ignored-Params header

Add middleware that detects unknown query params and JSON body fields,
logs a server-side warning, and returns an X-Ignored-Params response
header listing the ignored parameters. This surfaces silent parameter
ignoring (e.g. tag=source:slack on /memories/list) without breaking
forward compatibility between client and server versions.

Closes #792

* ci: report safe-to-test CI results on PR via status and comment

pull_request_target workflow runs are not linked to the PR by GitHub,
so the CI results are invisible on the PR page after adding safe-to-test.

Add a report-pr-status job that:
- Creates a commit status on the PR head SHA
- Posts/updates a summary comment with pass/fail counts and failed job names

* ci: skip secret-dependent jobs on fork pull_request events

Adds a has_secrets output to detect-changes that is false for fork PRs
via pull_request events. All 15 secret-dependent jobs now check this
output before running, avoiding guaranteed failures on fork PRs.

Fork contributors will see these jobs as skipped instead of failed,
and can use the safe-to-test label to run the full CI suite.
2026-03-31 11:09:24 +02:00
Nicolò Boschi cef42d8154 feat(api): warn on unknown request parameters via X-Ignored-Params header (#802)
Add middleware that detects unknown query params and JSON body fields,
logs a server-side warning, and returns an X-Ignored-Params response
header listing the ignored parameters. This surfaces silent parameter
ignoring (e.g. tag=source:slack on /memories/list) without breaking
forward compatibility between client and server versions.

Closes #792
2026-03-31 10:34:04 +02:00
Nicolò Boschi f8f62030e3 Add /code-review skill for automated code quality checks (#805)
* feat: add /code-review skill for automated code quality checks

Adds a Claude Code skill that reviews changes against project standards:
missing tests, dead code, type safety, lint, and CLAUDE.md conventions.
CLAUDE.md now instructs contributors to run /code-review after implementation.

* refactor: move code standards from CLAUDE.md into /code-review skill

Single source of truth for coding conventions (Python style, type safety,
TypeScript style) is now .claude/skills/code-review.md. CLAUDE.md points
to the skill for reading before coding and running after implementation.

* feat: add code comments convention to /code-review skill

Require comments explaining non-trivial technical decisions, with history
of previous approaches. Review step checks for missing reasoning comments,
stale comments, and undocumented approach changes.

* fix: move skill to directory structure for Claude Code discovery

Claude Code requires .claude/skills/<name>/SKILL.md, not loose .md files.

* feat: add branch hygiene checks to /code-review skill

Review step 1 now verifies branch is based on recent origin/main and
all commits are relevant to the feature. Unrelated commits flagged as
must-fix.
2026-03-31 10:21:25 +02:00
Nicolò Boschi 4768bf39ef fix(http): recall endpoint drops metadata in response (#797) (#803)
_fact_to_result was missing metadata=fact.metadata, so the HTTP recall
endpoint always returned metadata: null even though the engine preserved it.
2026-03-31 10:12:12 +02:00
Nicolò Boschi 865fb91298 fix(tests): use random port for pg0 in tests to avoid port conflicts (#801)
pg0 supports auto-assigning a free port when port=None. This avoids
test failures when port 5556 is already in use by another process.
2026-03-31 09:38:46 +02:00
Nicolò Boschi 1f5dc8bd15 ci: support running secret-dependent tests on fork PRs via safe-to-test label (#800)
Fork PRs don't have access to repository secrets, so integration tests
that need API keys (GCP, OpenAI, Cohere, etc.) are skipped. Maintainers
can now add the `safe-to-test` label after reviewing fork PR code to
trigger the full test suite with secrets via pull_request_target.
2026-03-31 09:31:39 +02:00
Nicolò Boschi d3d2684b11 fix(openclaw): add warn log and tests for CLI mode no-op in waitForReady (#799)
Follow-up to #764. Upgrades the silent debug log in waitForReady to
log.warn so unexpected calls before service.start() are visible, and
adds tests covering the CLI mode no-op path.
2026-03-31 09:25:52 +02:00
Kagura 41025c3b7c fix(openclaw): defer heavy init to service.start() to avoid CLI slowdown (#764)
OpenClaw loads plugins on every CLI command (status, models auth add,
config validate, etc.), not just gateway start. The plugin was starting
LLM detection, daemon initialization, and API health checks immediately
in the default export, causing unnecessary resource usage and terminal
noise on routine CLI operations.

Move all heavy initialization (detectLLMConfig, embedManager.start(),
checkExternalApiHealth, client creation) into service.start() which is
only called when the gateway starts. The default export now only does
lightweight config parsing and service/hook registration.

Hooks (before_prompt_build, agent_end) gracefully no-op when called
before service.start() via the waitForReady guard.

Closes #746
2026-03-31 09:10:13 +02:00
Volodymyr Prypeshniuk 1b5c262a8a fix(gemini): thought_signature read from wrong object and type in 3.1+ tool calls (#785) 2026-03-31 09:09:12 +02:00
Nicolò Boschi 0096115678 fix(engine): classify first-person agent experiences as 'experience' fact type (#775)
* fix(engine): classify first-person agent experiences as 'experience' fact type

The extraction prompt defined "assistant" too narrowly as only "interactions
with assistant (requests, recommendations)", causing the LLM to classify
first-person agent actions (code changes, debugging, discoveries) as "world".

Broadened the fact_type definition in the prompt and Pydantic model descriptions
to cover all first-person actions, experiences, and observations by the speaker.

* style: fix line length in fact_extraction.py
2026-03-31 09:06:36 +02:00
Nicolò Boschi b104bad02c fix(codex): merge new settings on upgrade instead of skipping (#780)
The installer skipped settings.json entirely if it already existed,
leaving version and new config keys stale. Now merges: updates version,
adds new upstream keys, preserves user customizations.

Also fixes pre-existing typo: RERANK_URL → rerank_url in ZeroEntropy
cross-encoder.
2026-03-31 09:05:49 +02:00
Chris Bartholomew 45ffc7fe90 SEO: add title and description to all integration pages (#787)
* SEO: add title and description to all integration pages

All 17 integration docs pages were missing title and description
frontmatter, causing Docusaurus to generate unhelpful titles like
"OpenClaw | Hindsight" and pull body text as meta descriptions.

- Add keyword-rich title and description frontmatter to all integration
  pages in both docs/ (current) and versioned_docs/version-0.4/
- Add scripts/check-integration-seo.mjs to enforce title + description
  on all future integration pages
- Wire the check into the build script so it runs locally and in CI

* Fix missing frontmatter on docs/sdks/integrations/openclaw.md

* Regenerate docs skill after integration page SEO updates
2026-03-30 17:53:43 -04:00
Chris Bartholomew 99122055f0 Improve OpenClaw post title, tags, and meta description
- Retitle to match search intent: "How to Add Persistent Memory to
  OpenClaw with Hindsight" targets openclaw memory/persistent memory queries
- Add intro paragraph before <!-- truncate --> so Docusaurus generates a
  proper meta description instead of "TL;DR"
- Expand tags from [openclaw] to include memory, agents, persistent-memory,
  knowledge-graph
2026-03-30 15:41:55 -04:00
Nicolò Boschi 75e2679cf1 release(llamaindex): v0.1.3 2026-03-30 18:34:25 +02:00
DK09876andClaude Opus 4.6 d93dfea8ce fix(llamaindex): document_id, memory API, and ReAct trace fixes (#777)
* fix(llamaindex): use uuid for document_id and sync version metadata

- Replace timestamp-based document_id with uuid4 hex to prevent
  collisions on rapid retains (timestamp_ms can duplicate in tight loops)
- Sync __version__ in __init__.py to match pyproject.toml (0.1.2)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix(docs): pass memory to run() instead of ReActAgent constructor

LlamaIndex 0.14.x ReActAgent does not accept a memory parameter in
its constructor — it's silently dropped via **kwargs. Memory must be
passed to agent.run(memory=...) where AgentWorkflow picks it up.

Also fixes the undefined `tools` variable (now `tools=[]`).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix(llamaindex): strip ReAct reasoning traces from retained assistant messages

HindsightMemory.put/aput now extracts only the final Answer: text from
assistant messages containing ReAct reasoning (Thought:/Action:/Observation:
prefixes), preventing internal reasoning traces from polluting long-term memory.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix(llamaindex): fix docstring example to pass memory to run()

The HindsightMemory class docstring showed the broken pattern of passing
memory= to the ReActAgent constructor, which silently drops it. Updated
to show the correct pattern: pass memory to agent.run().

Co-Authored-By: Claude Opus 4.6 <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-30 18:33:03 +02:00
Nicolò Boschi e5209b18b3 docs: add RERANKER_ZEROENTROPY_BASE_URL to configuration page (#779)
Document the new configurable base URL for the ZeroEntropy reranker
provider added in #766. Also fix a type error where RERANK_URL was
renamed to rerank_url but one usage was missed.
2026-03-30 17:51:04 +02:00
Nicolò Boschi a7adfbb0df release(codex): v0.2.0 2026-03-30 17:50:02 +02:00
Nicolò Boschi 3461398b52 feat(codex): add structured tool call retention from Codex rollout files (#778)
Parse all Codex rollout item types (function_call, local_shell_call,
exec_command_end, patch_apply_end, mcp_tool_call_end, web_search_call)
into structured JSON content blocks matching Claude Code's format.
Enabled by default via retainToolCalls setting.
2026-03-30 17:48:47 +02:00
Timur Iskhakov a915584e39 feat: add configurable base URL for ZeroEntropy reranker (#766) 2026-03-30 17:43:19 +02:00
Nicolò Boschi 2c72af5525 release(openclaw): v0.5.1 2026-03-30 16:58:18 +02:00
Nicolò Boschi 41bb6d710b Revert "Bump openclaw integration to v0.5.1"
This reverts commit a3e458ad43.
2026-03-30 16:57:39 +02:00
DK09876andClaude Opus 4.6 7af01e35e9 fix(docs): use tools=[] in BaseMemory example (#772)
The automatic memory example referenced an undefined `tools` variable.
Since HindsightMemory handles retain/recall transparently, no tools
are needed — use an empty list.

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-30 16:56:22 +02:00
Nicolò Boschi a3e458ad43 Bump openclaw integration to v0.5.1 2026-03-30 16:55:58 +02:00
Nicolò Boschi 704e41fa27 Fix trailing commas in openclaw.plugin.json and add JSON manifest CI tests (#774)
Fixes #771 — two trailing commas in openclaw.plugin.json caused OpenClaw's
strict JSON parser to reject the plugin manifest during installation.

Also adds JSON validation tests for both the openclaw plugin manifest and
the claude-code hooks.json so CI catches invalid JSON before release.
2026-03-30 16:55:06 +02:00
Ben f30ca3deda Fix blog homepage: Hindsight Cloud section always shows top 3 posts (#770)
* Fix blog homepage: show all posts so Cloud section always gets top 3
2026-03-30 10:40:57 -04:00
Ben d61517d502 Update MCP OAuth blog post date to 2026-03-30 (#769) 2026-03-30 13:53:41 +00:00
Ben df17570d8a blog: What's New in Hindsight Cloud — Native OAuth for MCP Clients (#731)
* Add MCP OAuth blog post
2026-03-30 09:33:58 -04:00
Nicolò Boschi 7a9e99998a docs: 0.4.21 release blog post and changelog (#765)
* docs: 0.4.21 release blog post and changelog

* fix(blog): align 0.4.21 code snippets with docs, add release image

* chore: regenerate docs skill references
2026-03-30 15:28:27 +02:00
Nicolò Boschi cc3cdc2f83 Release v0.4.21
- Update version to 0.4.21 in all components
- Regenerate OpenAPI spec and client SDKs
- Python packages: hindsight-api, hindsight-dev, hindsight-all, hindsight-embed
- Python client: hindsight-clients/python
- TypeScript client: hindsight-clients/typescript
- Rust CLI: hindsight-cli
- Control Plane: hindsight-control-plane
- Helm chart
- Sync documentation to version-0.4
2026-03-30 14:52:50 +02:00
Nicolò Boschi 88630f93d7 release(hermes): v0.5.0 2026-03-30 14:50:04 +02:00
Nicolò Boschi 73460fa4e6 chore: add hermes to release and changelog valid lists 2026-03-30 14:49:33 +02:00
Nicolò Boschi 6a28373ecf release(openclaw): v0.5.0 2026-03-30 14:47:01 +02:00
Nicolò Boschi 66071fec9d feat(scripts): support patch/minor/major bump keywords in release-integration 2026-03-30 14:43:43 +02:00
Nicolò Boschi b8b40e458b release(codex): v0.1.1 2026-03-30 14:40:58 +02:00
Nicolò Boschi e65bd361cc release(llamaindex): v0.1.2 2026-03-30 14:37:15 +02:00
Nicolò Boschi b8fa0e8cfc chore: add llamaindex and codex to changelog generator 2026-03-30 14:36:47 +02:00
Nicolò Boschi b739b9e36a chore: add llamaindex to release-integration valid list 2026-03-30 14:35:18 +02:00
Nicolò Boschi 5b16882a5e chore(llamaindex): bump version to 0.1.1 2026-03-30 14:32:10 +02:00
Nicolò Boschi 56489a58d1 refactor(llamaindex): merge into single hindsight-llamaindex package (#760)
* refactor(llamaindex): merge two packages into single hindsight-llamaindex

Merge `llama-index-tools-hindsight` and `llama-index-memory-hindsight` into
a single `hindsight-llamaindex` package following our naming convention.

- Rename package to `hindsight-llamaindex` (Python module: `hindsight_llamaindex`)
- Move HindsightToolSpec and HindsightMemory into the same package
- Delete `llamaindex-memory/` directory
- Add CI test job for llamaindex integration
- Update docs, blog post, and integrations.json

* fix(blog): update llamaindex blog post for merged package

- Move date to 2026-03-30
- Add HindsightMemory (automatic BaseMemory) pattern
- Fix "bank must exist first" pitfall — mission auto-creates
- Align all code examples with docs page
- Update architecture diagram to show both patterns

* fix(docs): add llamaindex/openai icons, rename Codex

- Add llamaindex.png and openai.png icons
- Rename "OpenAI Codex CLI" to "Codex" in integrations.json and docs
- Use openai.png icon for Codex integration
2026-03-30 14:30:48 +02:00
Nicolò Boschi a8a63818c7 feat(api): add duration_ms to audit log entries (#758)
* feat(api): add duration_ms to audit log entries

Server-computed duration in milliseconds (started_at → ended_at) on
the list audit logs endpoint. Null when ended_at is not set.

Closes #749

* feat(api): add duration_ms to audit log entries and type audit endpoints

- Add server-computed duration_ms (started_at → ended_at) to audit log
  list response. Null when ended_at is not set.
- Add typed Pydantic response models for both audit log endpoints
  (list and stats) so they appear in the OpenAPI spec.
- Regenerate OpenAPI spec and all client SDKs.

Closes #749

* chore: regenerate docs skill after audit log response models
2026-03-30 12:32:02 +02:00
Nicolò Boschi d8050387e4 fix(mcp): handle Claude Code GET probe and make stateless_http configurable (#757)
* fix(mcp): handle Claude Code GET probe and make stateless_http configurable (#751)

Claude Code v2.1.84+ sends a GET to /mcp/ before POST initialize,
which fails with 405 (stateless) or 400 (stateful). Intercept
sessionless GET requests in MCPMiddleware and return 200 OK so the
client proceeds to POST initialize.

Also make stateless_http configurable via HINDSIGHT_API_MCP_STATELESS
(default: false/stateful) instead of hardcoding true.

Closes #751

* docs: add HINDSIGHT_API_MCP_STATELESS to configuration reference
2026-03-30 12:15:06 +02:00
Nicolò Boschi 38e03e419d Convert codex tool_choice test to pytest style (#752)
* Convert codex tool_choice test to pytest style

Follow-up to #734: replace unittest.TestCase + manual sys.path
manipulation with idiomatic pytest + @pytest.mark.asyncio,
matching the rest of the test suite.

* Fix test_hierarchical_fields_categorization for new configurable fields

Update expected count from 20 to 21 and add assertions for fields
added by recent PRs: retain_default_strategy, retain_strategies,
max_observations_per_scope, reflect_source_facts_max_tokens,
llm_gemini_safety_settings, mcp_enabled_tools.

* Add LlamaIndex doc to v0.4 versioned docs and sidebars

The LlamaIndex integration doc was added to docs/ (next version) in
#672 but not to versioned_docs/version-0.4/, causing a broken link
on the /integrations page which resolves to the latest version.

* Regenerate docs skill references

Run generate-docs-skill.sh to pick up new integration pages
(codex, llamaindex) and updated configuration docs.

* Add Codex integration doc to v0.4 versioned docs and sidebar

Same issue as LlamaIndex: doc was added to docs/ (next) but not
versioned_docs/version-0.4/, causing broken link on /integrations.
2026-03-30 11:58:59 +02:00
Nicolò Boschi 6488c9bc77 fix: per-bank index creation respects HINDSIGHT_API_VECTOR_EXTENSION config (#755)
create_bank_hnsw_indexes() hardcoded USING hnsw regardless of the configured
vector extension, causing "column cannot have more than 2000 dimensions for
hnsw index" when using pgvectorscale or vchord with high-dimensional embeddings.

Now reads get_config().vector_extension and uses the appropriate index type:
- pgvector → USING hnsw
- pgvectorscale → USING diskann
- vchord → USING vchordrq

Closes #738
2026-03-30 11:37:26 +02:00
Nicolò Boschi d2965e64e6 fix(retain): inject retain_mission into verbose extraction mode (#745) (#754)
Verbose mode was the only extraction mode that skipped injecting the
retain_mission FOCUS section into its prompt template. Users who set a
retain_mission got no filtering when using verbose mode.
2026-03-30 11:36:36 +02:00
Nicolò Boschi ecf16ea1e6 fix(codex): cleanup dead code, add release lifecycle and docs (#753)
* fix(codex): cleanup dead code and add to release lifecycle

- Remove orphaned reflect() method from client.py (leftover from dropped auto-mode)
- Remove dead retainToolCalls config default (never wired through)
- Add codex to release-integration.sh valid integrations
- Add settings.json version fallback to release script
- Add codex CI test job in test.yml
- Add codex to integrations.json registry

* docs(codex): add changelog page and link from integration docs

* feat(codex): add hosted installer script (get-codex)

Add self-contained installer at hindsight.vectorize.io/get-codex that
downloads scripts from GitHub, configures hooks, and supports local/cloud
mode selection — no git clone required.

Update docs and README to use the one-liner install:
  curl -fsSL https://hindsight.vectorize.io/get-codex | bash

* chore(codex): remove install.sh in favor of hosted get-codex

* fix(docs): use /next/ prefix for codex changelog link

* fix(docs): use GitHub link for codex changelog back-link
2026-03-30 11:34:43 +02:00
Ben 0b17a67c70 feat: add Hindsight memory integration for OpenAI Codex CLI (#730)
* feat(codex): add Hindsight memory integration for OpenAI Codex CLI

Hooks-based integration that gives Codex CLI long-term memory via Hindsight.
Three hooks keep memory in sync: SessionStart (daemon pre-warm), UserPromptSubmit
(recall + context injection), Stop (retain conversation to memory).

Key differences from the Claude Code integration:
- Codex transcript format: JSONL with {msg: {type, message}} (user_message/agent_message)
- No CODEX_PLUGIN_ROOT env var — install.sh writes hooks.json with absolute paths
- State stored in ~/.hindsight/codex/state/ (not CLAUDE_PLUGIN_DATA)
- No async: true in hooks (not supported by Codex)
- No SessionEnd event
- hooks.json written to ~/.codex/hooks.json with codex_hooks = true in config.toml

* fix(codex): fix transcript parser for actual Codex disk format

Codex stores sessions as rollout-*.jsonl with response_item entries:
  User:      {type:response_item, payload:{type:message, role:user, content:[{type:input_text, text:...}]}}
  Assistant: {type:response_item, payload:{type:message, role:assistant, phase:final_answer, content:[{type:output_text, text:...}]}}

Previous parser expected an undocumented {msg:{type:user_message}} format from the Rust protocol spec
that does not match the actual on-disk storage format.

* feat(codex): add reflect mode to UserPromptSubmit hook

Add recallMode config option (default: 'recall') that switches the
UserPromptSubmit hook between:
- 'recall': existing behavior, fast raw facts list
- 'reflect': agentic synthesis loop, returns coherent prose answer

Also adds reflect() method to HindsightClient and HINDSIGHT_RECALL_MODE
env var override. Reflect uses a 25s timeout (vs 10s for recall).

* feat(codex): auto mode for recall/reflect selection

Add recallMode: 'auto' (new default) that picks the operation per-query:
- Synthesis patterns (what do you know, what's my, summarize, etc.) → reflect
- All other prompts → recall (fast, raw facts, better for code tasks)

* feat(codex): add automated test suite and finalize recall-only mode

* docs(codex): add docs page and sidebar entry for Codex CLI integration
2026-03-30 10:51:53 +02:00
Nicolò Boschi e7c9a6832d fix(hermes): sync lifecycle hooks for hermes-agent 0.5.0 (#741)
* fix(hermes): convert lifecycle hooks to sync for hermes-agent 0.5.0 compatibility

hermes-agent 0.5.0 calls plugin hooks synchronously via invoke_hook(),
but our pre_llm_call/post_llm_call were async — coroutines were never
awaited, so recall context injection and auto-retain silently did nothing.

Switch hooks to sync client methods and add integration tests using
the real hermes-agent PluginManager.

* fix(hermes): use proper hermes-agent dep with uv source override

Replace inline git URL with standard `hermes-agent>=0.5.0` version
constraint plus `[tool.uv.sources]` to resolve from the git tag until
0.5.0 lands on PyPI.
2026-03-30 10:44:54 +02:00
DK09876andClaude Opus 4.6 2d787c4ffd feat: add LlamaIndex integration (#672)
* feat: add LlamaIndex integration for Hindsight

Add hindsight-llamaindex package providing persistent memory tools for
LlamaIndex agents via the native BaseToolSpec pattern. Includes retain,
recall, and reflect tools, a convenience factory, global config, full
test suite, docs page, blog post, and integrations.json entry.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address PR review feedback for llamaindex integration

- Fix ReActAgent API: from_tools() → constructor, chat() → await run()
- Add create_bank step to all quickstart examples
- Add production patterns section to docs (tags, error handling, bank lifecycle)
- Add memory scoping recommendation to README
- Add when-not-to-use section to blog post
- Add LlamaIndex compatibility tests (agent acceptance, FunctionTool.call)
- Fix self-hosted auth wording in cookbook notebook

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: use async client methods and asyncio.run() for runnable examples

- Use await client.acreate_bank() instead of sync create_bank() to
  avoid "event loop already running" errors in notebooks and async contexts
- Wrap plain Python examples in async def main() + asyncio.run(main())
  so they are copy-paste runnable as scripts
- Add Jupyter notebook tip to docs showing top-level await pattern
- Bank lifecycle example in docs now uses async acreate_bank

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: add async tool methods to avoid event loop conflicts

HindsightToolSpec now provides both sync and async tool implementations
using LlamaIndex's (sync_fn, async_fn) tuple pattern in spec_functions.
Async agents (ReActAgent, etc.) use aretain/arecall/areflect natively,
avoiding the "Timeout context manager should be used inside a task"
error that occurred when sync _run_async() was called from within an
active event loop.

- Add aretain_memory, arecall_memory, areflect_on_memory async methods
- Extract shared kwargs builders (_retain_kwargs, _recall_kwargs, etc.)
- spec_functions now uses tuples: [("retain_memory", "aretain_memory"), ...]
- Tests verify tools have both sync fn and async fn set
- Notebook verified end-to-end with nbclient against local Hindsight

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* chore: remove blog post from integration PR

The blog post will be pulled in separately from its own PR.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Address PR review: add context label, document_id auto-gen, bank mission, graceful errors

- Add `retain_context` param (default: "llamaindex") as source label on retain ops
- Auto-generate `document_id` as `{session_id}-{timestamp_ms}` when not provided
- Add `retain_async` param (default: True) for non-blocking retain processing
- Add `mission` param for automatic bank creation/management on first use
- Change error handling from raising HindsightError to graceful log + return message
- Add per-operation timeout constants in _client.py
- Add `context` and `mission` fields to config.py and configure()
- Update docs: document as standalone package (not LlamaHub), new params, patterns
- Tests: 51 passing (up from 34), covering all new features

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Restructure to LlamaIndex namespace packages + add BaseMemory implementation

Tools package (llama-index-tools-hindsight):
- Restructured from hindsight_llamaindex/ to llama_index/tools/hindsight/
- Import: from llama_index.tools.hindsight import HindsightToolSpec
- Follows PEP 420 implicit namespace package convention
- Removed retain_async param (client.retain() doesn't support async_processing)

Memory package (llama-index-memory-hindsight):
- New package: llama_index/memory/hindsight/
- HindsightMemory(BaseMemory) for automatic memory
- put() auto-retains user/assistant messages to Hindsight
- get(input) auto-recalls relevant memories, prepends as system message
- Graceful error handling, bank mission management, document_id generation
- 28 unit tests passing

Both packages follow LlamaIndex community conventions for future LlamaHub submission.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-30 10:43:35 +02:00
111e8c70a2 fix(codex): don't crash on startup when quota is exhausted (429) (#744)
A 429 usage_limit_reached response during verify_connection() caused the
server to refuse to start entirely. Quota exhaustion is not a configuration
error — the server should start and serve retain/recall requests normally,
it just can't make LLM calls until the quota resets.

Co-authored-by: Marco Rutsch <[email protected]>
Co-authored-by: Claude Sonnet 4.6 <[email protected]>
2026-03-30 10:39:12 +02:00
d441ab814d feat(openclaw): configurable logging with structured output (#739)
* feat(openclaw): configurable logging with structured output

Replace raw console.log/warn/error spam with a structured logger.
New plugin settings: logLevel, logSummaryIntervalMs, logCompact.
Bank mission log demoted to verbose-only. Retain/recall batched
into periodic summaries. Each recall now shows memory count injected.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* use api.logger for framework-consistent output, show autoRecall/autoRetain on init

Route all log output through OpenClaw's api.logger instead of raw console
calls. Matches mem0 plugin style. Startup now shows mode + feature flags.
Dropped logCompact setting (framework handles formatting). Added subtle
slate-blue color to hindsight prefix for visual differentiation.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* add bank name to init and summary logs, fix singular/plural consistency

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* rename log levels to standard: off, error, warning, info, debug

Per review feedback — use standard level names instead of custom ones.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: billy <[email protected]>
Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
Co-authored-by: Nicolò Boschi <[email protected]>
2026-03-30 10:38:02 +02:00
Mr. Khachaturov f8285b7b90 feat(mcp): add filter_mcp_tools hook for per-user tool visibility (#737)
Add optional filter_mcp_tools() method to OperationValidatorExtension.
Called during tools/list after bank-level mcp_enabled_tools filtering.
Extensions can override to hide MCP tools per-user-per-bank based on
access policies. Default returns all tools unchanged.

- Add filter_mcp_tools to OperationValidatorExtension with default pass-through
- Wire into _get_enabled_tools in _apply_bank_tool_filtering
- Move _ALL_TOOLS to mcp_tools.py to avoid circular import (re-exported from mcp.py)
- Fail-open: if filter raises, log warning and return unfiltered tools
- Enforce ceiling: validator can narrow but never expand beyond bank config
- Add 8 tests: default, filtering, empty set, integration, composition,
  can't-add-tools, exception fail-open, no-validator passthrough
2026-03-30 10:33:09 +02:00
akhaterandAntoine Khater a209ef1ae2 fix: parse query params from base_url in OpenAI embeddings client (#735)
* fix: parse query params from base_url in OpenAI embeddings client

The OpenAI-compatible LLM provider already parses query parameters
(e.g. ?api-version=xxx for Azure OpenAI) from the base_url and passes
them as default_query to the OpenAI client. However, the OpenAI
embeddings provider did not do this, causing Azure OpenAI embeddings
to fail with 404 errors at runtime.

This applies the same URL parsing logic from the LLM provider to the
embeddings provider, enabling Azure OpenAI embeddings to work correctly.

* ci: add workflow to build fork Docker image

* ci: add slim image build (no local models)

* ci: remove fork build workflow per review request

---------

Co-authored-by: Antoine Khater <[email protected]>
2026-03-30 10:32:10 +02:00
Daoyang ShanandSapientropic 3573e53b1d Fix Codex named tool_choice in reflect (#734)
Co-authored-by: Sapientropic <[email protected]>
2026-03-30 10:31:17 +02:00
KaguraandClaude Opus 4.6 585ac76f39 fix(claude-code): implement tool_choice support for forced tool calls (#733)
* fix(claude-code): implement tool_choice support for forced tool calls

The call_with_tools() method now properly handles the tool_choice parameter
to force specific tool calls. Previously, the parameter was accepted but ignored,
causing the reflect agent to fail when trying to force specific tools on each
iteration.

Fixes #732

Changes:
- When tool_choice forces a specific function: filter allowed_tools to only
  that tool (with mcp prefix) and add a strong system prompt instruction
- When tool_choice is 'required': add instruction that model must call at
  least one tool
- When tool_choice is 'none': clear allowed_tools and mcp_servers to disable
  all tools
- When tool_choice is 'auto' (default): no change (existing behavior)

This matches the approach used in the OpenAI provider while adapting to the
Claude Agent SDK's lack of native tool_choice parameter by using allowed_tools
filtering and system prompt instructions.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* style: fix ruff formatting in alembic migration

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-30 10:30:53 +02:00
Nicolò Boschi b32767caa8 feat: add max_observations_per_scope bank config (#729)
* feat: add max_observations_per_scope bank config

Adds a configurable limit on the number of observations per tag scope.
When the limit is reached, consolidation only updates/deletes existing
observations — no new ones are created. Enforcement is done via a
constrained Pydantic response model (max_length on creates list) so the
LLM structurally cannot exceed the limit, plus prompt guidance.

- Config: HINDSIGHT_API_MAX_OBSERVATIONS_PER_SCOPE (-1 = unlimited)
- Reorder action execution: deletes → updates → creates
- Dynamic _ConsolidationBatchResponse with max_length constraint
- Prompt CAPACITY CONSTRAINT section when near/at limit
- Observations with no tags skip the limit entirely
- Control plane UI field + docs

* fix: strengthen max_observations tests with mock LLM + defensive truncation

- Rewrite integration tests to use MockLLM with deterministic responses
  (one observation per fact) instead of relying on real LLM behavior
- Add defensive truncation in _consolidate_batch_with_llm as belt-and-
  suspenders — catches LLM providers that ignore JSON schema max_length
- Tests now assert exact counts, not just upper bounds
2026-03-30 10:29:56 +02:00
cd4d449f8e fix(openclaw): add recallTimeoutMs config option for auto-recall (#736)
The auto-recall timeout was hardcoded to 10s but recall with budget=high
can take 13s+. This adds a configurable recallTimeoutMs option (default:
10000ms) so users can increase the timeout when using higher recall budgets.

Also adds recallInjectionPosition to the plugin schema (it was already
implemented in code but missing from the JSON schema validation, causing
config rejection).

Co-authored-by: Marco Rutsch <[email protected]>
Co-authored-by: Claude Sonnet 4.6 <[email protected]>
2026-03-28 17:59:23 +01:00
Nicolò Boschi 7a3dbc1958 refactor(embedded): replace UI programmatic API with constructor flags (#728)
Replace start_ui()/stop_ui()/is_ui_running() methods with declarative
constructor flags (ui, ui_port, ui_hostname). UI lifecycle now follows
the daemon automatically - starts in _ensure_started, stops in _cleanup.

Add integration test verifying UI starts and can reach the dataplane
via the control plane's /api/health endpoint. Add Node.js setup to
test-hindsight-all CI job to support the UI test.
2026-03-27 18:01:04 +01:00
a69bdbb55f How We Built a 4-Way Hybrid Search System That Actually Runs in Parallel (#708)
* Add blog: How We Built a 4-Way Parallel Hybrid Search System

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

* Add cover image for parallel hybrid search post

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

* Update parallel hybrid search post date to 2026-03-27

* Set author to chrislatimer

* Update recall docs link

* review: align blog post to actual retrieval code

- Reframe as evolutionary narrative (V1 asyncio.gather → connection sharing)
- Add missing reranker section (cross-encoder + multiplicative boost scoring)
- Replace MPFP references with LinkExpansion (3-signal CTE)
- Fix SQL to match actual UNION ALL approach, explain CTE planner issue
- Fix acquire_with_retry, index types (ivfflat→HNSW), fusion code
- Remove fabricated perf numbers
- Add alpha calibration rationale and connection contention insight

* add nicoloboschi and benfrank241 as co-authors

---------

Co-authored-by: Claude Sonnet 4.6 <[email protected]>
Co-authored-by: Nicolò Boschi <[email protected]>
2026-03-27 11:21:01 -04:00
Nicolò Boschi f50cc25dfb perf(stats): add bank_id to memory_links for direct filtering (#718)
The stats endpoint JOINs memory_links to memory_units just to filter
by bank_id.  With 8.2M+ links per bank this takes 18+ seconds, and
the control plane polls every 10s — perpetually blocking the server.

Add bank_id column directly to memory_links so the query can filter
on ml.bank_id instead of mu.bank_id, letting Postgres push the filter
down before the JOIN.

- Migration: add bank_id TEXT NOT NULL, backfill from memory_units
- All 4 INSERT paths (temporal, semantic, entity, causal) now write bank_id
- Stats query filters on ml.bank_id instead of mu.bank_id
2026-03-27 16:12:13 +01:00
Kagura 6e90df9818 fix(docker): add graceful shutdown handler to prevent pg0 data loss on restart (#698)
* fix(docker): add graceful shutdown handler to prevent pg0 data loss on restart (#675)

- Trap SIGTERM/SIGINT in start-all.sh to forward signals to child processes
- pg0 (embedded PostgreSQL) now gets a clean shutdown with WAL flush
- 30-second timeout before force-killing unresponsive processes
- Add startup data integrity check: warn if pg0 data dir exists but PG_VERSION missing
- Improve wait loop robustness: trigger cleanup when any child exits unexpectedly

Fixes #675

* fix: address review feedback — re-entrant guard, timeout docs, cleaner glob

- Add SHUTTING_DOWN guard to prevent concurrent cleanup runs
- Document Docker stop_grace_period mismatch (30s cleanup vs 10s default)
- Replace find subprocess with compgen glob for PG_VERSION check
- Add comment explaining wait -n && true idiom
2026-03-27 16:03:10 +01:00
Chris BartholomewandNicolò Boschi dffb87080f fix(migrations): bypass PgBouncer for advisory locks via MIGRATION_DATABASE_URL (#726)
* fix(migrations): use HINDSIGHT_API_MIGRATION_DATABASE_URL when set

Session-level advisory locks are broken when the database URL goes
through PgBouncer in transaction mode: the backend connection is
returned to the pool on COMMIT, orphaning the lock, so multiple pods
can simultaneously run migrations for the same schema.

When HINDSIGHT_API_MIGRATION_DATABASE_URL is set, use it for both
the advisory lock connection and the Alembic run.  Callers should
point this at the direct PostgreSQL endpoint (bypassing the pooler)
so the session-level lock is held for the full migration duration.

* refactor(migrations): move MIGRATION_DATABASE_URL to standard config

Wire HINDSIGHT_API_MIGRATION_DATABASE_URL through HindsightConfig
instead of reading os.getenv() directly in migrations.py. Add the
field to the dataclass, from_env(), log_config(), all call sites,
.env.example, and the configuration docs page.

* fix: update test mocks for migration_database_url kwarg and regenerate docs skill

---------

Co-authored-by: Nicolò Boschi <[email protected]>
2026-03-27 16:01:38 +01:00
Nicolò Boschi 1cac35728f fix: silence noisy google_genai.models INFO logging (#727)
* fix: silence noisy google_genai.models INFO logging

The google-genai SDK logs "AFC is enabled with max remote calls: 10"
at INFO level on every initialization. Set its logger to WARNING.

* fix: regenerate docs skill in release-integration script

The release script generates changelog/SDK pages but never re-ran
generate-docs-skill.sh, causing CI to fail with out-of-sync skill
files after every integration release. Now it regenerates the skill
and includes the output in the release commit.

Also adds the missing ag2 skill files from the latest release.
2026-03-27 16:01:23 +01:00
Chris Bartholomew 26e6877b53 fix(migration): use IF EXISTS when dropping chunk FK constraint (#725)
* fix(migration): use IF EXISTS when dropping chunk FK constraint

The migration unconditionally dropped memory_units_chunk_fkey, but
depending on the order in which migrations were applied the constraint
may not exist. Use raw SQL with IF EXISTS so the drop is safe regardless.

* fix(migration): make chunk FK add idempotent with DO block

The previous fix only handled the DROP side with IF EXISTS. The ADD side
could still fail with DuplicateObject when the FK already existed on a
schema that was provisioned after the base migration ran.

Wrap the ADD CONSTRAINT in a DO block to catch duplicate_object and
continue, making the migration fully idempotent in both directions.
2026-03-27 14:56:43 +01:00
1ac80bda6f fix(codex): resolve JSON serialization and logging exception propagation in codex_llm (#724)
Port fixes from #461 (claude_code_llm) to codex_llm:
- Replace json.dumps(result) with result.model_dump_json() for Pydantic models to fix TypeError during consolidation
- Wrap record_llm_call tracing block in try/except so logging failures never propagate to retry handler

Co-authored-by: Marco Rutsch <[email protected]>
Co-authored-by: Claude Sonnet 4.6 <[email protected]>
2026-03-27 13:37:30 +01:00
Nicolò Boschi 3c78b717b0 docs: add AG2 integration page (#723)
- Add AG2 integration doc with quick start, configuration, GroupChat example, and API reference
- Add to sidebar, versioned sidebar, and integrations hub
- Add AG2 icon
2026-03-27 10:58:42 +01:00
Nicolò Boschi 9321c59bf1 release(ag2): v0.1.1 2026-03-27 10:12:30 +01:00
Nicolò Boschi 696d99ca1e chore(dev): add ag2 package name mapping for changelog generator 2026-03-27 10:12:14 +01:00
Nicolò Boschi 4b584e4d0b chore(dev): add ag2 to changelog generator valid integrations 2026-03-27 10:11:21 +01:00
Nicolò Boschi e5c7e166c5 fix(ag2): code cleanup and CI/release integration (#721)
- Remove unnecessary `pass` in HindsightError
- Add `Callable` return type annotations to create/register functions
- Use lazy logger formatting instead of f-strings
- Add test-ag2-integration CI job in test.yml
- Add ag2 to release-integration.sh valid integrations
2026-03-27 10:10:14 +01:00
Nicolò Boschi 083295dc6f feat: add audit log for feature usage tracking (#717)
* feat: add audit log for feature usage tracking

Add full auditability for all mutating and core API operations across
HTTP, MCP, and system (worker) transports. Audit entries record raw
request/response as JSONB, timing (started_at/ended_at), action, and
transport type.

Backend:
- New audit_log table with JSONB columns for expandability without
  future migrations (merge migration of 3 existing heads)
- AuditLogger with fire-and-forget writes via asyncio.create_task
- @audited decorator on 28 HTTP route handlers
- MCP tool audit wrapping for 16 auditable tools
- Worker task execution wrapped with audit_context
- List endpoint with action, transport, date range filters + pagination
- Stats endpoint with per-day counts for charting
- Configurable retention sweep (concurrent-safe DELETE)

Config (env-only, static):
- HINDSIGHT_API_AUDIT_LOG_ENABLED (default: false)
- HINDSIGHT_API_AUDIT_LOG_ACTIONS (comma-separated allowlist, empty=all)
- HINDSIGHT_API_AUDIT_LOG_RETENTION_DAYS (default: -1, keep forever)

Control Plane:
- New "Audit Logs" tab on bank configuration page
- Line chart showing request volume (today/7d/30d) with action filter
- Filterable table with action, transport, date range filters
- Paginated list with detail dialog showing raw request/response JSON

Tests:
- 13 tests covering list, filters, pagination, stats, disabled mode,
  action allowlist, and ordering

* fix: split 3-way merge migration into two 2-way merges

Alembic doesn't support 3-parent merge migrations. Split into a no-op
merge of 2 heads (b1c2d3e4f5g6) followed by the audit_log table
migration merging the third head.

* fix: correct merge migration to merge actual 2 heads

The original analysis incorrectly identified 3 heads. There were only 2
(a3b4c5d6e7f8 and c8e5f2a3b4d1). Remove the unnecessary intermediate
merge migration and fix the audit_log migration to merge these 2 heads.

* fix: use 'heads' instead of 'head' in migration runner

Alembic's upgrade('head') fails when multiple heads exist (e.g. from
namespace package overlaps between hindsight-api and hindsight-api-slim).
Using 'heads' (plural) handles this gracefully by upgrading all branches.

* chore: regenerate OpenAPI spec with audit log endpoints

* chore: regenerate TypeScript client and docs skill OpenAPI spec

Python and Go clients still need regeneration (requires Docker).

* chore: regenerate all client SDKs (Python, Go, TypeScript)

Adds generated audit log API clients for Python (audit_api.py),
Go (api_audit.go), and TypeScript client type updates.
2026-03-27 09:52:03 +01:00
Faridun Mirzoev 731238707d feat(integrations): add AG2 framework integration (#720)
Add hindsight-ag2 package providing persistent memory tools for AG2 agents via retain/recall/reflect operations.
2026-03-27 09:41:37 +01:00
Ben 62c0992075 Teaching the Llama to Remember (#707)
Llama index blog
2026-03-26 15:44:18 -04:00
Nicolò Boschi 02b0f7799d docs: add Volcano Engine as supported LLM provider (#715)
* docs: add Volcano Engine as supported LLM provider

Follow-up to #714. Add Volcano Engine (ByteDance) to the documentation:
- LLM providers grid component
- Provider list in configuration docs
- Provider example with base URL and default model
- Default model table in models page

* chore: regenerate docs skill references
2026-03-26 18:21:34 +01:00
Nicolò Boschi 7c18723fd9 fix(python-client): expose all configurable fields in update_bank_config() (#712)
Add 10 missing bank-configurable fields to update_bank_config():
- entity_labels, entities_allow_free_form
- consolidation_llm_batch_size, consolidation_source_facts_max_tokens,
  consolidation_source_facts_max_tokens_per_observation
- retain_default_strategy, retain_strategies
- reflect_source_facts_max_tokens
- mcp_enabled_tools
- llm_gemini_safety_settings

Previously these could only be set via raw PATCH to /config.
All new params are keyword-only with None defaults (backwards compatible).
2026-03-26 17:20:10 +01:00
shun yiandyishun.eason 417fac61e4 feat: add support for ark and volcano LLM providers (#714)
- Add 'ark' and 'volcano' as valid LLM providers (both are aliases for Volcano Engine)
- Set default model to 'doubao-pro-32k' for both providers
- Add them to OpenAICompatibleLLM provider list
- Exclude from json_object response format support

Co-authored-by: yishun.eason <[email protected]>
2026-03-26 17:14:13 +01:00
Nicolò Boschi 105cdf1fbf fix(python-client): expose all configurable fields in update_bank_config() (#712)
Add 10 missing bank-configurable fields to update_bank_config():
- entity_labels, entities_allow_free_form
- consolidation_llm_batch_size, consolidation_source_facts_max_tokens,
  consolidation_source_facts_max_tokens_per_observation
- retain_default_strategy, retain_strategies
- reflect_source_facts_max_tokens
- mcp_enabled_tools
- llm_gemini_safety_settings

Previously these could only be set via raw PATCH to /config.
All new params are keyword-only with None defaults (backwards compatible).
2026-03-26 16:29:09 +01:00
Nicolò Boschi a0cea84d82 docs(python-client): async-first pydoc + low-level API access + missing params (#711)
* docs(python-client): improve pydoc strings for async-first usage and low-level API access

- Class docstring now clearly documents async-first pattern: a* methods
  preferred, sync wrappers for scripts/REPLs only
- Every sync method docstring points to its async counterpart
- Every async method docstring says "preferred"
- Expose 10 low-level API properties (documents, entities, operations,
  webhooks, monitoring, etc.) so agents/users can discover the full API
  surface without guessing at _-prefixed internals
- Add missing API parameters: tag_groups (recall/reflect), fact_types,
  exclude_mental_models, exclude_mental_model_ids (reflect),
  observation_scopes/strategy (retain items), background (create_bank)
- Fix areflect missing include_facts param that sync reflect already had
- Sync recall/reflect now delegate to async counterparts (no logic duplication)

* style(retain): format long function call arguments one-per-line
2026-03-26 16:09:58 +01:00
Nicolò Boschi 200bab233e feat(openclaw): add recallInjectionPosition config to preserve prompt cache (#710)
* feat(openclaw): add recallInjectionPosition config to preserve prompt cache

Add configurable injection position for recalled memories to avoid
breaking prefix-based prompt caching (Anthropic/Google) when agents
have large static system prompts.

Options: 'prepend' (default, current behavior), 'append' (end of
system prompt, preserves cache), 'user' (before user message).

Closes #703

* docs(openclaw): document all plugin config flags

Add missing config options to the OpenClaw docs: recallTopK,
recallTypes, recallContextTurns, recallMaxQueryChars,
recallPromptPreamble, recallInjectionPosition, recallRoles,
retainEveryNTurns, retainOverlapTurns, and debug.
2026-03-26 16:09:25 +01:00
Nicolò Boschi c9ff37dcbf fix(python-client): async=true silently ignored on retain (#709)
* docs(claude-code): tidy configuration reference and sync README

Add missing settings (retainMode, retainToolCalls, retainTags,
retainMetadata, embedPackagePath, llmApiKeyEnv, agentName, and
several recall options) that existed in code but not in docs.
Restructure config tables with prose introductions, clearer
descriptions, and consistent layout across both files.

* refactor(claude-code): remove recallTopK setting

Unused client-side cap — Hindsight server already controls result
count via recallBudget and recallMaxTokens.

* fix(python-client): async=true was silently ignored on retain calls

The hand-written client wrapper passed `async_=retain_async` to
RetainRequest, but the generated Pydantic model uses `var_async` as the
Python field name (with `alias="async"`). The `async_` kwarg didn't
match either the field name or the alias, so Pydantic silently ignored
it — every retain call ran synchronously regardless of the flag.

This has been broken since the client was first introduced (6073ac4f),
not a regression.

Also adds unit tests that verify the async field serializes correctly
in the request JSON, preventing future regressions.
2026-03-26 15:21:43 +01:00
Nicolò Boschi 91397190c0 docs(claude-code): tidy configuration reference and sync README (#706)
* docs(claude-code): tidy configuration reference and sync README

Add missing settings (retainMode, retainToolCalls, retainTags,
retainMetadata, embedPackagePath, llmApiKeyEnv, agentName, and
several recall options) that existed in code but not in docs.
Restructure config tables with prose introductions, clearer
descriptions, and consistent layout across both files.

* refactor(claude-code): remove recallTopK setting

Unused client-side cap — Hindsight server already controls result
count via recallBudget and recallMaxTokens.
2026-03-26 14:07:35 +01:00
Nicolò Boschi fd88c0efa5 feat(retain): delta retain — skip LLM for unchanged chunks on upsert (#701)
* feat(retain): delta retain — skip LLM re-extraction for unchanged chunks on upsert

When upserting a document (same document_id), instead of deleting all
facts and re-extracting from scratch, compare chunk content hashes
and only process changed/new chunks. Unchanged chunks keep their
existing facts, entities, and links.

- Add content_hash column to chunks table (migration b3c4d5e6f7a8)
- Add chunk delta comparison functions in chunk_storage.py
- Add delta_mode to fact_storage.handle_document_tracking (skip full delete)
- Add update_memory_units_tags for propagating tag changes to existing facts
- Refactor orchestrator into _try_delta_retain and _full_retain paths
- Automatic fallback to full retain for pre-migration data or all-changed scenarios
- Fix ty type error in metrics.py (resource module import on Windows)
- 16 new tests covering entities, links, tags, metadata, edge cases

* refactor(retain): deduplicate delta and full retain paths

Extract shared _insert_facts_and_links() and _extract_and_embed()
functions used by both the full retain and delta retain paths.
Remove delta_mode flag from handle_document_tracking — delta path
uses dedicated upsert_document_metadata() instead.

* chore: regenerate clients, openapi spec, and lockfile

* chore: regenerate docs skill
2026-03-26 13:50:55 +01:00
Nicolò Boschi ea4df8dbb5 fix: resolve remaining Dependabot security alerts (#705)
- python-multipart: pin >=0.0.22 (arbitrary file write via non-default config)
- requests: pin >=2.33.0 in litellm, langgraph, crewai integrations (insecure temp file reuse)

Remaining unfixable alerts: diskcache (<=5.6.3, no patch) and Pygments (<=2.19.2, no patch).
2026-03-26 13:43:27 +01:00
662 changed files with 69812 additions and 19906 deletions
+196
View File
@@ -0,0 +1,196 @@
---
name: code-review
description: Review changed code against project standards. Checks for missing tests, dead code, type safety, lint issues, and coding conventions. Run after completing any implementation work.
user_invocable: true
---
# Code Review
Review all changed code against the project's quality standards and coding conventions.
## Code Standards
Read and internalize these standards before writing code. The review steps below verify compliance.
### Python Style
- Python 3.11+, type hints required
- Async throughout (asyncpg, async FastAPI)
- Pydantic models for request/response
- Ruff for linting (line-length 120)
- No Python files at project root - maintain clean directory structure
- **Never use multi-item tuple return values** — not even for internal/private functions. Always use a dataclass or Pydantic model. No exceptions, no "it's just two values" shortcuts. If a function returns more than one value, define a named type for it.
### Type Safety with Pydantic Models
**NEVER use raw `dict` types for structured data** — this applies to all code, including internal helpers and private functions. If the dict has known keys, it must be a dataclass or Pydantic model:
- Use Pydantic `BaseModel` for all data structures passed between functions
- Use `@dataclass` for lightweight internal data containers when Pydantic validation isn't needed
- Add `@field_validator` for type coercion (e.g., ensuring datetimes are timezone-aware)
- Avoid `dict.get()` patterns - use typed model attributes instead
- Parse external data (JSON, API responses) into Pydantic models at the boundary
- This catches type errors at parse time, not deep in business logic
- The only acceptable `dict` usage is for truly dynamic/unknown keys (e.g., arbitrary metadata, JSON blobs with no fixed schema)
```python
# BAD - error-prone dict access
def process(data: dict) -> str:
return data.get("name", "") # No validation, silent failures
# GOOD - typed and validated
class UserData(BaseModel):
name: str
created_at: datetime
@field_validator("created_at", mode="before")
@classmethod
def ensure_tz_aware(cls, v):
if isinstance(v, str):
v = datetime.fromisoformat(v.replace("Z", "+00:00"))
if v.tzinfo is None:
return v.replace(tzinfo=timezone.utc)
return v
def process(data: UserData) -> str:
return data.name # Type-safe, validated at construction
```
### TypeScript Style
- Next.js App Router for control plane
- Tailwind CSS with shadcn/ui components
### Code Comments
- **Always comment non-trivial technical decisions** with the reasoning behind the choice. If someone would ask "why is it done this way?", there should be a comment.
- **Keep comments up to date with history** — when changing an approach, update the comment to explain what was tried before and why it was changed. Comments serve as a tracker of previous implementations that likely had problems.
- Don't comment obvious code — only where the "why" isn't self-evident from the code itself.
```python
# BAD - no context for future readers
results = await asyncio.gather(*tasks, return_exceptions=True)
# GOOD - explains the non-obvious choice
# Use return_exceptions=True to avoid cancelling sibling tasks on failure.
# Previously we used TaskGroup but it cancelled all tasks when one failed,
# causing partial writes that left orphaned entity links (see #412).
results = await asyncio.gather(*tasks, return_exceptions=True)
```
### Branch Hygiene
- **Always start new feature branches from `origin/main`** — rebase to ensure a clean base.
- **Only include commits relevant to the PR/branch/feature** — no unrelated changes. If the branch contains commits that don't belong, they must be removed before merging.
### General Principles
- Don't add features, refactor code, or make "improvements" beyond what was asked
- Don't add unnecessary error handling for impossible scenarios
- Don't create helpers or abstractions for one-time operations
- No backwards-compatibility hacks (unused vars, re-exports, "removed" comments)
- Three similar lines of code is better than a premature abstraction
## Review Steps
### 1. Check branch hygiene
- Run `git log --oneline main..HEAD` to list all commits on the branch.
- Verify every commit is relevant to the feature/PR. Flag any unrelated commits.
- Check the branch is based on a recent `origin/main` (no stale base).
### 2. Identify changed files
Run `git diff --name-only HEAD` (unstaged) and `git diff --cached --name-only` (staged) to get all changed files. If there are no local changes, diff against the base branch using `git diff main...HEAD --name-only` and `git diff main...HEAD` to review all commits on the current branch.
### 3. Run linters
```bash
./scripts/hooks/lint.sh
```
Report any failures. Do NOT fix them yourself — just report.
### 4. Check for dead code
For each changed Python file, check for:
- Unused imports (Ruff should catch these, but verify)
- Functions/methods/classes that were added but are never called from anywhere
- Variables assigned but never read
- Commented-out code blocks that should be removed
For each changed TypeScript file, check for:
- Unused imports
- Unused variables or functions
- Commented-out code
### 5. Check type safety (Python)
For each changed Python file, check for violations:
- **No raw `dict` for structured data** — must use Pydantic model or dataclass, even for internal/private functions (only exception: truly dynamic/unknown keys)
- **No multi-item tuple returns** — must use dataclass or Pydantic model, even for internal/private functions (no exceptions)
- **Missing type hints** on function parameters and return types
- **Missing `@field_validator`** for datetime fields that should be timezone-aware
### 6. Check for missing tests
For each new or significantly changed function/endpoint/class:
- Check if there is a corresponding test addition or update
- New API endpoints MUST have integration tests
- New utility functions MUST have unit tests
- Bug fixes SHOULD have a regression test
Flag any new logic that lacks test coverage.
### 7. Check API consistency
If any files in `hindsight-api-slim/hindsight_api/api/` were changed:
- Were the OpenAPI specs regenerated? (`./scripts/generate-openapi.sh`)
- Were the client SDKs regenerated? (`./scripts/generate-clients.sh`)
- Were the control plane proxy routes updated? (`hindsight-control-plane/src/app/api/`)
### 8. Check code comments
For each non-trivial change:
- **New non-obvious logic** — is there a comment explaining the reasoning?
- **Changed approach** — does the comment include what was done before and why it changed?
- **Stale comments** — do existing comments near the changed code still accurately describe the behavior?
### 9. Check integration completeness
If any files in `hindsight-integrations/` were added or changed, verify:
- **Tests exist** — the integration must have tests that simulate/exercise the external framework (not just pure unit tests of helpers). Check for a `tests/` directory with meaningful test files.
- **CI job exists** — check `.github/workflows/test.yml` for a corresponding `test-<name>-integration` job. If missing, flag it.
- **Release process** — check that the integration name is in the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh`. If missing, flag it.
- **Code standards** — the integration code must follow all Python style rules (type hints, no raw dicts, no tuple returns, etc.).
### 10. Review against other coding standards
Check the diff for violations of the standards listed above:
- Python files at project root (not allowed)
- Missing async patterns (should be async throughout)
- Pydantic models for request/response
- Line length > 120 chars
- New features/code beyond what was asked (over-engineering)
- Unnecessary error handling for impossible scenarios
- Premature abstractions or speculative helpers
- Backwards-compatibility hacks (unused vars, re-exports, "removed" comments)
### 11. Report findings
Present a clear summary organized by severity:
**Must fix** — issues that will break CI or violate hard project rules:
- Unrelated commits on the branch
- Lint failures
- Missing type hints on public functions
- Raw dict usage for structured data (including internal code)
- Multi-item tuple returns (including internal code)
- Missing tests for new endpoints
- New integration missing tests, CI job, or release-integration.sh entry
**Should fix** — issues that hurt code quality:
- Dead code / unused imports missed by linter
- Missing tests for non-trivial utility functions
- Over-engineering beyond the task scope
**Note** — observations that may or may not need action:
- API changes that might need client regeneration
- Patterns that deviate from nearby code style
For each finding, include the file path, line number, and a brief explanation.
Do NOT auto-fix any issues. Report all findings and let the user decide what to address. If there are no findings, confirm the code looks good.
+2 -1
View File
@@ -2,7 +2,7 @@
# Copy this file to .env and fill in your values
# LLM Configuration (Required)
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, volcano
HINDSIGHT_API_LLM_PROVIDER=openai
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
@@ -44,6 +44,7 @@ HINDSIGHT_API_LOG_LEVEL=info
# Database (Optional - uses embedded pg0 by default)
# HINDSIGHT_API_DATABASE_URL=postgresql://user:pass@host:5432/db
# HINDSIGHT_API_MIGRATION_DATABASE_URL= # Direct PostgreSQL URL for migrations (bypasses PgBouncer). Falls back to DATABASE_URL.
# HINDSIGHT_API_DATABASE_SCHEMA=public # PostgreSQL schema name (default: public)
# Vector Extension (Optional - uses pgvector by default)
+1 -1
View File
@@ -44,5 +44,5 @@ jobs:
runs-on: ubuntu-latest
needs: build
steps:
- uses: actions/deploy-pages@v4
- uses: actions/deploy-pages@v5
id: deployment
+1 -1
View File
@@ -382,7 +382,7 @@ jobs:
- uses: actions/checkout@v6
- name: Install Helm
uses: azure/setup-helm@v4
uses: azure/setup-helm@v5
with:
version: 'latest'
File diff suppressed because it is too large Load Diff
+2 -1
View File
@@ -50,7 +50,8 @@ hindsight-dev/benchmarks/perf/results/
benchmarks/results/
hindsight-cli/target
hindsight-clients/rust/target
.claude
.claude/*
!.claude/skills/
whats-next.md
TASK.md
# Changelog is now tracked in hindsight-docs/src/pages/changelog.md
+32 -53
View File
@@ -11,9 +11,15 @@ Hindsight is an agent memory system that provides long-term memory for AI agents
## Development Commands
### Local Development (API + UI)
```bash
# Start both API server and control plane UI
./scripts/dev/start.sh
```
### API Server (Python/FastAPI)
```bash
# Start API server (loads .env automatically)
# Start API server only (loads .env automatically)
./scripts/dev/start-api.sh
# Run all tests (parallelized with pytest-xdist)
@@ -73,17 +79,16 @@ cd hindsight-control-plane && npm run dev
### Monorepo Structure
- **hindsight-api-slim/**: Core FastAPI server with memory engine (Python, uv)
- **hindsight/**: Embedded Python bundle (hindsight-all package)
- **hindsight-control-plane/**: Admin UI (Next.js, npm)
- **hindsight-cli/**: CLI tool (Rust, cargo, uses progenitor for API client)
- **hindsight-clients/**: Generated SDK clients (Python, TypeScript, Rust)
- **hindsight-docs/**: Docusaurus documentation site
- **hindsight-integrations/**: Framework integrations (LiteLLM, OpenAI)
- **hindsight-integrations/**: Framework integrations (LiteLLM, CrewAI, LangGraph, Pydantic AI, AG2, Claude Code, etc.)
- **hindsight-dev/**: Development tools and benchmarks
### Core Engine (hindsight-api-slim/hindsight_api/engine/)
- `memory_engine.py`: Main orchestrator (~170KB) for retain/recall/reflect operations
- `llm_wrapper.py`: LLM abstraction supporting OpenAI, Anthropic, Gemini, Groq, MiniMax, Ollama, LM Studio
- `memory_engine.py`: Main orchestrator for retain/recall/reflect operations
- `llm_wrapper.py`: LLM abstraction supporting OpenAI, Anthropic, Gemini, VertexAI, Groq, MiniMax, Ollama, LM Studio, LiteLLM, Claude Code
- `embeddings.py`: Embedding generation (local sentence-transformers or TEI)
- `cross_encoder.py`: Reranking (local or TEI)
- `entity_resolver.py`: Entity extraction and normalization
@@ -96,13 +101,13 @@ cd hindsight-control-plane && npm run dev
**search/**: Multi-strategy retrieval
- `retrieval.py`: Main retrieval orchestrator
- `graph_retrieval.py`: Entity/relationship graph traversal
- `mpfp_retrieval.py`: Multi-Path Fact Propagation retrieval
- `graph_retrieval.py`: Graph retrieval abstract base class
- `link_expansion_retrieval.py`: Link expansion graph retrieval
- `fusion.py`: Reciprocal rank fusion for combining results
- `reranking.py`: Cross-encoder reranking
### API Layer (hindsight-api-slim/hindsight_api/api/)
- `http.py`: FastAPI HTTP routers (~80KB) for all REST endpoints
- `http.py`: FastAPI HTTP routers for all REST endpoints
- `mcp.py`: Model Context Protocol server implementation
Main operations:
@@ -164,11 +169,17 @@ Key tables: `banks`, `memory_units`, `documents`, `entities`, `entity_links`
## Key Conventions
### Code Quality
**Before writing code, read `.claude/skills/code-review/SKILL.md`** for the full coding standards (Python style, type safety, TypeScript style, general principles).
**Always run the lint script after making Python or TypeScript/Node changes:**
```bash
./scripts/hooks/lint.sh
```
This runs the same checks as the pre-commit hook (Ruff for Python, ESLint/Prettier for TypeScript).
**After completing any implementation work, run `/code-review`** to verify your changes against project standards (missing tests, dead code, type safety, etc.). Fix any "must fix" issues before considering the task done.
**MANDATORY: Run `/code-review` before pushing code or creating a pull request.** Do not push or create a PR until all "must fix" issues are resolved.
### Memory Banks
- Each bank is an isolated memory store (like a "brain" for one user/agent)
@@ -200,48 +211,16 @@ When adding or modifying parameters in the dataplane API (hindsight-api), you mu
- Update the client type definition in `lib/api.ts`
- Update any UI components that need to use the new parameter
### Python Style
- Python 3.11+, type hints required
- Async throughout (asyncpg, async FastAPI)
- Pydantic models for request/response
- Ruff for linting (line-length 120)
- No Python files at project root - maintain clean directory structure
- **Never use multi-item tuple return values** - prefer dataclass or Pydantic model for structured returns
### Adding New Integrations
### Type Safety with Pydantic Models
**NEVER use raw `dict` types for structured data.** Always use Pydantic models:
- Use Pydantic `BaseModel` for all data structures passed between functions
- Add `@field_validator` for type coercion (e.g., ensuring datetimes are timezone-aware)
- Avoid `dict.get()` patterns - use typed model attributes instead
- Parse external data (JSON, API responses) into Pydantic models at the boundary
- This catches type errors at parse time, not deep in business logic
Every new integration in `hindsight-integrations/` must satisfy all of the following before it can be merged:
```python
# BAD - error-prone dict access
def process(data: dict) -> str:
return data.get("name", "") # No validation, silent failures
1. **Tests are required** — tests must simulate or exercise the external system (mock the framework's interfaces and verify the integration actually calls Hindsight correctly). Pure unit tests of helper functions are not sufficient.
2. **CI job** — add a test job in `.github/workflows/test.yml` following the existing pattern (e.g., `test-crewai-integration`). The job must build, install deps, and run `uv run pytest tests -v`. Also add the integration to `detect-changes` outputs so it only runs when its files change.
3. **Release process** — add the integration name to the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh` so it can be released via the standard release workflow.
4. **Follow project code standards** — Python style, type safety, no raw dicts for structured data, no multi-item tuple returns (see `.claude/skills/code-review/SKILL.md`).
# GOOD - typed and validated
class UserData(BaseModel):
name: str
created_at: datetime
@field_validator("created_at", mode="before")
@classmethod
def ensure_tz_aware(cls, v):
if isinstance(v, str):
v = datetime.fromisoformat(v.replace("Z", "+00:00"))
if v.tzinfo is None:
return v.replace(tzinfo=timezone.utc)
return v
def process(data: UserData) -> str:
return data.name # Type-safe, validated at construction
```
### TypeScript Style
- Next.js App Router for control plane
- Tailwind CSS with shadcn/ui components
If any of these are missing, the integration is incomplete and must not be pushed or merged.
### Adding New API Configuration Flags
@@ -255,17 +234,17 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
- Add `ENV_*` constant for the environment variable name (e.g., `ENV_MY_SETTING = "HINDSIGHT_API_MY_SETTING"`)
- Add `DEFAULT_*` constant for the default value
- Add field to `HindsightConfig` dataclass with type annotation
- **Mark as hierarchical or static** by adding to `_HIERARCHICAL_FIELDS` set (hierarchical) or leaving it out (static)
- **Mark as configurable** by adding to `_CONFIGURABLE_FIELDS` set if the field should be overridable per-tenant/bank via API
- Add initialization in `from_env()` method
```python
# Hierarchical field (can be overridden per-bank)
_HIERARCHICAL_FIELDS = {
# Configurable field (can be overridden per-tenant/bank via API)
_CONFIGURABLE_FIELDS = {
...,
"my_setting", # Add here for hierarchical
"my_setting", # Add here for configurable
}
# Static field - just don't add to _HIERARCHICAL_FIELDS
# Static field - just don't add to _CONFIGURABLE_FIELDS
```
2. **main.py** (`hindsight-api-slim/hindsight_api/main.py`):
+97 -5
View File
@@ -1,6 +1,28 @@
#!/bin/bash
set -e
# =============================================================================
# Embedded pg0 data integrity check (#675)
#
# When using embedded pg0, check if the data directory has existing PostgreSQL
# data before starting. If the directory exists but appears empty/corrupt
# (e.g., missing PG_VERSION file), log a warning. This helps diagnose data
# loss scenarios where a container restart caused the data directory to be
# wiped despite a volume mount being present.
# =============================================================================
PG0_DATA_DIR="${HOME}/.pg0"
if [ -d "$PG0_DATA_DIR" ]; then
# Look for actual PostgreSQL data directories (pg0 creates subdirs per instance)
if compgen -G "$PG0_DATA_DIR"/*/PG_VERSION > /dev/null 2>&1; then
echo "✅ Existing pg0 data directory detected at $PG0_DATA_DIR"
elif [ "$(ls -A "$PG0_DATA_DIR" 2>/dev/null)" ]; then
echo "⚠️ WARNING: pg0 data directory exists at $PG0_DATA_DIR but no PG_VERSION found."
echo " This may indicate data corruption or an incomplete previous shutdown."
echo " If you see all migrations running from scratch after this, your data may have been lost."
echo " See: https://github.com/vectorize-io/hindsight/issues/675"
fi
fi
# Service flags (default to true if not set)
ENABLE_API="${HINDSIGHT_ENABLE_API:-true}"
ENABLE_CP="${HINDSIGHT_ENABLE_CP:-true}"
@@ -71,6 +93,63 @@ if [ "${HINDSIGHT_WAIT_FOR_DEPS:-false}" = "true" ]; then
done
fi
# =============================================================================
# Graceful shutdown handler (#675)
#
# Docker sends SIGTERM on `docker stop`/`docker restart`. Without a trap, child
# processes (hindsight-api + pg0, control-plane) are killed abruptly. For the
# embedded pg0 database this can cause data loss when the data directory is on
# a Docker volume that gets remounted after restart.
#
# The trap forwards SIGTERM to all tracked child PIDs so that:
# - hindsight-api receives the signal and can run its shutdown hooks
# - pg0 gets a clean PostgreSQL shutdown (checkpoint + WAL flush)
# - The control-plane Node.js process exits cleanly
# =============================================================================
# Guard against concurrent cleanup (e.g., child crash + SIGTERM arriving together)
SHUTTING_DOWN=false
cleanup() {
if $SHUTTING_DOWN; then return; fi
SHUTTING_DOWN=true
echo ""
echo "🛑 Received shutdown signal, stopping services gracefully..."
for pid in "${PIDS[@]}"; do
if kill -0 "$pid" 2>/dev/null; then
kill -TERM "$pid" 2>/dev/null
fi
done
# Give processes time to shut down cleanly (pg0 needs to flush WAL).
# NOTE: Docker's default stop_grace_period is 10s. If you use the default,
# either set stop_grace_period: 30s in your compose file / docker stop -t 30,
# or Docker will SIGKILL the container before this timeout expires.
local timeout=30
for ((i=1; i<=timeout; i++)); do
local all_stopped=true
for pid in "${PIDS[@]}"; do
if kill -0 "$pid" 2>/dev/null; then
all_stopped=false
break
fi
done
if $all_stopped; then
echo "✅ All services stopped cleanly"
exit 0
fi
sleep 1
done
# Force kill if still running after timeout
echo "⚠️ Timeout reached, forcing shutdown..."
for pid in "${PIDS[@]}"; do
if kill -0 "$pid" 2>/dev/null; then
kill -9 "$pid" 2>/dev/null
fi
done
exit 1
}
trap cleanup SIGTERM SIGINT
# Track PIDs for wait
PIDS=()
@@ -138,8 +217,21 @@ if [ ${#PIDS[@]} -eq 0 ]; then
exit 1
fi
# Wait for any process to exit
wait -n
# Exit with status of first exited process
exit $?
# Wait for any process to exit (use wait -n with trap-safe loop)
while true; do
# wait -n returns when any child exits; it also returns on signal delivery
# (the trap handler will run and exit, so this loop is just for robustness).
# `&& true` prevents `set -e` from killing the script when wait -n returns
# non-zero (child exited with error or no backgrounded children remain).
wait -n && true
# Check if any tracked PID has exited
for pid in "${PIDS[@]}"; do
if ! kill -0 "$pid" 2>/dev/null; then
wait "$pid" 2>/dev/null
exit_code=$?
echo "⚠️ Service (PID $pid) exited with code $exit_code"
# Trigger cleanup for remaining services
cleanup
fi
done
done
+2 -2
View File
@@ -2,8 +2,8 @@ apiVersion: v2
name: hindsight
description: Hindsight helm chart
type: application
version: 0.4.20
appVersion: "0.4.20"
version: 0.4.22
appVersion: "0.4.22"
keywords:
- ai
- memory
@@ -95,6 +95,27 @@ spec:
{{- toYaml .Values.api.readinessProbe | nindent 10 }}
resources:
{{- toYaml .Values.api.resources | nindent 10 }}
{{- if or .Values.api.persistence.modelCache.enabled .Values.api.extraVolumeMounts }}
volumeMounts:
{{- if .Values.api.persistence.modelCache.enabled }}
- name: model-cache
mountPath: /home/hindsight/.cache
{{- end }}
{{- with .Values.api.extraVolumeMounts }}
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}
{{- if or .Values.api.persistence.modelCache.enabled .Values.api.extraVolumes }}
volumes:
{{- if .Values.api.persistence.modelCache.enabled }}
- name: model-cache
persistentVolumeClaim:
claimName: {{ include "hindsight.fullname" . }}-api-model-cache
{{- end }}
{{- with .Values.api.extraVolumes }}
{{- toYaml . | nindent 6 }}
{{- end }}
{{- end }}
{{- with .Values.nodeSelector }}
nodeSelector:
{{- toYaml . | nindent 8 }}
@@ -0,0 +1,21 @@
{{- if and .Values.api.enabled .Values.api.persistence.modelCache.enabled }}
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
name: {{ include "hindsight.fullname" . }}-api-model-cache
labels:
{{- include "hindsight.api.labels" . | nindent 4 }}
{{- with .Values.api.persistence.modelCache.annotations }}
annotations:
{{- toYaml . | nindent 4 }}
{{- end }}
spec:
accessModes:
{{- toYaml .Values.api.persistence.modelCache.accessModes | nindent 4 }}
{{- if .Values.api.persistence.modelCache.storageClass }}
storageClassName: {{ .Values.api.persistence.modelCache.storageClass }}
{{- end }}
resources:
requests:
storage: {{ .Values.api.persistence.modelCache.size }}
{{- end }}
@@ -95,6 +95,16 @@ spec:
{{- toYaml .Values.worker.readinessProbe | nindent 10 }}
resources:
{{- toYaml .Values.worker.resources | nindent 10 }}
{{- if or .Values.worker.persistence.modelCache.enabled .Values.worker.extraVolumeMounts }}
volumeMounts:
{{- if .Values.worker.persistence.modelCache.enabled }}
- name: model-cache
mountPath: /home/hindsight/.cache
{{- end }}
{{- with .Values.worker.extraVolumeMounts }}
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}
{{- with .Values.nodeSelector }}
nodeSelector:
{{- toYaml . | nindent 8 }}
@@ -107,4 +117,26 @@ spec:
tolerations:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.worker.extraVolumes }}
volumes:
{{- toYaml . | nindent 6 }}
{{- end }}
{{- if .Values.worker.persistence.modelCache.enabled }}
volumeClaimTemplates:
- metadata:
name: model-cache
{{- with .Values.worker.persistence.modelCache.annotations }}
annotations:
{{- toYaml . | nindent 8 }}
{{- end }}
spec:
accessModes:
{{- toYaml .Values.worker.persistence.modelCache.accessModes | nindent 8 }}
{{- if .Values.worker.persistence.modelCache.storageClass }}
storageClassName: {{ .Values.worker.persistence.modelCache.storageClass }}
{{- end }}
resources:
requests:
storage: {{ .Values.worker.persistence.modelCache.size }}
{{- end }}
{{- end }}
+53
View File
@@ -67,6 +67,33 @@ api:
# Pod affinity/anti-affinity (overrides global affinity for this component)
# affinity: {}
# Persistent volume for local model cache (reranker, embeddings)
# Models are downloaded to /home/hindsight/.cache on first use.
# Without persistence, models are re-downloaded on every pod restart.
persistence:
modelCache:
enabled: false
size: 5Gi
storageClass: ""
accessModes:
- ReadWriteOnce
annotations: {}
# Extra volume mounts for the api container
# e.g.
# extraVolumeMounts:
# - name: my-volume
# mountPath: /mnt/my-volume
extraVolumeMounts: []
# Extra volumes for the api pod
# e.g.
# extraVolumes:
# - name: my-volume
# configMap:
# name: my-configmap
extraVolumes: []
# Environment variables
env:
#HINDSIGHT_API_LLM_PROVIDER: "groq"
@@ -140,6 +167,32 @@ worker:
# Pod affinity/anti-affinity (overrides global affinity for this component)
# affinity: {}
# Persistent volume for local model cache (reranker, embeddings)
# Uses volumeClaimTemplates since worker is a StatefulSet.
persistence:
modelCache:
enabled: false
size: 5Gi
storageClass: ""
accessModes:
- ReadWriteOnce
annotations: {}
# Extra volume mounts for the worker container
# e.g.
# extraVolumeMounts:
# - name: my-volume
# mountPath: /mnt/my-volume
extraVolumeMounts: []
# Extra volumes for the worker pod
# e.g.
# extraVolumes:
# - name: my-volume
# configMap:
# name: my-configmap
extraVolumes: []
# Secret environment variables (inherited from api.secrets if not specified)
secrets: {}
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "hindsight-all-slim"
version = "0.4.20"
version = "0.4.22"
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
readme = "README.md"
requires-python = ">=3.11"
+27 -37
View File
@@ -73,6 +73,9 @@ class HindsightEmbedded:
database_url: Optional database URL override (default: profile-specific pg0)
idle_timeout: Seconds before daemon auto-exits when idle (default: 300)
log_level: Daemon log level (default: "info")
ui: Whether to start the control plane web UI alongside the daemon (default: False)
ui_port: Port for the UI. Defaults to daemon_port + 10000.
ui_hostname: Hostname to bind the UI to. Defaults to "0.0.0.0".
"""
def __init__(
@@ -85,6 +88,9 @@ class HindsightEmbedded:
database_url: Optional[str] = None,
idle_timeout: int = 300,
log_level: str = "info",
ui: bool = False,
ui_port: Optional[int] = None,
ui_hostname: str = "0.0.0.0",
):
"""
Initialize the embedded client (daemon starts on first use).
@@ -98,6 +104,9 @@ class HindsightEmbedded:
database_url: Optional database URL override
idle_timeout: Seconds before daemon auto-exits when idle
log_level: Daemon log level
ui: Whether to start the control plane web UI alongside the daemon
ui_port: Port for the UI (defaults to daemon_port + 10000)
ui_hostname: Hostname to bind the UI to (defaults to "0.0.0.0")
"""
self.profile = profile
@@ -116,6 +125,10 @@ class HindsightEmbedded:
if database_url:
self.config["HINDSIGHT_EMBED_API_DATABASE_URL"] = database_url
self._ui = ui
self._ui_port = ui_port
self._ui_hostname = ui_hostname
self._client: Optional[Hindsight] = None
self._lock = threading.Lock()
self._started = False
@@ -157,6 +170,15 @@ class HindsightEmbedded:
self._started = True
logger.info(f"Connected to daemon at {daemon_url}")
# Start UI if requested
if self._ui:
logger.info(f"Starting UI for profile '{self.profile}'...")
ui_started = self._manager.start_ui(
self.profile, self._ui_port, self._ui_hostname
)
if not ui_started:
logger.warning(f"Failed to start UI for profile '{self.profile}'")
def _cleanup(self, stop_daemon_on_close: bool = False):
"""
Cleanup client resources (idempotent).
@@ -176,6 +198,11 @@ class HindsightEmbedded:
self._client.close()
self._client = None
# Stop UI if it was started
if self._ui and self._started:
logger.info(f"Stopping UI for profile '{self.profile}'...")
self._manager.stop_ui(self.profile, self._ui_port)
# Optionally stop daemon (daemon has idle timeout, so not required)
if stop_daemon_on_close and self._started:
logger.info(f"Stopping daemon for profile '{self.profile}'...")
@@ -379,43 +406,6 @@ class HindsightEmbedded:
"""Check if the client is initialized."""
return self._started and not self._closed and self._client is not None
def start_ui(self, ui_port: int | None = None, hostname: str = "0.0.0.0") -> bool:
"""Start the control plane web UI.
The daemon is started automatically if not already running.
Args:
ui_port: Port for the UI. Defaults to daemon_port + 10000.
hostname: Hostname to bind to. Defaults to 0.0.0.0.
Returns:
True if UI started successfully.
"""
self._ensure_started()
return self._manager.start_ui(self.profile, ui_port, hostname)
def stop_ui(self, ui_port: int | None = None) -> bool:
"""Stop the control plane web UI.
Args:
ui_port: Port the UI is running on. Defaults to daemon_port + 10000.
Returns:
True if stopped successfully.
"""
return self._manager.stop_ui(self.profile, ui_port)
def is_ui_running(self, ui_port: int | None = None) -> bool:
"""Check if the control plane web UI is running.
Args:
ui_port: Port to check. Defaults to daemon_port + 10000.
Returns:
True if UI is running and responsive.
"""
return self._manager.is_ui_running(self.profile, ui_port)
@property
def ui_url(self) -> str:
"""Get the UI URL for this profile."""
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "hindsight-all"
version = "0.4.20"
version = "0.4.22"
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
readme = "README.md"
requires-python = ">=3.11"
+82 -17
View File
@@ -15,6 +15,8 @@ import os
import uuid
import pytest
import urllib.request
import json
from hindsight import HindsightEmbedded
@@ -23,12 +25,20 @@ from hindsight import HindsightEmbedded
def llm_config():
"""Get LLM configuration from environment (session-scoped)."""
# Try both naming conventions
provider = os.getenv("HINDSIGHT_API_LLM_PROVIDER") or os.getenv("HINDSIGHT_LLM_PROVIDER", "groq")
api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY") or os.getenv("HINDSIGHT_LLM_API_KEY", "")
model = os.getenv("HINDSIGHT_API_LLM_MODEL") or os.getenv("HINDSIGHT_LLM_MODEL", "openai/gpt-oss-120b")
provider = os.getenv("HINDSIGHT_API_LLM_PROVIDER") or os.getenv(
"HINDSIGHT_LLM_PROVIDER", "groq"
)
api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY") or os.getenv(
"HINDSIGHT_LLM_API_KEY", ""
)
model = os.getenv("HINDSIGHT_API_LLM_MODEL") or os.getenv(
"HINDSIGHT_LLM_MODEL", "openai/gpt-oss-120b"
)
if not api_key:
pytest.skip("LLM API key not configured. Set HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_LLM_API_KEY.")
pytest.skip(
"LLM API key not configured. Set HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_LLM_API_KEY."
)
return {
"llm_provider": provider,
@@ -78,7 +88,9 @@ def test_embedded_context_manager(llm_config):
# Recall memory
recall_results = client.recall(bank_id=bank_id, query="context")
assert isinstance(recall_results.results, list), "Recall should return results list"
assert isinstance(recall_results.results, list), (
"Recall should return results list"
)
# Server should be stopped after context exit
# Note: We can't check client.is_running here as client is out of scope
@@ -105,7 +117,9 @@ def test_embedded_complete_workflow(llm_config):
# Step 1: Create a memory bank
print(f"\n1. Creating memory bank: {bank_id}")
bank_response = client.create_bank(
bank_id=bank_id, name="Test Assistant", mission="Help with programming tasks"
bank_id=bank_id,
name="Test Assistant",
mission="Help with programming tasks",
)
assert bank_response.bank_id == bank_id
@@ -126,7 +140,9 @@ def test_embedded_complete_workflow(llm_config):
items=[
{"content": "User works with pandas and numpy."},
{"content": "User likes matplotlib for visualization."},
{"content": "User is interested in machine learning with scikit-learn."},
{
"content": "User is interested in machine learning with scikit-learn."
},
],
)
assert batch_response.success
@@ -134,7 +150,9 @@ def test_embedded_complete_workflow(llm_config):
# Step 4: Recall memories
print("\n4. Recalling memories...")
recall_response = client.recall(bank_id=bank_id, query="What tools does the user prefer?", max_tokens=2000)
recall_response = client.recall(
bank_id=bank_id, query="What tools does the user prefer?", max_tokens=2000
)
assert isinstance(recall_response.results, list)
assert len(recall_response.results) > 0
print(f" Found {len(recall_response.results)} relevant memories")
@@ -152,7 +170,9 @@ def test_embedded_complete_workflow(llm_config):
# Verify answer mentions relevant tools
answer_lower = reflect_response.text.lower()
assert any(term in answer_lower for term in ["python", "pandas", "numpy", "data"])
assert any(
term in answer_lower for term in ["python", "pandas", "numpy", "data"]
)
# Step 6: List memories
print("\n6. Listing memories...")
@@ -215,7 +235,9 @@ def test_embedded_method_proxying(llm_config):
assert bank.bank_id == bank_id
# Test mission setting
mission_response = client.set_mission(bank_id=bank_id, mission="Test mission for proxying")
mission_response = client.set_mission(
bank_id=bank_id, mission="Test mission for proxying"
)
assert mission_response.bank_id == bank_id
# Test retain
@@ -264,7 +286,9 @@ def test_embedded_multiple_banks(llm_config):
# Create second bank and store data
client.create_bank(bank_id=bank2_id, name="Bank 2")
client.retain(bank_id=bank2_id, content="Bob uses JavaScript for web development")
client.retain(
bank_id=bank2_id, content="Bob uses JavaScript for web development"
)
# Recall from both banks
results1 = client.recall(bank_id=bank1_id, query="programming language")
@@ -275,9 +299,9 @@ def test_embedded_multiple_banks(llm_config):
# Verify banks are isolated (each should only see their own content)
# This is a basic check - content isolation is tested more thoroughly in other tests
assert results1.results[0].text != results2.results[0].text or len(results1.results) != len(
results2.results
)
assert results1.results[0].text != results2.results[0].text or len(
results1.results
) != len(results2.results)
finally:
client.close()
@@ -296,10 +320,14 @@ def test_embedded_profile_isolation(llm_config):
try:
# Store data in profile1
client1.retain(bank_id=bank_id, content="User likes TypeScript for frontend development")
client1.retain(
bank_id=bank_id, content="User likes TypeScript for frontend development"
)
# Store different data in profile2
client2.retain(bank_id=bank_id, content="User prefers Rust for systems programming")
client2.retain(
bank_id=bank_id, content="User prefers Rust for systems programming"
)
# Each profile should only see its own data
results1 = client1.recall(bank_id=bank_id, query="programming preference")
@@ -334,5 +362,42 @@ def test_embedded_error_after_close(llm_config):
assert not client.is_running
# Trying to use it after close should raise an error
with pytest.raises(RuntimeError, match="Cannot use HindsightEmbedded after it has been closed"):
with pytest.raises(
RuntimeError, match="Cannot use HindsightEmbedded after it has been closed"
):
client.retain(bank_id=bank_id, content="This should fail")
def test_embedded_ui_flag(llm_config):
"""
Test that ui=True starts the control plane UI alongside the daemon,
and that the UI's health endpoint reports a connected dataplane.
"""
profile = f"test_ui_{uuid.uuid4().hex[:8]}"
bank_id = f"bank_{uuid.uuid4().hex[:8]}"
client = HindsightEmbedded(profile=profile, log_level="info", ui=True, **llm_config)
try:
# First use triggers daemon + UI startup
result = client.retain(bank_id=bank_id, content="UI integration test content")
assert result.success, "Retain should succeed"
assert client.is_running, "Daemon should be running"
# Verify UI is reachable and reports connected dataplane
ui_url = client.ui_url
assert ui_url, "ui_url should be set"
health_url = f"{ui_url}/api/health"
with urllib.request.urlopen(health_url, timeout=10) as resp:
health = json.loads(resp.read().decode())
assert health["status"] == "ok", (
f"UI health status should be 'ok', got: {health['status']}"
)
assert health["dataplane"]["status"] == "connected", (
f"Dataplane should be connected, got: {health['dataplane']}"
)
finally:
client.close()
+1 -1
View File
@@ -46,4 +46,4 @@ __all__ = [
"RemoteTEICrossEncoder",
"LLMConfig",
]
__version__ = "0.4.20"
__version__ = "0.4.22"
@@ -249,7 +249,7 @@ async def _run_migration(
schemas = list(dict.fromkeys(schemas))
for schema in schemas:
run_migrations(resolved_url, schema=schema)
run_migrations(resolved_url, schema=schema, migration_database_url=config.migration_database_url)
if embedding_dimension is not None:
for schema in schemas:
@@ -0,0 +1,45 @@
"""Recreate entities trigram index on LOWER(canonical_name) for case-insensitive matching
The previous GIN trigram index on canonical_name was case-sensitive, causing
"Alice" and "alice" to have different trigram sets. This recreates it on
LOWER(canonical_name) so the % operator matches case-insensitively.
Revision ID: d6e7f8a9b0c1
Revises: c5d6e7f8a9b0
Create Date: 2026-03-31
"""
from collections.abc import Sequence
from alembic import context, op
revision: str = "d6e7f8a9b0c1"
down_revision: str | Sequence[str] | None = "c5d6e7f8a9b0"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def upgrade() -> None:
schema = _get_schema_prefix()
# Drop the old case-sensitive trigram index
op.execute("DROP INDEX IF EXISTS entities_canonical_name_trgm_idx")
# Create case-insensitive trigram index on LOWER(canonical_name)
op.execute(
f"CREATE INDEX IF NOT EXISTS entities_canonical_name_lower_trgm_idx "
f"ON {schema}entities USING GIN (LOWER(canonical_name) gin_trgm_ops)"
)
def downgrade() -> None:
op.execute("DROP INDEX IF EXISTS entities_canonical_name_lower_trgm_idx")
schema = _get_schema_prefix()
# Restore original case-sensitive index
op.execute(
f"CREATE INDEX IF NOT EXISTS entities_canonical_name_trgm_idx "
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
)
@@ -0,0 +1,142 @@
"""Fix per-bank vector indexes to match configured extension
Revision ID: a4b5c6d7e8f9
Revises: d6e7f8a9b0c1
Create Date: 2026-04-01
Migration d5e6f7a8b9c0 hardcoded HNSW when creating per-bank partial vector
indexes, ignoring HINDSIGHT_API_VECTOR_EXTENSION. Banks that existed when that
migration ran got HNSW indexes even when pgvectorscale (DiskANN) or vchord
was configured.
This migration detects the mismatch and recreates the affected indexes with
the correct type. Skipped entirely when the configured extension is pgvector
(the default), since those indexes are already correct.
"""
import os
from collections.abc import Sequence
from alembic import context, op
from sqlalchemy import text
revision: str = "a4b5c6d7e8f9"
down_revision: str | Sequence[str] | None = "d6e7f8a9b0c1"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_FACT_TYPES: dict[str, str] = {
"world": "worl",
"experience": "expr",
"observation": "obsv",
}
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _target_index_type() -> str | None:
"""Return the target index type, or None if pgvector (no fix needed)."""
ext = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
if ext == "pgvectorscale":
return "diskann"
elif ext == "vchord":
return "vchordrq"
return None
def _vector_index_using_clause() -> str:
"""Return the USING clause based on the configured vector extension."""
ext = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
if ext == "pgvectorscale":
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
elif ext == "vchord":
return "USING vchordrq (embedding vector_l2_ops)"
else:
return "USING hnsw (embedding vector_cosine_ops)"
def upgrade() -> None:
target = _target_index_type()
if target is None:
# pgvector — indexes are already HNSW, nothing to fix
return
bind = op.get_bind()
schema_name = context.config.get_main_option("target_schema")
schema = _get_schema_prefix()
table_ref = f'"{schema_name}".memory_units' if schema_name else "memory_units"
banks_ref = f'"{schema_name}".banks' if schema_name else "banks"
using_clause = _vector_index_using_clause()
pg_schema = schema_name or "public"
rows = bind.execute(text(f"SELECT bank_id, internal_id FROM {banks_ref}")).fetchall() # noqa: S608
for row in rows:
bank_id = row[0]
internal_id = str(row[1]).replace("-", "")[:16]
escaped_bank_id = bank_id.replace("'", "''")
for ft, ft_short in _FACT_TYPES.items():
idx_name = f"idx_mu_emb_{ft_short}_{internal_id}"
# Check if this index exists and what type it is
idx_info = bind.execute(
text("SELECT indexdef FROM pg_indexes WHERE schemaname = :schema AND indexname = :idx"),
{"schema": pg_schema, "idx": idx_name},
).fetchone()
if idx_info is None:
# Index doesn't exist — create it with the correct type
bind.execute(
text(
f"CREATE INDEX IF NOT EXISTS {idx_name} "
f"ON {table_ref} {using_clause} "
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped_bank_id}'"
)
)
continue
indexdef = idx_info[0].lower()
if target in indexdef:
# Already the correct type
continue
# Wrong type — drop and recreate
bind.execute(text(f"DROP INDEX IF EXISTS {schema}{idx_name}"))
bind.execute(
text(
f"CREATE INDEX IF NOT EXISTS {idx_name} "
f"ON {table_ref} {using_clause} "
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped_bank_id}'"
)
)
def downgrade() -> None:
# Downgrade recreates indexes as HNSW (the original hardcoded behavior)
target = _target_index_type()
if target is None:
return
bind = op.get_bind()
schema_name = context.config.get_main_option("target_schema")
schema = _get_schema_prefix()
table_ref = f'"{schema_name}".memory_units' if schema_name else "memory_units"
banks_ref = f'"{schema_name}".banks' if schema_name else "banks"
rows = bind.execute(text(f"SELECT bank_id, internal_id FROM {banks_ref}")).fetchall() # noqa: S608
for row in rows:
bank_id = row[0]
internal_id = str(row[1]).replace("-", "")[:16]
escaped_bank_id = bank_id.replace("'", "''")
for ft, ft_short in _FACT_TYPES.items():
idx_name = f"idx_mu_emb_{ft_short}_{internal_id}"
bind.execute(text(f"DROP INDEX IF EXISTS {schema}{idx_name}"))
bind.execute(
text(
f"CREATE INDEX IF NOT EXISTS {idx_name} "
f"ON {table_ref} USING hnsw (embedding vector_cosine_ops) "
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped_bank_id}'"
)
)
@@ -0,0 +1,32 @@
"""add content_hash to chunks table for delta retain
Revision ID: b3c4d5e6f7a8
Revises: a3b4c5d6e7f8
Create Date: 2026-03-25
"""
from collections.abc import Sequence
from alembic import context, op
revision: str = "b3c4d5e6f7a8"
down_revision: str | Sequence[str] | None = "a3b4c5d6e7f8"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def upgrade() -> None:
schema = _get_schema_prefix()
# Add content_hash column to chunks table for delta comparison
op.execute(f"ALTER TABLE {schema}chunks ADD COLUMN IF NOT EXISTS content_hash TEXT")
def downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}chunks DROP COLUMN IF EXISTS content_hash")
@@ -0,0 +1,61 @@
"""Add audit_log table for feature usage tracking.
Merge migration that combines the two existing heads (a3b4c5d6e7f8 + c8e5f2a3b4d1).
Stores raw request/response as JSONB for expandability without future migrations.
The metadata JSONB column allows adding arbitrary fields in the future.
Revision ID: c2d3e4f5g6h7
Revises: a3b4c5d6e7f8, c8e5f2a3b4d1
Create Date: 2026-03-26
"""
from collections.abc import Sequence
from alembic import context, op
revision: str = "c2d3e4f5g6h7"
down_revision: str | Sequence[str] | None = ("a3b4c5d6e7f8", "c8e5f2a3b4d1")
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def upgrade() -> None:
schema = _get_schema_prefix()
op.execute(
f"""
CREATE TABLE IF NOT EXISTS {schema}audit_log (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
action TEXT NOT NULL,
transport TEXT NOT NULL,
bank_id TEXT,
started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
ended_at TIMESTAMPTZ,
request JSONB,
response JSONB,
metadata JSONB DEFAULT '{{}}'::jsonb
)
"""
)
op.execute(
f"CREATE INDEX IF NOT EXISTS idx_audit_log_action_started ON {schema}audit_log (action, started_at DESC)"
)
op.execute(f"CREATE INDEX IF NOT EXISTS idx_audit_log_bank_started ON {schema}audit_log (bank_id, started_at DESC)")
op.execute(f"CREATE INDEX IF NOT EXISTS idx_audit_log_started ON {schema}audit_log (started_at DESC)")
def downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"DROP INDEX IF EXISTS {schema}idx_audit_log_started")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_audit_log_bank_started")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_audit_log_action_started")
op.execute(f"DROP TABLE IF EXISTS {schema}audit_log")
@@ -0,0 +1,48 @@
"""Add bank_id column to memory_links for direct filtering
The stats endpoint JOINs memory_links to memory_units just to filter by
bank_id. With millions of links this takes 18+ seconds. Adding bank_id
directly to memory_links lets Postgres push the filter down before the JOIN.
Revision ID: c5d6e7f8a9b0
Revises: b3c4d5e6f7a8
Create Date: 2026-03-26
"""
from collections.abc import Sequence
from alembic import context, op
revision: str = "c5d6e7f8a9b0"
down_revision: str | Sequence[str] | None = "b3c4d5e6f7a8"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def upgrade() -> None:
schema = _get_schema_prefix()
# 1. Add nullable column
op.execute(f"ALTER TABLE {schema}memory_links ADD COLUMN IF NOT EXISTS bank_id TEXT")
# 2. Backfill from memory_units
op.execute(f"""
UPDATE {schema}memory_links ml
SET bank_id = mu.bank_id
FROM {schema}memory_units mu
WHERE ml.from_unit_id = mu.id
AND ml.bank_id IS NULL
""")
# 3. Set NOT NULL
op.execute(f"ALTER TABLE {schema}memory_links ALTER COLUMN bank_id SET NOT NULL")
def downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"ALTER TABLE {schema}memory_links DROP COLUMN IF EXISTS bank_id")
@@ -1,4 +1,4 @@
"""Add internal_id to banks and per-(bank, fact_type) partial HNSW indexes
"""Add internal_id to banks and per-(bank, fact_type) partial vector indexes
Revision ID: d5e6f7a8b9c0
Revises: a3b4c5d6e7f8
@@ -6,25 +6,20 @@ Create Date: 2026-03-11
This migration:
1. Adds internal_id UUID column to banks (stable identifier for index naming)
2. Drops the global HNSW index (competes with per-bank partial indexes)
3. Creates per-(bank_id, fact_type) partial HNSW indexes for all existing banks
(new banks get indexes created at bank-creation time via bank_utils.create_bank_hnsw_indexes)
2. Drops the global vector index (competes with per-bank partial indexes)
3. Creates per-(bank_id, fact_type) partial vector indexes for all existing banks
using the configured vector extension (HNSW for pgvector, DiskANN for
pgvectorscale, vchordrq for vchord).
(new banks get indexes created at bank-creation time via bank_utils.create_bank_vector_indexes)
Why per-(bank, fact_type) indexes:
- fact_type-only partial indexes are never chosen by the planner when bank_id is in the WHERE
clause, because the idx_memory_units_bank_id B-tree index always wins at planning time.
- Per-(bank, fact_type) partial indexes have both predicates matching → planner selects them.
- The global HNSW index competes for larger partitions (world, observation) and must be dropped.
For large deployments, create indexes CONCURRENTLY before running this migration:
SELECT internal_id, bank_id FROM banks;
-- for each bank and each fact_type in (world, experience, observation):
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_mu_emb_{ft}_{uid16}
ON memory_units USING hnsw (embedding vector_cosine_ops)
WHERE fact_type = '{ft}' AND bank_id = '{bank_id}';
DROP INDEX CONCURRENTLY IF EXISTS idx_memory_units_embedding;
- The global vector index competes for larger partitions (world, observation) and must be dropped.
"""
import os
from collections.abc import Sequence
from alembic import context, op
@@ -35,7 +30,7 @@ down_revision: str | Sequence[str] | None = "c3d4e5f6g7h8"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_HNSW_FACT_TYPES: dict[str, str] = {
_FACT_TYPES: dict[str, str] = {
"world": "worl",
"experience": "expr",
"observation": "obsv",
@@ -47,6 +42,17 @@ def _get_schema_prefix() -> str:
return f'"{schema}".' if schema else ""
def _vector_index_using_clause() -> str:
"""Return the USING clause based on the configured vector extension."""
ext = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
if ext == "pgvectorscale":
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
elif ext == "vchord":
return "USING vchordrq (embedding vector_l2_ops)"
else:
return "USING hnsw (embedding vector_cosine_ops)"
def upgrade() -> None:
schema = _get_schema_prefix()
@@ -56,33 +62,35 @@ def upgrade() -> None:
)
op.execute(f"ALTER TABLE {schema}banks ADD CONSTRAINT banks_internal_id_unique UNIQUE (internal_id)")
# 2. Drop any fact_type-only partial HNSW indexes that may exist from prior migrations
# 2. Drop any fact_type-only partial indexes that may exist from prior migrations
# (bank_id B-tree always wins over them when bank_id is in the WHERE clause)
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mu_emb_world")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mu_emb_observation")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mu_emb_experience")
# 4. Drop global HNSW index (competes with per-bank partial indexes)
# 4. Drop global vector index (competes with per-bank partial indexes)
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_embedding")
# 5. Create per-(bank, fact_type) partial HNSW indexes for all existing banks
# 5. Create per-(bank, fact_type) partial vector indexes for all existing banks
# using the configured extension (HNSW / DiskANN / vchordrq)
bind = op.get_bind()
schema_name = context.config.get_main_option("target_schema")
table_ref = f'"{schema_name}".memory_units' if schema_name else "memory_units"
banks_ref = f'"{schema_name}".banks' if schema_name else "banks"
using_clause = _vector_index_using_clause()
rows = bind.execute(text(f"SELECT bank_id, internal_id FROM {banks_ref}")).fetchall() # noqa: S608
for row in rows:
bank_id = row[0]
internal_id = str(row[1]).replace("-", "")[:16]
escaped_bank_id = bank_id.replace("'", "''")
for ft, ft_short in _HNSW_FACT_TYPES.items():
for ft, ft_short in _FACT_TYPES.items():
idx_name = f"idx_mu_emb_{ft_short}_{internal_id}"
# Index name is schema-unqualified (indexes live in the schema of their table)
bind.execute(
text(
f"CREATE INDEX IF NOT EXISTS {idx_name} "
f"ON {table_ref} USING hnsw (embedding vector_cosine_ops) "
f"ON {table_ref} {using_clause} "
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped_bank_id}'"
)
)
@@ -5,7 +5,7 @@ Revises: e0a1b2c3d4e5
Create Date: 2025-01-12
Add composite index on memory_links (from_unit_id, link_type, weight DESC)
to optimize MPFP graph traversal queries that need top-k edges per type.
to optimize graph traversal queries that need top-k edges per type.
"""
from collections.abc import Sequence
@@ -26,7 +26,7 @@ def _get_schema_prefix() -> str:
def upgrade() -> None:
"""Add composite index for efficient MPFP edge loading."""
"""Add composite index for efficient graph retrieval edge loading."""
schema = _get_schema_prefix()
# Create composite index for efficient top-k per (from_node, link_type) queries
# This enables LATERAL joins to use index-only scans with early termination
@@ -24,9 +24,28 @@ def upgrade() -> None:
the memory_units rows survived with chunk_id = NULL, leaving ghost records.
Switching to CASCADE ensures they are removed together with their chunk.
"""
op.drop_constraint("memory_units_chunk_fkey", "memory_units", type_="foreignkey")
op.create_foreign_key(
"memory_units_chunk_fkey", "memory_units", "chunks", ["chunk_id"], ["chunk_id"], ondelete="CASCADE"
from alembic import context
schema = context.config.get_main_option("target_schema")
schema_prefix = f'"{schema}".' if schema else ""
# Use raw SQL with IF EXISTS so this is safe on schemas where the FK was
# already dropped or never existed under this name.
op.execute(f"ALTER TABLE {schema_prefix}memory_units DROP CONSTRAINT IF EXISTS memory_units_chunk_fkey")
# Use a DO block so the ADD is also idempotent: if the FK already exists (e.g.
# the schema was provisioned after the base migration already added it) the
# duplicate_object exception is swallowed rather than failing the migration.
op.execute(
f"""
DO $$ BEGIN
ALTER TABLE {schema_prefix}memory_units
ADD CONSTRAINT memory_units_chunk_fkey
FOREIGN KEY (chunk_id)
REFERENCES {schema_prefix}chunks (chunk_id)
ON DELETE CASCADE;
EXCEPTION
WHEN duplicate_object THEN NULL;
END $$;
"""
)
@@ -0,0 +1,83 @@
"""remove_opinion_fact_type
Revision ID: g2h3i4j5k6l7
Revises: f1a2b3c4d5e6
Create Date: 2026-04-02
Remove the deprecated 'opinion' fact type: drop opinion-specific indexes,
update CHECK constraints, delete any remaining opinion rows, and drop the
confidence_score column (was only used for opinions, always NULL otherwise).
"""
from collections.abc import Sequence
from alembic import context, op
# revision identifiers, used by Alembic.
revision: str = "g2h3i4j5k6l7"
down_revision: str | Sequence[str] | None = "f1a2b3c4d5e6"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def upgrade() -> None:
schema = _get_schema_prefix()
# 1. Delete any remaining opinion rows
op.execute(f"DELETE FROM {schema}memory_units WHERE fact_type = 'opinion'")
# 2. Drop opinion-specific indexes
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_opinion_confidence")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_opinion_date")
# 3. Drop confidence_score constraints and column (only used for opinions, always NULL otherwise)
op.execute(f"ALTER TABLE {schema}memory_units DROP CONSTRAINT IF EXISTS confidence_score_fact_type_check")
op.execute(f"ALTER TABLE {schema}memory_units DROP CONSTRAINT IF EXISTS memory_units_confidence_score_check")
op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS confidence_score")
# 4. Replace fact_type CHECK constraint
op.execute(f"ALTER TABLE {schema}memory_units DROP CONSTRAINT IF EXISTS memory_units_fact_type_check")
op.execute(
f"ALTER TABLE {schema}memory_units ADD CONSTRAINT memory_units_fact_type_check "
f"CHECK (fact_type IN ('world', 'experience', 'observation'))"
)
def downgrade() -> None:
schema = _get_schema_prefix()
# Restore confidence_score column
op.execute(f"ALTER TABLE {schema}memory_units ADD COLUMN IF NOT EXISTS confidence_score float")
op.execute(
f"ALTER TABLE {schema}memory_units ADD CONSTRAINT memory_units_confidence_score_check "
f"CHECK (confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0))"
)
op.execute(
f"ALTER TABLE {schema}memory_units ADD CONSTRAINT confidence_score_fact_type_check "
f"CHECK ((fact_type = 'opinion' AND confidence_score IS NOT NULL) OR "
f"(fact_type = 'observation') OR "
f"(fact_type NOT IN ('opinion', 'observation') AND confidence_score IS NULL))"
)
# Restore original fact_type CHECK constraint (with opinion)
op.execute(f"ALTER TABLE {schema}memory_units DROP CONSTRAINT IF EXISTS memory_units_fact_type_check")
op.execute(
f"ALTER TABLE {schema}memory_units ADD CONSTRAINT memory_units_fact_type_check "
f"CHECK (fact_type IN ('world', 'experience', 'opinion', 'observation'))"
)
# Recreate opinion indexes
op.execute(
f"CREATE INDEX idx_memory_units_opinion_confidence ON {schema}memory_units "
f"(bank_id, confidence_score DESC) WHERE fact_type = 'opinion'"
)
op.execute(
f"CREATE INDEX idx_memory_units_opinion_date ON {schema}memory_units "
f"(bank_id, event_date DESC) WHERE fact_type = 'opinion'"
)
File diff suppressed because it is too large Load Diff
+137 -45
View File
@@ -12,44 +12,9 @@ from hindsight_api.config import _get_raw_config
from hindsight_api.engine.memory_engine import _current_schema
from hindsight_api.extensions import MCPExtension, load_extension
from hindsight_api.extensions.tenant import AuthenticationError
from hindsight_api.mcp_tools import MCPToolsConfig, register_mcp_tools
from hindsight_api.mcp_tools import _ALL_TOOLS, MCPToolsConfig, register_mcp_tools
from hindsight_api.models import RequestContext
# All tools available in the system (explicit list — no wildcards)
_ALL_TOOLS: frozenset[str] = frozenset(
{
"retain",
"recall",
"reflect",
"list_banks",
"create_bank",
"list_mental_models",
"get_mental_model",
"create_mental_model",
"update_mental_model",
"delete_mental_model",
"refresh_mental_model",
"list_directives",
"create_directive",
"delete_directive",
"list_memories",
"get_memory",
"delete_memory",
"list_documents",
"get_document",
"delete_document",
"list_operations",
"get_operation",
"cancel_operation",
"list_tags",
"get_bank",
"get_bank_stats",
"update_bank",
"delete_bank",
"clear_memories",
}
)
# Configure logging from HINDSIGHT_API_LOG_LEVEL environment variable
_log_level_str = os.environ.get("HINDSIGHT_API_LOG_LEVEL", "info").lower()
_log_level_map = {
@@ -191,24 +156,65 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
return mcp
def _get_mcp_tools(mcp: FastMCP) -> dict:
"""Get tool name→object mapping, compatible with FastMCP 2.x and 3.x."""
# FastMCP 2.x: _tool_manager._tools
if hasattr(mcp, "_tool_manager"):
return mcp._tool_manager._tools # type: ignore[union-attr]
# FastMCP 3.x: _local_provider._components with "tool:" prefix
if hasattr(mcp, "_local_provider"):
return {
k.split(":")[1].split("@")[0]: v
for k, v in mcp._local_provider._components.items() # type: ignore[union-attr]
if k.startswith("tool:")
}
msg = "Cannot locate tools on FastMCP instance"
raise AttributeError(msg)
def _make_tools_tolerant(mcp: FastMCP) -> None:
"""Wrap all tool run methods to strip unknown arguments before validation.
"""Wrap all tool run methods to strip unknown arguments and coerce string-encoded JSON.
LLMs frequently add extra fields like "explanation" or "reasoning" to tool calls.
FastMCP's Pydantic TypeAdapter rejects these with "Unexpected keyword argument".
This wraps each tool's run() to filter arguments to only known parameters.
LLMs also frequently serialize list/dict arguments as JSON strings instead of native
types (e.g., tags='["a","b"]' instead of tags=["a","b"]). This auto-coerces them.
This wraps each tool's run() to apply both fixes before validation.
"""
try:
for name, tool in mcp._tool_manager._tools.items():
tools = _get_mcp_tools(mcp)
for name, tool in tools.items():
if hasattr(tool, "parameters") and tool.parameters:
allowed = set(tool.parameters.get("properties", {}).keys())
properties = tool.parameters.get("properties", {})
allowed = set(properties.keys())
# Build sets of parameter names that expect array or object types.
# Handles both direct types {"type": "array"} and anyOf/oneOf unions
# like {"anyOf": [{"type": "array", ...}, {"type": "null"}]}.
array_params: set[str] = set()
object_params: set[str] = set()
for param_name, param_schema in properties.items():
_collect_coercible_types(param_schema, param_name, array_params, object_params)
original_run = tool.run
async def _tolerant_run(arguments, _allowed=allowed, _orig=original_run):
async def _tolerant_run(
arguments,
_allowed=allowed,
_orig=original_run,
_array_params=array_params,
_object_params=object_params,
):
extra_keys = set(arguments.keys()) - _allowed
if extra_keys:
logger.debug(f"Stripping unknown arguments from tool call: {extra_keys}")
arguments = {k: v for k, v in arguments.items() if k in _allowed}
# Coerce string-encoded JSON for list/dict parameters
arguments = _coerce_string_json(arguments, _array_params, _object_params)
return await _orig(arguments)
# FunctionTool is a Pydantic model with extra='forbid', so use
@@ -218,6 +224,59 @@ def _make_tools_tolerant(mcp: FastMCP) -> None:
logger.warning(f"Could not make tools tolerant of extra arguments: {e}")
def _collect_coercible_types(schema: dict, param_name: str, array_params: set[str], object_params: set[str]) -> None:
"""Check a JSON Schema property and add param_name to array_params/object_params if applicable."""
# Direct type
schema_type = schema.get("type")
if schema_type == "array":
array_params.add(param_name)
return
if schema_type == "object":
object_params.add(param_name)
return
# anyOf / oneOf unions (e.g., list[str] | None → {"anyOf": [{"type": "array"}, {"type": "null"}]})
for variant in schema.get("anyOf", []) + schema.get("oneOf", []):
variant_type = variant.get("type")
if variant_type == "array":
array_params.add(param_name)
return
if variant_type == "object":
object_params.add(param_name)
return
def _coerce_string_json(arguments: dict, array_params: set[str], object_params: set[str]) -> dict:
"""Auto-coerce string-encoded JSON arrays/objects to native types.
LLM agents frequently serialize list and dict tool arguments as JSON strings.
This is backward-compatible: native arrays/objects pass through unchanged.
"""
for param_name in array_params:
val = arguments.get(param_name)
if isinstance(val, str):
try:
parsed = json.loads(val)
if isinstance(parsed, list):
arguments = {**arguments, param_name: parsed}
logger.debug(f"Coerced string to list for parameter '{param_name}'")
except (json.JSONDecodeError, TypeError):
pass
for param_name in object_params:
val = arguments.get(param_name)
if isinstance(val, str):
try:
parsed = json.loads(val)
if isinstance(parsed, dict):
arguments = {**arguments, param_name: parsed}
logger.debug(f"Coerced string to dict for parameter '{param_name}'")
except (json.JSONDecodeError, TypeError):
pass
return arguments
class MCPMiddleware:
"""ASGI middleware that intercepts MCP requests and routes to appropriate MCP server.
@@ -281,10 +340,12 @@ class MCPMiddleware:
self.single_bank_server = single_bank_server
else:
# Create servers internally (for direct construction / tests)
global_config = _get_raw_config()
stateless = global_config.mcp_stateless
self.multi_bank_server = create_mcp_server(memory, multi_bank=True)
self.multi_bank_app = self.multi_bank_server.http_app(path="/", stateless_http=True)
self.multi_bank_app = self.multi_bank_server.http_app(path="/", stateless_http=stateless)
self.single_bank_server = create_mcp_server(memory, multi_bank=False)
self.single_bank_app = self.single_bank_server.http_app(path="/", stateless_http=True)
self.single_bank_app = self.single_bank_server.http_app(path="/", stateless_http=stateless)
def _get_header(self, scope: dict, name: str) -> str | None:
"""Extract a header value from ASGI scope."""
@@ -307,6 +368,17 @@ class MCPMiddleware:
await self.app(scope, receive, send)
return
# Handle GET-before-POST gracefully (Claude Code v2.1.84+ sends GET probe before POST initialize).
# Without a valid Mcp-Session-Id, GET has no meaningful response — return 200 OK so
# the client proceeds to POST initialize instead of marking the server as failed.
method = scope.get("method", "")
if method == "GET":
session_id = self._get_header(scope, "Mcp-Session-Id")
if not session_id:
logger.debug("MCP GET without session ID (client probe) — returning 200 OK")
await self._send_ok(send)
return
# Strip prefix from path
path = path[len(self.prefix) :] or "/"
@@ -436,6 +508,22 @@ class MCPMiddleware:
if schema_token is not None:
_current_schema.reset(schema_token)
async def _send_ok(self, send):
"""Send a 200 OK response with empty body (used for GET probes without session)."""
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [(b"content-type", b"application/json")],
}
)
await send(
{
"type": "http.response.body",
"body": b"{}",
}
)
async def _send_error(self, send, status: int, message: str, extra_headers: dict[str, str] | None = None):
"""Send an error response."""
body = json.dumps({"error": message}).encode()
@@ -466,10 +554,14 @@ def create_mcp_servers(memory: MemoryEngine):
Returns:
Tuple of (multi_bank_server, single_bank_server, multi_bank_app, single_bank_app)
"""
global_config = _get_raw_config()
stateless = global_config.mcp_stateless
multi_bank_server = create_mcp_server(memory, multi_bank=True)
multi_bank_app = multi_bank_server.http_app(path="/", stateless_http=True)
multi_bank_app = multi_bank_server.http_app(path="/", stateless_http=stateless)
single_bank_server = create_mcp_server(memory, multi_bank=False)
single_bank_app = single_bank_server.http_app(path="/", stateless_http=True)
single_bank_app = single_bank_server.http_app(path="/", stateless_http=stateless)
logger.info(f"MCP servers created (stateless_http={stateless})")
return multi_bank_server, single_bank_server, multi_bank_app, single_bank_app
+126 -5
View File
@@ -118,6 +118,7 @@ def normalize_config_dict(config: dict[str, Any]) -> dict[str, Any]:
# Environment variable names
ENV_DATABASE_URL = "HINDSIGHT_API_DATABASE_URL"
ENV_MIGRATION_DATABASE_URL = "HINDSIGHT_API_MIGRATION_DATABASE_URL"
ENV_DATABASE_SCHEMA = "HINDSIGHT_API_DATABASE_SCHEMA"
ENV_LLM_PROVIDER = "HINDSIGHT_API_LLM_PROVIDER"
ENV_LLM_API_KEY = "HINDSIGHT_API_LLM_API_KEY"
@@ -130,10 +131,12 @@ ENV_LLM_MAX_BACKOFF = "HINDSIGHT_API_LLM_MAX_BACKOFF"
ENV_LLM_TIMEOUT = "HINDSIGHT_API_LLM_TIMEOUT"
ENV_LLM_GROQ_SERVICE_TIER = "HINDSIGHT_API_LLM_GROQ_SERVICE_TIER"
ENV_LLM_OPENAI_SERVICE_TIER = "HINDSIGHT_API_LLM_OPENAI_SERVICE_TIER"
ENV_LLM_EXTRA_BODY = "HINDSIGHT_API_LLM_EXTRA_BODY"
# Defaults for service tiers
DEFAULT_LLM_GROQ_SERVICE_TIER = "auto" # "on_demand", "flex", or "auto"
DEFAULT_LLM_OPENAI_SERVICE_TIER = None # None (default) or "flex" (50% cheaper)
DEFAULT_LLM_EXTRA_BODY = None # None = no extra body params; JSON dict merged into OpenAI extra_body
# Per-operation LLM configuration (optional, falls back to global LLM config)
ENV_RETAIN_LLM_PROVIDER = "HINDSIGHT_API_RETAIN_LLM_PROVIDER"
@@ -175,6 +178,14 @@ ENV_EMBEDDINGS_OPENAI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY"
ENV_EMBEDDINGS_OPENAI_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL"
ENV_EMBEDDINGS_OPENAI_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL"
# Gemini/Vertex AI embeddings configuration
ENV_EMBEDDINGS_GEMINI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_GEMINI_API_KEY"
ENV_EMBEDDINGS_GEMINI_MODEL = "HINDSIGHT_API_EMBEDDINGS_GEMINI_MODEL"
ENV_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY = "HINDSIGHT_API_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY"
ENV_EMBEDDINGS_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI_PROJECT_ID"
ENV_EMBEDDINGS_VERTEXAI_REGION = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI_REGION"
ENV_EMBEDDINGS_VERTEXAI_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI_SERVICE_ACCOUNT_KEY"
# Cohere configuration (separate for embeddings and reranker)
ENV_EMBEDDINGS_COHERE_API_KEY = "HINDSIGHT_API_EMBEDDINGS_COHERE_API_KEY"
ENV_EMBEDDINGS_COHERE_MODEL = "HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL"
@@ -199,6 +210,7 @@ ENV_RERANKER_LITELLM_MAX_TOKENS_PER_DOC = "HINDSIGHT_API_RERANKER_LITELLM_MAX_TO
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_KEY"
ENV_EMBEDDINGS_LITELLM_SDK_MODEL = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL"
ENV_EMBEDDINGS_LITELLM_SDK_API_BASE = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_BASE"
ENV_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS"
ENV_RERANKER_LITELLM_SDK_API_KEY = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY"
ENV_RERANKER_LITELLM_SDK_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL"
ENV_RERANKER_LITELLM_SDK_API_BASE = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_BASE"
@@ -225,6 +237,12 @@ ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR"
# ZeroEntropy configuration (reranker only)
ENV_RERANKER_ZEROENTROPY_API_KEY = "HINDSIGHT_API_RERANKER_ZEROENTROPY_API_KEY"
ENV_RERANKER_ZEROENTROPY_MODEL = "HINDSIGHT_API_RERANKER_ZEROENTROPY_MODEL"
ENV_RERANKER_ZEROENTROPY_BASE_URL = "HINDSIGHT_API_RERANKER_ZEROENTROPY_BASE_URL"
# Google Discovery Engine reranker configuration
ENV_RERANKER_GOOGLE_MODEL = "HINDSIGHT_API_RERANKER_GOOGLE_MODEL"
ENV_RERANKER_GOOGLE_PROJECT_ID = "HINDSIGHT_API_RERANKER_GOOGLE_PROJECT_ID"
ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY"
ENV_VECTOR_EXTENSION = "HINDSIGHT_API_VECTOR_EXTENSION"
ENV_TEXT_SEARCH_EXTENSION = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION"
@@ -237,9 +255,9 @@ ENV_LOG_FORMAT = "HINDSIGHT_API_LOG_FORMAT"
ENV_WORKERS = "HINDSIGHT_API_WORKERS"
ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
ENV_MCP_ENABLED_TOOLS = "HINDSIGHT_API_MCP_ENABLED_TOOLS"
ENV_MCP_STATELESS = "HINDSIGHT_API_MCP_STATELESS"
ENV_ENABLE_BANK_CONFIG_API = "HINDSIGHT_API_ENABLE_BANK_CONFIG_API"
ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER"
ENV_MPFP_TOP_K_NEIGHBORS = "HINDSIGHT_API_MPFP_TOP_K_NEIGHBORS"
ENV_RECALL_MAX_CONCURRENT = "HINDSIGHT_API_RECALL_MAX_CONCURRENT"
ENV_RECALL_CONNECTION_BUDGET = "HINDSIGHT_API_RECALL_CONNECTION_BUDGET"
ENV_RECALL_MAX_QUERY_TOKENS = "HINDSIGHT_API_RECALL_MAX_QUERY_TOKENS"
@@ -251,6 +269,7 @@ ENV_OTEL_EXPORTER_OTLP_ENDPOINT = "HINDSIGHT_API_OTEL_EXPORTER_OTLP_ENDPOINT"
ENV_OTEL_EXPORTER_OTLP_HEADERS = "HINDSIGHT_API_OTEL_EXPORTER_OTLP_HEADERS"
ENV_OTEL_SERVICE_NAME = "HINDSIGHT_API_OTEL_SERVICE_NAME"
ENV_OTEL_DEPLOYMENT_ENVIRONMENT = "HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT"
ENV_METRICS_INCLUDE_BANK_ID = "HINDSIGHT_API_METRICS_INCLUDE_BANK_ID"
# Vertex AI configuration
ENV_LLM_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID"
@@ -272,6 +291,7 @@ ENV_RETAIN_BATCH_TOKENS = "HINDSIGHT_API_RETAIN_BATCH_TOKENS"
ENV_RETAIN_ENTITY_LOOKUP = "HINDSIGHT_API_RETAIN_ENTITY_LOOKUP"
ENV_RETAIN_BATCH_ENABLED = "HINDSIGHT_API_RETAIN_BATCH_ENABLED"
ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS = "HINDSIGHT_API_RETAIN_BATCH_POLL_INTERVAL_SECONDS"
ENV_RETAIN_CHUNK_BATCH_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_BATCH_SIZE"
# File storage configuration
ENV_FILE_STORAGE_TYPE = "HINDSIGHT_API_FILE_STORAGE_TYPE"
@@ -304,6 +324,7 @@ ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION = (
"HINDSIGHT_API_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION"
)
ENV_OBSERVATIONS_MISSION = "HINDSIGHT_API_OBSERVATIONS_MISSION"
ENV_MAX_OBSERVATIONS_PER_SCOPE = "HINDSIGHT_API_MAX_OBSERVATIONS_PER_SCOPE"
ENV_ENABLE_OBSERVATION_HISTORY = "HINDSIGHT_API_ENABLE_OBSERVATION_HISTORY"
ENV_ENABLE_MENTAL_MODEL_HISTORY = "HINDSIGHT_API_ENABLE_MENTAL_MODEL_HISTORY"
@@ -334,6 +355,7 @@ ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES"
ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT"
ENV_WORKER_MAX_SLOTS = "HINDSIGHT_API_WORKER_MAX_SLOTS"
ENV_WORKER_CONSOLIDATION_MAX_SLOTS = "HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLOTS"
ENV_RETAIN_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_MAX_CONCURRENT"
# Reflect agent settings
ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS"
@@ -342,6 +364,11 @@ ENV_REFLECT_WALL_TIMEOUT = "HINDSIGHT_API_REFLECT_WALL_TIMEOUT"
ENV_REFLECT_MISSION = "HINDSIGHT_API_REFLECT_MISSION"
ENV_REFLECT_SOURCE_FACTS_MAX_TOKENS = "HINDSIGHT_API_REFLECT_SOURCE_FACTS_MAX_TOKENS"
# Audit log settings
ENV_AUDIT_LOG_ENABLED = "HINDSIGHT_API_AUDIT_LOG_ENABLED"
ENV_AUDIT_LOG_ACTIONS = "HINDSIGHT_API_AUDIT_LOG_ACTIONS"
ENV_AUDIT_LOG_RETENTION_DAYS = "HINDSIGHT_API_AUDIT_LOG_RETENTION_DAYS"
# Disposition settings
ENV_DISPOSITION_SKEPTICISM = "HINDSIGHT_API_DISPOSITION_SKEPTICISM"
ENV_DISPOSITION_LITERALISM = "HINDSIGHT_API_DISPOSITION_LITERALISM"
@@ -368,6 +395,7 @@ PROVIDER_DEFAULT_MODELS = {
"none": "none",
"litellm": "gpt-4o-mini",
"bedrock": "us.amazon.nova-2-lite-v1:0",
"volcano": "doubao-pro-32k",
}
DEFAULT_LLM_MODEL = "gpt-4o-mini" # Fallback if provider not in table
DEFAULT_LLM_MAX_CONCURRENT = 32
@@ -389,6 +417,8 @@ DEFAULT_EMBEDDINGS_LOCAL_MODEL = "BAAI/bge-small-en-v1.5"
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU = False # Force CPU mode for local embeddings (avoids MPS/XPC issues on macOS)
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE = False # Security: disabled by default, required for some models
DEFAULT_EMBEDDINGS_OPENAI_MODEL = "text-embedding-3-small"
DEFAULT_EMBEDDINGS_GEMINI_MODEL = "gemini-embedding-001"
DEFAULT_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY = 768
DEFAULT_EMBEDDING_DIMENSION = 384
DEFAULT_RERANKER_PROVIDER = "local"
@@ -412,6 +442,8 @@ DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
DEFAULT_RERANKER_ZEROENTROPY_MODEL = "zerank-2"
DEFAULT_RERANKER_GOOGLE_MODEL = "semantic-ranker-default-004"
# Vector extension (pgvector, vchord, or pgvectorscale)
DEFAULT_VECTOR_EXTENSION = "pgvector" # Options: "pgvector", "vchord", "pgvectorscale"
@@ -436,9 +468,9 @@ DEFAULT_LOG_FORMAT = "text" # Options: "text", "json"
DEFAULT_WORKERS = 1
DEFAULT_MCP_ENABLED = True
DEFAULT_MCP_ENABLED_TOOLS: list[str] | None = None # None = all tools enabled
DEFAULT_MCP_STATELESS = False # False = stateful (supports SSE/GET); True = stateless (POST-only)
DEFAULT_ENABLE_BANK_CONFIG_API = True
DEFAULT_GRAPH_RETRIEVER = "link_expansion" # Options: "link_expansion", "mpfp", "bfs"
DEFAULT_MPFP_TOP_K_NEIGHBORS = 20 # Fan-out limit per node in MPFP graph traversal
DEFAULT_GRAPH_RETRIEVER = "link_expansion"
DEFAULT_RECALL_MAX_CONCURRENT = 32 # Max concurrent recall operations per worker
DEFAULT_RECALL_CONNECTION_BUDGET = 4 # Max concurrent DB connections per recall operation
DEFAULT_RECALL_MAX_QUERY_TOKENS = 500 # Maximum tokens allowed in recall query
@@ -454,6 +486,9 @@ DEFAULT_RETAIN_MISSION = None # Declarative spec of what to retain (injected in
DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS = None # Custom extraction guidelines (only used when mode="custom")
DEFAULT_RETAIN_DEFAULT_STRATEGY = None # Default strategy name (None = no strategy override)
DEFAULT_RETAIN_STRATEGIES: dict | None = None # Named retain strategies (dict of name → config overrides)
DEFAULT_RETAIN_CHUNK_BATCH_SIZE = (
100 # Max chunks per streaming batch. Each chunk produces ~17 facts, so 100 chunks = ~1700 facts/batch.
)
DEFAULT_RETAIN_BATCH_TOKENS = 10_000 # ~40KB of text # Max chars per sub-batch for async retain auto-splitting
DEFAULT_RETAIN_ENTITY_LOOKUP = "trigram" # "full" or "trigram"
DEFAULT_RETAIN_BATCH_ENABLED = False # Use LLM Batch API for fact extraction (only when async=True)
@@ -482,6 +517,7 @@ DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION = (
256 # Max tokens of source facts per observation in consolidation prompt (-1 = unlimited)
)
DEFAULT_OBSERVATIONS_MISSION = None # Declarative spec of what observations are for this bank
DEFAULT_MAX_OBSERVATIONS_PER_SCOPE = -1 # Max observations per tag scope (-1 = unlimited)
# Database migrations
DEFAULT_RUN_MIGRATIONS_ON_STARTUP = True
@@ -500,6 +536,7 @@ DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed
DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health
DEFAULT_WORKER_MAX_SLOTS = 10 # Total concurrent tasks per worker
DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS = 2 # Max concurrent consolidation tasks per worker
DEFAULT_RETAIN_MAX_CONCURRENT = 4 # Max concurrent retain DB phases (HNSW reads + writes). Limits I/O contention.
# Reflect agent settings
DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response
@@ -516,6 +553,12 @@ DEFAULT_DISPOSITION_EMPATHY = None
DEFAULT_OTEL_TRACES_ENABLED = False # Disabled by default for backward compatibility
DEFAULT_OTEL_SERVICE_NAME = "hindsight-api"
DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT = "development"
DEFAULT_METRICS_INCLUDE_BANK_ID = False # Disabled by default to avoid high-cardinality OTel metric growth
# Audit log defaults
DEFAULT_AUDIT_LOG_ENABLED = False # Disabled by default
DEFAULT_AUDIT_LOG_ACTIONS = "" # Empty = audit all eligible actions
DEFAULT_AUDIT_LOG_RETENTION_DAYS = -1 # -1 = keep forever
# Default MCP tool descriptions (can be customized via env vars)
DEFAULT_MCP_RETAIN_DESCRIPTION = """Store important information to long-term memory.
@@ -605,6 +648,7 @@ class HindsightConfig:
# Database
database_url: str
migration_database_url: str | None
database_schema: str
vector_extension: str # "pgvector" or "vchord"
text_search_extension: str # "native" or "vchord"
@@ -621,6 +665,9 @@ class HindsightConfig:
llm_timeout: float
llm_groq_service_tier: str # Groq: "on_demand", "flex", or "auto"
llm_openai_service_tier: str | None # OpenAI: None (default) or "flex" (50% cheaper)
llm_extra_body: (
dict | None
) # Extra body params merged into OpenAI-compatible API calls (e.g. {"chat_template_kwargs": {"enable_thinking": true}})
# Vertex AI configuration
llm_vertexai_project_id: str | None
@@ -677,6 +724,14 @@ class HindsightConfig:
embeddings_litellm_sdk_api_key: str | None
embeddings_litellm_sdk_model: str
embeddings_litellm_sdk_api_base: str | None
embeddings_litellm_sdk_output_dimensions: int | None
# Gemini/Vertex AI embeddings
embeddings_gemini_api_key: str | None
embeddings_gemini_model: str
embeddings_gemini_output_dimensionality: int | None
embeddings_vertexai_project_id: str | None
embeddings_vertexai_region: str | None
embeddings_vertexai_service_account_key: str | None
# Reranker
reranker_provider: str
@@ -703,6 +758,10 @@ class HindsightConfig:
reranker_litellm_sdk_api_base: str | None
reranker_zeroentropy_api_key: str | None
reranker_zeroentropy_model: str
reranker_zeroentropy_base_url: str | None
reranker_google_model: str
reranker_google_project_id: str | None
reranker_google_service_account_key: str | None
# Server
host: str
@@ -712,11 +771,11 @@ class HindsightConfig:
log_format: str
mcp_enabled: bool
mcp_enabled_tools: list[str] | None # None = all tools; explicit list = allowlist
mcp_stateless: bool # True = stateless HTTP (POST-only); False = stateful (supports GET/SSE)
enable_bank_config_api: bool
# Recall
graph_retriever: str
mpfp_top_k_neighbors: int
recall_max_concurrent: int
recall_connection_budget: int
recall_max_query_tokens: int
@@ -735,6 +794,7 @@ class HindsightConfig:
retain_batch_enabled: bool
retain_batch_poll_interval_seconds: int
retain_entity_lookup: str # "full" or "trigram"
retain_chunk_batch_size: int # Max chunks per streaming batch (0 = disabled)
# File storage (static - server-level only)
file_storage_type: str # "native" (PostgreSQL) or "s3" (S3-compatible)
@@ -767,6 +827,7 @@ class HindsightConfig:
consolidation_source_facts_max_tokens: int
consolidation_source_facts_max_tokens_per_observation: int
observations_mission: str | None
max_observations_per_scope: int
# Entity labels (controlled vocabulary of key:value classification labels extracted at retain time)
# List of label group dicts: [{key, description, type, optional, values: [{value, description}]}]
@@ -805,6 +866,7 @@ class HindsightConfig:
worker_http_port: int
worker_max_slots: int
worker_consolidation_max_slots: int
retain_max_concurrent: int
# Reflect agent settings
reflect_max_iterations: int
@@ -817,6 +879,12 @@ class HindsightConfig:
otel_exporter_otlp_headers: str | None
otel_service_name: str
otel_deployment_environment: str
metrics_include_bank_id: bool
# Audit log configuration (static - server-level only)
audit_log_enabled: bool # Master switch for audit logging
audit_log_actions: list[str] # Allowlist of action types (empty = all)
audit_log_retention_days: int # -1 = keep forever, >0 = delete after N days
# Webhook configuration (static - server-level only, not per-bank)
webhook_url: str | None # Global webhook URL (None = disabled)
@@ -841,8 +909,13 @@ class HindsightConfig:
"embeddings_tei_base_url",
"reranker_tei_base_url",
"reranker_cohere_base_url",
"reranker_zeroentropy_base_url",
# Service Account Keys
"llm_vertexai_service_account_key",
"embeddings_vertexai_service_account_key",
"reranker_google_service_account_key",
# Embeddings API keys
"embeddings_gemini_api_key",
# File storage credentials
"file_storage_s3_access_key_id",
"file_storage_s3_secret_access_key",
@@ -865,6 +938,7 @@ class HindsightConfig:
"retain_custom_instructions",
"retain_default_strategy",
"retain_strategies",
"retain_chunk_batch_size",
# Entity labels (controlled vocabulary for entity classification)
"entity_labels",
"entities_allow_free_form",
@@ -874,6 +948,7 @@ class HindsightConfig:
"consolidation_source_facts_max_tokens",
"consolidation_source_facts_max_tokens_per_observation",
"observations_mission",
"max_observations_per_scope",
# Reflect settings
"reflect_mission",
"reflect_source_facts_max_tokens",
@@ -993,6 +1068,7 @@ class HindsightConfig:
config = cls(
# Database
database_url=os.getenv(ENV_DATABASE_URL, DEFAULT_DATABASE_URL),
migration_database_url=os.getenv(ENV_MIGRATION_DATABASE_URL) or None,
database_schema=os.getenv(ENV_DATABASE_SCHEMA, DEFAULT_DATABASE_SCHEMA),
vector_extension=os.getenv(ENV_VECTOR_EXTENSION, DEFAULT_VECTOR_EXTENSION).lower(),
text_search_extension=os.getenv(ENV_TEXT_SEARCH_EXTENSION, DEFAULT_TEXT_SEARCH_EXTENSION).lower(),
@@ -1008,6 +1084,7 @@ class HindsightConfig:
llm_timeout=float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT))),
llm_groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
llm_openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
llm_extra_body=json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null")),
# Vertex AI
llm_vertexai_project_id=os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID) or DEFAULT_LLM_VERTEXAI_PROJECT_ID,
llm_vertexai_region=os.getenv(ENV_LLM_VERTEXAI_REGION, DEFAULT_LLM_VERTEXAI_REGION),
@@ -1114,6 +1191,23 @@ class HindsightConfig:
ENV_EMBEDDINGS_LITELLM_SDK_MODEL, DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL
),
embeddings_litellm_sdk_api_base=os.getenv(ENV_EMBEDDINGS_LITELLM_SDK_API_BASE) or None,
embeddings_litellm_sdk_output_dimensions=int(v)
if (v := os.getenv(ENV_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS))
else None,
# Gemini/Vertex AI embeddings (with fallback to LLM keys)
embeddings_gemini_api_key=os.getenv(ENV_EMBEDDINGS_GEMINI_API_KEY) or os.getenv(ENV_LLM_API_KEY),
embeddings_gemini_model=os.getenv(ENV_EMBEDDINGS_GEMINI_MODEL, DEFAULT_EMBEDDINGS_GEMINI_MODEL),
embeddings_gemini_output_dimensionality=int(
os.getenv(
ENV_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY,
str(DEFAULT_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY),
)
),
embeddings_vertexai_project_id=os.getenv(ENV_EMBEDDINGS_VERTEXAI_PROJECT_ID)
or os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID),
embeddings_vertexai_region=os.getenv(ENV_EMBEDDINGS_VERTEXAI_REGION) or os.getenv(ENV_LLM_VERTEXAI_REGION),
embeddings_vertexai_service_account_key=os.getenv(ENV_EMBEDDINGS_VERTEXAI_SERVICE_ACCOUNT_KEY)
or os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY),
# Reranker
reranker_provider=os.getenv(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER),
reranker_local_model=os.getenv(ENV_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_LOCAL_MODEL),
@@ -1162,6 +1256,13 @@ class HindsightConfig:
# ZeroEntropy reranker
reranker_zeroentropy_api_key=os.getenv(ENV_RERANKER_ZEROENTROPY_API_KEY),
reranker_zeroentropy_model=os.getenv(ENV_RERANKER_ZEROENTROPY_MODEL, DEFAULT_RERANKER_ZEROENTROPY_MODEL),
reranker_zeroentropy_base_url=os.getenv(ENV_RERANKER_ZEROENTROPY_BASE_URL) or None,
# Google Discovery Engine reranker (with fallback to LLM Vertex AI keys)
reranker_google_model=os.getenv(ENV_RERANKER_GOOGLE_MODEL, DEFAULT_RERANKER_GOOGLE_MODEL),
reranker_google_project_id=os.getenv(ENV_RERANKER_GOOGLE_PROJECT_ID)
or os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID),
reranker_google_service_account_key=os.getenv(ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY)
or os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY),
# Server
host=os.getenv(ENV_HOST, DEFAULT_HOST),
port=int(os.getenv(ENV_PORT, DEFAULT_PORT)),
@@ -1172,11 +1273,11 @@ class HindsightConfig:
mcp_enabled_tools=[t.strip() for t in os.getenv(ENV_MCP_ENABLED_TOOLS).split(",") if t.strip()]
if os.getenv(ENV_MCP_ENABLED_TOOLS)
else DEFAULT_MCP_ENABLED_TOOLS,
mcp_stateless=os.getenv(ENV_MCP_STATELESS, str(DEFAULT_MCP_STATELESS)).lower() == "true",
enable_bank_config_api=os.getenv(ENV_ENABLE_BANK_CONFIG_API, str(DEFAULT_ENABLE_BANK_CONFIG_API)).lower()
== "true",
# Recall
graph_retriever=os.getenv(ENV_GRAPH_RETRIEVER, DEFAULT_GRAPH_RETRIEVER),
mpfp_top_k_neighbors=int(os.getenv(ENV_MPFP_TOP_K_NEIGHBORS, str(DEFAULT_MPFP_TOP_K_NEIGHBORS))),
recall_max_concurrent=int(os.getenv(ENV_RECALL_MAX_CONCURRENT, str(DEFAULT_RECALL_MAX_CONCURRENT))),
recall_connection_budget=int(
os.getenv(ENV_RECALL_CONNECTION_BUDGET, str(DEFAULT_RECALL_CONNECTION_BUDGET))
@@ -1211,6 +1312,7 @@ class HindsightConfig:
retain_batch_poll_interval_seconds=int(
os.getenv(ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS, str(DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS))
),
retain_chunk_batch_size=int(os.getenv(ENV_RETAIN_CHUNK_BATCH_SIZE, str(DEFAULT_RETAIN_CHUNK_BATCH_SIZE))),
# File storage
file_storage_type=os.getenv(ENV_FILE_STORAGE_TYPE, DEFAULT_FILE_STORAGE_TYPE),
file_storage_s3_bucket=os.getenv(ENV_FILE_STORAGE_S3_BUCKET) or None,
@@ -1270,6 +1372,9 @@ class HindsightConfig:
)
),
observations_mission=os.getenv(ENV_OBSERVATIONS_MISSION) or DEFAULT_OBSERVATIONS_MISSION,
max_observations_per_scope=int(
os.getenv(ENV_MAX_OBSERVATIONS_PER_SCOPE, str(DEFAULT_MAX_OBSERVATIONS_PER_SCOPE))
),
entity_labels=None,
entities_allow_free_form=True,
# Database migrations
@@ -1289,6 +1394,7 @@ class HindsightConfig:
worker_consolidation_max_slots=int(
os.getenv(ENV_WORKER_CONSOLIDATION_MAX_SLOTS, str(DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS))
),
retain_max_concurrent=int(os.getenv(ENV_RETAIN_MAX_CONCURRENT, str(DEFAULT_RETAIN_MAX_CONCURRENT))),
# Reflect agent settings
reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))),
reflect_max_context_tokens=int(
@@ -1316,6 +1422,16 @@ class HindsightConfig:
otel_exporter_otlp_headers=os.getenv(ENV_OTEL_EXPORTER_OTLP_HEADERS) or None,
otel_service_name=os.getenv(ENV_OTEL_SERVICE_NAME, DEFAULT_OTEL_SERVICE_NAME),
otel_deployment_environment=os.getenv(ENV_OTEL_DEPLOYMENT_ENVIRONMENT, DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT),
metrics_include_bank_id=os.getenv(ENV_METRICS_INCLUDE_BANK_ID, str(DEFAULT_METRICS_INCLUDE_BANK_ID)).lower()
in ("true", "1", "yes"),
# Audit log configuration (static, server-level only)
audit_log_enabled=os.getenv(ENV_AUDIT_LOG_ENABLED, str(DEFAULT_AUDIT_LOG_ENABLED)).lower() == "true",
audit_log_actions=[
a.strip() for a in os.getenv(ENV_AUDIT_LOG_ACTIONS, DEFAULT_AUDIT_LOG_ACTIONS).split(",") if a.strip()
],
audit_log_retention_days=int(
os.getenv(ENV_AUDIT_LOG_RETENTION_DAYS, str(DEFAULT_AUDIT_LOG_RETENTION_DAYS))
),
# Webhook configuration (static, server-level only)
webhook_url=os.getenv(ENV_WEBHOOK_URL) or DEFAULT_WEBHOOK_URL,
webhook_secret=os.getenv(ENV_WEBHOOK_SECRET) or DEFAULT_WEBHOOK_SECRET,
@@ -1385,9 +1501,14 @@ class HindsightConfig:
root_logger.addHandler(handler)
# Silence noisy third-party loggers
logging.getLogger("google_genai.models").setLevel(logging.WARNING)
def log_config(self) -> None:
"""Log the current configuration (without sensitive values)."""
logger.info(f"Database: {self.database_url} (schema: {self.database_schema})")
if self.migration_database_url:
logger.info(f"Migration database: {self.migration_database_url}")
logger.info(f"LLM: provider={self.llm_provider}, model={self.llm_model}")
if self.retain_llm_provider or self.retain_llm_model:
retain_provider = self.retain_llm_provider or self.llm_provider
@@ -0,0 +1,209 @@
"""Audit logging for feature usage tracking.
Provides fire-and-forget audit logging of all mutating and core operations
(retain, recall, reflect, bank CRUD, etc.) across HTTP, MCP, and system transports.
"""
from __future__ import annotations
import asyncio
import json
import logging
import uuid
from collections.abc import Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
import asyncpg
from ..engine.db_utils import acquire_with_retry
logger = logging.getLogger(__name__)
@dataclass
class AuditEntry:
"""A single audit log entry."""
action: str
transport: str # "http", "mcp", "system"
bank_id: str | None = None
started_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
ended_at: datetime | None = None
request: dict[str, Any] | None = None
response: dict[str, Any] | None = None
metadata: dict[str, Any] = field(default_factory=dict)
def _json_default(obj: Any) -> str:
"""JSON serializer for objects not serializable by default."""
if isinstance(obj, datetime):
return obj.isoformat()
if isinstance(obj, uuid.UUID):
return str(obj)
if isinstance(obj, bytes):
return "<bytes>"
if isinstance(obj, set):
return list(obj)
return str(obj)
def _safe_json(data: Any) -> str | None:
"""Serialize data to JSON string, returning None on failure."""
if data is None:
return None
try:
return json.dumps(data, default=_json_default)
except Exception:
logger.debug("Failed to serialize audit data", exc_info=True)
return None
_SWEEP_INTERVAL_SECONDS = 3600 # Run retention sweep every hour
class AuditLogger:
"""Fire-and-forget audit log writer with optional retention sweep."""
def __init__(
self,
pool_getter: Callable[[], asyncpg.Pool | None],
schema_getter: Callable[[], str],
enabled: bool,
allowed_actions: list[str],
retention_days: int = -1,
) -> None:
self._pool_getter = pool_getter
self._schema_getter = schema_getter
self._enabled = enabled
self._allowed_actions: frozenset[str] | None = frozenset(allowed_actions) if allowed_actions else None
self._retention_days = retention_days
self._sweep_task: asyncio.Task | None = None
def is_enabled(self, action: str) -> bool:
"""Check if audit logging is enabled for this action."""
if not self._enabled:
return False
if self._allowed_actions is not None:
return action in self._allowed_actions
return True
def log_fire_and_forget(self, entry: AuditEntry) -> None:
"""Schedule an audit write as a background task."""
if not self.is_enabled(entry.action):
return
try:
asyncio.create_task(self._safe_log(entry))
except RuntimeError:
# No running event loop (e.g. during shutdown)
logger.debug("Cannot schedule audit log write: no running event loop")
async def _safe_log(self, entry: AuditEntry) -> None:
"""Write audit entry to DB. Errors are logged, never raised."""
pool = self._pool_getter()
if pool is None:
logger.debug("Audit log skipped: pool not available")
return
try:
schema = self._schema_getter()
table = f"{schema}.audit_log"
async with acquire_with_retry(pool, max_retries=1) as conn:
await conn.execute(
f"""
INSERT INTO {table}
(id, action, transport, bank_id, started_at, ended_at, request, response, metadata)
VALUES
($1, $2, $3, $4, $5, $6, $7::jsonb, $8::jsonb, $9::jsonb)
""",
uuid.uuid4(),
entry.action,
entry.transport,
entry.bank_id,
entry.started_at,
entry.ended_at,
_safe_json(entry.request),
_safe_json(entry.response),
_safe_json(entry.metadata) or "{}",
)
except Exception as e:
logger.warning(f"Audit log write failed for action={entry.action}: {e}")
def start_retention_sweep(self) -> None:
"""Start the periodic retention sweep if retention is configured."""
if self._retention_days <= 0 or not self._enabled:
return
try:
self._sweep_task = asyncio.create_task(self._sweep_loop())
except RuntimeError:
logger.debug("Cannot start retention sweep: no running event loop")
async def stop_retention_sweep(self) -> None:
"""Stop the periodic retention sweep."""
if self._sweep_task and not self._sweep_task.done():
self._sweep_task.cancel()
try:
await self._sweep_task
except asyncio.CancelledError:
pass
self._sweep_task = None
async def _sweep_loop(self) -> None:
"""Periodically delete audit log entries older than retention_days."""
while True:
await self._run_sweep()
await asyncio.sleep(_SWEEP_INTERVAL_SECONDS)
async def _run_sweep(self) -> None:
"""Delete expired audit log entries. Concurrent-safe via row-level deletes."""
pool = self._pool_getter()
if pool is None:
return
try:
schema = self._schema_getter()
table = f"{schema}.audit_log"
async with acquire_with_retry(pool, max_retries=1) as conn:
result = await conn.execute(
f"DELETE FROM {table} WHERE started_at < NOW() - INTERVAL '{self._retention_days} days'"
)
if result and result != "DELETE 0":
logger.info(f"Audit log retention sweep: {result}")
except Exception as e:
logger.warning(f"Audit log retention sweep failed: {e}")
@asynccontextmanager
async def audit_context(
audit_logger: AuditLogger | None,
action: str,
transport: str,
bank_id: str | None = None,
request: dict[str, Any] | None = None,
metadata: dict[str, Any] | None = None,
):
"""Async context manager that times the operation and writes audit on exit.
Usage:
async with audit_context(logger, "retain", "http", bank_id, request_dict) as entry:
result = await do_work()
entry.response = result_dict
"""
if audit_logger is None or not audit_logger.is_enabled(action):
entry = AuditEntry(action=action, transport=transport, bank_id=bank_id)
yield entry
return
entry = AuditEntry(
action=action,
transport=transport,
bank_id=bank_id,
started_at=datetime.now(timezone.utc),
request=request,
metadata=metadata or {},
)
try:
yield entry
finally:
entry.ended_at = datetime.now(timezone.utc)
audit_logger.log_fire_and_forget(entry)
@@ -119,6 +119,39 @@ def _aggregate_source_fields(source_mems: list[dict[str, Any]], tags: list[str]
)
async def _count_observations_for_scope(
conn: "Connection",
bank_id: str,
tags: list[str],
) -> int:
"""Count existing observations matching the given tag scope.
Returns the count of observations whose tags contain all specified tags.
Observations with no tags are not counted (the limit does not apply to them).
"""
return await conn.fetchval(
f"SELECT COUNT(*) FROM {fq_table('memory_units')} "
f"WHERE bank_id = $1 AND fact_type = 'observation' AND tags @> $2::varchar[]",
bank_id,
tags,
)
def _build_response_model(max_creates: int | None = None) -> type[_ConsolidationBatchResponse]:
"""Build a response model, optionally constraining max creates via JSON schema."""
if max_creates is None or max_creates < 0:
return _ConsolidationBatchResponse
from pydantic import Field as PydanticField
clamped = max(max_creates, 0)
class _ConstrainedConsolidationBatchResponse(_ConsolidationBatchResponse):
creates: list[_CreateAction] = PydanticField(default=[], max_length=clamped)
return _ConstrainedConsolidationBatchResponse
class ConsolidationPerfLog:
"""Performance logging for consolidation operations."""
@@ -698,24 +731,6 @@ async def _process_memory_batch(
if recall_result.source_facts:
union_source_facts.update(recall_result.source_facts)
# 3. Single LLM call
t0 = time.time()
llm_result = await _consolidate_batch_with_llm(
llm_config=llm_config,
memories=memories,
union_observations=union_observations,
union_source_facts=union_source_facts,
config=config,
)
if perf:
perf.record_timing("llm", time.time() - t0)
perf.record_llm_call(llm_result.obs_count, llm_result.prompt_chars)
# 4. Sequential execution of creates / updates / deletes
# Track which memory indices participated so we can build per-memory results for stats
per_memory_created: set[str] = set()
per_memory_updated: set[str] = set()
# Determine effective tag scope for observations.
# When obs_tags_override is set, use it; otherwise use the memory's own tags.
if obs_tags_override is not None:
@@ -724,28 +739,52 @@ async def _process_memory_batch(
# All memories in the batch share the same tag set (enforced by batching)
fact_tags = memories[0].get("tags") or [] if memories else []
# 2b. Compute remaining observation slots for this scope (if limit configured)
max_obs = config.max_observations_per_scope if config is not None else -1
remaining_observation_slots: int | None = None
if max_obs > 0 and fact_tags:
current_count = await _count_observations_for_scope(conn, bank_id, fact_tags)
remaining_observation_slots = max(max_obs - current_count, 0)
if remaining_observation_slots == 0:
logger.info(
f"[CONSOLIDATION] bank={bank_id} scope={fact_tags} at observation limit "
f"({current_count}/{max_obs}), only updates/deletes allowed"
)
# 3. Single LLM call
t0 = time.time()
llm_result = await _consolidate_batch_with_llm(
llm_config=llm_config,
memories=memories,
union_observations=union_observations,
union_source_facts=union_source_facts,
config=config,
remaining_observation_slots=remaining_observation_slots,
max_observations_per_scope=max_obs,
)
if perf:
perf.record_timing("llm", time.time() - t0)
perf.record_llm_call(llm_result.obs_count, llm_result.prompt_chars)
# 4. Sequential execution of deletes / updates / creates
# Deletes run first to free observation slots before creates consume them.
# Track which memory indices participated so we can build per-memory results for stats
per_memory_created: set[str] = set()
per_memory_updated: set[str] = set()
mem_by_id = {str(m["id"]): m for m in memories}
for create in llm_result.creates:
source_mems = [mem_by_id[fid] for fid in create.source_fact_ids if fid in mem_by_id]
if not source_mems:
# Execute deletes first to free observation slots before creates consume them
deleted_count = 0
for delete in llm_result.deletes:
# Security: the observation must be present in the unioned recall
if not any(str(obs.id) == delete.observation_id for obs in union_observations):
logger.debug(
f"Batch consolidation: rejected delete — observation {delete.observation_id} not in unioned recall"
)
continue
agg = _aggregate_source_fields(source_mems, tags=fact_tags)
await _execute_create_action(
conn=conn,
memory_engine=memory_engine,
bank_id=bank_id,
source_memory_ids=[m["id"] for m in source_mems],
text=create.text,
source_fact_tags=agg.tags,
event_date=agg.event_date,
occurred_start=agg.occurred_start,
occurred_end=agg.occurred_end,
mentioned_at=agg.mentioned_at,
perf=perf,
)
for m in source_mems:
per_memory_created.add(str(m["id"]))
await _execute_delete_action(conn=conn, bank_id=bank_id, observation_id=delete.observation_id)
deleted_count += 1
for update in llm_result.updates:
source_mems = [mem_by_id[fid] for fid in update.source_fact_ids if fid in mem_by_id]
@@ -776,16 +815,26 @@ async def _process_memory_batch(
for m in source_mems:
per_memory_updated.add(str(m["id"]))
deleted_count = 0
for delete in llm_result.deletes:
# Security: the observation must be present in the unioned recall
if not any(str(obs.id) == delete.observation_id for obs in union_observations):
logger.debug(
f"Batch consolidation: rejected delete — observation {delete.observation_id} not in unioned recall"
)
for create in llm_result.creates:
source_mems = [mem_by_id[fid] for fid in create.source_fact_ids if fid in mem_by_id]
if not source_mems:
continue
await _execute_delete_action(conn=conn, bank_id=bank_id, observation_id=delete.observation_id)
deleted_count += 1
agg = _aggregate_source_fields(source_mems, tags=fact_tags)
await _execute_create_action(
conn=conn,
memory_engine=memory_engine,
bank_id=bank_id,
source_memory_ids=[m["id"] for m in source_mems],
text=create.text,
source_fact_tags=agg.tags,
event_date=agg.event_date,
occurred_start=agg.occurred_start,
occurred_end=agg.occurred_end,
mentioned_at=agg.mentioned_at,
perf=perf,
)
for m in source_mems:
per_memory_created.add(str(m["id"]))
# Build per-memory result dicts for the stats tracker in the outer loop
results: list[dict[str, Any]] = []
@@ -1083,6 +1132,8 @@ async def _consolidate_batch_with_llm(
union_observations: "list[MemoryFact]",
union_source_facts: "dict[str, MemoryFact]",
config: Any = None,
remaining_observation_slots: int | None = None,
max_observations_per_scope: int = -1,
) -> _BatchLLMResult:
"""Single LLM call for a batch of facts against a pooled set of observations."""
if union_observations:
@@ -1106,24 +1157,51 @@ async def _consolidate_batch_with_llm(
facts_lines = "\n".join(_fact_line(m) for m in memories)
# Build capacity note for the prompt when observation limit is configured
observation_capacity_note: str | None = None
if remaining_observation_slots is not None and max_observations_per_scope > 0:
if remaining_observation_slots == 0:
observation_capacity_note = (
f"OBSERVATION LIMIT REACHED ({max_observations_per_scope}/{max_observations_per_scope}). "
"Only UPDATE or DELETE existing observations. Do NOT create new ones — "
"merge new knowledge into existing observations via UPDATE."
)
elif remaining_observation_slots <= len(memories):
observation_capacity_note = (
f"This scope has {remaining_observation_slots} observation slot(s) remaining "
f"(out of {max_observations_per_scope}). Prefer UPDATE over CREATE when possible."
)
observations_mission = config.observations_mission if config is not None else None
prompt_template = build_batch_consolidation_prompt(observations_mission)
prompt_template = build_batch_consolidation_prompt(observations_mission, observation_capacity_note)
prompt = prompt_template.format(
facts_text=facts_lines,
observations_text=observations_text,
)
# Use a constrained response model when observation limit is active
response_model = _build_response_model(max_creates=remaining_observation_slots)
max_attempts = 3
last_exc: Exception | None = None
for attempt in range(1, max_attempts + 1):
try:
response: _ConsolidationBatchResponse = await llm_config.call(
messages=[{"role": "user", "content": prompt}],
response_format=_ConsolidationBatchResponse,
response_format=response_model,
scope="consolidation",
)
# Defensive truncation: some LLM providers may not enforce JSON schema max_length
creates = response.creates
if remaining_observation_slots is not None and remaining_observation_slots >= 0:
if len(creates) > remaining_observation_slots:
logger.info(
f"[CONSOLIDATION] Truncating {len(creates)} creates to {remaining_observation_slots} "
f"(max_observations_per_scope={max_observations_per_scope})"
)
creates = creates[:remaining_observation_slots]
return _BatchLLMResult(
creates=response.creates,
creates=creates,
updates=response.updates,
deletes=response.deletes,
obs_count=len(union_observations),
@@ -5,10 +5,24 @@ _DEFAULT_MISSION = "Track every detail: names, numbers, dates, places, and relat
# Processing rules — always present regardless of mission
_PROCESSING_RULES = """Processing rules (always apply):
- REDUNDANT: same info worded differently → UPDATE the existing observation.
- CONTRADICTION/UPDATE: capture both states with temporal markers ("used to X, now Y").
- RESOLVE REFERENCES: when a new fact provides a concrete value resolving a vague placeholder in an existing observation (e.g. "home country", "hometown", "birthplace", "native language", "her ex", "that city"), UPDATE the observation to embed the resolved value explicitly. Example: new fact says "grandma in Sweden" + existing observation says "moved from her home country" → update to "home country is Sweden".
- NEVER merge observations about different people or unrelated topics."""
1. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), etc. Never merge different facets into one observation.
2. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
3. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
4. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
5. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
6. SAME FACET → UPDATE, NOT CREATE: a new count supersedes the old count — UPDATE the existing count observation, don't create a second one. If there's an existing observation for the same specific facet, always UPDATE it rather than creating a duplicate.
7. PRESERVE HISTORY: observations that record significant events (sold, died, moved, changed) are important history — never DELETE them. Only delete an observation when it is restated identically or truly meaningless. Be very conservative with deletes.
8. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country""Sweden"), UPDATE to embed the resolved value.
9. NEVER merge observations about different people or unrelated topics."""
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
_BATCH_DATA_SECTION = """
@@ -26,8 +40,8 @@ Each observation includes:
- source_memories: array of supporting facts with their text and dates
Compare the facts against existing observations:
- Same topic as an existing observation → UPDATE it (observation_id + source_fact_ids)
- New topic with durable knowledge → CREATE a new observation (source_fact_ids)
- Same facet as an existing observation → UPDATE it (observation_id + source_fact_ids)
- New facet with durable knowledge → CREATE a new observation (source_fact_ids)
- Cross-reference facts within the batch: a later fact may resolve a vague reference in an earlier one
- Purely ephemeral facts → omit them unless the MISSION above explicitly targets such data (e.g. timestamped events, session state, screen content)"""
@@ -66,7 +80,10 @@ Rules:
- Return {{"creates": [], "updates": [], "deletes": []}} if nothing durable is found."""
def build_batch_consolidation_prompt(observations_mission: str | None = None) -> str:
def build_batch_consolidation_prompt(
observations_mission: str | None = None,
observation_capacity_note: str | None = None,
) -> str:
"""
Build the consolidation prompt for batch mode (multiple facts per LLM call).
@@ -75,9 +92,13 @@ def build_batch_consolidation_prompt(observations_mission: str | None = None) ->
"""
mission = observations_mission or _DEFAULT_MISSION
capacity_section = ""
if observation_capacity_note:
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n{observation_capacity_note}"
return (
"You are a memory consolidation system. Synthesize facts into observations "
"and merge with existing observations when appropriate.\n\n"
f"## MISSION\n{mission}\n\n"
f"## MISSION\n{mission}{capacity_section}\n\n"
f"{_PROCESSING_RULES}" + _BATCH_DATA_SECTION + _BATCH_OUTPUT_FORMAT
)
@@ -20,6 +20,7 @@ from ..config import (
DEFAULT_RERANKER_COHERE_MODEL,
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
DEFAULT_RERANKER_FLASHRANK_MODEL,
DEFAULT_RERANKER_GOOGLE_MODEL,
DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
DEFAULT_RERANKER_LITELLM_MODEL,
DEFAULT_RERANKER_LITELLM_SDK_MODEL,
@@ -36,6 +37,7 @@ from ..config import (
ENV_RERANKER_COHERE_MODEL,
ENV_RERANKER_FLASHRANK_CACHE_DIR,
ENV_RERANKER_FLASHRANK_MODEL,
ENV_RERANKER_GOOGLE_PROJECT_ID,
ENV_RERANKER_LITELLM_SDK_API_KEY,
ENV_RERANKER_LOCAL_FORCE_CPU,
ENV_RERANKER_LOCAL_MAX_CONCURRENT,
@@ -544,6 +546,7 @@ class CohereCrossEncoder(CrossEncoderModel):
self.base_url = base_url
self.timeout = timeout
self._client = None
self._httpx_client: httpx.Client | None = None
@property
def provider_name(self) -> str:
@@ -551,23 +554,32 @@ class CohereCrossEncoder(CrossEncoderModel):
async def initialize(self) -> None:
"""Initialize the Cohere client."""
if self._client is not None:
if self._client is not None or self._httpx_client is not None:
return
try:
import cohere
except ImportError:
raise ImportError("cohere is required for CohereCrossEncoder. Install it with: pip install cohere")
base_url_msg = f" at {self.base_url}" if self.base_url else ""
logger.info(f"Reranker: initializing Cohere provider with model {self.model}{base_url_msg}")
# Build client kwargs, only including base_url if set (for Azure or custom endpoints)
client_kwargs = {"api_key": self.api_key, "timeout": self.timeout}
if self.base_url:
client_kwargs["base_url"] = self.base_url
self._client = cohere.Client(**client_kwargs)
logger.info("Reranker: Cohere provider initialized")
# For custom endpoints (Azure AI Foundry), use httpx directly to avoid SDK path appending
# Azure endpoints already include the full path (e.g., /models/.../invoke)
self._httpx_client = httpx.Client(
timeout=self.timeout,
headers={
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
},
)
logger.info("Reranker: Cohere provider initialized (using httpx for custom endpoint)")
else:
# For native Cohere API, use the official SDK
try:
import cohere
except ImportError:
raise ImportError("cohere is required for CohereCrossEncoder. Install it with: pip install cohere")
self._client = cohere.Client(api_key=self.api_key, timeout=self.timeout)
logger.info("Reranker: Cohere provider initialized")
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
"""
@@ -579,7 +591,7 @@ class CohereCrossEncoder(CrossEncoderModel):
Returns:
List of relevance scores
"""
if self._client is None:
if self._client is None and self._httpx_client is None:
raise RuntimeError("Reranker not initialized. Call initialize() first.")
if not pairs:
@@ -605,18 +617,40 @@ class CohereCrossEncoder(CrossEncoderModel):
texts = [text for _, text in indexed_texts]
indices = [idx for idx, _ in indexed_texts]
response = self._client.rerank(
query=query,
documents=texts,
model=self.model,
return_documents=False,
)
if self._httpx_client:
# Direct HTTP request for custom endpoints (Azure AI Foundry)
response = self._httpx_client.post(
self.base_url,
json={
"model": self.model,
"query": query,
"documents": texts,
"return_documents": False,
},
)
response.raise_for_status()
result = response.json()
# Map scores back to original positions
for result in response.results:
original_idx = result.index
score = result.relevance_score
all_scores[indices[original_idx]] = score
# Map scores back to original positions
# Azure Cohere response format: {"results": [{"index": 0, "relevance_score": 0.9}, ...]}
for item in result.get("results", []):
original_idx = item["index"]
score = item["relevance_score"]
all_scores[indices[original_idx]] = score
else:
# Native Cohere SDK for standard API
response = self._client.rerank(
query=query,
documents=texts,
model=self.model,
return_documents=False,
)
# Map scores back to original positions
for result in response.results:
original_idx = result.index
score = result.relevance_score
all_scores[indices[original_idx]] = score
return all_scores
@@ -629,12 +663,14 @@ class ZeroEntropyCrossEncoder(CrossEncoderModel):
See: https://docs.zeroentropy.dev/models
"""
RERANK_URL = "https://api.zeroentropy.dev/v1/models/rerank"
DEFAULT_BASE_URL = "https://api.zeroentropy.dev"
RERANK_PATH = "/v1/models/rerank"
def __init__(
self,
api_key: str,
model: str = DEFAULT_RERANKER_ZEROENTROPY_MODEL,
base_url: str | None = None,
timeout: float = 60.0,
):
"""
@@ -643,10 +679,13 @@ class ZeroEntropyCrossEncoder(CrossEncoderModel):
Args:
api_key: ZeroEntropy API key
model: ZeroEntropy rerank model name (default: zerank-2)
base_url: Custom base URL for ZeroEntropy-compatible API (e.g., mock server or proxy)
timeout: Request timeout in seconds (default: 60.0)
"""
self.api_key = api_key
self.model = model
self.base_url = base_url.rstrip("/") if base_url else self.DEFAULT_BASE_URL
self.rerank_url = f"{self.base_url}{self.RERANK_PATH}"
self.timeout = timeout
self._async_client: httpx.AsyncClient | None = None
@@ -699,7 +738,7 @@ class ZeroEntropyCrossEncoder(CrossEncoderModel):
indices = [idx for idx, _ in indexed_texts]
response = await self._async_client.post(
self.RERANK_URL,
self.rerank_url,
json={
"model": self.model,
"query": query,
@@ -1229,6 +1268,164 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
return await loop.run_in_executor(None, self._predict_sync, pairs)
class GoogleCrossEncoder(CrossEncoderModel):
"""
Google Discovery Engine cross-encoder using the Ranking REST API.
Uses httpx + google-auth for lightweight REST calls (no gRPC/protobuf).
Supports ADC (Application Default Credentials) or service account key file.
Available models:
- semantic-ranker-default-004: Best quality, 1024 tokens/record (recommended)
- semantic-ranker-fast-004: Lower latency, 1024 tokens/record
Max 200 records per API request. Location is always "global".
"""
MAX_RECORDS_PER_REQUEST = 200
API_BASE = "https://discoveryengine.googleapis.com/v1"
SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
def __init__(
self,
project_id: str,
model: str = DEFAULT_RERANKER_GOOGLE_MODEL,
service_account_key: str | None = None,
location: str = "global",
timeout: float = 60.0,
):
"""
Initialize Google Discovery Engine cross-encoder.
Args:
project_id: Google Cloud project ID
model: Ranking model name (default: semantic-ranker-default-004)
service_account_key: Path to service account JSON key file.
If None, uses Application Default Credentials (ADC).
location: API location (default: "global")
timeout: Request timeout in seconds (default: 60.0)
"""
self.project_id = project_id
self.model = model
self.service_account_key = service_account_key
self.location = location
self.timeout = timeout
self._credentials = None
self._client: httpx.Client | None = None
self._rank_url: str | None = None
@property
def provider_name(self) -> str:
return "google"
def _get_auth_headers(self) -> dict[str, str]:
"""Get Authorization header with a fresh access token."""
import google.auth.transport.requests
if not self._credentials.valid:
self._credentials.refresh(google.auth.transport.requests.Request())
return {"Authorization": f"Bearer {self._credentials.token}"}
async def initialize(self) -> None:
"""Initialize credentials and HTTP client."""
if self._client is not None:
return
auth_method = "ADC" if not self.service_account_key else "service_account"
logger.info(
f"Reranker: initializing Google Discovery Engine provider "
f"(project={self.project_id}, model={self.model}, auth={auth_method})"
)
if self.service_account_key:
try:
from google.oauth2 import service_account
except ImportError:
raise ImportError(
"google-auth is required for GoogleCrossEncoder. Install it with: pip install google-auth"
)
self._credentials = service_account.Credentials.from_service_account_file(
self.service_account_key,
scopes=self.SCOPES,
)
else:
try:
import google.auth
except ImportError:
raise ImportError(
"google-auth is required for GoogleCrossEncoder. Install it with: pip install google-auth"
)
self._credentials, _ = google.auth.default(scopes=self.SCOPES)
ranking_config = f"projects/{self.project_id}/locations/{self.location}/rankingConfigs/default_ranking_config"
self._rank_url = f"{self.API_BASE}/{ranking_config}:rank"
self._client = httpx.Client(timeout=self.timeout)
logger.info("Reranker: Google Discovery Engine provider initialized")
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
"""Synchronous predict via REST API."""
if not pairs:
return []
# Group pairs by query
query_groups: dict[str, list[tuple[int, str]]] = {}
for idx, (query, text) in enumerate(pairs):
if query not in query_groups:
query_groups[query] = []
query_groups[query].append((idx, text))
all_scores = [0.0] * len(pairs)
for query, indexed_texts in query_groups.items():
texts = [text for _, text in indexed_texts]
indices = [idx for idx, _ in indexed_texts]
# Process in batches of MAX_RECORDS_PER_REQUEST
for batch_start in range(0, len(texts), self.MAX_RECORDS_PER_REQUEST):
batch_texts = texts[batch_start : batch_start + self.MAX_RECORDS_PER_REQUEST]
batch_indices = indices[batch_start : batch_start + self.MAX_RECORDS_PER_REQUEST]
records = [{"id": str(i), "content": text} for i, text in enumerate(batch_texts)]
response = self._client.post(
self._rank_url,
headers=self._get_auth_headers(),
json={
"model": self.model,
"query": query,
"records": records,
"topN": len(records),
},
)
response.raise_for_status()
result = response.json()
for record in result.get("records", []):
local_idx = int(record["id"])
all_scores[batch_indices[local_idx]] = record["score"]
return all_scores
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
"""
Score query-document pairs using Google Discovery Engine Ranking API.
Args:
pairs: List of (query, document) tuples to score
Returns:
List of relevance scores (0-1, higher = more relevant)
"""
if self._client is None:
raise RuntimeError("Reranker not initialized. Call initialize() first.")
if not pairs:
return []
loop = asyncio.get_event_loop()
return await loop.run_in_executor(None, self._predict_sync, pairs)
def create_cross_encoder_from_env() -> CrossEncoderModel:
"""
Create a CrossEncoderModel instance based on configuration.
@@ -1304,11 +1501,23 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
api_key=api_key,
model=config.reranker_zeroentropy_model,
)
elif provider == "google":
project_id = config.reranker_google_project_id
if not project_id:
raise ValueError(
f"{ENV_RERANKER_GOOGLE_PROJECT_ID} (or HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID) "
f"is required when {ENV_RERANKER_PROVIDER} is 'google'"
)
return GoogleCrossEncoder(
project_id=project_id,
model=config.reranker_google_model,
service_account_key=config.reranker_google_service_account_key,
)
elif provider == "rrf":
return RRFPassthroughCrossEncoder()
elif provider == "jina-mlx":
return JinaMLXCrossEncoder()
else:
raise ValueError(
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
)
@@ -13,11 +13,13 @@ import logging
import os
import warnings
from abc import ABC, abstractmethod
from urllib.parse import parse_qs, urlparse, urlunparse
import httpx
from ..config import (
DEFAULT_EMBEDDINGS_COHERE_MODEL,
DEFAULT_EMBEDDINGS_GEMINI_MODEL,
DEFAULT_EMBEDDINGS_LITELLM_MODEL,
DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU,
@@ -27,6 +29,7 @@ from ..config import (
DEFAULT_EMBEDDINGS_PROVIDER,
DEFAULT_LITELLM_API_BASE,
ENV_EMBEDDINGS_COHERE_API_KEY,
ENV_EMBEDDINGS_GEMINI_API_KEY,
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY,
ENV_EMBEDDINGS_LOCAL_FORCE_CPU,
ENV_EMBEDDINGS_LOCAL_MODEL,
@@ -426,9 +429,19 @@ class OpenAIEmbeddings(Embeddings):
logger.info(f"Embeddings: initializing OpenAI provider with model {self.model}{base_url_msg}")
# Build client kwargs, only including base_url if set (for Azure or custom endpoints)
# Parse query parameters from base_url (e.g. ?api-version=xxx for Azure OpenAI)
# and pass them as default_query so they're included in every request.
client_kwargs = {"api_key": self.api_key, "max_retries": self.max_retries}
if self.base_url:
client_kwargs["base_url"] = self.base_url
parsed = urlparse(self.base_url)
if parsed.query:
clean_url = urlunparse(parsed._replace(query=""))
client_kwargs["base_url"] = clean_url
default_query = {k: v[0] for k, v in parse_qs(parsed.query).items()}
client_kwargs["default_query"] = default_query
self.base_url = clean_url
else:
client_kwargs["base_url"] = self.base_url
self._client = OpenAI(**client_kwargs)
# Try to get dimension from known models, otherwise do a test embedding
@@ -741,6 +754,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
api_key: str,
model: str = DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
api_base: str | None = None,
output_dimensions: int | None = None,
batch_size: int = 100,
timeout: float = 60.0,
):
@@ -751,12 +765,14 @@ class LiteLLMSDKEmbeddings(Embeddings):
api_key: API key for the embedding provider
model: Model name with provider prefix (e.g., "cohere/embed-english-v3.0")
api_base: Custom base URL for API (optional)
output_dimensions: Optional output embedding dimensions (provider-dependent)
batch_size: Maximum batch size for embedding requests (default: 100)
timeout: Request timeout in seconds (default: 60.0)
"""
self.api_key = api_key
self.model = model
self.api_base = api_base
self.output_dimensions = output_dimensions
self.batch_size = batch_size
self.timeout = timeout
self._litellm = None # Will be set during initialization
@@ -798,6 +814,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
}
if self.api_base:
embed_kwargs["api_base"] = self.api_base
if self.output_dimensions is not None:
embed_kwargs["dimensions"] = self.output_dimensions
# Use async embedding method (standard in litellm)
response = await self._litellm.aembedding(**embed_kwargs)
@@ -845,6 +863,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
}
if self.api_base:
embed_kwargs["api_base"] = self.api_base
if self.output_dimensions is not None:
embed_kwargs["dimensions"] = self.output_dimensions
# Use sync embedding (litellm doesn't have async in thread-safe way)
response = self._litellm.embedding(**embed_kwargs)
@@ -866,6 +886,179 @@ class LiteLLMSDKEmbeddings(Embeddings):
return all_embeddings
class GeminiEmbeddings(Embeddings):
"""
Google embeddings via the google.genai SDK.
Supports both:
1. Gemini API (api.generativeai.google.com) with API key authentication
2. Vertex AI with service account or Application Default Credentials (ADC)
Uses the embed_content API: client.models.embed_content(model, contents)
"""
def __init__(
self,
model: str = DEFAULT_EMBEDDINGS_GEMINI_MODEL,
api_key: str | None = None,
vertexai_project_id: str | None = None,
vertexai_region: str | None = None,
vertexai_service_account_key: str | None = None,
output_dimensionality: int | None = None,
batch_size: int = 100,
):
self.model = model
self.api_key = api_key
self.vertexai_project_id = vertexai_project_id
self.vertexai_region = vertexai_region or "us-central1"
self.vertexai_service_account_key = vertexai_service_account_key
self.output_dimensionality = output_dimensionality
self.batch_size = batch_size
self._client = None
self._dimension: int | None = None
self._is_vertexai = vertexai_project_id is not None
self._embed_config = None # EmbedContentConfig, built during initialize()
@property
def provider_name(self) -> str:
return "google"
@property
def dimension(self) -> int:
if self._dimension is None:
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
return self._dimension
async def initialize(self) -> None:
"""Initialize the Google genai client and detect embedding dimension."""
if self._client is not None:
return
from google import genai
from google.genai import types as genai_types
if self._is_vertexai:
self._init_vertexai(genai)
else:
self._init_gemini(genai)
# Build EmbedContentConfig if output_dimensionality is set
if self.output_dimensionality is not None:
self._embed_config = genai_types.EmbedContentConfig(
output_dimensionality=self.output_dimensionality,
)
# Detect dimension via a test embedding (respects output_dimensionality)
embed_kwargs = {"model": self.model, "contents": ["test"]}
if self._embed_config is not None:
embed_kwargs["config"] = self._embed_config
result = self._client.models.embed_content(**embed_kwargs) # type: ignore[union-attr]
if result.embeddings and len(result.embeddings) > 0:
self._dimension = len(result.embeddings[0].values)
auth_mode = "vertex_ai" if self._is_vertexai else "api_key"
logger.info(
f"Embeddings: google provider initialized (auth: {auth_mode}, model: {self.model}, dim: {self._dimension})"
)
def _init_gemini(self, genai) -> None:
"""Initialize Gemini API client with API key."""
if not self.api_key:
raise ValueError("Gemini embeddings provider requires an API key")
self._client = genai.Client(api_key=self.api_key)
logger.info(f"Embeddings: initializing Gemini provider with model {self.model}")
def _init_vertexai(self, genai) -> None:
"""Initialize Vertex AI client with project, region, and credentials."""
if not self.vertexai_project_id:
raise ValueError(
"HINDSIGHT_API_EMBEDDINGS_VERTEXAI_PROJECT_ID (or HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID) "
"is required for Vertex AI embeddings provider."
)
auth_method = "ADC"
credentials = None
if self.vertexai_service_account_key:
try:
from google.oauth2 import service_account
except ImportError:
raise ImportError(
"Vertex AI service account auth requires 'google-auth' package. "
"Install with: pip install google-auth"
)
credentials = service_account.Credentials.from_service_account_file(
self.vertexai_service_account_key,
scopes=["https://www.googleapis.com/auth/cloud-platform"],
)
auth_method = "service_account"
logger.info(f"Embeddings: Vertex AI using service account key: {self.vertexai_service_account_key}")
# Strip google/ prefix from model name — native SDK uses bare names
if self.model.startswith("google/"):
self.model = self.model[len("google/") :]
client_kwargs = {
"vertexai": True,
"project": self.vertexai_project_id,
"location": self.vertexai_region,
}
if credentials is not None:
client_kwargs["credentials"] = credentials
self._client = genai.Client(**client_kwargs)
logger.info(
f"Embeddings: initializing Vertex AI provider "
f"(project={self.vertexai_project_id}, region={self.vertexai_region}, "
f"model={self.model}, auth={auth_method})"
)
def encode(self, texts: list[str]) -> list[list[float]]:
"""
Generate embeddings using the Google genai SDK.
Args:
texts: List of text strings to encode
Returns:
List of embedding vectors
"""
if self._client is None:
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
if not texts:
return []
all_embeddings = []
# Process in batches
for i in range(0, len(texts), self.batch_size):
batch = texts[i : i + self.batch_size]
embed_kwargs = {"model": self.model, "contents": batch}
if self._embed_config is not None:
embed_kwargs["config"] = self._embed_config
result = self._client.models.embed_content(**embed_kwargs)
all_embeddings.extend([emb.values for emb in result.embeddings])
# L2-normalize when output_dimensionality is set — Gemini only returns
# normalized vectors at full 3072 dims; truncated dims need re-normalization
# for accurate cosine similarity.
if self.output_dimensionality is not None:
import numpy as np
arr = np.array(all_embeddings)
norms = np.linalg.norm(arr, axis=1, keepdims=True)
norms[norms == 0] = 1
all_embeddings = (arr / norms).tolist()
return all_embeddings
def create_embeddings_from_env() -> Embeddings:
"""
Create an Embeddings instance based on configuration.
@@ -927,9 +1120,29 @@ def create_embeddings_from_env() -> Embeddings:
api_key=api_key,
model=config.embeddings_litellm_sdk_model,
api_base=config.embeddings_litellm_sdk_api_base,
output_dimensions=config.embeddings_litellm_sdk_output_dimensions,
)
elif provider == "google":
vertexai_project_id = config.embeddings_vertexai_project_id
if vertexai_project_id:
api_key = None # Vertex AI uses ADC or service account
else:
api_key = config.embeddings_gemini_api_key
if not api_key:
raise ValueError(
f"{ENV_EMBEDDINGS_GEMINI_API_KEY} or {ENV_LLM_API_KEY} is required "
f"when {ENV_EMBEDDINGS_PROVIDER} is 'google' (set VERTEXAI_PROJECT_ID for Vertex AI auth instead)"
)
return GeminiEmbeddings(
model=config.embeddings_gemini_model,
api_key=api_key,
vertexai_project_id=vertexai_project_id,
vertexai_region=config.embeddings_vertexai_region,
vertexai_service_account_key=config.embeddings_vertexai_service_account_key,
output_dimensionality=config.embeddings_gemini_output_dimensionality,
)
else:
raise ValueError(
f"Unknown embeddings provider: {provider}. "
f"Supported: 'local', 'tei', 'openai', 'cohere', 'litellm', 'litellm-sdk'"
f"Supported: 'local', 'tei', 'openai', 'cohere', 'google', 'litellm', 'litellm-sdk'"
)
@@ -317,8 +317,13 @@ class EntityResolver:
entity_texts = list(set(e["text"] for e in entities_data))
# Fetch candidates for all unique entity texts in a single batched query.
# The trigram % operator uses the GIN index; the substring conditions cover
# exact prefix/suffix matches that trigrams might miss at low similarity.
# Uses the GIN trigram index on LOWER(canonical_name) for case-insensitive
# similarity lookup. Previous version also had LIKE '%...' substring fallbacks,
# but those forced full sequential scans of the entities table and caused
# TimeoutErrors on banks with 10k+ entities. Lowering the similarity threshold
# to 0.15 (from default 0.3) catches most substring relationships while
# staying fully index-based.
await conn.execute("SET pg_trgm.similarity_threshold = 0.15")
rows = await conn.fetch(
f"""
SELECT DISTINCT ON (e.id)
@@ -327,16 +332,13 @@ class EntityResolver:
FROM unnest($2::text[]) AS q(query_text)
JOIN {fq_table("entities")} e ON (
e.bank_id = $1
AND (
e.canonical_name % q.query_text
OR LOWER(e.canonical_name) LIKE '%' || LOWER(q.query_text) || '%'
OR LOWER(q.query_text) LIKE '%' || LOWER(e.canonical_name) || '%'
)
AND LOWER(e.canonical_name) % LOWER(q.query_text)
)
""",
bank_id,
entity_texts,
)
await conn.execute("RESET pg_trgm.similarity_threshold")
# Group candidates by query_text
all_candidates: dict[str, list] = {t: [] for t in entity_texts}
@@ -808,14 +810,19 @@ class EntityResolver:
return await self._link_units_to_entities_batch_impl(conn, unit_entity_pairs)
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(
# Sorted bulk insert to prevent deadlocks from inconsistent lock ordering
# across concurrent transactions on the unit_entities unique index.
sorted_pairs = sorted(unit_entity_pairs)
unit_ids = [p[0] for p in sorted_pairs]
entity_ids = [p[1] for p in sorted_pairs]
await conn.execute(
f"""
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
VALUES ($1, $2)
SELECT u, e FROM unnest($1::uuid[], $2::uuid[]) AS t(u, e)
ON CONFLICT DO NOTHING
""",
unit_entity_pairs,
unit_ids,
entity_ids,
)
# Build map of unit -> entities for co-occurrence calculation
@@ -240,6 +240,7 @@ class MemoryEngineInterface(ABC):
bank_id: str,
*,
fact_type: str | None = None,
delete_bank_profile: bool = True,
request_context: "RequestContext",
) -> dict[str, int]:
"""
@@ -248,6 +249,8 @@ class MemoryEngineInterface(ABC):
Args:
bank_id: The memory bank ID.
fact_type: If specified, only delete memories of this type.
delete_bank_profile: If True, also delete the bank profile row itself.
If False, only delete memories/entities/documents but preserve the bank.
request_context: Request context for authentication.
Returns:
@@ -146,6 +146,7 @@ def create_llm_provider(
reasoning_effort: str,
groq_service_tier: str | None = None,
openai_service_tier: str | None = None,
extra_body: dict[str, Any] | None = None,
vertexai_project_id: str | None = None,
vertexai_region: str | None = None,
vertexai_credentials: Any = None,
@@ -162,6 +163,7 @@ def create_llm_provider(
reasoning_effort: Reasoning effort level for supported providers.
groq_service_tier: Groq service tier (for Groq provider) - "on_demand", "flex", or "auto".
openai_service_tier: OpenAI service tier (for OpenAI provider) - None (default) or "flex" (50% cheaper).
extra_body: Extra body params merged into OpenAI-compatible API calls.
vertexai_project_id: Vertex AI project ID (for VertexAI provider).
vertexai_region: Vertex AI region (for VertexAI provider).
vertexai_credentials: Vertex AI credentials object (for VertexAI provider).
@@ -261,7 +263,7 @@ def create_llm_provider(
reasoning_effort=reasoning_effort,
)
elif provider_lower in ("openai", "groq", "ollama", "lmstudio", "minimax"):
elif provider_lower in ("openai", "groq", "ollama", "lmstudio", "minimax", "volcano"):
return OpenAICompatibleLLM(
provider=provider,
api_key=api_key,
@@ -270,6 +272,7 @@ def create_llm_provider(
reasoning_effort=reasoning_effort,
groq_service_tier=groq_service_tier,
openai_service_tier=openai_service_tier,
extra_body=extra_body,
)
else:
@@ -293,6 +296,7 @@ class LLMProvider:
groq_service_tier: str | None = None,
openai_service_tier: str | None = None,
gemini_safety_settings: list | None = None,
extra_body: dict[str, Any] | None = None,
):
"""
Initialize LLM provider.
@@ -306,6 +310,7 @@ class LLMProvider:
groq_service_tier: Groq service tier ("on_demand", "flex", "auto") - from config.
openai_service_tier: OpenAI service tier (None or "flex") - from config.
gemini_safety_settings: Safety settings for Gemini/VertexAI providers.
extra_body: Extra body params merged into OpenAI-compatible API calls.
"""
self.provider = provider.lower()
self.api_key = api_key
@@ -317,6 +322,8 @@ class LLMProvider:
self.openai_service_tier = openai_service_tier
# Gemini safety settings (instance default; can be overridden per-request via context var)
self.gemini_safety_settings = gemini_safety_settings
# Extra body params for OpenAI-compatible providers (e.g. chat_template_kwargs)
self.extra_body = extra_body
# Validate provider
valid_providers = [
@@ -334,6 +341,7 @@ class LLMProvider:
"minimax",
"litellm",
"bedrock",
"volcano",
]
if self.provider not in valid_providers:
raise ValueError(f"Invalid LLM provider: {self.provider}. Must be one of: {', '.join(valid_providers)}")
@@ -412,6 +420,7 @@ class LLMProvider:
reasoning_effort=self.reasoning_effort,
groq_service_tier=self.groq_service_tier,
openai_service_tier=self.openai_service_tier,
extra_body=self.extra_body,
vertexai_project_id=vertexai_project_id,
vertexai_region=vertexai_region,
vertexai_credentials=vertexai_credentials,
@@ -706,65 +715,40 @@ class LLMProvider:
pass
@classmethod
def for_memory(cls) -> "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", "")
def from_env(cls) -> "LLMProvider":
"""Create provider from environment variables using config.py constants."""
from ..config import (
DEFAULT_LLM_MODEL,
DEFAULT_LLM_PROVIDER,
ENV_LLM_API_KEY,
ENV_LLM_BASE_URL,
ENV_LLM_EXTRA_BODY,
ENV_LLM_MODEL,
ENV_LLM_PROVIDER,
)
provider = os.getenv(ENV_LLM_PROVIDER, DEFAULT_LLM_PROVIDER)
api_key = os.getenv(ENV_LLM_API_KEY, "")
# API key not needed for openai-codex (uses OAuth), claude-code (uses Keychain OAuth),
# ollama (local), vertexai (uses GCP service account credentials),
# or litellm (uses provider-specific auth, e.g. AWS credentials for Bedrock)
if not api_key and not requires_api_key(provider):
pass # Provider handles its own auth
elif not api_key:
raise ValueError(
"HINDSIGHT_API_LLM_API_KEY environment variable is required (unless using openai-codex, claude-code, or litellm)"
f"{ENV_LLM_API_KEY} environment variable is required (unless using openai-codex, claude-code, or litellm)"
)
base_url = os.getenv("HINDSIGHT_API_LLM_BASE_URL", "")
model = os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b")
base_url = os.getenv(ENV_LLM_BASE_URL, "")
model = os.getenv(ENV_LLM_MODEL, DEFAULT_LLM_MODEL)
extra_body = json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null"))
return cls(provider=provider, api_key=api_key, base_url=base_url, model=model, reasoning_effort="low")
@classmethod
def for_answer_generation(cls) -> "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", ""))
# API key not needed for providers with their own auth mechanisms
if not api_key and not requires_api_key(provider):
pass # Provider handles its own auth
elif not api_key:
raise ValueError(
"HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_API_ANSWER_LLM_API_KEY environment variable is required "
"(unless using openai-codex, claude-code, or litellm)"
)
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"))
return cls(provider=provider, api_key=api_key, base_url=base_url, model=model, reasoning_effort="high")
@classmethod
def for_judge(cls) -> "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", ""))
# API key not needed for providers with their own auth mechanisms
if not api_key and not requires_api_key(provider):
pass # Provider handles its own auth
elif not api_key:
raise ValueError(
"HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_API_JUDGE_LLM_API_KEY environment variable is required "
"(unless using openai-codex, claude-code, or litellm)"
)
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"))
return cls(provider=provider, api_key=api_key, base_url=base_url, model=model, reasoning_effort="high")
return cls(
provider=provider,
api_key=api_key,
base_url=base_url,
model=model,
reasoning_effort="low",
extra_body=extra_body,
)
class ConfiguredLLMProvider:
@@ -16,6 +16,7 @@ import logging
import time
import uuid
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any
@@ -28,6 +29,7 @@ from ..metrics import get_metrics_collector
from ..tracing import create_operation_span
from ..utils import mask_network_location
from ..worker.exceptions import RetryTaskAt
from .audit import AuditLogger, audit_context
from .db_budget import budgeted_operation
from .operation_metadata import (
BatchRetainChildMetadata,
@@ -225,6 +227,42 @@ def _get_tiktoken_encoding():
return _TIKTOKEN_ENCODING
@dataclass(frozen=True)
class RefreshTagFiltering:
"""Resolved tag filtering parameters for mental model refresh."""
tags: list[str] | None
tags_match: TagsMatch
tag_groups: list[TagGroup] | None
def _resolve_refresh_tag_filtering(
model_tags: list[str] | None,
trigger_data: dict[str, Any],
) -> RefreshTagFiltering:
"""Resolve tag filtering parameters for mental model refresh.
Takes raw trigger dict from DB (JSONB with no fixed schema guarantee)
and resolves the tag filtering to use during reflect.
Priority:
- If trigger has tag_groups, use those (overrides flat tags entirely)
- If trigger has tags_match, use model's tags with that match mode
- Otherwise default to all_strict when tags present (security isolation)
"""
trigger_tag_groups = trigger_data.get("tag_groups")
if trigger_tag_groups is not None:
from pydantic import TypeAdapter
adapter = TypeAdapter(TagGroup)
parsed = [adapter.validate_python(tg) for tg in trigger_tag_groups]
return RefreshTagFiltering(tags=None, tags_match="any", tag_groups=parsed)
trigger_tags_match = trigger_data.get("tags_match")
tags_match: TagsMatch = trigger_tags_match if trigger_tags_match else ("all_strict" if model_tags else "any")
return RefreshTagFiltering(tags=model_tags, tags_match=tags_match, tag_groups=None)
class MemoryEngine(MemoryEngineInterface):
"""
Advanced memory system using temporal and semantic linking with PostgreSQL.
@@ -232,7 +270,7 @@ class MemoryEngine(MemoryEngineInterface):
This class provides:
- Embedding generation for semantic search
- Entity, temporal, and semantic link creation
- Think operations for formulating answers with opinions
- Think operations for formulating answers with observations
- bank profile and disposition management
"""
@@ -395,6 +433,7 @@ class MemoryEngine(MemoryEngineInterface):
api_key=memory_llm_api_key,
base_url=memory_llm_base_url,
model=memory_llm_model,
extra_body=config.llm_extra_body,
)
# Store client and model for convenience (deprecated: use _llm_config.call() instead)
@@ -421,6 +460,7 @@ class MemoryEngine(MemoryEngineInterface):
api_key=retain_api_key,
base_url=retain_base_url,
model=retain_model,
extra_body=config.llm_extra_body,
)
# Reflect LLM config - for think/observe operations (can use lighter models)
@@ -442,6 +482,7 @@ class MemoryEngine(MemoryEngineInterface):
api_key=reflect_api_key,
base_url=reflect_base_url,
model=reflect_model,
extra_body=config.llm_extra_body,
)
# Consolidation LLM config - for mental model consolidation (can use efficient models)
@@ -463,6 +504,7 @@ class MemoryEngine(MemoryEngineInterface):
api_key=consolidation_api_key,
base_url=consolidation_base_url,
model=consolidation_model,
extra_body=config.llm_extra_body,
)
# Initialize cross-encoder reranker (cached for performance)
@@ -476,14 +518,25 @@ class MemoryEngine(MemoryEngineInterface):
schema_getter=get_current_schema,
)
# Audit logger for feature usage tracking
config = get_config()
self._audit_logger = AuditLogger(
pool_getter=lambda: self._pool,
schema_getter=get_current_schema,
enabled=config.audit_log_enabled,
allowed_actions=config.audit_log_actions,
retention_days=config.audit_log_retention_days,
)
# Backpressure mechanism: limit concurrent searches to prevent overwhelming the database
# Configurable via HINDSIGHT_API_RECALL_MAX_CONCURRENT (default: 50)
self._search_semaphore = asyncio.Semaphore(get_config().recall_max_concurrent)
# Backpressure for put operations: limit concurrent puts to prevent database contention
# Each put_batch holds a connection for the entire transaction, so we limit to 5
# concurrent puts to avoid connection pool exhaustion and reduce write contention
self._put_semaphore = asyncio.Semaphore(5)
# Backpressure for retain DB writes: limit concurrent transactions to prevent contention
# on entity/link tables. Acquired in the orchestrator *after* LLM extraction completes,
# so LLM calls run in full parallelism while only the DB-heavy phase is throttled.
# Configurable via HINDSIGHT_API_RETAIN_MAX_CONCURRENT (default: 4).
self._put_semaphore = asyncio.Semaphore(get_config().retain_max_concurrent)
# initialize encoding eagerly to avoid delaying the first time
_get_tiktoken_encoding()
@@ -498,6 +551,11 @@ class MemoryEngine(MemoryEngineInterface):
tenant_extension = DefaultTenantExtension(config={})
self._tenant_extension = tenant_extension
@property
def audit_logger(self) -> AuditLogger:
"""The audit logger for feature usage tracking."""
return self._audit_logger
@property
def tenant_extension(self) -> "TenantExtension | None":
"""The configured tenant extension, if any."""
@@ -888,26 +946,23 @@ class MemoryEngine(MemoryEngineInterface):
source_query = mental_model["source_query"]
# SECURITY: If the mental model has tags, pass them to reflect with "all_strict" matching
# to ensure it can only access other mental models/memories with the SAME tags.
# This prevents cross-tenant/cross-user information leakage by excluding untagged content.
tags = mental_model.get("tags")
tags_match = "all_strict" if tags else "any"
# Read reflect options from trigger (if stored)
trigger_data = mental_model.get("trigger") or {}
fact_types = trigger_data.get("fact_types")
exclude_mental_models = trigger_data.get("exclude_mental_models", False)
stored_exclude_ids: list[str] = trigger_data.get("exclude_mental_model_ids") or []
tag_filtering = _resolve_refresh_tag_filtering(mental_model.get("tags"), trigger_data)
# Run reflect to generate new content, excluding the mental model being refreshed
# Always add self to excluded IDs to prevent circular reference
reflect_result = await self.reflect_async(
bank_id=bank_id,
query=source_query,
request_context=internal_context,
tags=tags,
tags_match=tags_match,
tags=tag_filtering.tags,
tags_match=tag_filtering.tags_match,
tag_groups=tag_filtering.tag_groups,
fact_types=fact_types,
exclude_mental_models=exclude_mental_models,
exclude_mental_model_ids=list({*stored_exclude_ids, mental_model_id}),
@@ -1030,72 +1085,78 @@ class MemoryEngine(MemoryEngineInterface):
# Continue with processing if we can't check status
consolidation_result: dict | None = None
try:
if task_type == "batch_retain":
await self._handle_batch_retain(task_dict)
elif task_type == "file_convert_retain":
await self._handle_file_convert_retain(task_dict)
elif task_type == "consolidation":
consolidation_result = await self._handle_consolidation(task_dict)
elif task_type == "refresh_mental_model":
await self._handle_refresh_mental_model(task_dict)
elif task_type == "webhook_delivery":
await self._handle_webhook_delivery(task_dict)
else:
logger.error(f"Unknown task type: {task_type}")
# Don't retry unknown task types
if operation_id:
await self._delete_operation_record(operation_id)
return
# Task succeeded - mark operation as completed
# file_convert_retain marks itself as completed in a transaction, skip double-marking
if operation_id and task_type not in ("file_convert_retain",):
if task_type == "consolidation":
# Atomically mark completed AND queue webhook delivery in one transaction
await self._mark_operation_completed_and_fire_webhook(
operation_id=operation_id,
bank_id=task_dict.get("bank_id", ""),
status="completed",
result=consolidation_result,
schema=schema,
)
bank_id = task_dict.get("bank_id")
async with audit_context(
self._audit_logger, task_type or "unknown", "system", bank_id, request=task_dict
) as audit_entry:
try:
if task_type == "batch_retain":
await self._handle_batch_retain(task_dict)
elif task_type == "file_convert_retain":
await self._handle_file_convert_retain(task_dict)
elif task_type == "consolidation":
consolidation_result = await self._handle_consolidation(task_dict)
elif task_type == "refresh_mental_model":
await self._handle_refresh_mental_model(task_dict)
elif task_type == "webhook_delivery":
await self._handle_webhook_delivery(task_dict)
else:
await self._mark_operation_completed(operation_id)
logger.error(f"Unknown task type: {task_type}")
# Don't retry unknown task types
if operation_id:
await self._delete_operation_record(operation_id)
return
except RetryTaskAt:
# Task-owned retry: let the poller handle scheduling
raise
except Exception as e:
logger.error(f"Task execution failed: {task_type}, error: {e}")
import traceback
# Task succeeded - mark operation as completed
# file_convert_retain marks itself as completed in a transaction, skip double-marking
if operation_id and task_type not in ("file_convert_retain",):
if task_type == "consolidation":
# Atomically mark completed AND queue webhook delivery in one transaction
await self._mark_operation_completed_and_fire_webhook(
operation_id=operation_id,
bank_id=task_dict.get("bank_id", ""),
status="completed",
result=consolidation_result,
schema=schema,
)
else:
await self._mark_operation_completed(operation_id)
error_traceback = traceback.format_exc()
traceback.print_exc()
audit_entry.response = {"status": "completed", "operation_id": operation_id}
if task_type == "file_convert_retain":
# Non-retryable: mark as failed immediately.
# Conversion failures won't improve on retry (missing OCR, corrupted file, etc.)
logger.error(f"Not retrying task {task_type} (non-retryable), marking as failed")
if operation_id:
await self._mark_operation_failed(operation_id, str(e), error_traceback)
else:
if task_type == "consolidation" and operation_id:
# Fire failure webhook (non-transactional — operation not yet marked failed;
# poller will mark it failed after this raise)
await self._fire_consolidation_webhook(
bank_id=task_dict.get("bank_id", ""),
operation_id=operation_id,
status="failed",
result=None,
error_message=str(e),
schema=schema,
)
# Retryable: use RetryTaskAt if under the retry limit, else re-raise (poller marks failed)
retry_count = task_dict.get("_retry_count", 0)
if retry_count < 3:
raise RetryTaskAt(retry_at=datetime.now(UTC) + timedelta(seconds=60), message=str(e))
except RetryTaskAt:
# Task-owned retry: let the poller handle scheduling
raise
except Exception as e:
logger.error(f"Task execution failed: {task_type}, error: {e}")
import traceback
error_traceback = traceback.format_exc()
traceback.print_exc()
if task_type == "file_convert_retain":
# Non-retryable: mark as failed immediately.
# Conversion failures won't improve on retry (missing OCR, corrupted file, etc.)
logger.error(f"Not retrying task {task_type} (non-retryable), marking as failed")
if operation_id:
await self._mark_operation_failed(operation_id, str(e), error_traceback)
else:
if task_type == "consolidation" and operation_id:
# Fire failure webhook (non-transactional — operation not yet marked failed;
# poller will mark it failed after this raise)
await self._fire_consolidation_webhook(
bank_id=task_dict.get("bank_id", ""),
operation_id=operation_id,
status="failed",
result=None,
error_message=str(e),
schema=schema,
)
# Retryable: use RetryTaskAt if under the retry limit, else re-raise (poller marks failed)
retry_count = task_dict.get("_retry_count", 0)
if retry_count < 3:
raise RetryTaskAt(retry_at=datetime.now(UTC) + timedelta(seconds=60), message=str(e))
raise
async def _fire_consolidation_webhook(
self,
@@ -1659,18 +1720,16 @@ class MemoryEngine(MemoryEngineInterface):
# Migrate all schemas from the tenant extension
# The tenant extension is the single source of truth for which schemas exist
logger.info("Running database migrations...")
config = get_config()
tenants = await self._tenant_extension.list_tenants()
if tenants:
logger.info(f"Running migrations on {len(tenants)} schema(s)...")
for tenant in tenants:
schema = tenant.schema
if schema:
run_migrations(self.db_url, schema=schema)
run_migrations(self.db_url, schema=schema, migration_database_url=config.migration_database_url)
logger.info("Schema migrations completed")
# Get config for vector extension setting
config = get_config()
# Ensure embedding column dimension matches the model's dimension
# This is done after migrations and after embeddings.initialize()
for tenant in tenants:
@@ -1791,6 +1850,9 @@ class MemoryEngine(MemoryEngineInterface):
self._task_backend.set_executor(self.execute_task)
await self._task_backend.initialize()
# Start audit log retention sweep (if configured)
self._audit_logger.start_retention_sweep()
self._initialized = True
logger.info("Memory system initialized (pool and task backend started)")
@@ -1843,6 +1905,9 @@ class MemoryEngine(MemoryEngineInterface):
"""Close the connection pool and shutdown background workers."""
logger.info("close() started")
# Stop audit log retention sweep
await self._audit_logger.stop_retention_sweep()
# Shutdown task backend
await self._task_backend.shutdown()
@@ -1938,7 +2003,6 @@ class MemoryEngine(MemoryEngineInterface):
event_date: datetime | None = None,
document_id: str | None = None,
fact_type_override: str | None = None,
confidence_score: float | None = None,
*,
request_context: "RequestContext",
) -> list[str]:
@@ -1954,7 +2018,6 @@ class MemoryEngine(MemoryEngineInterface):
event_date: When the event occurred (defaults to now)
document_id: Optional document ID for tracking (always upserts if document already exists)
fact_type_override: Override fact type ('world', 'experience')
confidence_score: Confidence score (0.0 to 1.0)
request_context: Request context for authentication.
Returns:
@@ -1973,7 +2036,6 @@ class MemoryEngine(MemoryEngineInterface):
contents=[content_dict],
request_context=request_context,
fact_type_override=fact_type_override,
confidence_score=confidence_score,
)
# Return the first (and only) list of unit IDs
@@ -1987,7 +2049,6 @@ class MemoryEngine(MemoryEngineInterface):
request_context: "RequestContext",
document_id: str | None = None,
fact_type_override: str | None = None,
confidence_score: float | None = None,
document_tags: list[str] | None = None,
return_usage: bool = False,
operation_id: str | None = None,
@@ -2013,7 +2074,6 @@ class MemoryEngine(MemoryEngineInterface):
document_id: **DEPRECATED** - Use "document_id" key in each content dict instead.
Applies the same document_id to ALL content items that don't specify their own.
fact_type_override: Override fact type for all facts ('world', 'experience')
confidence_score: Confidence score (0.0 to 1.0)
return_usage: If True, returns tuple of (unit_ids, TokenUsage). Default False for backward compatibility.
Returns:
@@ -2063,7 +2123,6 @@ class MemoryEngine(MemoryEngineInterface):
request_context=request_context,
document_id=document_id,
fact_type_override=fact_type_override,
confidence_score=confidence_score,
)
result = await self._validate_operation(self._operation_validator.validate_retain(ctx))
if result and result.contents is not None:
@@ -2149,7 +2208,6 @@ class MemoryEngine(MemoryEngineInterface):
document_id=document_id,
is_first_batch=i == 1, # Only upsert on first batch
fact_type_override=fact_type_override,
confidence_score=confidence_score,
document_tags=document_tags,
operation_id=operation_id,
strategy=strategy,
@@ -2174,7 +2232,6 @@ class MemoryEngine(MemoryEngineInterface):
document_id=document_id,
is_first_batch=True,
fact_type_override=fact_type_override,
confidence_score=confidence_score,
document_tags=document_tags,
operation_id=operation_id,
strategy=strategy,
@@ -2191,7 +2248,6 @@ class MemoryEngine(MemoryEngineInterface):
request_context=request_context,
document_id=document_id,
fact_type_override=fact_type_override,
confidence_score=confidence_score,
unit_ids=result,
success=True,
error=None,
@@ -2226,7 +2282,6 @@ class MemoryEngine(MemoryEngineInterface):
document_id: str | None = None,
is_first_batch: bool = True,
fact_type_override: str | None = None,
confidence_score: float | None = None,
document_tags: list[str] | None = None,
operation_id: str | None = None,
outbox_callback: "Callable[[asyncpg.Connection], Awaitable[None]] | None" = None,
@@ -2247,54 +2302,51 @@ class MemoryEngine(MemoryEngineInterface):
document_id: Optional document ID (always upserts if exists)
is_first_batch: Whether this is the first batch (for chunked operations, only delete on first batch)
fact_type_override: Override fact type for all facts
confidence_score: Confidence score for opinions
document_tags: Tags applied to all items in this batch
Returns:
Tuple of (unit ID lists, token usage for fact extraction)
"""
# Backpressure: limit concurrent retains to prevent database contention
async with self._put_semaphore:
# Use the new modular orchestrator
from .retain import orchestrator
# Use the new modular orchestrator
from .retain import orchestrator
pool = await self._get_pool()
pool = await self._get_pool()
# Resolve bank-specific config for this operation
resolved_config = await self._config_resolver.resolve_full_config(bank_id, request_context)
# Resolve bank-specific config for this operation
resolved_config = await self._config_resolver.resolve_full_config(bank_id, request_context)
# Force chunks mode when LLM provider is "none" (no LLM available for fact extraction)
if self._llm_config.provider == "none":
resolved_config.retain_extraction_mode = "chunks"
resolved_config.enable_observations = False
# Force chunks mode when LLM provider is "none" (no LLM available for fact extraction)
if self._llm_config.provider == "none":
resolved_config.retain_extraction_mode = "chunks"
resolved_config.enable_observations = False
# Apply strategy overrides: explicit strategy > bank default strategy
from hindsight_api.config_resolver import apply_strategy
# Apply strategy overrides: explicit strategy > bank default strategy
from hindsight_api.config_resolver import apply_strategy
effective_strategy = strategy or resolved_config.retain_default_strategy
if effective_strategy:
resolved_config = apply_strategy(resolved_config, effective_strategy)
effective_strategy = strategy or resolved_config.retain_default_strategy
if effective_strategy:
resolved_config = apply_strategy(resolved_config, effective_strategy)
# Create parent span for retain operation
with create_operation_span("retain", bank_id):
return await orchestrator.retain_batch(
pool=pool,
embeddings_model=self.embeddings,
llm_config=self._retain_llm_config.with_config(resolved_config),
entity_resolver=self.entity_resolver,
format_date_fn=self._format_readable_date,
bank_id=bank_id,
contents_dicts=contents,
document_id=document_id,
is_first_batch=is_first_batch,
fact_type_override=fact_type_override,
confidence_score=confidence_score,
document_tags=document_tags,
config=resolved_config,
operation_id=operation_id,
schema=_current_schema.get(),
outbox_callback=outbox_callback,
)
# Create parent span for retain operation
with create_operation_span("retain", bank_id):
return await orchestrator.retain_batch(
pool=pool,
embeddings_model=self.embeddings,
llm_config=self._retain_llm_config.with_config(resolved_config),
entity_resolver=self.entity_resolver,
format_date_fn=self._format_readable_date,
bank_id=bank_id,
contents_dicts=contents,
document_id=document_id,
is_first_batch=is_first_batch,
fact_type_override=fact_type_override,
document_tags=document_tags,
config=resolved_config,
operation_id=operation_id,
schema=_current_schema.get(),
outbox_callback=outbox_callback,
db_semaphore=self._put_semaphore,
)
def recall(
self,
@@ -2314,7 +2366,7 @@ class MemoryEngine(MemoryEngineInterface):
Args:
bank_id: bank ID to recall for
query: Recall query
fact_type: Required filter for fact type ('world', 'experience', or 'opinion')
fact_type: Required filter for fact type ('world' or 'experience')
budget: Budget level for graph traversal (low=100, mid=300, high=600 units)
max_tokens: Maximum tokens to return (counts only 'text' field, default 4096)
enable_trace: If True, returns detailed trace object
@@ -2408,8 +2460,10 @@ class MemoryEngine(MemoryEngineInterface):
if fact_type is None:
fact_type = list(VALID_RECALL_FACT_TYPES)
# Filter out 'opinion' early (deprecated, silently ignore)
# Filter out 'opinion' (removed fact type, silently ignore for backwards compat)
fact_type = [ft for ft in fact_type if ft != "opinion"]
if not fact_type:
return RecallResultModel(results=[], entities={}, chunks={})
# Validate fact types
invalid_types = set(fact_type) - VALID_RECALL_FACT_TYPES
@@ -2418,9 +2472,6 @@ class MemoryEngine(MemoryEngineInterface):
f"Invalid fact type(s): {', '.join(sorted(invalid_types))}. "
f"Must be one of: {', '.join(sorted(VALID_RECALL_FACT_TYPES))}"
)
if not fact_type:
# All requested types were opinions - return empty result
return RecallResultModel(results=[], entities={}, chunks={})
# Validate operation if validator is configured
if self._operation_validator:
@@ -2768,7 +2819,7 @@ class MemoryEngine(MemoryEngineInterface):
"temporal": 0.0,
"temporal_extraction": 0.0,
}
all_mpfp_timings = []
all_graph_timings = []
detected_temporal_constraint = None
max_conn_wait = multi_result.max_conn_wait
@@ -2830,25 +2881,25 @@ class MemoryEngine(MemoryEngineInterface):
)
# Log graph retriever timing breakdown if available
if all_mpfp_timings:
if all_graph_timings:
retriever_name = get_default_graph_retriever().name.upper()
mpfp_total = all_mpfp_timings[0] # Take first fact type's timing as representative
mpfp_parts = [
f"db_queries={mpfp_total.db_queries}",
f"edge_load={mpfp_total.edge_load_time:.3f}s",
f"edges={mpfp_total.edge_count}",
f"patterns={mpfp_total.pattern_count}",
graph_total = all_graph_timings[0] # Take first fact type's timing as representative
graph_parts = [
f"db_queries={graph_total.db_queries}",
f"edge_load={graph_total.edge_load_time:.3f}s",
f"edges={graph_total.edge_count}",
f"patterns={graph_total.pattern_count}",
]
if mpfp_total.seeds_time > 0.01:
mpfp_parts.append(f"seeds={mpfp_total.seeds_time:.3f}s")
if mpfp_total.fusion > 0.001:
mpfp_parts.append(f"fusion={mpfp_total.fusion:.3f}s")
if mpfp_total.fetch > 0.001:
mpfp_parts.append(f"fetch={mpfp_total.fetch:.3f}s")
log_buffer.append(f" [{retriever_name}] {', '.join(mpfp_parts)}")
if graph_total.seeds_time > 0.01:
graph_parts.append(f"seeds={graph_total.seeds_time:.3f}s")
if graph_total.fusion > 0.001:
graph_parts.append(f"fusion={graph_total.fusion:.3f}s")
if graph_total.fetch > 0.001:
graph_parts.append(f"fetch={graph_total.fetch:.3f}s")
log_buffer.append(f" [{retriever_name}] {', '.join(graph_parts)}")
# Log detailed hop timing for debugging slow queries
if mpfp_total.hop_details:
for hd in mpfp_total.hop_details:
if graph_total.hop_details:
for hd in graph_total.hop_details:
log_buffer.append(
f" hop{hd['hop']}: exec={hd.get('exec_time', 0) * 1000:.0f}ms, "
f"uncached={hd.get('uncached_after_filter', 0)}, "
@@ -3483,11 +3534,13 @@ class MemoryEngine(MemoryEngineInterface):
doc = await conn.fetchrow(
f"""
SELECT d.id, d.bank_id, d.original_text, d.content_hash,
d.created_at, d.updated_at, d.tags, COUNT(mu.id) as unit_count
d.created_at, d.updated_at, d.tags, d.retain_params,
COUNT(mu.id) as unit_count
FROM {fq_table("documents")} d
LEFT JOIN {fq_table("memory_units")} mu ON mu.document_id = d.id
WHERE d.id = $1 AND d.bank_id = $2
GROUP BY d.id, d.bank_id, d.original_text, d.content_hash, d.created_at, d.updated_at, d.tags
GROUP BY d.id, d.bank_id, d.original_text, d.content_hash,
d.created_at, d.updated_at, d.tags, d.retain_params
""",
document_id,
bank_id,
@@ -3496,6 +3549,14 @@ class MemoryEngine(MemoryEngineInterface):
if not doc:
return None
retain_params_raw = doc["retain_params"]
retain_params_parsed = (
json.loads(retain_params_raw) if isinstance(retain_params_raw, str) else retain_params_raw
)
# document_metadata is sourced from retain_params.metadata
document_metadata = retain_params_parsed.get("metadata") if retain_params_parsed else None
return {
"id": doc["id"],
"bank_id": doc["bank_id"],
@@ -3505,6 +3566,8 @@ class MemoryEngine(MemoryEngineInterface):
"created_at": doc["created_at"].isoformat() if doc["created_at"] else None,
"updated_at": doc["updated_at"].isoformat() if doc["updated_at"] else None,
"tags": list(doc["tags"]) if doc["tags"] else [],
"document_metadata": document_metadata or None,
"retain_params": retain_params_parsed or None,
}
async def delete_document(
@@ -3562,7 +3625,10 @@ class MemoryEngine(MemoryEngineInterface):
}
if invalidated_obs > 0:
await self.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
try:
await self.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
except Exception as e:
logger.warning(f"Failed to submit consolidation after document deletion for bank {bank_id}: {e}")
return result
@@ -3696,7 +3762,10 @@ class MemoryEngine(MemoryEngineInterface):
)
if invalidated_obs > 0:
await self.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
try:
await self.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
except Exception as e:
logger.warning(f"Failed to submit consolidation after document update for bank {bank_id}: {e}")
return True
@@ -3758,7 +3827,14 @@ class MemoryEngine(MemoryEngineInterface):
}
if bank_id_for_consolidation:
await self.submit_async_consolidation(bank_id=bank_id_for_consolidation, request_context=request_context)
try:
await self.submit_async_consolidation(
bank_id=bank_id_for_consolidation, request_context=request_context
)
except Exception as e:
logger.warning(
f"Failed to submit consolidation after memory deletion for bank {bank_id_for_consolidation}: {e}"
)
return result
@@ -3767,6 +3843,7 @@ class MemoryEngine(MemoryEngineInterface):
bank_id: str,
fact_type: str | None = None,
*,
delete_bank_profile: bool = True,
request_context: "RequestContext",
) -> dict[str, int]:
"""
@@ -3782,7 +3859,7 @@ class MemoryEngine(MemoryEngineInterface):
Args:
bank_id: bank ID to delete
fact_type: Optional fact type filter (world, experience, opinion). If provided, only deletes memories of that type.
fact_type: Optional fact type filter (world, experience). If provided, only deletes memories of that type.
request_context: Request context for authentication.
Returns:
@@ -3853,31 +3930,35 @@ class MemoryEngine(MemoryEngineInterface):
# Delete entities (cascades to unit_entities, entity_cooccurrences, memory_links with entity_id)
await conn.execute(f"DELETE FROM {fq_table('entities')} WHERE bank_id = $1", bank_id)
# Delete the bank profile and retrieve internal_id for HNSW index cleanup
internal_id = await conn.fetchval(
f"DELETE FROM {fq_table('banks')} WHERE bank_id = $1 RETURNING internal_id", bank_id
)
if internal_id:
bank_internal_id = str(internal_id)
result = {
"memory_units_deleted": units_count,
"entities_deleted": entities_count,
"documents_deleted": documents_count,
"bank_deleted": True,
}
if delete_bank_profile:
# Delete the bank profile and retrieve internal_id for HNSW index cleanup
internal_id = await conn.fetchval(
f"DELETE FROM {fq_table('banks')} WHERE bank_id = $1 RETURNING internal_id", bank_id
)
if internal_id:
bank_internal_id = str(internal_id)
result["bank_deleted"] = True
except Exception as e:
raise Exception(f"Failed to delete agent data: {str(e)}")
# Drop per-bank HNSW indexes AFTER the transaction commits to avoid
# Drop per-bank vector indexes AFTER the transaction commits to avoid
# AccessExclusiveLock deadlocks with concurrent bank deletions.
# (DROP INDEX on memory_units conflicts with RowExclusiveLock from DELETE inside tx)
if bank_internal_id:
await bank_utils.drop_bank_hnsw_indexes(conn, bank_internal_id)
await bank_utils.drop_bank_vector_indexes(conn, bank_internal_id)
if invalidated_obs > 0:
await self.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
try:
await self.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
except Exception as e:
logger.warning(f"Failed to submit consolidation after bank deletion for bank {bank_id}: {e}")
return result
@@ -4100,7 +4181,7 @@ class MemoryEngine(MemoryEngineInterface):
Args:
bank_id: Filter by bank ID
fact_type: Filter by fact type (world, experience, opinion)
fact_type: Filter by fact type (world, experience)
limit: Maximum number of items to return (default: 1000)
q: Full-text search query (searches text and context fields)
tags: Filter by tags
@@ -4186,23 +4267,27 @@ class MemoryEngine(MemoryEngineInterface):
source_memory_ids.extend(unit["source_memory_ids"])
source_memory_ids = list(set(source_memory_ids)) # Deduplicate
# Fetch links involving both visible units AND source memories
# Fetch links where BOTH endpoints are in the visible set (or source memories)
# Cap at 10k edges — the UI can't usefully render more, and uncapped queries
# on highly-connected graphs (e.g. 1000 nodes with 500k+ edges) are too slow.
max_edges = 10000
all_relevant_ids = unit_ids + source_memory_ids
if all_relevant_ids:
links = await conn.fetch(
f"""
SELECT DISTINCT ON (LEAST(ml.from_unit_id, ml.to_unit_id), GREATEST(ml.from_unit_id, ml.to_unit_id), ml.link_type, COALESCE(ml.entity_id, '00000000-0000-0000-0000-000000000000'::uuid))
ml.from_unit_id,
ml.to_unit_id,
ml.link_type,
ml.weight,
e.canonical_name as entity_name
SELECT ml.from_unit_id,
ml.to_unit_id,
ml.link_type,
ml.weight,
e.canonical_name as entity_name
FROM {fq_table("memory_links")} ml
LEFT JOIN {fq_table("entities")} e ON ml.entity_id = e.id
WHERE ml.from_unit_id = ANY($1::uuid[]) OR ml.to_unit_id = ANY($1::uuid[])
ORDER BY LEAST(ml.from_unit_id, ml.to_unit_id), GREATEST(ml.from_unit_id, ml.to_unit_id), ml.link_type, COALESCE(ml.entity_id, '00000000-0000-0000-0000-000000000000'::uuid), ml.weight DESC
WHERE ml.from_unit_id = ANY($1::uuid[]) AND ml.to_unit_id = ANY($1::uuid[])
ORDER BY ml.weight DESC NULLS LAST
LIMIT $2
""",
all_relevant_ids,
max_edges,
)
else:
links = []
@@ -4263,13 +4348,23 @@ class MemoryEngine(MemoryEngineInterface):
link for link in links if link["from_unit_id"] in unit_id_set and link["to_unit_id"] in unit_id_set
]
# Get entity information
unit_entities = await conn.fetch(f"""
SELECT ue.unit_id, e.canonical_name
FROM {fq_table("unit_entities")} ue
JOIN {fq_table("entities")} e ON ue.entity_id = e.id
ORDER BY ue.unit_id
""")
# Get entity information — only for visible units
# Fetch entities for visible units AND their source memories
# (so observations can inherit entities from source memories)
entity_lookup_ids = unit_ids + source_memory_ids
if entity_lookup_ids:
unit_entities = await conn.fetch(
f"""
SELECT ue.unit_id, e.canonical_name
FROM {fq_table("unit_entities")} ue
JOIN {fq_table("entities")} e ON ue.entity_id = e.id
WHERE ue.unit_id = ANY($1::uuid[])
ORDER BY ue.unit_id
""",
entity_lookup_ids,
)
else:
unit_entities = []
# Build entity mapping
entity_map = {}
@@ -4467,7 +4562,7 @@ class MemoryEngine(MemoryEngineInterface):
Args:
bank_id: Filter by bank ID
fact_type: Filter by fact type (world, experience, opinion)
fact_type: Filter by fact type (world, experience)
search_query: Full-text search query (searches text and context fields)
limit: Maximum number of results to return
offset: Offset for pagination
@@ -4924,6 +5019,14 @@ class MemoryEngine(MemoryEngineInterface):
bank_id_val = row["bank_id"]
unit_count = count_map.get((doc_id, bank_id_val), 0)
retain_params_val = row["retain_params"]
retain_params_val = (
json.loads(retain_params_val) if isinstance(retain_params_val, str) else retain_params_val
)
# document_metadata is sourced from retain_params.metadata
document_metadata = retain_params_val.get("metadata") if retain_params_val else None
items.append(
{
"id": doc_id,
@@ -4933,7 +5036,8 @@ class MemoryEngine(MemoryEngineInterface):
"updated_at": row["updated_at"].isoformat() if row["updated_at"] else "",
"text_length": row["text_length"] or 0,
"memory_unit_count": unit_count,
"retain_params": row["retain_params"] if row["retain_params"] else None,
"retain_params": retain_params_val or None,
"document_metadata": document_metadata or None,
"tags": row["tags"] if row["tags"] else [],
}
)
@@ -5826,14 +5930,16 @@ class MemoryEngine(MemoryEngineInterface):
bank_id,
)
# Single query for all link stats — avoids triple join on memory_links (can be 21M+ rows).
# link_counts and link_counts_by_fact_type are derived in Python from the breakdown.
# Link stats — filter on ml.bank_id (indexed) instead of joining through mu.bank_id.
# With the idx_memory_links_bank_link_type index this turns a full-table hash join
# into an indexed scan + PK lookups. link_counts and link_counts_by_fact_type are
# derived in Python from the breakdown.
link_breakdown_stats = await conn.fetch(
f"""
SELECT mu.fact_type, ml.link_type, COUNT(*) as count
FROM {fq_table("memory_links")} ml
JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
WHERE mu.bank_id = $1
WHERE ml.bank_id = $1
GROUP BY mu.fact_type, ml.link_type
""",
bank_id,
@@ -6255,6 +6361,7 @@ class MemoryEngine(MemoryEngineInterface):
*,
tags: list[str] | None = None,
tags_match: str = "any",
detail: str = "full",
limit: int = 100,
offset: int = 0,
request_context: "RequestContext",
@@ -6265,6 +6372,7 @@ class MemoryEngine(MemoryEngineInterface):
bank_id: Bank identifier
tags: Optional tags to filter by
tags_match: How to match tags - 'any', 'all', or 'exact'
detail: Detail level - 'metadata', 'content', or 'full'
limit: Maximum number of results
offset: Offset for pagination
request_context: Request context for authentication
@@ -6306,13 +6414,14 @@ class MemoryEngine(MemoryEngineInterface):
*params,
)
return [self._row_to_mental_model(row) for row in rows]
return [self._row_to_mental_model(row, detail=detail) for row in rows]
async def get_mental_model(
self,
bank_id: str,
mental_model_id: str,
*,
detail: str = "full",
request_context: "RequestContext",
) -> dict[str, Any] | None:
"""Get a single pinned mental model by ID.
@@ -6320,6 +6429,7 @@ class MemoryEngine(MemoryEngineInterface):
Args:
bank_id: Bank identifier
mental_model_id: Pinned mental model UUID
detail: Detail level - 'metadata', 'content', or 'full'
request_context: Request context for authentication
Returns:
@@ -6353,7 +6463,7 @@ class MemoryEngine(MemoryEngineInterface):
mental_model_id,
)
result = self._row_to_mental_model(row) if row else None
result = self._row_to_mental_model(row, detail=detail) if row else None
# Post-operation hook (usage recording)
if result and self._operation_validator:
@@ -6528,26 +6638,23 @@ class MemoryEngine(MemoryEngineInterface):
# Create parent span for mental model refresh operation
with create_operation_span("mental_model_refresh", bank_id):
# SECURITY: If the mental model has tags, pass them to reflect with "all_strict" matching
# to ensure it can only access other mental models/memories with the SAME tags.
# This prevents cross-tenant/cross-user information leakage by excluding untagged content.
tags = mental_model.get("tags")
tags_match = "all_strict" if tags else "any"
# Read reflect options from trigger (if stored)
trigger_data = mental_model.get("trigger") or {}
fact_types = trigger_data.get("fact_types")
exclude_mental_models = trigger_data.get("exclude_mental_models", False)
stored_exclude_ids: list[str] = trigger_data.get("exclude_mental_model_ids") or []
tag_filtering = _resolve_refresh_tag_filtering(mental_model.get("tags"), trigger_data)
# Run reflect with the source query, excluding the mental model being refreshed
# Skip creating a nested "hindsight.reflect" span since we already have "hindsight.mental_model_refresh"
reflect_result = await self.reflect_async(
bank_id=bank_id,
query=mental_model["source_query"],
request_context=request_context,
tags=tags,
tags_match=tags_match,
tags=tag_filtering.tags,
tags_match=tag_filtering.tags_match,
tag_groups=tag_filtering.tag_groups,
fact_types=fact_types,
exclude_mental_models=exclude_mental_models,
exclude_mental_model_ids=list({*stored_exclude_ids, mental_model_id}),
@@ -6754,34 +6861,45 @@ class MemoryEngine(MemoryEngineInterface):
return result == "DELETE 1"
def _row_to_mental_model(self, row) -> dict[str, Any]:
"""Convert a database row to a mental model dict."""
reflect_response = row.get("reflect_response")
# Parse JSON string to dict if needed (asyncpg may return JSONB as string)
if isinstance(reflect_response, str):
try:
reflect_response = json.loads(reflect_response)
except json.JSONDecodeError:
reflect_response = None
def _row_to_mental_model(self, row, *, detail: str = "full") -> dict[str, Any]:
"""Convert a database row to a mental model dict.
Args:
row: Database row
detail: Detail level - 'metadata', 'content', or 'full'
"""
result: dict[str, Any] = {
"id": str(row["id"]),
"bank_id": row["bank_id"],
"name": row["name"],
"tags": row["tags"] or [],
"last_refreshed_at": row["last_refreshed_at"].isoformat() if row["last_refreshed_at"] else None,
"created_at": row["created_at"].isoformat() if row["created_at"] else None,
}
if detail == "metadata":
return result
trigger = row.get("trigger")
if isinstance(trigger, str):
try:
trigger = json.loads(trigger)
except json.JSONDecodeError:
trigger = None
return {
"id": str(row["id"]),
"bank_id": row["bank_id"],
"name": row["name"],
"source_query": row["source_query"],
"content": row["content"],
"tags": row["tags"] or [],
"max_tokens": row.get("max_tokens"),
"trigger": trigger,
"last_refreshed_at": row["last_refreshed_at"].isoformat() if row["last_refreshed_at"] else None,
"created_at": row["created_at"].isoformat() if row["created_at"] else None,
"reflect_response": reflect_response,
}
result["source_query"] = row["source_query"]
result["content"] = row["content"]
result["max_tokens"] = row.get("max_tokens")
result["trigger"] = trigger
if detail == "full":
reflect_response = row.get("reflect_response")
if isinstance(reflect_response, str):
try:
reflect_response = json.loads(reflect_response)
except json.JSONDecodeError:
reflect_response = None
result["reflect_response"] = reflect_response
return result
# =========================================================================
# Directives - Hard rules injected into prompts
@@ -331,7 +331,11 @@ class ClaudeCodeLLM(LLMInterface):
max_retries: Maximum retry attempts.
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
tool_choice: How to choose tools (not used by Claude Agent SDK).
tool_choice: How to choose tools - "auto", "none", "required", or specific function dict.
- "auto": Model decides whether to call tools (default)
- "required": Model must call at least one tool
- "none": Model must not call any tools
- {"type": "function", "function": {"name": "..."}}: Force specific tool call
Returns:
LLMToolCallResult with content and/or tool_calls.
@@ -410,16 +414,57 @@ class ClaudeCodeLLM(LLMInterface):
tool_call_id = msg.get("tool_call_id", "")
user_content += f"\n\n[Tool result for {tool_call_id}: {content}]"
# Handle tool_choice parameter to filter tools and adjust instructions
# The Claude Agent SDK doesn't have a native tool_choice parameter, so we
# enforce it via allowed_tools filtering and system prompt instructions.
# Format tool names for SDK MCP servers: mcp__{server_name}__{tool_name}
# This is required by the Claude Agent SDK for MCP server tools
allowed_tool_names = [f"mcp__hindsight_tools__{name}" for name in tool_names]
mcp_servers_config = {"hindsight_tools": mcp_server} if sdk_tools else {}
# Process tool_choice
if isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
# Force a specific tool: filter allowed_tools to only that tool and add instruction
forced_name = tool_choice.get("function", {}).get("name")
if forced_name:
# Filter to only the forced tool (with MCP prefix)
forced_tool_mcp_name = f"mcp__hindsight_tools__{forced_name}"
if forced_tool_mcp_name in allowed_tool_names:
allowed_tool_names = [forced_tool_mcp_name]
# Add strong instruction to system prompt
force_instruction = (
f"\n\nIMPORTANT: You MUST call the '{forced_name}' tool. Do not respond with text only."
)
system_prompt += force_instruction
logger.debug(f"Claude Code: Forcing tool call to '{forced_name}'")
else:
logger.warning(f"Claude Code: Forced tool '{forced_name}' not found in available tools")
elif tool_choice == "required":
# Must call at least one tool
tool_instruction = (
"\n\nIMPORTANT: You MUST call at least one of the available tools. Do not respond with text only."
)
system_prompt += tool_instruction
logger.debug("Claude Code: Tool call required")
elif tool_choice == "none":
# No tools should be called - disable all tools
allowed_tool_names = []
mcp_servers_config = {}
logger.debug("Claude Code: Tools disabled (tool_choice=none)")
# else: tool_choice == "auto" or unspecified - use default behavior (no changes needed)
# Configure SDK options with MCP server
# tools=[] disables built-in CLI tools (Read, Write, Bash, ToolSearch, etc.)
# Without this, Claude Code CLI defers MCP tools when too many built-in tools
# are loaded, forcing Claude to use ToolSearch first — which wastes the max_turns
# budget and prevents direct MCP tool calls.
options = ClaudeAgentOptions(
system_prompt=system_prompt if system_prompt else None,
max_turns=1, # Single-turn for API-style interactions
mcp_servers={"hindsight_tools": mcp_server} if sdk_tools else {},
allowed_tools=allowed_tool_names if allowed_tool_names else [],
tools=[], # Disable built-in tools so MCP tools load eagerly
max_turns=2, # Allow tool call + tool result round-trip
mcp_servers=mcp_servers_config,
allowed_tools=allowed_tool_names,
)
# Call Claude Agent SDK with retry logic
@@ -126,6 +126,32 @@ class CodexLLM(LLMInterface):
}
return mapping.get(effort.lower(), "auto")
def _normalize_tool_choice(self, tool_choice: str | dict[str, Any]) -> str | dict[str, Any]:
"""Normalize forced function tool choice for the Codex Responses API.
Older agent paths may still pass OpenAI chat-completions style named
tool choice payloads such as:
{"type": "function", "function": {"name": "recall"}}
Codex Responses expects the named function at the top level instead:
{"type": "function", "name": "recall"}
"""
if not isinstance(tool_choice, dict):
return tool_choice
if str(tool_choice.get("type") or "").strip() != "function":
return tool_choice
function_payload = tool_choice.get("function")
if isinstance(function_payload, dict):
function_name = str(function_payload.get("name") or "").strip()
if function_name:
return {"type": "function", "name": function_name}
function_name = str(tool_choice.get("name") or "").strip()
if function_name:
return {"type": "function", "name": function_name}
return tool_choice
async def verify_connection(self) -> None:
"""Verify Codex connection by making a simple test call."""
try:
@@ -140,6 +166,10 @@ class CodexLLM(LLMInterface):
)
logger.info(f"Codex LLM verified: {self.model}")
except Exception as e:
# 429 means quota exhausted, not a configuration error — warn but allow startup
if "429" in str(e) or "usage_limit_reached" in str(e):
logger.warning(f"Codex LLM quota exhausted for {self.model}, continuing startup: {e}")
return
raise RuntimeError(f"Codex LLM connection verification failed for {self.model}: {e}") from e
async def call(
@@ -263,24 +293,27 @@ class CodexLLM(LLMInterface):
)
# Record trace span
from hindsight_api.tracing import get_span_recorder
try:
from hindsight_api.tracing import get_span_recorder
# Estimate tokens for tracing
estimated_input = sum(len(m.get("content", "")) for m in messages) // 4
estimated_output = len(content) // 4
span_recorder = get_span_recorder()
span_recorder.record_llm_call(
provider=self.provider,
model=self.model,
scope=scope,
messages=messages,
response_content=result if isinstance(result, str) else json.dumps(result),
input_tokens=estimated_input,
output_tokens=estimated_output,
duration=duration,
finish_reason=None,
error=None,
)
# Estimate tokens for tracing
estimated_input = sum(len(m.get("content", "")) for m in messages) // 4
estimated_output = len(content) // 4
span_recorder = get_span_recorder()
span_recorder.record_llm_call(
provider=self.provider,
model=self.model,
scope=scope,
messages=messages,
response_content=result if isinstance(result, str) else result.model_dump_json(),
input_tokens=estimated_input,
output_tokens=estimated_output,
duration=duration,
finish_reason=None,
error=None,
)
except Exception:
pass # logging failure must never affect the operation
if return_usage:
# Codex doesn't provide token counts, estimate based on content
@@ -422,7 +455,7 @@ class CodexLLM(LLMInterface):
max_retries: Maximum retry attempts.
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
tool_choice: How to choose tools - "auto", "none", "required", or specific function.
tool_choice: How to choose tools - "auto", "none", "required", or a specific function.
Returns:
LLMToolCallResult with content and/or tool_calls.
@@ -479,7 +512,7 @@ class CodexLLM(LLMInterface):
"instructions": system_instruction,
"input": user_messages,
"tools": codex_tools,
"tool_choice": tool_choice,
"tool_choice": self._normalize_tool_choice(tool_choice),
"parallel_tool_calls": True,
"reasoning": {"summary": reasoning_summary},
"store": False,
@@ -526,26 +559,31 @@ class CodexLLM(LLMInterface):
)
# Record OpenTelemetry span
from hindsight_api.tracing import get_span_recorder
try:
from hindsight_api.tracing import get_span_recorder
span_recorder = get_span_recorder()
# Convert LLMToolCall objects to dicts for span recording
tool_calls_dict = (
[{"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in tool_calls] if tool_calls else None
)
span_recorder.record_llm_call(
provider=self.provider,
model=self.model,
scope=scope,
messages=messages,
response_content=content,
input_tokens=0, # Codex doesn't provide token counts
output_tokens=0,
duration=duration,
finish_reason="tool_calls" if tool_calls else "stop",
error=None,
tool_calls=tool_calls_dict,
)
span_recorder = get_span_recorder()
# Convert LLMToolCall objects to dicts for span recording
tool_calls_dict = (
[{"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in tool_calls]
if tool_calls
else None
)
span_recorder.record_llm_call(
provider=self.provider,
model=self.model,
scope=scope,
messages=messages,
response_content=content,
input_tokens=0, # Codex doesn't provide token counts
output_tokens=0,
duration=duration,
finish_reason="tool_calls" if tool_calls else "stop",
error=None,
tool_calls=tool_calls_dict,
)
except Exception:
pass # logging failure must never affect the operation
return LLMToolCallResult(
content=content,
@@ -7,6 +7,7 @@ This provider supports both:
"""
import asyncio
import base64
import json
import logging
import os
@@ -472,9 +473,10 @@ class GeminiLLM(LLMInterface):
fn_args = parse_llm_json(fn_args_str)
thought_signature = tc.get("thought_signature")
fc_kwargs: dict[str, Any] = {"name": fn_name, "args": fn_args}
part_kwargs: dict[str, Any] = {"function_call": genai_types.FunctionCall(**fc_kwargs)}
if thought_signature:
fc_kwargs["thought_signature"] = thought_signature
parts.append(genai_types.Part(function_call=genai_types.FunctionCall(**fc_kwargs)))
part_kwargs["thought_signature"] = base64.b64decode(thought_signature)
parts.append(genai_types.Part(**part_kwargs))
gemini_contents.append(genai_types.Content(role="model", parts=parts))
else:
gemini_contents.append(genai_types.Content(role="model", parts=[genai_types.Part(text=content)]))
@@ -547,7 +549,10 @@ class GeminiLLM(LLMInterface):
content = part.text
if hasattr(part, "function_call") and part.function_call:
fc = part.function_call
thought_signature = getattr(fc, "thought_signature", None)
_raw_ts = getattr(part, "thought_signature", None)
thought_signature = (
base64.b64encode(_raw_ts).decode("ascii") if isinstance(_raw_ts, bytes) else _raw_ts
)
tool_calls.append(
LLMToolCall(
id=f"gemini_{len(tool_calls)}",
@@ -80,6 +80,7 @@ class OpenAICompatibleLLM(LLMInterface):
reasoning_effort: str = "low",
timeout: float | None = None,
groq_service_tier: str | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs: Any,
):
"""
@@ -93,12 +94,13 @@ class OpenAICompatibleLLM(LLMInterface):
reasoning_effort: Reasoning effort level for supported models ("low", "medium", "high").
timeout: Request timeout in seconds (uses env var or 300s default).
groq_service_tier: Groq service tier ("on_demand", "flex", "auto").
extra_body: Extra body params merged into every API call.
**kwargs: Additional provider-specific parameters.
"""
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
# Validate provider
valid_providers = ["openai", "groq", "ollama", "lmstudio", "minimax"]
valid_providers = ["openai", "groq", "ollama", "lmstudio", "minimax", "volcano"]
if self.provider not in valid_providers:
raise ValueError(f"OpenAICompatibleLLM only supports: {', '.join(valid_providers)}. Got: {self.provider}")
@@ -124,6 +126,8 @@ class OpenAICompatibleLLM(LLMInterface):
# Service tier configuration (from config, not env vars)
self.groq_service_tier = groq_service_tier
self.openai_service_tier = kwargs.get("openai_service_tier")
# User-configured extra body params (merged into every API call)
self._config_extra_body = extra_body or {}
# Get timeout config
self.timeout = timeout or float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT)))
@@ -187,6 +191,23 @@ class OpenAICompatibleLLM(LLMInterface):
return None
def _max_tokens_param_name(self) -> str:
"""Return the correct parameter name for limiting response tokens.
Native OpenAI and Groq accept 'max_completion_tokens'. Mistral and other
OpenAI-compatible endpoints that haven't adopted the newer parameter name
require 'max_tokens'. Using a custom base_url with the openai provider
signals a third-party compatible API, so fall back to 'max_tokens'.
"""
# Native OpenAI (no custom base URL) and Groq use max_completion_tokens
if self.provider == "groq":
return "max_completion_tokens"
if self.provider == "openai" and not self.base_url:
return "max_completion_tokens"
# openai with custom base_url, ollama, lmstudio, minimax, volcano —
# use the widely-supported max_tokens
return "max_tokens"
async def call(
self,
messages: list[dict[str, str]],
@@ -259,9 +280,7 @@ class OpenAICompatibleLLM(LLMInterface):
# For reasoning models, enforce minimum to ensure space for reasoning + output
if is_reasoning_model and max_completion_tokens < 16000:
max_completion_tokens = 16000
call_params["max_completion_tokens"] = max_completion_tokens
# Temperature - reasoning models don't support custom temperature
call_params[self._max_tokens_param_name()] = max_completion_tokens
if temperature is not None and not is_reasoning_model:
# MiniMax requires temperature in (0.0, 1.0] — clamp accordingly
if self.provider == "minimax":
@@ -273,17 +292,17 @@ class OpenAICompatibleLLM(LLMInterface):
call_params["reasoning_effort"] = self.reasoning_effort
# Provider-specific parameters
extra_body: dict[str, Any] = {**self._config_extra_body}
if self.provider == "groq":
call_params["seed"] = DEFAULT_LLM_SEED
extra_body: dict[str, Any] = {}
# Add service_tier if configured
if self.groq_service_tier:
extra_body["service_tier"] = self.groq_service_tier
# Add reasoning parameters for reasoning models
if is_reasoning_model:
extra_body["include_reasoning"] = False
if extra_body:
call_params["extra_body"] = extra_body
if extra_body:
call_params["extra_body"] = extra_body
# Prepare response format ONCE before retry loop
if response_format is not None:
@@ -316,8 +335,8 @@ class OpenAICompatibleLLM(LLMInterface):
first_msg = call_params["messages"][0]
if isinstance(first_msg, dict) and isinstance(first_msg.get("content"), str):
first_msg["content"] = schema_msg + "\n\n" + first_msg["content"]
if self.provider not in ("lmstudio", "ollama"):
# LM Studio and Ollama don't support json_object response format reliably
if self.provider not in ("lmstudio", "ollama", "volcano"):
# LM Studio, Ollama and Volcano don't support json_object response format reliably
call_params["response_format"] = {"type": "json_object"}
last_exception = None
@@ -573,7 +592,7 @@ class OpenAICompatibleLLM(LLMInterface):
}
if max_completion_tokens is not None:
call_params["max_completion_tokens"] = max_completion_tokens
call_params[self._max_tokens_param_name()] = max_completion_tokens
if temperature is not None:
# MiniMax requires temperature in (0.0, 1.0] — clamp accordingly
if self.provider == "minimax":
@@ -581,8 +600,11 @@ class OpenAICompatibleLLM(LLMInterface):
call_params["temperature"] = temperature
# Provider-specific parameters
extra_body: dict[str, Any] = {**self._config_extra_body}
if self.provider == "groq":
call_params["seed"] = DEFAULT_LLM_SEED
if extra_body:
call_params["extra_body"] = extra_body
last_exception = None
@@ -137,7 +137,21 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
"RETURN_AS_TIMEZONE_AWARE": False,
}
results = self._search_dates(query, settings=settings)
# Wrap dateparser in a defensive try/except. dateparser has been
# observed to crash with internal errors (e.g., IndexError from
# locale.translate_search) on certain query inputs. A parser bug
# should not bring down the whole search/consolidation pipeline —
# treat any failure as "no temporal constraint found" so the caller
# can fall back to non-temporal retrieval.
try:
results = self._search_dates(query, settings=settings)
except Exception as e:
logger.warning(
"dateparser raised %s on query (treating as no temporal constraint): %s",
type(e).__name__,
e,
)
return QueryAnalysis(temporal_constraint=None)
if not results:
return QueryAnalysis(temporal_constraint=None)
@@ -10,7 +10,6 @@ from typing import Any
from pydantic import BaseModel, ConfigDict, Field, field_validator
# Valid fact types for recall operations (excludes 'opinion' which is deprecated)
VALID_RECALL_FACT_TYPES = frozenset(["world", "experience", "observation"])
@@ -10,32 +10,47 @@ from typing import TypedDict
from pydantic import BaseModel, Field
from ...config import get_config
from ..db_utils import acquire_with_retry
from ..memory_engine import fq_table, get_current_schema
from ..response_models import DispositionTraits
logger = logging.getLogger(__name__)
# Fact types that get per-bank partial HNSW indexes, mapped to their 4-char index suffix.
_HNSW_FACT_TYPES: dict[str, str] = {
# Fact types that get per-bank partial vector indexes, mapped to their 4-char index suffix.
_BANK_INDEX_FACT_TYPES: dict[str, str] = {
"world": "worl",
"experience": "expr",
"observation": "obsv",
}
def _hnsw_index_name(ft: str, internal_id: str) -> str:
"""Deterministic, schema-safe HNSW index name for a (bank, fact_type) pair.
def _bank_index_name(ft: str, internal_id: str) -> str:
"""Deterministic, schema-safe vector index name for a (bank, fact_type) pair.
Uses the first 16 hex chars of internal_id (8 bytes of entropy) — unique
enough in practice, fits comfortably within PostgreSQL's 63-char identifier limit.
"""
uid = str(internal_id).replace("-", "")[:16]
return f"idx_mu_emb_{_HNSW_FACT_TYPES[ft]}_{uid}"
return f"idx_mu_emb_{_BANK_INDEX_FACT_TYPES[ft]}_{uid}"
async def create_bank_hnsw_indexes(conn, bank_id: str, internal_id: str) -> None:
"""Create per-(bank, fact_type) partial HNSW indexes for a newly created bank.
def _vector_index_clause() -> str:
"""Return the USING clause for vector index creation based on the configured extension."""
ext = get_config().vector_extension
if ext == "pgvectorscale":
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
elif ext == "vchord":
return "USING vchordrq (embedding vector_l2_ops)"
else: # pgvector (default)
return "USING hnsw (embedding vector_cosine_ops)"
async def create_bank_vector_indexes(conn, bank_id: str, internal_id: str) -> None:
"""Create per-(bank, fact_type) partial vector indexes for a newly created bank.
Respects the HINDSIGHT_API_VECTOR_EXTENSION config to use the appropriate
index type (HNSW for pgvector, DiskANN for pgvectorscale, vchordrq for vchord).
Called immediately after the bank row is first inserted. Safe on empty banks
(index build is instant). Idempotent via CREATE INDEX IF NOT EXISTS.
@@ -43,24 +58,25 @@ async def create_bank_hnsw_indexes(conn, bank_id: str, internal_id: str) -> None
"""
table = fq_table("memory_units")
escaped = bank_id.replace("'", "''")
for ft in _HNSW_FACT_TYPES:
idx = _hnsw_index_name(ft, internal_id)
using_clause = _vector_index_clause()
for ft in _BANK_INDEX_FACT_TYPES:
idx = _bank_index_name(ft, internal_id)
await conn.execute(
f"CREATE INDEX IF NOT EXISTS {idx} "
f"ON {table} USING hnsw (embedding vector_cosine_ops) "
f"ON {table} {using_clause} "
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped}'"
)
async def drop_bank_hnsw_indexes(conn, internal_id: str) -> None:
"""Drop per-(bank, fact_type) partial HNSW indexes for a bank being deleted.
async def drop_bank_vector_indexes(conn, internal_id: str) -> None:
"""Drop per-(bank, fact_type) partial vector indexes for a bank being deleted.
Called before the bank row is deleted so internal_id is still known.
Idempotent via DROP INDEX IF EXISTS.
"""
schema = get_current_schema()
for ft in _HNSW_FACT_TYPES:
idx = _hnsw_index_name(ft, internal_id)
for ft in _BANK_INDEX_FACT_TYPES:
idx = _bank_index_name(ft, internal_id)
await conn.execute(f"DROP INDEX IF EXISTS {schema}.{idx}")
@@ -121,7 +137,7 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
# Bank doesn't exist, create with defaults.
# Generate internal_id here so we control the value and can use it
# immediately for HNSW index creation without a RETURNING round-trip.
# immediately for vector index creation without a RETURNING round-trip.
internal_id = uuid.uuid4()
inserted = await conn.fetchval(
f"""
@@ -138,8 +154,8 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
)
if inserted:
# Fresh insert — create per-bank HNSW indexes (instant on empty bank)
await create_bank_hnsw_indexes(conn, bank_id, str(internal_id))
# Fresh insert — create per-bank vector indexes (instant on empty bank)
await create_bank_vector_indexes(conn, bank_id, str(internal_id))
return BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission="")
@@ -4,7 +4,9 @@ Chunk storage for retain pipeline.
Handles storage of document chunks in the database.
"""
import hashlib
import logging
from dataclasses import dataclass
from ..memory_engine import fq_table
from .types import ChunkMetadata
@@ -12,6 +14,61 @@ from .types import ChunkMetadata
logger = logging.getLogger(__name__)
def compute_chunk_hash(chunk_text: str) -> str:
"""Compute SHA256 hash of chunk text for delta comparison."""
return hashlib.sha256(chunk_text.encode()).hexdigest()
@dataclass
class ExistingChunk:
"""Represents a chunk already stored in the database."""
chunk_id: str
chunk_index: int
content_hash: str | None
async def load_existing_chunks(conn, bank_id: str, document_id: str) -> list[ExistingChunk]:
"""
Load existing chunk metadata for a document.
Returns list of ExistingChunk with chunk_id, chunk_index, and content_hash.
"""
rows = await conn.fetch(
f"""
SELECT chunk_id, chunk_index, content_hash
FROM {fq_table("chunks")}
WHERE document_id = $1 AND bank_id = $2
ORDER BY chunk_index
""",
document_id,
bank_id,
)
return [
ExistingChunk(
chunk_id=row["chunk_id"],
chunk_index=row["chunk_index"],
content_hash=row["content_hash"],
)
for row in rows
]
async def delete_chunks_by_ids(conn, chunk_ids: list[str]) -> None:
"""
Delete specific chunks by their IDs.
This cascades to memory_units (via FK with CASCADE delete)
and their links.
"""
if not chunk_ids:
return
await conn.execute(
f"DELETE FROM {fq_table('chunks')} WHERE chunk_id = ANY($1::text[])",
chunk_ids,
)
async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[ChunkMetadata]) -> dict[int, str]:
"""
Store document chunks in the database.
@@ -32,6 +89,7 @@ async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[
chunk_ids = []
chunk_texts = []
chunk_indices = []
content_hashes = []
chunk_id_map = {}
for chunk in chunks:
@@ -39,19 +97,21 @@ async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[
chunk_ids.append(chunk_id)
chunk_texts.append(chunk.chunk_text)
chunk_indices.append(chunk.chunk_index)
content_hashes.append(compute_chunk_hash(chunk.chunk_text))
chunk_id_map[chunk.chunk_index] = chunk_id
# Batch insert all chunks
await conn.execute(
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[])
INSERT INTO {fq_table("chunks")} (chunk_id, document_id, bank_id, chunk_text, chunk_index, content_hash)
SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[], $5::integer[], $6::text[])
""",
chunk_ids,
[document_id] * len(chunk_texts),
[bank_id] * len(chunk_texts),
chunk_texts,
chunk_indices,
content_hashes,
)
return chunk_id_map
@@ -12,61 +12,27 @@ from .types import EntityLink, ProcessedFact
logger = logging.getLogger(__name__)
async def process_entities_batch(
entity_resolver,
conn,
bank_id: str,
unit_ids: list[str],
def _prepare_facts_for_entity_processing(
facts: list[ProcessedFact],
log_buffer: list[str] = None,
user_entities_per_content: dict[int, list[dict]] = None,
entity_labels: list | None = None,
) -> list[EntityLink]:
user_entities_per_content: dict[int, list[dict]] | None = None,
) -> tuple[list[str], list, list[list[dict]]]:
"""
Process entities for all facts and create entity links.
This function:
1. Extracts entity mentions from fact texts
2. Merges user-provided entities with LLM-extracted entities
3. Resolves entity names to canonical entities
4. Creates entity records in the database
5. Returns entity links ready for insertion
Args:
entity_resolver: EntityResolver instance for entity resolution
conn: Database connection
bank_id: Bank identifier
unit_ids: List of unit IDs (same length as facts)
facts: List of ProcessedFact objects
log_buffer: Optional buffer for detailed logging
user_entities_per_content: Dict mapping content_index to list of user-provided entities
Extract fact texts, dates, and merged entity lists from ProcessedFact objects.
Returns:
List of EntityLink objects for batch insertion
Tuple of (fact_texts, fact_dates, entities_per_fact)
"""
if not unit_ids or not facts:
return []
if len(unit_ids) != len(facts):
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})")
user_entities_per_content = user_entities_per_content or {}
# Extract data for link_utils function
fact_texts = [fact.fact_text for fact in facts]
# Use occurred_start if available, otherwise use mentioned_at for entity timestamps
fact_dates = [fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at for fact in facts]
# Convert EntityRef objects to dict format and merge with user-provided entities
entities_per_fact = []
for fact in facts:
# Start with LLM-extracted entities
llm_entities = [{"text": entity.name, "type": "CONCEPT"} for entity in (fact.entities or [])]
# Get user entities for this content (use content_index from fact)
user_entities = user_entities_per_content.get(fact.content_index, [])
# Merge with case-insensitive deduplication
seen_texts = {e["text"].lower() for e in llm_entities}
for user_entity in user_entities:
if user_entity["text"].lower() not in seen_texts:
@@ -80,8 +46,48 @@ async def process_entities_batch(
entities_per_fact.append(llm_entities)
# Use existing link_utils function for entity processing
entity_links = await link_utils.extract_entities_batch_optimized(
return fact_texts, fact_dates, entities_per_fact
async def resolve_entities(
entity_resolver,
conn,
bank_id: str,
unit_ids: list[str],
facts: list[ProcessedFact],
log_buffer: list[str] = None,
user_entities_per_content: dict[int, list[dict]] = None,
entity_labels: list | None = None,
) -> tuple[list[str], list[tuple], dict[str, list[str]]]:
"""
Phase 1: Resolve entity names to canonical IDs (read-heavy).
Should be called on a SEPARATE connection OUTSIDE the main write transaction
to avoid holding the transaction open during expensive trigram scans.
Args:
entity_resolver: EntityResolver instance
conn: Database connection (separate from the main write transaction)
bank_id: Bank identifier
unit_ids: Placeholder unit IDs (used only for grouping)
facts: List of ProcessedFact objects
log_buffer: Optional buffer for detailed logging
user_entities_per_content: Dict mapping content_index to user-provided entities
entity_labels: Optional entity label taxonomy
Returns:
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids)
to pass to build_entity_links().
"""
if not unit_ids or not facts:
return [], [], {}
if len(unit_ids) != len(facts):
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})")
fact_texts, fact_dates, entities_per_fact = _prepare_facts_for_entity_processing(facts, user_entities_per_content)
return await link_utils.resolve_entities_only(
entity_resolver,
conn,
bank_id,
@@ -90,22 +96,67 @@ async def process_entities_batch(
"", # context (not used in current implementation)
fact_dates,
entities_per_fact,
log_buffer, # Pass log_buffer for detailed logging
log_buffer,
entity_labels=entity_labels,
)
return entity_links
async def build_entity_links(
entity_resolver,
conn,
bank_id: str,
unit_ids: list[str],
resolved_entity_ids: list[str],
entity_to_unit: list[tuple],
unit_to_entity_ids: dict[str, list[str]],
log_buffer: list[str] = None,
skip_unit_entities_insert: bool = False,
) -> list[EntityLink]:
"""
Build entity links for UI graph visualization.
Queries unit_entities to find shared entities between new and existing units,
then generates EntityLink objects. When called from Phase 3 (post-transaction),
set skip_unit_entities_insert=True since unit_entities were already inserted
in Phase 2.
Args:
entity_resolver: EntityResolver instance
conn: Database connection
bank_id: Bank identifier
unit_ids: Actual unit IDs (must already be inserted in the DB)
resolved_entity_ids: From resolve_entities()
entity_to_unit: From resolve_entities()
unit_to_entity_ids: From resolve_entities()
log_buffer: Optional buffer for detailed logging
skip_unit_entities_insert: Skip unit_entities INSERT (already done in Phase 2)
Returns:
List of EntityLink objects for batch insertion
"""
return await link_utils.build_entity_links_from_resolved(
entity_resolver,
conn,
bank_id,
unit_ids,
resolved_entity_ids,
entity_to_unit,
unit_to_entity_ids,
log_buffer,
skip_unit_entities_insert=skip_unit_entities_insert,
)
async def insert_entity_links_batch(conn, entity_links: list[EntityLink]) -> None:
async def insert_entity_links_batch(conn, entity_links: list[EntityLink], bank_id: str) -> None:
"""
Insert entity links in batch.
Args:
conn: Database connection
entity_links: List of EntityLink objects
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
"""
if not entity_links:
return
await link_utils.insert_entity_links_batch(conn, entity_links)
await link_utils.insert_entity_links_batch(conn, entity_links, bank_id)
@@ -87,7 +87,7 @@ class Fact(BaseModel):
# Required fields
fact: str = Field(description="Combined fact text: what | when | where | who | why")
fact_type: Literal["world", "experience", "opinion"] = Field(description="Perspective: world/experience/opinion")
fact_type: Literal["world", "experience"] = Field(description="Perspective: world/experience")
# Optional temporal fields
occurred_start: str | None = None
@@ -159,7 +159,9 @@ class ExtractedFact(BaseModel):
fact_kind: str = Field(default="conversation", description="'event' or 'conversation'")
occurred_start: str | None = Field(default=None, description="ISO timestamp for events")
occurred_end: str | None = Field(default=None, description="ISO timestamp for event end")
fact_type: Literal["world", "assistant"] = Field(description="'world' or 'assistant'")
fact_type: Literal["world", "assistant"] = Field(
description="'world' = objective/external facts. 'assistant' = first-person actions, experiences, or observations by the speaker."
)
entities: list[Entity] | None = Field(default=None, description="People, places, concepts")
causal_relations: list[FactCausalRelation] | None = Field(
default=None, description="Links to previous facts (target_index < this fact's index)"
@@ -261,7 +263,7 @@ class ExtractedFactVerbose(BaseModel):
)
fact_type: Literal["world", "assistant"] = Field(
description="'world' = about the user/others (background, experiences). 'assistant' = experience with the assistant."
description="'world' = objective/external facts about other people, events, general knowledge. 'assistant' = first-person actions, experiences, or observations by the speaker (e.g., 'I changed X', 'I discovered Y')."
)
entities: list[Entity] | None = Field(
@@ -352,7 +354,9 @@ class VerbatimExtractedFact(BaseModel):
fact_kind: str = Field(default="conversation", description="'event' or 'conversation'")
occurred_start: str | None = Field(default=None, description="ISO timestamp for events")
occurred_end: str | None = Field(default=None, description="ISO timestamp for event end")
fact_type: Literal["world", "assistant"] = Field(description="'world' or 'assistant'")
fact_type: Literal["world", "assistant"] = Field(
description="'world' = objective/external facts. 'assistant' = first-person actions, experiences, or observations by the speaker."
)
entities: list[Entity] | None = Field(default=None, description="People, places, concepts")
@field_validator("entities", mode="before")
@@ -499,8 +503,8 @@ fact_kind:
- "conversation": Ongoing state, preference, trait (no dates)
fact_type:
- "world": About user's life, other people, external events
- "assistant": Interactions with assistant (requests, recommendations)
- "world": About other people, external events, general knowledge, objective facts
- "assistant": First-person actions, experiences, or observations by the speaker/author (e.g., "I changed X", "I discovered Y", "I debugged Z"). Also includes interactions with the user (requests, recommendations). If the narrator describes something they did, tried, learned, or decided — use "assistant".
══════════════════════════════════════════════════════════════════════════
TEMPORAL HANDLING
@@ -616,7 +620,7 @@ VERBOSE_FACT_EXTRACTION_PROMPT = """Extract facts from text into structured form
LANGUAGE: MANDATORY — Detect the language of the input text and produce ALL output in that EXACT same language. You are STRICTLY FORBIDDEN from translating or switching to any other language. Every single word of your output must be in the same language as the input. Do NOT output in a different language under any circumstance.
══════════════════════════════════════════════════════════════════════════
{retain_mission_section}══════════════════════════════════════════════════════════════════════════
FACT FORMAT - ALL FIVE DIMENSIONS REQUIRED - MAXIMUM VERBOSITY
══════════════════════════════════════════════════════════════════════════
@@ -827,7 +831,9 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
custom_instructions=config.retain_custom_instructions,
)
elif extraction_mode == "verbose":
prompt = VERBOSE_FACT_EXTRACTION_PROMPT
prompt = VERBOSE_FACT_EXTRACTION_PROMPT.format(
retain_mission_section=retain_mission_section,
)
elif extraction_mode == "verbatim":
prompt = VERBATIM_FACT_EXTRACTION_PROMPT.format(
retain_mission_section=retain_mission_section,
@@ -1049,7 +1055,7 @@ async def _extract_facts_from_chunk(
f"LLM response missing 'facts' field or returned empty list. "
f"Response: {extraction_response_json}. "
f"Input: "
f"date: {event_date.isoformat()}, "
f"date: {event_date.isoformat() if event_date else 'unset'}, "
f"context: {context if context else 'none'}, "
f"text: {chunk}"
)
@@ -1458,28 +1464,76 @@ async def extract_facts_from_text(
f"chunk_size={config.retain_chunk_size:,}) - starting parallel LLM extraction"
)
tasks = [
_extract_facts_with_auto_split(
chunk=chunk,
chunk_index=i,
total_chunks=len(chunks),
event_date=event_date,
context=context,
llm_config=llm_config,
config=config,
agent_name=agent_name,
metadata=metadata,
)
for i, chunk in enumerate(chunks)
]
chunk_results = await asyncio.gather(*tasks)
# Per-chunk retry wrapper: each chunk gets up to MAX_CHUNK_RETRIES attempts.
# This handles transient LLM failures (timeouts, rate limits, malformed responses)
# without discarding the entire batch. If a chunk still fails after all retries,
# the ENTIRE retain fails — we do not accept partial extraction.
MAX_CHUNK_RETRIES = 3
CHUNK_RETRY_BASE_DELAY = 2.0 # seconds, doubles each retry
async def _extract_chunk_with_retry(chunk: str, chunk_index: int) -> tuple:
"""Extract facts from a single chunk with retries on failure."""
last_exception = None
for attempt in range(MAX_CHUNK_RETRIES):
try:
return await _extract_facts_with_auto_split(
chunk=chunk,
chunk_index=chunk_index,
total_chunks=len(chunks),
event_date=event_date,
context=context,
llm_config=llm_config,
config=config,
agent_name=agent_name,
metadata=metadata,
)
except Exception as e:
last_exception = e
if attempt < MAX_CHUNK_RETRIES - 1:
delay = CHUNK_RETRY_BASE_DELAY * (2**attempt)
logger.warning(
f"Chunk {chunk_index}/{len(chunks)} extraction failed "
f"(attempt {attempt + 1}/{MAX_CHUNK_RETRIES}): "
f"{type(e).__name__}. Retrying in {delay:.0f}s..."
)
await asyncio.sleep(delay)
else:
logger.error(
f"Chunk {chunk_index}/{len(chunks)} extraction failed after "
f"{MAX_CHUNK_RETRIES} attempts: {type(e).__name__}: {e}"
)
raise last_exception
tasks = [_extract_chunk_with_retry(chunk, i) for i, chunk in enumerate(chunks)]
# return_exceptions=True so we can collect all results even if some chunks
# exhausted their retries. We check for failures below and fail the retain
# if ANY chunk could not be extracted — partial extraction is not acceptable.
chunk_results = await asyncio.gather(*tasks, return_exceptions=True)
all_facts = []
chunk_metadata = [] # [(chunk_text, fact_count), ...]
total_usage = TokenUsage()
for chunk, (chunk_facts, chunk_usage) in zip(chunks, chunk_results):
failed_chunks = []
for i, (chunk, result) in enumerate(zip(chunks, chunk_results)):
if isinstance(result, Exception):
failed_chunks.append((i, result))
continue
chunk_facts, chunk_usage = result
all_facts.extend(chunk_facts)
chunk_metadata.append((chunk, len(chunk_facts)))
total_usage = total_usage + chunk_usage
if failed_chunks:
# Fail the entire retain — partial extraction is not acceptable.
# All successfully extracted facts are discarded because the transaction
# hasn't committed yet. The worker poller will retry the entire task.
failed_summary = ", ".join(f"chunk {idx}: {type(err).__name__}" for idx, err in failed_chunks[:5])
raise RuntimeError(
f"Fact extraction failed: {len(failed_chunks)}/{len(chunks)} chunks failed "
f"after {MAX_CHUNK_RETRIES} retries each. First failures: {failed_summary}"
)
return all_facts, chunk_metadata, total_usage
@@ -1913,7 +1967,7 @@ async def extract_facts_from_contents_batch_api(
for fact_from_llm in chunk_facts:
extracted_fact = ExtractedFactType(
fact_text=fact_from_llm.fact,
fact_type="experience" if fact_from_llm.fact_type == "assistant" else "world",
fact_type=fact_from_llm.fact_type,
entities=[e.text for e in (fact_from_llm.entities or [])],
occurred_start=_parse_datetime(fact_from_llm.occurred_start) if fact_from_llm.occurred_start else None,
occurred_end=_parse_datetime(fact_from_llm.occurred_end) if fact_from_llm.occurred_end else None,
@@ -2049,8 +2103,9 @@ async def extract_facts_from_contents(
)
fact_extraction_tasks.append(task)
# Step 2: Wait for all fact extractions to complete
all_fact_results = await asyncio.gather(*fact_extraction_tasks)
# Step 2: Wait for all fact extractions to complete.
# Use return_exceptions=True so one content item failure doesn't discard the rest.
all_fact_results = await asyncio.gather(*fact_extraction_tasks, return_exceptions=True)
# Step 3: Flatten and convert to typed objects
extracted_facts: list[ExtractedFactType] = []
@@ -2060,9 +2115,16 @@ async def extract_facts_from_contents(
global_chunk_idx = 0
global_fact_idx = 0
for content_index, (content, (facts_from_llm, chunks_from_llm, content_usage)) in enumerate(
zip(contents, all_fact_results)
):
# Filter out failed content items
valid_results = []
for content, result in zip(contents, all_fact_results):
if isinstance(result, Exception):
logger.warning(f"Content extraction failed (skipping): {type(result).__name__}: {result}")
valid_results.append((content, ([], [], TokenUsage())))
else:
valid_results.append((content, result))
for content_index, (content, (facts_from_llm, chunks_from_llm, content_usage)) in enumerate(valid_results):
total_usage = total_usage + content_usage
chunk_start_idx = global_chunk_idx
@@ -2090,7 +2152,7 @@ async def extract_facts_from_contents(
# mentioned_at is always the event_date (when the conversation/document occurred)
extracted_fact = ExtractedFactType(
fact_text=fact_from_llm.fact,
fact_type="experience" if fact_from_llm.fact_type == "assistant" else "world",
fact_type=fact_from_llm.fact_type,
entities=[e.text for e in (fact_from_llm.entities or [])],
# occurred_start/end: from LLM only, leave None if not provided
occurred_start=_parse_datetime(fact_from_llm.occurred_start)
@@ -10,7 +10,7 @@ import uuid
from ...config import get_config
from ..memory_engine import fq_table
from .bank_utils import DEFAULT_DISPOSITION, create_bank_hnsw_indexes
from .bank_utils import DEFAULT_DISPOSITION, create_bank_vector_indexes
from .fact_extraction import _sanitize_text
from .types import ProcessedFact
@@ -44,7 +44,6 @@ async def insert_facts_batch(
mentioned_ats = []
contexts = []
fact_types = []
confidence_scores = []
metadata_jsons = []
chunk_ids = []
document_ids = []
@@ -64,8 +63,6 @@ async def insert_facts_batch(
mentioned_ats.append(fact.mentioned_at)
contexts.append(_sanitize_text(fact.context))
fact_types.append(fact.fact_type)
# confidence_score is only for opinion facts
confidence_scores.append(1.0 if fact.fact_type == "opinion" else None)
metadata_jsons.append(json.dumps(fact.metadata))
chunk_ids.append(fact.chunk_id)
# Use per-fact document_id if available, otherwise fallback to batch-level document_id
@@ -103,18 +100,18 @@ async def insert_facts_batch(
WITH input_data AS (
SELECT * FROM unnest(
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
$8::text[], $9::text[], $10::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[], $15::jsonb[], $16::text[]
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags_json,
context, fact_type, metadata, chunk_id, document_id, tags_json,
observation_scopes_json, text_signals)
)
INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags,
context, fact_type, metadata, chunk_id, document_id, tags,
observation_scopes, text_signals, search_vector)
SELECT
$1,
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, confidence_score, metadata, chunk_id, document_id,
context, fact_type, metadata, chunk_id, document_id,
COALESCE(
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
'{{}}'::varchar[]
@@ -135,18 +132,18 @@ async def insert_facts_batch(
WITH input_data AS (
SELECT * FROM unnest(
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
$8::text[], $9::text[], $10::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[], $15::jsonb[], $16::text[]
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags_json,
context, fact_type, metadata, chunk_id, document_id, tags_json,
observation_scopes_json, text_signals)
)
INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags,
context, fact_type, metadata, chunk_id, document_id, tags,
observation_scopes, text_signals)
SELECT
$1,
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, confidence_score, metadata, chunk_id, document_id,
context, fact_type, metadata, chunk_id, document_id,
COALESCE(
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
'{{}}'::varchar[]
@@ -168,7 +165,6 @@ async def insert_facts_batch(
mentioned_ats,
contexts,
fact_types,
confidence_scores,
metadata_jsons,
chunk_ids,
document_ids,
@@ -207,8 +203,8 @@ async def ensure_bank_exists(conn, bank_id: str) -> None:
internal_id,
)
if inserted:
# Fresh insert — create per-bank HNSW indexes
await create_bank_hnsw_indexes(conn, bank_id, str(internal_id))
# Fresh insert — create per-bank vector indexes
await create_bank_vector_indexes(conn, bank_id, str(internal_id))
async def handle_document_tracking(
@@ -221,7 +217,10 @@ async def handle_document_tracking(
document_tags: list[str] | None = None,
) -> None:
"""
Handle document tracking in the database.
Handle document tracking in the database (full-replace mode).
Deletes the existing document (cascading to all units and links) on the
first batch, then inserts the new document record.
Args:
conn: Database connection
@@ -238,22 +237,58 @@ async def handle_document_tracking(
combined_content = _sanitize_text(combined_content) or ""
content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
# Always delete old document first if it exists (cascades to units and links)
# Delete old document first (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(
f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING id", document_id, bank_id
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 _upsert_document_row(conn, bank_id, document_id, combined_content, content_hash, retain_params, document_tags)
async def upsert_document_metadata(
conn,
bank_id: str,
document_id: str,
combined_content: str,
retain_params: dict | None = None,
document_tags: list[str] | None = None,
) -> None:
"""
Update document metadata without deleting existing facts/chunks.
Used by delta retain: the document row is upserted but chunks and
memory_units are managed separately at the chunk level.
"""
import hashlib
combined_content = _sanitize_text(combined_content) or ""
content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
await _upsert_document_row(conn, bank_id, document_id, combined_content, content_hash, retain_params, document_tags)
async def _upsert_document_row(
conn,
bank_id: str,
document_id: str,
combined_content: str,
content_hash: str,
retain_params: dict | None = None,
document_tags: list[str] | None = None,
) -> None:
"""Insert or update a document row."""
await conn.execute(
f"""
INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, metadata, retain_params, tags)
VALUES ($1, $2, $3, $4, $5, $6, $7)
INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, retain_params, tags)
VALUES ($1, $2, $3, $4, $5, $6)
ON CONFLICT (id, bank_id) DO UPDATE
SET original_text = EXCLUDED.original_text,
content_hash = EXCLUDED.content_hash,
metadata = EXCLUDED.metadata,
retain_params = EXCLUDED.retain_params,
tags = EXCLUDED.tags,
updated_at = NOW()
@@ -262,7 +297,37 @@ async def handle_document_tracking(
bank_id,
combined_content,
content_hash,
json.dumps({}), # Empty metadata dict
json.dumps(retain_params) if retain_params else None,
document_tags or [],
)
async def update_memory_units_tags(
conn,
bank_id: str,
document_id: str,
tags: list[str],
) -> int:
"""
Update tags on all memory_units belonging to a document.
Used during delta retain to propagate tag changes to unchanged facts.
Returns:
Number of memory units updated.
"""
result = await conn.execute(
f"""
UPDATE {fq_table("memory_units")}
SET tags = $3, updated_at = NOW()
WHERE bank_id = $1 AND document_id = $2
""",
bank_id,
document_id,
tags or [],
)
# result is a status string like "UPDATE 5"
try:
return int(result.split()[-1])
except (ValueError, IndexError):
return 0
@@ -32,17 +32,26 @@ async def create_temporal_links_batch(conn, bank_id: str, unit_ids: list[str]) -
return await link_utils.create_temporal_links_batch_per_fact(conn, bank_id, unit_ids, log_buffer=[])
async def create_semantic_links_batch(conn, bank_id: str, unit_ids: list[str], embeddings: list[list[float]]) -> int:
async def create_semantic_links_batch(
conn,
bank_id: str,
unit_ids: list[str],
embeddings: list[list[float]],
pre_computed_ann_links: list[tuple] | None = None,
) -> int:
"""
Create semantic links between facts.
Links facts that are semantically similar based on embeddings.
When pre_computed_ann_links are provided (from Phase 1), they are used
instead of running ANN queries inside the transaction.
Args:
conn: Database connection
bank_id: Bank identifier
unit_ids: List of unit IDs to create links for
embeddings: List of embedding vectors (same length as unit_ids)
pre_computed_ann_links: Pre-computed ANN results from Phase 1
Returns:
Number of semantic links created
@@ -53,10 +62,12 @@ async def create_semantic_links_batch(conn, bank_id: str, unit_ids: list[str], e
if len(unit_ids) != len(embeddings):
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and embeddings ({len(embeddings)})")
return await link_utils.create_semantic_links_batch(conn, bank_id, unit_ids, embeddings, log_buffer=[])
return await link_utils.create_semantic_links_batch(
conn, bank_id, unit_ids, embeddings, log_buffer=[], pre_computed_ann_links=pre_computed_ann_links
)
async def create_causal_links_batch(conn, unit_ids: list[str], facts: list[ProcessedFact]) -> int:
async def create_causal_links_batch(conn, bank_id: str, unit_ids: list[str], facts: list[ProcessedFact]) -> int:
"""
Create causal links between facts.
@@ -94,6 +105,6 @@ async def create_causal_links_batch(conn, unit_ids: list[str], facts: list[Proce
else:
causal_relations_per_fact.append([])
link_count = await link_utils.create_causal_links_batch(conn, unit_ids, causal_relations_per_fact)
link_count = await link_utils.create_causal_links_batch(conn, bank_id, unit_ids, causal_relations_per_fact)
return link_count
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -107,7 +107,7 @@ class ExtractedFact:
"""
fact_text: str
fact_type: str # "world", "experience", "opinion", "observation"
fact_type: str # "world", "experience", "observation"
entities: list[str] = field(default_factory=list)
occurred_start: datetime | None = None
occurred_end: datetime | None = None
@@ -221,6 +221,45 @@ class ProcessedFact:
)
@dataclass
class Phase3Context:
"""
Data passed from Phase 2 to Phase 3 for entity link building.
Contains the unit IDs and entity resolution data needed to build
entity links for UI graph visualization after the write transaction commits.
"""
unit_ids: list[str] = field(default_factory=list)
resolved_entity_ids: list[str] = field(default_factory=list)
entity_to_unit: list[tuple] = field(default_factory=list)
unit_to_entity_ids: dict[str, list[str]] = field(default_factory=dict)
@dataclass
class EntityResolutionResult:
"""
Result of Phase 1 entity resolution.
Contains resolved entity IDs and the mapping data needed to remap
placeholder unit IDs to real IDs after fact insertion in Phase 2.
"""
resolved_entity_ids: list[str]
entity_to_unit: list[tuple]
unit_to_entity_ids: dict[str, list[str]]
@dataclass
class Phase1Result:
"""
Full result of Phase 1 (entity resolution + optional semantic ANN).
"""
entities: EntityResolutionResult
semantic_ann_links: list[tuple]
@dataclass
class EntityLink:
"""
@@ -248,7 +287,6 @@ class RetainBatch:
contents: list[RetainContent]
document_id: str | None = None
fact_type_override: str | None = None
confidence_score: float | None = None
document_tags: list[str] = field(default_factory=list) # Tags applied to all items
# Extracted data (populated during processing)
@@ -3,12 +3,11 @@ Search module for memory retrieval.
Provides modular search architecture:
- Retrieval: 4-way parallel (semantic + BM25 + graph + temporal)
- Graph retrieval: Pluggable strategies (BFS, PPR)
- Graph retrieval: Link expansion strategy
- Reranking: Pluggable strategies (heuristic, cross-encoder)
"""
from .graph_retrieval import BFSGraphRetriever, GraphRetriever
from .mpfp_retrieval import MPFPGraphRetriever
from .graph_retrieval import GraphRetriever
from .reranking import CrossEncoderReranker
from .retrieval import (
ParallelRetrievalResult,
@@ -21,7 +20,5 @@ __all__ = [
"set_default_graph_retriever",
"ParallelRetrievalResult",
"GraphRetriever",
"BFSGraphRetriever",
"MPFPGraphRetriever",
"CrossEncoderReranker",
]
@@ -2,17 +2,15 @@
Graph retrieval strategies for memory recall.
This module provides an abstraction for graph-based memory retrieval,
allowing different algorithms (BFS spreading activation, PPR, etc.) to be
swapped without changing the rest of the recall pipeline.
allowing different algorithms to be swapped without changing the rest
of the recall pipeline.
"""
import logging
from abc import ABC, abstractmethod
from ..db_utils import acquire_with_retry
from ..memory_engine import fq_table
from .tags import TagGroup, TagsMatch, filter_results_by_tag_groups, filter_results_by_tags
from .types import MPFPTimings, RetrievalResult
from .tags import TagGroup, TagsMatch
from .types import GraphRetrievalTimings, RetrievalResult
logger = logging.getLogger(__name__)
@@ -29,7 +27,7 @@ class GraphRetriever(ABC):
@property
@abstractmethod
def name(self) -> str:
"""Return identifier for this retrieval strategy (e.g., 'bfs', 'mpfp')."""
"""Return identifier for this retrieval strategy (e.g., 'link_expansion')."""
pass
@abstractmethod
@@ -47,7 +45,7 @@ class GraphRetriever(ABC):
tags: list[str] | None = None, # Visibility scope tags for filtering
tags_match: TagsMatch = "any", # How to match tags: 'any' (OR) or 'all' (AND)
tag_groups: list[TagGroup] | None = None, # Compound boolean tag filter groups
) -> tuple[list[RetrievalResult], MPFPTimings | None]:
) -> tuple[list[RetrievalResult], GraphRetrievalTimings | None]:
"""
Retrieve relevant facts via graph traversal.
@@ -55,228 +53,15 @@ class GraphRetriever(ABC):
pool: Database connection pool
query_embedding_str: Query embedding as string (for finding entry points)
bank_id: Memory bank identifier
fact_type: Fact type to filter ('world', 'experience', 'opinion', 'observation')
fact_type: Fact type to filter ('world', 'experience', 'observation')
budget: Maximum number of nodes to explore/return
query_text: Original query text (optional, for some strategies)
semantic_seeds: Pre-computed semantic entry points (from semantic retrieval)
temporal_seeds: Pre-computed temporal entry points (from temporal retrieval)
adjacency: Pre-loaded typed adjacency graph (optional, for MPFP)
adjacency: Pre-loaded typed adjacency graph (optional)
tags: Optional list of tags for visibility filtering (OR matching)
Returns:
Tuple of (List of RetrievalResult with activation scores, optional timing info)
"""
pass
class BFSGraphRetriever(GraphRetriever):
"""
Graph retrieval using BFS-style spreading activation.
Starting from semantic entry points, spreads activation through
the memory graph (entity, temporal, causal links) using breadth-first
traversal with decaying activation.
This is the original Hindsight graph retrieval algorithm.
"""
def __init__(
self,
entry_point_limit: int = 5,
entry_point_threshold: float = 0.5,
activation_decay: float = 0.8,
min_activation: float = 0.1,
batch_size: int = 20,
):
"""
Initialize BFS graph retriever.
Args:
entry_point_limit: Maximum number of entry points to start from
entry_point_threshold: Minimum semantic similarity for entry points
activation_decay: Decay factor per hop (activation *= decay)
min_activation: Minimum activation to continue spreading
batch_size: Number of nodes to process per batch (for neighbor fetching)
"""
self.entry_point_limit = entry_point_limit
self.entry_point_threshold = entry_point_threshold
self.activation_decay = activation_decay
self.min_activation = min_activation
self.batch_size = batch_size
@property
def name(self) -> str:
return "bfs"
async def retrieve(
self,
pool,
query_embedding_str: str,
bank_id: str,
fact_type: str,
budget: int,
query_text: str | None = None,
semantic_seeds: list[RetrievalResult] | None = None,
temporal_seeds: list[RetrievalResult] | None = None,
adjacency=None, # Not used by BFS
tags: list[str] | None = None,
tags_match: TagsMatch = "any",
tag_groups: list[TagGroup] | None = None,
) -> tuple[list[RetrievalResult], MPFPTimings | None]:
"""
Retrieve facts using BFS spreading activation.
Algorithm:
1. Find entry points (top semantic matches above threshold)
2. BFS traversal: visit neighbors, propagate decaying activation
3. Boost causal links (causes, enables, prevents)
4. Return visited nodes up to budget
Note: BFS finds its own entry points via embedding search.
The semantic_seeds, temporal_seeds, and adjacency parameters are accepted
for interface compatibility but not used.
"""
async with acquire_with_retry(pool) as conn:
results = await self._retrieve_with_conn(
conn,
query_embedding_str,
bank_id,
fact_type,
budget,
tags=tags,
tags_match=tags_match,
tag_groups=tag_groups,
)
return results, None
async def _retrieve_with_conn(
self,
conn,
query_embedding_str: str,
bank_id: str,
fact_type: str,
budget: int,
tags: list[str] | None = None,
tags_match: TagsMatch = "any",
tag_groups: list[TagGroup] | None = None,
) -> list[RetrievalResult]:
"""Internal implementation with connection."""
from .tags import build_tag_groups_where_clause, build_tags_where_clause_simple
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
tag_groups_param_start = 6 + (1 if tags else 0)
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
params = [query_embedding_str, bank_id, fact_type, self.entry_point_threshold, self.entry_point_limit]
if tags:
params.append(tags)
params.extend(groups_params)
# Step 1: Find entry points
entry_points = await conn.fetch(
f"""
SELECT id, text, context, event_date, occurred_start, occurred_end,
mentioned_at, fact_type, document_id, chunk_id, tags,
1 - (embedding <=> $1::vector) AS similarity
FROM {fq_table("memory_units")}
WHERE bank_id = $2
AND embedding IS NOT NULL
AND fact_type = $3
AND (1 - (embedding <=> $1::vector)) >= $4
{tags_clause}
{groups_clause}
ORDER BY embedding <=> $1::vector
LIMIT $5
""",
*params,
)
if not entry_points:
logger.debug(
f"[BFS] No entry points found for fact_type={fact_type} (tags={tags}, tags_match={tags_match})"
)
return []
logger.debug(
f"[BFS] Found {len(entry_points)} entry points for fact_type={fact_type} "
f"(tags={tags}, tags_match={tags_match})"
)
# Step 2: BFS spreading activation
visited = set()
results = []
queue = [(RetrievalResult.from_db_row(dict(r)), r["similarity"]) for r in entry_points]
budget_remaining = budget
while queue and budget_remaining > 0:
# Collect a batch of nodes to process
batch_nodes = []
batch_activations = {}
while queue and len(batch_nodes) < self.batch_size and budget_remaining > 0:
current, activation = queue.pop(0)
unit_id = current.id
if unit_id not in visited:
visited.add(unit_id)
budget_remaining -= 1
current.activation = activation
results.append(current)
batch_nodes.append(current.id)
batch_activations[unit_id] = activation
# Batch fetch neighbors
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.fact_type,
mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
ml.weight, ml.link_type, ml.from_unit_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
ORDER BY ml.weight DESC
LIMIT $4
""",
batch_nodes,
self.min_activation,
fact_type,
max_neighbors,
)
for n in neighbors:
neighbor_id = str(n["id"])
if neighbor_id not in visited:
parent_id = str(n["from_unit_id"])
parent_activation = batch_activations.get(parent_id, 0.5)
# Boost causal links
link_type = n["link_type"]
base_weight = n["weight"]
if link_type in ("causes", "caused_by"):
causal_boost = 2.0
elif link_type in ("enables", "prevents"):
causal_boost = 1.5
else:
causal_boost = 1.0
effective_weight = base_weight * causal_boost
new_activation = parent_activation * effective_weight * self.activation_decay
if new_activation > self.min_activation:
neighbor_result = RetrievalResult.from_db_row(dict(n))
queue.append((neighbor_result, new_activation))
# Apply tags filtering (BFS may traverse into memories that don't match tags criteria)
if tags:
results = filter_results_by_tags(results, tags, match=tags_match)
# Apply compound tag group filtering (post-traversal)
if tag_groups:
results = filter_results_by_tag_groups(results, tag_groups)
return results
@@ -4,9 +4,9 @@ Link Expansion graph retrieval.
Expands from semantic/temporal seeds through three parallel, first-class signals
stored in memory_links:
1. Entity links precomputed co-occurrence graph (created at retain time, bounded to
MAX_LINKS_PER_ENTITY per entity). Score = number of distinct shared
entities between the seed set and each candidate.
1. Entity links query-time self-join through unit_entities. Score = number of distinct
shared entities between the seed set and each candidate, computed via
COUNT(DISTINCT entity_id). More accurate than precomputed entity links.
2. Semantic links precomputed kNN graph (each new fact linked to its top-5 most
similar existing facts at insert time, similarity >= 0.7). Checked
in both directions since the graph is not symmetric. Score = weight.
@@ -29,7 +29,7 @@ from ..db_utils import acquire_with_retry
from ..memory_engine import fq_table
from .graph_retrieval import GraphRetriever
from .tags import TagGroup, TagsMatch, filter_results_by_tag_groups, filter_results_by_tags
from .types import MPFPTimings, RetrievalResult
from .types import GraphRetrievalTimings, RetrievalResult
logger = logging.getLogger(__name__)
@@ -59,7 +59,7 @@ async def _find_semantic_seeds(
rows = await conn.fetch(
f"""
SELECT id, text, context, event_date, occurred_start, occurred_end,
mentioned_at, fact_type, document_id, chunk_id, tags,
mentioned_at, fact_type, document_id, chunk_id, tags, proof_count,
1 - (embedding <=> $1::vector) AS similarity
FROM {fq_table("memory_units")}
WHERE bank_id = $2
@@ -116,7 +116,7 @@ class LinkExpansionRetriever(GraphRetriever):
tags: list[str] | None = None,
tags_match: TagsMatch = "any",
tag_groups: list[TagGroup] | None = None,
) -> tuple[list[RetrievalResult], MPFPTimings | None]:
) -> tuple[list[RetrievalResult], GraphRetrievalTimings | None]:
"""
Retrieve facts by expanding links from seeds.
@@ -136,7 +136,7 @@ class LinkExpansionRetriever(GraphRetriever):
Tuple of (results, timings)
"""
start_time = time.time()
timings = MPFPTimings(fact_type=fact_type)
timings = GraphRetrievalTimings(fact_type=fact_type)
async with acquire_with_retry(pool) as conn:
# Find seeds if not provided
@@ -264,29 +264,33 @@ class LinkExpansionRetriever(GraphRetriever):
"""
ml = fq_table("memory_links")
mu = fq_table("memory_units")
all_rows = await conn.fetch(
f"""
WITH entity_expanded AS (
-- Entity co-occurrence: seeds their precomputed entity-link neighbors.
-- Score = distinct shared entities (bounded at retain time to
-- MAX_LINKS_PER_ENTITY=50). GROUP BY mu.id is sufficient because mu.id
-- is the primary key and functionally determines all other mu columns.
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
COUNT(DISTINCT ml.entity_id)::float AS score,
'entity'::text AS source
FROM {ml} ml
JOIN {mu} mu ON mu.id = ml.to_unit_id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type = 'entity'
ue = fq_table("unit_entities")
entity_cte = f"""
entity_expanded AS (
-- Entity co-occurrence via unit_entities self-join.
-- Finds units sharing entities with seeds at query time more accurate
-- than precomputed entity links (no stale 50-neighbor cap).
-- Score = COUNT(DISTINCT shared entities), mapped to [0,1] via tanh.
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
COUNT(DISTINCT ue_seed.entity_id)::float AS score,
'entity'::text AS source
FROM {ue} ue_seed
JOIN {ue} ue_target ON ue_seed.entity_id = ue_target.entity_id
JOIN {mu} mu ON mu.id = ue_target.unit_id
WHERE ue_seed.unit_id = ANY($1::uuid[])
AND ue_target.unit_id != ALL($1::uuid[])
AND mu.fact_type = $2
AND mu.id != ALL($1::uuid[])
GROUP BY mu.id
ORDER BY score DESC
LIMIT $3
),
)"""
all_rows = await conn.fetch(
f"""
WITH {entity_cte},
semantic_expanded AS (
-- Semantic kNN: both outgoing (seeds their kNN at insert time) and
-- incoming (facts inserted after seeds that found seeds as kNN).
@@ -294,14 +298,14 @@ class LinkExpansionRetriever(GraphRetriever):
SELECT
id, text, context, event_date, occurred_start,
occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags,
fact_type, document_id, chunk_id, tags, proof_count,
MAX(weight) AS score,
'semantic'::text AS source
FROM (
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
ml.weight
FROM {ml} ml
JOIN {mu} mu ON mu.id = ml.to_unit_id
@@ -313,7 +317,7 @@ class LinkExpansionRetriever(GraphRetriever):
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
ml.weight
FROM {ml} ml
JOIN {mu} mu ON mu.id = ml.from_unit_id
@@ -324,7 +328,7 @@ class LinkExpansionRetriever(GraphRetriever):
) sem_raw
GROUP BY id, text, context, event_date, occurred_start,
occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags
fact_type, document_id, chunk_id, tags, proof_count
ORDER BY score DESC
LIMIT $3
),
@@ -335,7 +339,7 @@ class LinkExpansionRetriever(GraphRetriever):
SELECT DISTINCT ON (mu.id)
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
ml.weight AS score,
'causal'::text AS source
FROM {ml} ml
@@ -397,6 +401,19 @@ class LinkExpansionRetriever(GraphRetriever):
f"{len(source_ids_found)} source_memory_ids found"
)
ue = fq_table("unit_entities")
connected_sources_cte = f"""
connected_sources AS (
-- Find sources sharing entities with seed observation sources
-- via unit_entities self-join (query-time, no precomputed links needed).
SELECT DISTINCT ue_target.unit_id AS source_id
FROM seed_sources ss
JOIN {ue} ue_seed ON ue_seed.unit_id = ss.source_id
JOIN {ue} ue_target ON ue_seed.entity_id = ue_target.entity_id
WHERE ue_target.unit_id != ss.source_id
)"""
entity_rows = await conn.fetch(
f"""
WITH seed_sources AS (
@@ -405,22 +422,14 @@ class LinkExpansionRetriever(GraphRetriever):
WHERE id = ANY($1::uuid[])
AND source_memory_ids IS NOT NULL
),
connected_sources AS (
-- Mirror the non-observation entity expansion: follow pre-bounded entity
-- links in memory_links (capped to MAX_LINKS_PER_ENTITY=50 at retain time).
-- Score = number of distinct shared entities, same as the non-obs path.
SELECT DISTINCT ml.to_unit_id AS source_id
FROM seed_sources ss
JOIN {fq_table("memory_links")} ml ON ml.from_unit_id = ss.source_id
WHERE ml.link_type = 'entity'
),
{connected_sources_cte},
connected_array AS (
SELECT array_agg(source_id) AS source_ids FROM connected_sources
)
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
(SELECT COUNT(DISTINCT s) FROM unnest(mu.source_memory_ids) s WHERE s = ANY(ca.source_ids))::float AS score
FROM {fq_table("memory_units")} mu, connected_array ca
WHERE mu.fact_type = 'observation'
@@ -444,13 +453,13 @@ class LinkExpansionRetriever(GraphRetriever):
SELECT
id, text, context, event_date, occurred_start,
occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags,
fact_type, document_id, chunk_id, tags, proof_count,
MAX(weight) AS score,
'semantic'::text AS source
FROM (
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, ml.weight
mu.chunk_id, mu.tags, mu.proof_count, ml.weight
FROM {ml} ml JOIN {mu} mu ON mu.id = ml.to_unit_id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
@@ -458,21 +467,21 @@ class LinkExpansionRetriever(GraphRetriever):
UNION ALL
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, ml.weight
mu.chunk_id, mu.tags, mu.proof_count, ml.weight
FROM {ml} ml JOIN {mu} mu ON mu.id = ml.from_unit_id
WHERE ml.to_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
) sem_raw
GROUP BY id, text, context, event_date, occurred_start, occurred_end,
mentioned_at, fact_type, document_id, chunk_id, tags
mentioned_at, fact_type, document_id, chunk_id, tags, proof_count
ORDER BY score DESC LIMIT $2
),
causal_expanded AS (
SELECT DISTINCT ON (mu.id)
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, ml.weight AS score, 'causal'::text AS source
mu.chunk_id, mu.tags, mu.proof_count, ml.weight AS score, 'causal'::text AS source
FROM {ml} ml JOIN {mu} mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
@@ -1,702 +0,0 @@
"""
Meta-Path Forward Push (MPFP) graph retrieval.
A sublinear graph traversal algorithm for memory retrieval over heterogeneous
graphs with multiple edge types (semantic, temporal, causal, entity).
Combines meta-path patterns from HIN literature with Forward Push local
propagation from Approximate PPR.
Key properties:
- Sublinear in graph size (threshold pruning bounds active nodes)
- Lazy edge loading: only loads edges for frontier nodes, not entire graph
- Predefined patterns capture different retrieval intents
- All patterns run in parallel, results fused via RRF
- No LLM in the loop during traversal
"""
import asyncio
import logging
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 .tags import TagGroup, TagsMatch
from .types import MPFPTimings, RetrievalResult
logger = logging.getLogger(__name__)
# -----------------------------------------------------------------------------
# Data Classes
# -----------------------------------------------------------------------------
@dataclass
class EdgeTarget:
"""A neighbor node with its edge weight."""
node_id: str
weight: float
@dataclass
class EdgeCache:
"""
Cache for lazily-loaded edges.
Grows per-hop as edges are loaded for frontier nodes.
Shared across patterns to avoid redundant loads.
Loads ALL edge types at once to minimize DB queries.
Thread-safe via asyncio lock to prevent redundant concurrent loads.
"""
# edge_type -> from_node_id -> list of EdgeTarget
graphs: dict[str, dict[str, list[EdgeTarget]]] = field(default_factory=dict)
# Track which nodes have been fully loaded (all edge types)
_fully_loaded: set[str] = field(default_factory=set)
# Timing stats
db_queries: int = 0
edge_load_time: float = 0.0
# Detailed hop timing for debugging
hop_details: list[dict] = field(default_factory=list)
# Lock to prevent redundant concurrent loads
_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
def get_neighbors(self, edge_type: str, node_id: str) -> list[EdgeTarget]:
"""Get neighbors for a node via a specific edge type."""
return self.graphs.get(edge_type, {}).get(node_id, [])
def get_normalized_neighbors(self, edge_type: str, node_id: str, top_k: int) -> list[EdgeTarget]:
"""Get top-k neighbors with weights normalized to sum to 1."""
neighbors = self.get_neighbors(edge_type, node_id)[:top_k]
if not neighbors:
return []
total = sum(n.weight for n in neighbors)
if total == 0:
return []
return [EdgeTarget(node_id=n.node_id, weight=n.weight / total) for n in neighbors]
def is_fully_loaded(self, node_id: str) -> bool:
"""Check if all edges for this node have been loaded."""
return node_id in self._fully_loaded
def get_uncached(self, node_ids: list[str]) -> list[str]:
"""Get node IDs that haven't been fully loaded yet."""
return [n for n in node_ids if not self.is_fully_loaded(n)]
def add_all_edges(self, edges_by_type: dict[str, dict[str, list[EdgeTarget]]], all_queried: list[str]):
"""
Add loaded edges to the cache (all edge types at once).
Args:
edges_by_type: Dict mapping edge_type -> from_node_id -> list of EdgeTarget
all_queried: All node IDs that were queried (marks them as fully loaded)
"""
for edge_type, edges in edges_by_type.items():
if edge_type not in self.graphs:
self.graphs[edge_type] = {}
for node_id, neighbors in edges.items():
self.graphs[edge_type][node_id] = neighbors
# Mark all queried nodes as fully loaded (even if they have no edges)
self._fully_loaded.update(all_queried)
@dataclass
class PatternResult:
"""Result from a single pattern traversal."""
pattern: list[str]
scores: dict[str, float] # node_id -> accumulated mass
@dataclass
class MPFPConfig:
"""Configuration for MPFP algorithm."""
alpha: float = 0.15 # teleport/keep probability
threshold: float = 1e-6 # mass pruning threshold (lower = explore more)
top_k_neighbors: int = 20 # fan-out limit per node
# Patterns from semantic seeds
patterns_semantic: list[list[str]] = field(
default_factory=lambda: [
["semantic", "semantic"], # topic expansion
["entity", "temporal"], # entity timeline
["semantic", "causes"], # reasoning chains (forward)
["semantic", "caused_by"], # reasoning chains (backward)
["entity", "semantic"], # entity context
]
)
# Patterns from temporal seeds
patterns_temporal: list[list[str]] = field(
default_factory=lambda: [
["temporal", "semantic"], # what was happening then
["temporal", "entity"], # who was involved then
]
)
@dataclass
class SeedNode:
"""An entry point node with its initial score."""
node_id: str
score: float # initial mass (e.g., similarity score)
# -----------------------------------------------------------------------------
# Lazy Edge Loading
# -----------------------------------------------------------------------------
async def load_all_edges_for_frontier(
pool,
node_ids: list[str],
top_k_per_type: int = 20,
) -> dict[str, dict[str, list[EdgeTarget]]]:
"""
Load top-k edges per (node, edge_type) for frontier nodes.
Uses a LATERAL join to efficiently fetch only the top-k edges per type,
avoiding loading hundreds of entity edges when only 20 are needed.
Requires composite index: (from_unit_id, link_type, weight DESC)
Args:
pool: Database connection pool
node_ids: Frontier node IDs to load edges for
top_k_per_type: Max edges to load per (node, link_type) pair
Returns:
Dict mapping edge_type -> from_node_id -> list of EdgeTarget
"""
if not node_ids:
return {}
async with acquire_with_retry(pool) as conn:
# Use LATERAL join to get top-k per (from_node, link_type)
# This leverages the composite index for efficient early termination
rows = await conn.fetch(
f"""
WITH frontier(node_id) AS (SELECT unnest($1::uuid[]))
SELECT f.node_id as from_unit_id, lt.link_type, edges.to_unit_id, edges.weight
FROM frontier f
CROSS JOIN (VALUES ('semantic'), ('temporal'), ('entity'), ('causes'), ('caused_by')) AS lt(link_type)
CROSS JOIN LATERAL (
SELECT ml.to_unit_id, ml.weight
FROM {fq_table("memory_links")} ml
WHERE ml.from_unit_id = f.node_id
AND ml.link_type = lt.link_type
AND ml.weight >= 0.1
ORDER BY ml.weight DESC
LIMIT $2
) edges
""",
node_ids,
top_k_per_type,
)
# Group by edge_type -> from_node -> neighbors
result: dict[str, dict[str, list[EdgeTarget]]] = defaultdict(lambda: defaultdict(list))
for row in rows:
edge_type = row["link_type"]
from_id = str(row["from_unit_id"])
to_id = str(row["to_unit_id"])
weight = row["weight"]
result[edge_type][from_id].append(EdgeTarget(node_id=to_id, weight=weight))
# Convert nested defaultdicts to regular dicts
return {edge_type: dict(edges) for edge_type, edges in result.items()}
# -----------------------------------------------------------------------------
# Core Algorithm (Async with Lazy Loading)
# -----------------------------------------------------------------------------
@dataclass
class PatternState:
"""State for a pattern traversal between hops."""
pattern: list[str]
hop_index: int
scores: dict[str, float]
frontier: dict[str, float]
def _init_pattern_state(seeds: list[SeedNode], pattern: list[str]) -> PatternState:
"""Initialize pattern state from seeds."""
if not seeds:
return PatternState(pattern=pattern, hop_index=0, scores={}, frontier={})
total_seed_score = sum(s.score for s in seeds)
if total_seed_score == 0:
total_seed_score = len(seeds)
frontier = {s.node_id: s.score / total_seed_score for s in seeds}
return PatternState(pattern=pattern, hop_index=0, scores={}, frontier=frontier)
def _execute_hop(state: PatternState, cache: EdgeCache, config: MPFPConfig) -> set[str]:
"""
Execute ONE hop of traversal, return frontier nodes for next hop.
This is a pure function that uses cached edges (no DB access).
Returns set of uncached nodes needed for next hop.
"""
if state.hop_index >= len(state.pattern):
return set()
edge_type = state.pattern[state.hop_index]
# Collect active nodes above threshold
active_nodes = [node_id for node_id, mass in state.frontier.items() if mass >= config.threshold]
if not active_nodes:
state.frontier = {}
return set()
# Propagate mass using cached edges
next_frontier: dict[str, float] = {}
uncached_for_next: set[str] = set()
for node_id, mass in state.frontier.items():
if mass < config.threshold:
continue
# Keep α portion for this node
state.scores[node_id] = state.scores.get(node_id, 0) + config.alpha * mass
# Push (1-α) to neighbors
push_mass = (1 - config.alpha) * mass
neighbors = cache.get_normalized_neighbors(edge_type, node_id, config.top_k_neighbors)
for neighbor in neighbors:
next_frontier[neighbor.node_id] = next_frontier.get(neighbor.node_id, 0) + push_mass * neighbor.weight
# Track if we'll need edges for this node in the next hop
if not cache.is_fully_loaded(neighbor.node_id):
uncached_for_next.add(neighbor.node_id)
state.frontier = next_frontier
state.hop_index += 1
return uncached_for_next
def _finalize_pattern(state: PatternState, config: MPFPConfig) -> PatternResult:
"""Finalize pattern by adding remaining frontier mass to scores."""
for node_id, mass in state.frontier.items():
if mass >= config.threshold:
state.scores[node_id] = state.scores.get(node_id, 0) + mass
return PatternResult(pattern=state.pattern, scores=state.scores)
async def mpfp_traverse_hop_synchronized(
pool,
pattern_jobs: list[tuple[list[SeedNode], list[str]]],
config: MPFPConfig,
cache: EdgeCache,
) -> list[PatternResult]:
"""
Execute ALL patterns with hop-synchronized edge loading.
Instead of running each pattern independently (causing multiple DB queries),
this function:
1. Runs hop 1 for ALL patterns (using pre-warmed seed edges)
2. Collects ALL unique hop-2 frontier nodes across patterns
3. Pre-warms hop-2 edges in ONE query
4. Runs hop 2 for ALL patterns
This reduces DB queries from O(patterns * hops) to O(hops).
Args:
pool: Database connection pool
pattern_jobs: List of (seeds, pattern) tuples
config: Algorithm parameters
cache: Shared edge cache (should be pre-warmed with seed edges)
Returns:
List of PatternResult for each pattern
"""
import time
# Initialize all pattern states
states = [_init_pattern_state(seeds, pattern) for seeds, pattern in pattern_jobs]
# Determine max hops (all patterns should be same length, but be safe)
max_hops = max((len(p) for _, p in pattern_jobs), default=0)
# Detailed timing for debugging
hop_times: list[dict] = []
# Execute hop-by-hop across ALL patterns
for hop in range(max_hops):
hop_start = time.time()
hop_timing = {"hop": hop, "patterns_executed": 0, "uncached_count": 0, "load_time": 0.0}
# Execute this hop for all patterns, collect uncached nodes for next hop
all_uncached: set[str] = set()
exec_start = time.time()
for state in states:
if state.hop_index < len(state.pattern):
uncached = _execute_hop(state, cache, config)
all_uncached.update(uncached)
hop_timing["patterns_executed"] += 1
hop_timing["exec_time"] = time.time() - exec_start
# Pre-warm edges for ALL uncached nodes before next hop
hop_timing["uncached_count"] = len(all_uncached)
if all_uncached:
uncached_list = list(all_uncached - cache._fully_loaded)
hop_timing["uncached_after_filter"] = len(uncached_list)
if uncached_list:
load_start = time.time()
edges_by_type = await load_all_edges_for_frontier(pool, uncached_list, config.top_k_neighbors)
hop_timing["load_time"] = time.time() - load_start
cache.edge_load_time += hop_timing["load_time"]
cache.db_queries += 1
cache.add_all_edges(edges_by_type, uncached_list)
hop_timing["edges_loaded"] = sum(
len(neighbors) for edges in edges_by_type.values() for neighbors in edges.values()
)
hop_timing["total_time"] = time.time() - hop_start
hop_times.append(hop_timing)
# Store hop timing details in cache for logging
cache.hop_details = hop_times
# Finalize all patterns
return [_finalize_pattern(state, config) for state in states]
async def mpfp_traverse_async(
pool,
seeds: list[SeedNode],
pattern: list[str],
config: MPFPConfig,
cache: EdgeCache,
) -> PatternResult:
"""
Async Forward Push traversal with lazy edge loading.
NOTE: For better performance with multiple patterns, use mpfp_traverse_hop_synchronized().
This function is kept for single-pattern use cases.
"""
if not seeds:
return PatternResult(pattern=pattern, scores={})
results = await mpfp_traverse_hop_synchronized(pool, [(seeds, pattern)], config, cache)
return results[0] if results else PatternResult(pattern=pattern, scores={})
def rrf_fusion(
results: list[PatternResult],
k: int = 60,
top_k: int = 50,
) -> list[tuple[str, float]]:
"""
Reciprocal Rank Fusion to combine pattern results.
Args:
results: List of pattern results
k: RRF constant (higher = more uniform weighting)
top_k: Number of results to return
Returns:
List of (node_id, fused_score) tuples, sorted by score descending
"""
fused: dict[str, float] = {}
for result in results:
if not result.scores:
continue
# Rank nodes by their score in this pattern
ranked = sorted(result.scores.keys(), key=lambda n: result.scores[n], reverse=True)
for rank, node_id in enumerate(ranked):
fused[node_id] = fused.get(node_id, 0) + 1.0 / (k + rank + 1)
# Sort by fused score and return top-k
sorted_results = sorted(fused.items(), key=lambda x: x[1], reverse=True)
return sorted_results[:top_k]
# -----------------------------------------------------------------------------
# Database Loading
# -----------------------------------------------------------------------------
async def fetch_memory_units_by_ids(
pool,
node_ids: list[str],
fact_type: str,
) -> list[RetrievalResult]:
"""Fetch full memory unit details for a list of node IDs."""
if not node_ids:
return []
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, fact_type, document_id, chunk_id, tags, metadata
FROM {fq_table("memory_units")}
WHERE id = ANY($1::uuid[])
AND fact_type = $2
""",
node_ids,
fact_type,
)
return [RetrievalResult.from_db_row(dict(r)) for r in rows]
# -----------------------------------------------------------------------------
# Graph Retriever Implementation
# -----------------------------------------------------------------------------
class MPFPGraphRetriever(GraphRetriever):
"""
Graph retrieval using Meta-Path Forward Push with lazy edge loading.
Runs predefined patterns in parallel from semantic and temporal seeds,
loading edges on-demand per hop instead of loading entire graph upfront.
"""
def __init__(self, config: MPFPConfig | None = None):
"""
Initialize MPFP retriever.
Args:
config: Algorithm configuration (uses defaults if None)
"""
if config is None:
# Read top_k_neighbors from global config
from ...config import get_config
global_config = get_config()
config = MPFPConfig(top_k_neighbors=global_config.mpfp_top_k_neighbors)
self.config = config
@property
def name(self) -> str:
return "mpfp"
async def retrieve(
self,
pool,
query_embedding_str: str,
bank_id: str,
fact_type: str,
budget: int,
query_text: str | None = None,
semantic_seeds: list[RetrievalResult] | None = None,
temporal_seeds: list[RetrievalResult] | None = None,
adjacency=None, # Ignored - kept for interface compatibility
tags: list[str] | None = None,
tags_match: TagsMatch = "any",
tag_groups: list[TagGroup] | None = None,
) -> tuple[list[RetrievalResult], MPFPTimings | None]:
"""
Retrieve facts using MPFP algorithm with lazy edge loading.
Args:
pool: Database connection pool
query_embedding_str: Query embedding (used for fallback seed finding)
bank_id: Memory bank ID
fact_type: Fact type to filter
budget: Maximum results to return
query_text: Original query text (optional)
semantic_seeds: Pre-computed semantic entry points
temporal_seeds: Pre-computed temporal entry points
adjacency: Ignored (kept for interface compatibility)
tags: Optional list of tags for visibility filtering (OR matching)
Returns:
Tuple of (List of RetrievalResult with activation scores, MPFPTimings)
"""
import time
timings = MPFPTimings(fact_type=fact_type)
# Convert seeds to SeedNode format
semantic_seed_nodes = self._convert_seeds(semantic_seeds, "similarity")
temporal_seed_nodes = self._convert_seeds(temporal_seeds, "temporal_score")
# If no semantic seeds provided, fall back to finding our own
if not semantic_seed_nodes:
seeds_start = time.time()
semantic_seed_nodes = await self._find_semantic_seeds(
pool,
query_embedding_str,
bank_id,
fact_type,
tags=tags,
tags_match=tags_match,
tag_groups=tag_groups,
)
timings.seeds_time = time.time() - seeds_start
logger.debug(
f"[MPFP] Found {len(semantic_seed_nodes)} semantic seeds for fact_type={fact_type} (tags={tags}, tags_match={tags_match})"
)
# Collect all pattern jobs
pattern_jobs = []
# Patterns from semantic seeds
for pattern in self.config.patterns_semantic:
if semantic_seed_nodes:
pattern_jobs.append((semantic_seed_nodes, pattern))
# Patterns from temporal seeds
for pattern in self.config.patterns_temporal:
if temporal_seed_nodes:
pattern_jobs.append((temporal_seed_nodes, pattern))
if not pattern_jobs:
logger.debug(
f"[MPFP] No pattern jobs (semantic_seeds={len(semantic_seed_nodes)}, temporal_seeds={len(temporal_seed_nodes)})"
)
return [], timings
timings.pattern_count = len(pattern_jobs)
# Shared edge cache across all patterns
cache = EdgeCache()
# Pre-warm cache with ALL seed node edges BEFORE running patterns
# This prevents redundant DB queries at hop 1
all_seed_ids = list({s.node_id for seeds, _ in pattern_jobs for s in seeds})
if all_seed_ids:
import time as time_module
prewarm_start = time_module.time()
edges_by_type = await load_all_edges_for_frontier(pool, all_seed_ids, self.config.top_k_neighbors)
cache.edge_load_time += time_module.time() - prewarm_start
cache.db_queries += 1
cache.add_all_edges(edges_by_type, all_seed_ids)
# Run all patterns with HOP-SYNCHRONIZED edge loading
# This batches hop-2 edge loads across ALL patterns into ONE query
# Reduces DB queries from O(patterns * hops) to O(hops)
step_start = time.time()
pattern_results = await mpfp_traverse_hop_synchronized(pool, pattern_jobs, self.config, cache)
timings.traverse = time.time() - step_start
# Record edge loading stats from cache
timings.edge_count = sum(len(neighbors) for g in cache.graphs.values() for neighbors in g.values())
timings.db_queries = cache.db_queries
timings.edge_load_time = cache.edge_load_time
timings.hop_details = cache.hop_details
# Fuse results
step_start = time.time()
fused = rrf_fusion(pattern_results, top_k=budget)
timings.fusion = time.time() - step_start
if not fused:
logger.debug(f"[MPFP] No fused results after RRF fusion (pattern_count={len(pattern_results)})")
return [], timings
# Get top result IDs
result_ids = [node_id for node_id, score in fused][:budget]
# Fetch full details
step_start = time.time()
results = await fetch_memory_units_by_ids(pool, result_ids, fact_type)
timings.fetch = time.time() - step_start
# Filter results by tags (graph traversal may have picked up unfiltered memories)
if tags:
from .tags import filter_results_by_tags
results = filter_results_by_tags(results, tags, match=tags_match)
# Apply compound tag group filtering (post-traversal)
if tag_groups:
from .tags import filter_results_by_tag_groups
results = filter_results_by_tag_groups(results, tag_groups)
timings.result_count = len(results)
# Add activation scores from fusion
score_map = {node_id: score for node_id, score in fused}
for result in results:
result.activation = score_map.get(result.id, 0.0)
# Sort by activation
results.sort(key=lambda r: r.activation or 0, reverse=True)
return results, timings
def _convert_seeds(
self,
seeds: list[RetrievalResult] | None,
score_attr: str,
) -> list[SeedNode]:
"""Convert RetrievalResult seeds to SeedNode format."""
if not seeds:
return []
result = []
for seed in seeds:
score = getattr(seed, score_attr, None)
if score is None:
score = seed.activation or seed.similarity or 1.0
result.append(SeedNode(node_id=seed.id, score=score))
return result
async def _find_semantic_seeds(
self,
pool,
query_embedding_str: str,
bank_id: str,
fact_type: str,
limit: int = 20,
threshold: float = 0.3,
tags: list[str] | None = None,
tags_match: TagsMatch = "any",
tag_groups: list[TagGroup] | None = None,
) -> list[SeedNode]:
"""Fallback: find semantic seeds via embedding search."""
from .tags import build_tag_groups_where_clause, build_tags_where_clause_simple
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
tag_groups_param_start = 6 + (1 if tags else 0)
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
params = [query_embedding_str, bank_id, fact_type, threshold, limit]
if tags:
params.append(tags)
params.extend(groups_params)
async with acquire_with_retry(pool) as conn:
rows = await conn.fetch(
f"""
SELECT id, 1 - (embedding <=> $1::vector) AS similarity
FROM {fq_table("memory_units")}
WHERE bank_id = $2
AND embedding IS NOT NULL
AND fact_type = $3
AND (1 - (embedding <=> $1::vector)) >= $4
{tags_clause}
{groups_clause}
ORDER BY embedding <=> $1::vector
LIMIT $5
""",
*params,
)
return [SeedNode(node_id=str(r["id"]), score=r["similarity"]) for r in rows]
@@ -2,6 +2,7 @@
Cross-encoder neural reranking for search results.
"""
import math
from datetime import datetime, timezone
from .types import MergedCandidate, ScoredResult
@@ -13,6 +14,7 @@ UTC = timezone.utc
# so the max combined boost is (1 + alpha/2)^2 ≈ +21% and min is (1 - alpha/2)^2 ≈ -19%.
_RECENCY_ALPHA: float = 0.2
_TEMPORAL_ALPHA: float = 0.2
_PROOF_COUNT_ALPHA: float = 0.1 # Conservative: max ±5% for evidence strength
def apply_combined_scoring(
@@ -20,28 +22,40 @@ def apply_combined_scoring(
now: datetime,
recency_alpha: float = _RECENCY_ALPHA,
temporal_alpha: float = _TEMPORAL_ALPHA,
proof_count_alpha: float = _PROOF_COUNT_ALPHA,
) -> None:
"""Apply combined scoring to a list of ScoredResults in-place.
Uses the cross-encoder score as the primary relevance signal, with recency
and temporal proximity applied as multiplicative boosts. This ensures the
influence of these secondary signals is always proportional to the base
relevance score, regardless of the cross-encoder model's score calibration.
Uses the cross-encoder score as the primary relevance signal, with recency,
temporal proximity, and proof count applied as multiplicative boosts. This
ensures the influence of these secondary signals is always proportional to
the base relevance score, regardless of the cross-encoder model's score
calibration.
Formula::
recency_boost = 1 + recency_alpha * (recency - 0.5) # in [1-α/2, 1+α/2]
temporal_boost = 1 + temporal_alpha * (temporal - 0.5) # in [1-α/2, 1+α/2]
combined_score = cross_encoder_score_normalized * recency_boost * temporal_boost
recency_boost = 1 + recency_alpha * (recency - 0.5) # in [1-α/2, 1+α/2]
temporal_boost = 1 + temporal_alpha * (temporal - 0.5) # in [1-α/2, 1+α/2]
proof_count_boost = 1 + proof_count_alpha * (proof_norm - 0.5) # in [1-α/2, 1+α/2]
combined_score = CE_normalized * recency_boost * temporal_boost * proof_count_boost
proof_norm maps proof_count using a smooth logarithmic curve centered at 0.5,
clamped to [0, 1]:
proof_count=1 0.5 + 0 = 0.5 (neutral multiplier)
proof_count=150 clamped to 1.0 (max +5% boost)
Temporal proximity is treated as neutral (0.5) when not set by temporal retrieval,
so temporal_boost collapses to 1.0 for non-temporal queries.
Proof count is treated as neutral (0.5) when not available (non-observation facts),
so proof_count_boost collapses to 1.0 for world/experience/opinion facts.
Args:
scored_results: Results from the cross-encoder reranker. Mutated in place.
now: Current UTC datetime for recency calculation.
recency_alpha: Max relative recency adjustment (default 0.2 ±10%).
temporal_alpha: Max relative temporal adjustment (default 0.2 ±10%).
proof_count_alpha: Max relative proof count adjustment (default 0.1 ±5%).
"""
if now.tzinfo is None:
now = now.replace(tzinfo=UTC)
@@ -59,13 +73,23 @@ def apply_combined_scoring(
# Temporal proximity: meaningful only for temporal queries; neutral otherwise.
sr.temporal = sr.retrieval.temporal_proximity if sr.retrieval.temporal_proximity is not None else 0.5
# Proof count: log-normalized evidence strength; neutral for non-observations.
proof_count = sr.retrieval.proof_count
if proof_count is not None and proof_count >= 1:
# Clamp to [0, 1] so extreme counts stay within documented ±5% range
proof_norm = min(1.0, max(0.0, 0.5 + (math.log(proof_count) / 10.0)))
else:
# Neutral baseline is precisely 0.5, ensuring neutral multiplier (1.0)
proof_norm = 0.5
# RRF: kept at 0.0 for trace continuity but excluded from scoring.
# RRF is batch-relative (min-max normalised) and redundant after reranking.
sr.rrf_normalized = 0.0
recency_boost = 1.0 + recency_alpha * (sr.recency - 0.5)
temporal_boost = 1.0 + temporal_alpha * (sr.temporal - 0.5)
sr.combined_score = sr.cross_encoder_score_normalized * recency_boost * temporal_boost
proof_count_boost = 1.0 + proof_count_alpha * (proof_norm - 0.5)
sr.combined_score = sr.cross_encoder_score_normalized * recency_boost * temporal_boost * proof_count_boost
sr.weight = sr.combined_score
@@ -18,11 +18,10 @@ 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 .graph_retrieval import GraphRetriever
from .link_expansion_retrieval import LinkExpansionRetriever
from .mpfp_retrieval import MPFPGraphRetriever
from .tags import TagGroup, TagsMatch, build_tag_groups_where_clause, build_tags_where_clause_simple
from .types import MPFPTimings, RetrievalResult
from .types import GraphRetrievalTimings, RetrievalResult
logger = logging.getLogger(__name__)
@@ -46,7 +45,9 @@ class ParallelRetrievalResult:
temporal: list[RetrievalResult] | None
timings: dict[str, float] = field(default_factory=dict)
temporal_constraint: tuple | None = None # (start_date, end_date)
mpfp_timings: list[MPFPTimings] = field(default_factory=list) # MPFP sub-step timings per fact type
graph_timings: list[GraphRetrievalTimings] = field(
default_factory=list
) # Graph retrieval sub-step timings per fact type
max_conn_wait: float = 0.0 # Maximum connection acquisition wait time across all methods
@@ -72,15 +73,7 @@ def get_default_graph_retriever() -> GraphRetriever:
if _default_graph_retriever is None:
config = get_config()
retriever_type = config.graph_retriever.lower()
if retriever_type == "mpfp":
_default_graph_retriever = MPFPGraphRetriever()
logger.info(
f"Using MPFP graph retriever (top_k_neighbors={_default_graph_retriever.config.top_k_neighbors})"
)
elif retriever_type == "bfs":
_default_graph_retriever = BFSGraphRetriever()
logger.info("Using BFS graph retriever")
elif retriever_type == "link_expansion":
if retriever_type == "link_expansion":
_default_graph_retriever = LinkExpansionRetriever()
logger.info("Using LinkExpansion graph retriever")
else:
@@ -148,7 +141,7 @@ async def retrieve_semantic_bm25_combined(
cols = (
"id, text, context, event_date, occurred_start, occurred_end, mentioned_at, "
"fact_type, document_id, chunk_id, tags, metadata"
"fact_type, document_id, chunk_id, tags, metadata, proof_count"
)
table = fq_table("memory_units")
@@ -343,7 +336,7 @@ async def retrieve_temporal_combined(
{groups_clause}
),
sim_ranked AS (
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.proof_count, mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
1 - (mu.embedding <=> $1::vector) AS similarity,
ROW_NUMBER() OVER (PARTITION BY mu.fact_type ORDER BY mu.embedding <=> $1::vector) AS sim_rn
FROM date_ranked dr
@@ -351,7 +344,7 @@ async def retrieve_temporal_combined(
WHERE dr.rn <= 50
AND (1 - (mu.embedding <=> $1::vector)) >= $6
)
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, metadata, similarity
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, proof_count, document_id, chunk_id, tags, metadata, similarity
FROM sim_ranked
WHERE sim_rn <= 10
""",
@@ -627,9 +620,11 @@ async def retrieve_all_fact_types_parallel(
timings["temporal_combined"] = temporal_time
# Step 3: Run graph retrieval for each fact type in parallel
async def run_graph_for_fact_type(ft: str) -> tuple[str, list[RetrievalResult], float, MPFPTimings | None]:
async def run_graph_for_fact_type(
ft: str,
) -> tuple[str, list[RetrievalResult], float, GraphRetrievalTimings | None]:
graph_start = time.time()
results, mpfp_timing = await retriever.retrieve(
results, graph_timing = await retriever.retrieve(
pool=pool,
query_embedding_str=query_embedding_str,
bank_id=bank_id,
@@ -642,7 +637,7 @@ async def retrieve_all_fact_types_parallel(
tags_match=tags_match,
tag_groups=tag_groups,
)
return ft, results, time.time() - graph_start, mpfp_timing
return ft, results, time.time() - graph_start, graph_timing
# Run graph for all fact types in parallel
graph_tasks = [run_graph_for_fact_type(ft) for ft in fact_types]
@@ -651,7 +646,7 @@ async def retrieve_all_fact_types_parallel(
# Organize results by fact type
results_by_fact_type: dict[str, ParallelRetrievalResult] = {}
max_conn_wait = conn_wait # Single connection for semantic+bm25+temporal
all_mpfp_timings: list[MPFPTimings] = []
all_graph_timings: list[GraphRetrievalTimings] = []
for ft in fact_types:
# Get semantic + bm25 results for this fact type
@@ -660,14 +655,14 @@ async def retrieve_all_fact_types_parallel(
# Find graph results for this fact type
graph_results = []
graph_time = 0.0
mpfp_timing = None
graph_timing = None
for gr in graph_results_list:
if gr[0] == ft:
graph_results = gr[1]
graph_time = gr[2]
mpfp_timing = gr[3]
if mpfp_timing:
all_mpfp_timings.append(mpfp_timing)
graph_timing = gr[3]
if graph_timing:
all_graph_timings.append(graph_timing)
break
# Get temporal results for this fact type from combined result
@@ -688,7 +683,7 @@ async def retrieve_all_fact_types_parallel(
"temporal_extraction": temporal_extraction_time,
},
temporal_constraint=temporal_constraint,
mpfp_timings=[mpfp_timing] if mpfp_timing else [],
graph_timings=[graph_timing] if graph_timing else [],
max_conn_wait=max_conn_wait,
)
@@ -110,11 +110,7 @@ def build_think_prompt(
context: str | None = None,
entity_summaries_text: str | None = None,
) -> str:
"""Build the think prompt for the LLM.
Note: opinion_facts_text parameter removed - opinions are now stored as mental models
and included via entity_summaries_text.
"""
"""Build the think prompt for the LLM."""
disposition_desc = build_disposition_description(disposition)
name_section = f"""
@@ -131,7 +131,7 @@ class RetrievalResult(BaseModel):
text: str = Field(description="Memory unit text content")
context: str = Field(default="", description="Memory unit context")
event_date: datetime | None = Field(default=None, description="When the memory occurred")
fact_type: str | None = Field(default=None, description="Fact type (world, experience, opinion)")
fact_type: str | None = Field(default=None, description="Fact type (world, experience)")
score: float = Field(description="Score from this retrieval method")
score_name: str = Field(description="Name of the score (e.g., 'similarity', 'bm25_score', 'activation')")
@@ -140,9 +140,7 @@ class RetrievalMethodResults(BaseModel):
"""Results from a single retrieval method."""
method_name: Literal["semantic", "bm25", "graph", "temporal"] = Field(description="Name of retrieval method")
fact_type: str | None = Field(
default=None, description="Fact type this retrieval was for (world, experience, opinion)"
)
fact_type: str | None = Field(default=None, description="Fact type this retrieval was for (world, experience)")
results: list[RetrievalResult] = Field(description="Retrieved results with ranks")
duration_seconds: float = Field(description="Time taken for this retrieval")
metadata: dict[str, Any] = Field(default_factory=dict, description="Method-specific metadata")
@@ -319,7 +319,7 @@ class SearchTracer:
duration_seconds: Time taken for this retrieval
score_field: Field name containing the score in data dict
metadata: Optional metadata about this retrieval method
fact_type: Fact type this retrieval was for (world, experience, opinion)
fact_type: Fact type this retrieval was for (world, experience)
"""
retrieval_results = []
for rank, (doc_id, data) in enumerate(results, start=1):
@@ -11,8 +11,8 @@ from typing import Any
@dataclass
class MPFPTimings:
"""Timing breakdown for a single MPFP retrieval call."""
class GraphRetrievalTimings:
"""Timing breakdown for a single graph retrieval call."""
fact_type: str
edge_count: int = 0 # Total edges loaded
@@ -48,6 +48,7 @@ class RetrievalResult:
chunk_id: str | None = None
tags: list[str] | None = None # Visibility scope tags
metadata: dict[str, str] | None = None # User-provided metadata
proof_count: int | None = None # Number of supporting memories (observations only)
# Retrieval-specific scores (only one will be set depending on retrieval method)
similarity: float | None = None # Semantic retrieval
@@ -72,6 +73,7 @@ class RetrievalResult:
chunk_id=row.get("chunk_id"),
tags=row.get("tags"),
metadata=row.get("metadata"),
proof_count=row.get("proof_count"),
similarity=row.get("similarity"),
bm25_score=row.get("bm25_score"),
activation=row.get("activation"),
@@ -82,20 +82,16 @@ class TaskBackend(ABC):
Args:
task_dict: Task dictionary to execute
Raises:
Exception: Re-raised from executor on failure.
"""
if self._executor is None:
task_type = task_dict.get("type", "unknown")
logger.warning(f"No executor registered, skipping task {task_type}")
return
try:
await self._executor(task_dict)
except Exception as e:
task_type = task_dict.get("type", "unknown")
logger.error(f"Error executing task {task_type}: {e}")
import traceback
traceback.print_exc()
await self._executor(task_dict)
class SyncTaskBackend(TaskBackend):
@@ -120,7 +120,9 @@ class DefaultExtensionContext(ExtensionContext):
# CREATE INDEX CONCURRENTLY inside the migration waits for those transactions
# forever — a deadlock.
config = get_config()
await asyncio.to_thread(run_migrations, db_url, schema=schema)
await asyncio.to_thread(
run_migrations, db_url, schema=schema, migration_database_url=config.migration_database_url
)
# Ensure embedding column dimension matches the model's dimension
# This is needed because migrations create columns with default dimension
@@ -96,7 +96,6 @@ class RetainContext:
request_context: "RequestContext"
document_id: str | None = None
fact_type_override: str | None = None
confidence_score: float | None = None
@dataclass
@@ -169,7 +168,6 @@ class RetainResult:
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
@@ -402,7 +400,6 @@ class OperationValidatorExtension(Extension, ABC):
- 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.
@@ -722,3 +719,28 @@ class OperationValidatorExtension(Extension, ABC):
BankListResult with the filtered list of banks.
"""
return BankListResult(banks=ctx.banks)
async def filter_mcp_tools(
self,
bank_id: str,
request_context: "RequestContext",
tools: frozenset[str],
) -> frozenset[str]:
"""
Filter MCP tools visible to this user on this bank.
Called during tools/list after bank-level mcp_enabled_tools filtering.
The input set is already narrowed by bank config this method can only
remove tools, never add ones the bank config excluded.
Default: return all tools unchanged (no per-user filtering).
Args:
bank_id: Target bank ID (from URL path or header).
request_context: Authenticated context with tenant_id set.
tools: Tools remaining after bank config filtering.
Returns:
Subset of tools this user should see.
"""
return tools
+185 -9
View File
@@ -8,7 +8,7 @@ This module provides the core tool logic used by both:
import json
import logging
from dataclasses import dataclass
from datetime import datetime
from datetime import datetime, timezone
from typing import Any, Callable
from fastmcp import FastMCP
@@ -18,11 +18,48 @@ from hindsight_api.config import (
DEFAULT_MCP_RECALL_DESCRIPTION,
DEFAULT_MCP_RETAIN_DESCRIPTION,
)
from hindsight_api.engine.audit import AuditEntry, AuditLogger
from hindsight_api.engine.memory_engine import Budget
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
from hindsight_api.extensions import OperationValidationError
from hindsight_api.models import RequestContext
# All tools available in the system (explicit list — no wildcards).
# Defined here (shared module) to avoid circular imports with api/mcp.py.
_ALL_TOOLS: frozenset[str] = frozenset(
{
"retain",
"recall",
"reflect",
"list_banks",
"create_bank",
"list_mental_models",
"get_mental_model",
"create_mental_model",
"update_mental_model",
"delete_mental_model",
"refresh_mental_model",
"list_directives",
"create_directive",
"delete_directive",
"list_memories",
"get_memory",
"delete_memory",
"list_documents",
"get_document",
"delete_document",
"list_operations",
"get_operation",
"cancel_operation",
"list_tags",
"get_bank",
"get_bank_stats",
"update_bank",
"delete_bank",
"clear_memories",
}
)
logger = logging.getLogger(__name__)
@@ -289,6 +326,7 @@ def register_mcp_tools(
_register_clear_memories(mcp, memory, config)
_apply_bank_tool_filtering(mcp, memory, config)
_apply_audit_logging(mcp, memory, config)
def _apply_bank_tool_filtering(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
@@ -303,11 +341,29 @@ def _apply_bank_tool_filtering(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
if not bank_id:
return None
request_context = _get_request_context(config)
# Layer 1: bank config filter (existing)
bank_cfg = await memory._config_resolver.get_bank_config(bank_id, request_context)
enabled: list[str] | None = bank_cfg.get("mcp_enabled_tools")
if enabled is None:
return None
return set(enabled)
bank_tools: list[str] | None = bank_cfg.get("mcp_enabled_tools")
enabled: set[str] | None = set(bank_tools) if bank_tools is not None else None
# Layer 2: operation validator filter
validator = memory._operation_validator
if validator is not None:
candidate = frozenset(enabled) if enabled is not None else _ALL_TOOLS
try:
filtered = await validator.filter_mcp_tools(bank_id, request_context, candidate)
except Exception:
logger.warning("filter_mcp_tools raised, returning unfiltered tools", exc_info=True)
return enabled
if filtered != candidate:
# Validator can only narrow, never expand beyond the bank config ceiling.
if bank_tools is not None:
enabled = set(filtered) & set(bank_tools)
else:
enabled = set(filtered)
return enabled
if hasattr(mcp, "list_tools"):
# FastMCP 3.x: wrap list_tools() and get_tool() on the instance
@@ -361,6 +417,112 @@ def _apply_bank_tool_filtering(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
logger.warning("Could not apply bank tool filtering: unknown FastMCP version")
_AUDITABLE_MCP_TOOLS: frozenset[str] = frozenset(
{
"retain",
"recall",
"reflect",
"create_bank",
"update_bank",
"delete_bank",
"clear_memories",
"create_mental_model",
"update_mental_model",
"delete_mental_model",
"refresh_mental_model",
"create_directive",
"delete_directive",
"delete_memory",
"delete_document",
"cancel_operation",
}
)
def _apply_audit_logging(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
"""Wrap auditable MCP tool run methods with audit logging."""
audit_logger: AuditLogger = memory.audit_logger
def _wrap_tool_run(tool_name: str, original_run):
"""Create an audited wrapper for a tool's run method."""
async def _audited_run(arguments, _name=tool_name, _orig=original_run):
if not audit_logger.is_enabled(_name):
return await _orig(arguments)
bank_id = None
if isinstance(arguments, dict):
bank_id = arguments.get("bank_id") or (config.bank_id_resolver() if config.bank_id_resolver else None)
elif hasattr(arguments, "get"):
bank_id = arguments.get("bank_id")
entry = AuditEntry(
action=_name,
transport="mcp",
bank_id=bank_id,
started_at=datetime.now(timezone.utc),
request=dict(arguments) if isinstance(arguments, dict) else {"raw": str(arguments)},
)
try:
result = await _orig(arguments)
if isinstance(result, dict):
entry.response = result
elif isinstance(result, list):
entry.response = {"items": result}
elif isinstance(result, str):
entry.response = {"text": result}
return result
finally:
entry.ended_at = datetime.now(timezone.utc)
audit_logger.log_fire_and_forget(entry)
return _audited_run
if hasattr(mcp, "_tool_manager"):
# FastMCP 2.x
try:
for name, tool in mcp._tool_manager._tools.items(): # type: ignore[unresolved-attribute] # FastMCP 2.x internal; guarded by hasattr
if name in _AUDITABLE_MCP_TOOLS:
object.__setattr__(tool, "run", _wrap_tool_run(name, tool.run))
except (AttributeError, KeyError) as e:
logger.warning(f"Could not apply MCP audit logging (v2): {e}")
elif hasattr(mcp, "get_tool"):
# FastMCP 3.x: wrap call_tool
original_call_tool = getattr(mcp, "call_tool", None)
if original_call_tool:
async def _audited_call_tool(name, arguments=None, **kwargs):
if name not in _AUDITABLE_MCP_TOOLS or not audit_logger.is_enabled(name):
return await original_call_tool(name, arguments, **kwargs)
bank_id = None
if isinstance(arguments, dict):
bank_id = arguments.get("bank_id") or (
config.bank_id_resolver() if config.bank_id_resolver else None
)
entry = AuditEntry(
action=name,
transport="mcp",
bank_id=bank_id,
started_at=datetime.now(timezone.utc),
request=dict(arguments) if isinstance(arguments, dict) else {},
)
try:
result = await original_call_tool(name, arguments, **kwargs)
entry.response = {"result": str(result)[:4096]}
return result
finally:
entry.ended_at = datetime.now(timezone.utc)
audit_logger.log_fire_and_forget(entry)
object.__setattr__(mcp, "call_tool", _audited_call_tool)
else:
logger.warning("Could not apply MCP audit logging: unknown FastMCP version")
def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
"""Register the retain tool."""
description = config.retain_description or DEFAULT_MCP_RETAIN_DESCRIPTION
@@ -840,6 +1002,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
@mcp.tool()
async def list_mental_models(
tags: list[str] | None = None,
detail: str = "full",
bank_id: str | None = None,
) -> str:
"""
@@ -851,6 +1014,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
Args:
tags: Optional tags to filter by (returns models matching any tag)
detail: Detail level - 'metadata' (names/tags only), 'content' (adds content/config), 'full' (includes reflect_response). Default: 'full'
bank_id: Optional bank to list from (defaults to session bank). Use for cross-bank operations.
"""
try:
@@ -861,6 +1025,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
models = await memory.list_mental_models(
bank_id=target_bank,
tags=tags,
detail=detail,
request_context=_get_request_context(config),
)
return json.dumps({"items": models}, indent=2, default=str)
@@ -876,6 +1041,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
@mcp.tool()
async def list_mental_models(
tags: list[str] | None = None,
detail: str = "full",
) -> dict:
"""
List mental models (pinned reflections) for this memory bank.
@@ -886,6 +1052,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
Args:
tags: Optional tags to filter by (returns models matching any tag)
detail: Detail level - 'metadata' (names/tags only), 'content' (adds content/config), 'full' (includes reflect_response). Default: 'full'
"""
try:
target_bank = config.bank_id_resolver()
@@ -895,6 +1062,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
models = await memory.list_mental_models(
bank_id=target_bank,
tags=tags,
detail=detail,
request_context=_get_request_context(config),
)
return {"items": models}
@@ -914,16 +1082,18 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
@mcp.tool()
async def get_mental_model(
mental_model_id: str,
detail: str = "full",
bank_id: str | None = None,
) -> str:
"""
Get a specific mental model by ID.
Returns the full mental model including its generated content, source query,
and metadata. Use list_mental_models first to discover available model IDs.
Returns the mental model with the requested detail level. Use list_mental_models
first to discover available model IDs.
Args:
mental_model_id: The ID of the mental model to retrieve
detail: Detail level - 'metadata' (names/tags only), 'content' (adds content/config), 'full' (includes reflect_response). Default: 'full'
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
"""
try:
@@ -934,6 +1104,7 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
model = await memory.get_mental_model(
bank_id=target_bank,
mental_model_id=mental_model_id,
detail=detail,
request_context=_get_request_context(config),
)
if model is None:
@@ -951,15 +1122,17 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
@mcp.tool()
async def get_mental_model(
mental_model_id: str,
detail: str = "full",
) -> dict:
"""
Get a specific mental model by ID.
Returns the full mental model including its generated content, source query,
and metadata. Use list_mental_models first to discover available model IDs.
Returns the mental model with the requested detail level. Use list_mental_models
first to discover available model IDs.
Args:
mental_model_id: The ID of the mental model to retrieve
detail: Detail level - 'metadata' (names/tags only), 'content' (adds content/config), 'full' (includes reflect_response). Default: 'full'
"""
try:
target_bank = config.bank_id_resolver()
@@ -969,6 +1142,7 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
model = await memory.get_mental_model(
bank_id=target_bank,
mental_model_id=mental_model_id,
detail=detail,
request_context=_get_request_context(config),
)
if model is None:
@@ -2711,6 +2885,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
result = await memory.delete_bank(
target_bank,
fact_type=type,
delete_bank_profile=False,
request_context=_get_request_context(config),
)
return json.dumps({"status": "cleared", "bank_id": target_bank, **result}, default=str)
@@ -2743,6 +2918,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
result = await memory.delete_bank(
target_bank,
fact_type=type,
delete_bank_profile=False,
request_context=_get_request_context(config),
)
return {"status": "cleared", "bank_id": target_bank, **result}
+11 -10
View File
@@ -11,14 +11,11 @@ This module provides metrics for:
- Database connection pool metrics
"""
import importlib
import logging
import os
import types
try:
import resource
except ImportError:
resource: types.ModuleType | None = None # Windows doesn't have resource module
_resource_mod = importlib.import_module("resource") if importlib.util.find_spec("resource") else None
import threading
import time
from contextlib import contextmanager
@@ -255,6 +252,9 @@ class MetricsCollector(MetricsCollectorBase):
def __init__(self):
self.meter = get_meter()
from .config import get_config
self._include_bank_id = get_config().metrics_include_bank_id
# Operation latency histogram (in seconds)
# Records duration of retain, recall, reflect operations
@@ -335,10 +335,11 @@ class MetricsCollector(MetricsCollectorBase):
start_time = time.time()
attributes = {
"operation": operation,
"bank_id": bank_id,
"source": source,
"tenant": _get_tenant(),
}
if self._include_bank_id:
attributes["bank_id"] = bank_id
if budget:
attributes["budget"] = budget
if max_tokens:
@@ -460,13 +461,13 @@ class MetricsCollector(MetricsCollectorBase):
def _setup_process_metrics(self):
"""Set up observable gauges for process metrics."""
if resource is None:
if _resource_mod is None:
return # Skip process metrics on Windows
def get_cpu_times(_options):
"""Get process CPU times."""
try:
rusage = resource.getrusage(resource.RUSAGE_SELF)
rusage = _resource_mod.getrusage(_resource_mod.RUSAGE_SELF)
yield metrics.Observation(rusage.ru_utime, {"type": "user"})
yield metrics.Observation(rusage.ru_stime, {"type": "system"})
except Exception:
@@ -475,7 +476,7 @@ class MetricsCollector(MetricsCollectorBase):
def get_memory_usage(_options):
"""Get process memory usage in bytes."""
try:
rusage = resource.getrusage(resource.RUSAGE_SELF)
rusage = _resource_mod.getrusage(_resource_mod.RUSAGE_SELF)
# ru_maxrss is in kilobytes on Linux, bytes on macOS
max_rss = rusage.ru_maxrss
if os.uname().sysname == "Linux":
@@ -493,7 +494,7 @@ class MetricsCollector(MetricsCollectorBase):
yield metrics.Observation(count)
else:
# Fallback: use resource limits
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
soft, hard = _resource_mod.getrlimit(_resource_mod.RLIMIT_NOFILE)
yield metrics.Observation(soft, {"limit": "soft"})
except Exception:
pass
+12 -3
View File
@@ -157,7 +157,7 @@ def _run_migrations_internal(database_url: str, script_location: str, schema: st
# calls from different threads corrupt each other's context.
try:
with _alembic_lock:
command.upgrade(alembic_cfg, "head")
command.upgrade(alembic_cfg, "heads")
except ResolutionError as e:
# This happens during rolling deployments when a newer version of the code
# has already run migrations, and this older replica doesn't have the new
@@ -176,6 +176,7 @@ def run_migrations(
database_url: str,
script_location: str | None = None,
schema: str | None = None,
migration_database_url: str | None = None,
) -> None:
"""
Run database migrations to the latest version using programmatic Alembic configuration.
@@ -213,6 +214,14 @@ def run_migrations(
script_location="/path/to/copied/_alembic"
)
"""
# Prefer a dedicated migration URL that bypasses connection poolers (e.g.
# PgBouncer in transaction mode). Session-level advisory locks don't
# survive a PgBouncer transaction-mode cycle, so the distributed lock is
# ineffective when the app URL goes through a pooler. Configure
# HINDSIGHT_API_MIGRATION_DATABASE_URL to the direct PostgreSQL endpoint
# (e.g. hindsight-pg-rw) to restore correct locking behaviour.
migration_url = migration_database_url or database_url
try:
# Determine script location
if script_location is None:
@@ -249,7 +258,7 @@ def run_migrations(
# 2. After acquiring the lock, COMMIT the transaction on the advisory-lock
# connection itself before running migrations. pg_advisory_lock is
# session-level, so the lock survives the COMMIT.
engine = create_engine(database_url)
engine = create_engine(migration_url)
with engine.connect() as conn:
logger.debug(f"Acquiring migration advisory lock for schema '{schema_name}' (id={lock_id})...")
while True:
@@ -394,7 +403,7 @@ def run_migrations(
conn.commit()
# Run migrations while holding the lock
_run_migrations_internal(database_url, script_location, schema=schema)
_run_migrations_internal(migration_url, script_location, schema=schema)
finally:
# Explicitly release the lock (also released on connection close)
conn.execute(text(f"SELECT pg_advisory_unlock({lock_id})"))
+1 -23
View File
@@ -97,7 +97,6 @@ class MemoryUnit(Base):
occurred_end: Mapped[datetime | None] = mapped_column(TIMESTAMP(timezone=True)) # When fact occurred (range end)
mentioned_at: Mapped[datetime | None] = mapped_column(TIMESTAMP(timezone=True)) # When fact was mentioned
fact_type: Mapped[str] = mapped_column(Text, nullable=False, server_default="world")
confidence_score: Mapped[float | None] = mapped_column(Float)
unit_metadata: Mapped[dict] = mapped_column(
"metadata", JSONB, server_default=sql_text("'{}'::jsonb")
) # User-defined metadata (str->str)
@@ -121,14 +120,7 @@ class MemoryUnit(Base):
name="memory_units_document_fkey",
ondelete="CASCADE",
),
CheckConstraint("fact_type IN ('world', 'experience', 'opinion', 'observation')"),
CheckConstraint("confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0)"),
CheckConstraint(
"(fact_type = 'opinion' AND confidence_score IS NOT NULL) OR "
"(fact_type = 'observation') OR "
"(fact_type NOT IN ('opinion', 'observation') AND confidence_score IS NULL)",
name="confidence_score_fact_type_check",
),
CheckConstraint("fact_type IN ('world', 'experience', 'observation')"),
Index("idx_memory_units_bank_id", "bank_id"),
Index("idx_memory_units_document_id", "document_id"),
Index("idx_memory_units_event_date", "event_date", postgresql_ops={"event_date": "DESC"}),
@@ -142,20 +134,6 @@ class MemoryUnit(Base):
"event_date",
postgresql_ops={"event_date": "DESC"},
),
Index(
"idx_memory_units_opinion_confidence",
"bank_id",
"confidence_score",
postgresql_where=sql_text("fact_type = 'opinion'"),
postgresql_ops={"confidence_score": "DESC"},
),
Index(
"idx_memory_units_opinion_date",
"bank_id",
"event_date",
postgresql_where=sql_text("fact_type = 'opinion'"),
postgresql_ops={"event_date": "DESC"},
),
Index(
"idx_memory_units_observation_date",
"bank_id",
+7 -4
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "hindsight-api-slim"
version = "0.4.20"
version = "0.4.22"
description = "Hindsight: Agent Memory That Works Like Human Memory"
readme = "README.md"
requires-python = ">=3.11"
@@ -26,7 +26,7 @@ dependencies = [
"tiktoken>=0.12.0",
"httpx>=0.27.0",
"PyJWT[crypto]>=2.8.0",
"fastmcp>=2.14.0", # CVE-2025-66416
"fastmcp>=3.2.0", # SSRF/path traversal, OAuth confused deputy, command injection fixes
"python-dateutil>=2.8.0",
"opentelemetry-api>=1.20.0",
"opentelemetry-sdk>=1.20.0",
@@ -48,17 +48,19 @@ dependencies = [
# Transitive dependency security fixes
"pyasn1>=0.6.3", # DoS vulnerability fix
"urllib3>=2.6.3", # Decompression-bomb safeguards bypass fix
"langchain-core>=1.2.11", # Serialization injection + SSRF vulnerability fix
"langchain-core>=1.2.22", # Path traversal in legacy load_prompt functions fix
"langsmith>=0.6.3", # SSRF via tracing header injection fix
"protobuf>=6.33.5", # JSON recursion depth bypass fix
"pillow>=12.1.1", # Out-of-bounds write in PSD image loading fix
"cryptography>=46.0.5", # Subgroup attack vulnerability fix
"cryptography>=46.0.6", # Incomplete DNS name constraint enforcement fix
"filelock>=3.20.1", # TOCTOU race condition fix
"authlib>=1.6.9", # Account takeover/JWS header injection vulnerability fix
"pyjwt>=2.12.0", # Accepts unknown crit header extensions fix
"orjson>=3.11.6", # Unbounded recursion DoS fix
"python-multipart>=0.0.22", # Arbitrary file write via non-default configuration fix
"tornado>=6.5.5", # DoS multipart/incomplete cookie validation fix
"aiohttp>=3.13.3", # Multiple DoS vulnerabilities
"pygments>=2.20.0", # ReDoS via inefficient GUID regex fix
"claude-agent-sdk>=0.1.27",
"boto3>=1.42.74",
]
@@ -134,6 +136,7 @@ dev = [
"pytest-asyncio>=1.3.0",
"pytest-timeout>=2.4.0",
"pytest-xdist>=3.8.0",
"pytest-rerunfailures>=15.0",
"python-dotenv>=1.2.1",
"filelock>=3.20.1", # TOCTOU race condition fix
"ruff>=0.8.0",
+2 -2
View File
@@ -17,7 +17,7 @@ from hindsight_api.pg0 import EmbeddedPostgres
# Default pg0 instance configuration for tests
DEFAULT_PG0_INSTANCE_NAME = "hindsight-test"
DEFAULT_PG0_PORT = 5556
DEFAULT_PG0_PORT = int(os.environ.get("HINDSIGHT_TEST_PG_PORT", "5556"))
# Load environment variables from .env at the start of test session
@@ -126,7 +126,7 @@ def llm_config():
Provide LLM configuration for tests.
This can be used by tests that need to call LLM directly without memory system.
"""
return LLMConfig.for_memory()
return LLMConfig.from_env()
@pytest.fixture(scope="session")
@@ -314,7 +314,7 @@ async def test_run_migration_without_schema_discovers_and_deduplicates_schemas(m
async def fake_resolve_database_url(db_url: str) -> str:
return f"resolved::{db_url}"
def fake_run_migrations(database_url: str, schema: str | None = None) -> None:
def fake_run_migrations(database_url: str, schema: str | None = None, **kwargs) -> None:
calls["run_migrations"].append((database_url, schema))
def fake_ensure_vector_extension(
@@ -376,7 +376,7 @@ async def test_run_migration_without_schema_runs_optional_post_migration_hooks(m
async def fake_resolve_database_url(db_url: str) -> str:
return f"resolved::{db_url}"
def fake_run_migrations(database_url: str, schema: str | None = None) -> None:
def fake_run_migrations(database_url: str, schema: str | None = None, **kwargs) -> None:
calls["run_migrations"].append((database_url, schema))
def fake_ensure_embedding_dimension(
@@ -453,7 +453,7 @@ async def test_run_migration_with_schema_only_runs_requested_schema(monkeypatch)
async def fake_resolve_database_url(db_url: str) -> str:
return f"resolved::{db_url}"
def fake_run_migrations(database_url: str, schema: str | None = None) -> None:
def fake_run_migrations(database_url: str, schema: str | None = None, **kwargs) -> None:
calls["run_migrations"].append((database_url, schema))
def fake_ensure_vector_extension(
+449
View File
@@ -0,0 +1,449 @@
"""
Tests for the audit log feature.
Tests the audit log list, stats, filtering, and pagination endpoints.
Verifies that audit entries are created for operations when audit logging is enabled.
"""
import asyncio
import httpx
import pytest
import pytest_asyncio
from hindsight_api.api import create_app
from hindsight_api.config import get_config
@pytest_asyncio.fixture
async def audit_api_client(memory):
"""Create a test client with audit logging enabled."""
# Enable audit logging on the memory engine's audit logger
memory._audit_logger._enabled = True
memory._audit_logger._allowed_actions = None # All actions
app = create_app(memory, initialize_memory=False)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
yield client
@pytest.fixture
def bank_id():
"""Provide a unique bank ID for audit tests."""
from datetime import datetime
return f"audit_test_{datetime.now().timestamp()}"
@pytest.mark.asyncio
async def test_audit_log_list_empty(audit_api_client, bank_id):
"""Test listing audit logs for a bank with no entries returns empty."""
# Create the bank first
await audit_api_client.put(
f"/v1/default/banks/{bank_id}",
json={"name": "Audit Test Bank"},
)
# Small delay for fire-and-forget audit writes
await asyncio.sleep(0.5)
response = await audit_api_client.get(f"/v1/default/banks/{bank_id}/audit-logs")
assert response.status_code == 200
data = response.json()
assert data["bank_id"] == bank_id
assert "total" in data
assert "items" in data
assert "limit" in data
assert "offset" in data
assert isinstance(data["items"], list)
@pytest.mark.asyncio
async def test_audit_log_created_for_retain(audit_api_client, bank_id):
"""Test that a retain operation creates an audit log entry."""
# Create bank
await audit_api_client.put(
f"/v1/default/banks/{bank_id}",
json={"name": "Audit Test Bank"},
)
# Perform a retain
response = await audit_api_client.post(
f"/v1/default/banks/{bank_id}/memories",
json={
"items": [{"content": "Alice likes cats", "context": "preferences"}],
},
)
assert response.status_code == 200
# Wait for fire-and-forget audit writes
await asyncio.sleep(1.0)
# List audit logs - should have entries for create_bank and retain
response = await audit_api_client.get(f"/v1/default/banks/{bank_id}/audit-logs")
assert response.status_code == 200
data = response.json()
assert data["total"] >= 1
actions = [item["action"] for item in data["items"]]
assert "retain" in actions, f"Expected 'retain' in audit actions, got: {actions}"
@pytest.mark.asyncio
async def test_audit_log_entry_fields(audit_api_client, bank_id):
"""Test that audit log entries have all expected fields."""
# Create bank + recall to generate entries
await audit_api_client.put(
f"/v1/default/banks/{bank_id}",
json={"name": "Audit Test Bank"},
)
await audit_api_client.post(
f"/v1/default/banks/{bank_id}/memories/recall",
json={"query": "test query"},
)
await asyncio.sleep(1.0)
response = await audit_api_client.get(f"/v1/default/banks/{bank_id}/audit-logs")
assert response.status_code == 200
data = response.json()
assert data["total"] >= 1
# Check the recall entry has all fields
recall_entries = [item for item in data["items"] if item["action"] == "recall"]
assert len(recall_entries) >= 1, f"Expected recall entry, got actions: {[i['action'] for i in data['items']]}"
entry = recall_entries[0]
assert entry["id"] is not None
assert entry["action"] == "recall"
assert entry["transport"] == "http"
assert entry["bank_id"] == bank_id
assert entry["started_at"] is not None
assert entry["ended_at"] is not None
# Request should contain the recall parameters
assert entry["request"] is not None
assert "query" in entry["request"]
# Response should contain the recall results
assert entry["response"] is not None
@pytest.mark.asyncio
async def test_audit_log_filter_by_action(audit_api_client, bank_id):
"""Test filtering audit logs by action type."""
# Create bank and do retain + recall
await audit_api_client.put(
f"/v1/default/banks/{bank_id}",
json={"name": "Audit Test Bank"},
)
await audit_api_client.post(
f"/v1/default/banks/{bank_id}/memories",
json={"items": [{"content": "test content", "context": "test"}]},
)
await audit_api_client.post(
f"/v1/default/banks/{bank_id}/memories/recall",
json={"query": "test"},
)
await asyncio.sleep(1.0)
# Filter by retain only
response = await audit_api_client.get(
f"/v1/default/banks/{bank_id}/audit-logs",
params={"action": "retain"},
)
assert response.status_code == 200
data = response.json()
for item in data["items"]:
assert item["action"] == "retain"
# Filter by recall only
response = await audit_api_client.get(
f"/v1/default/banks/{bank_id}/audit-logs",
params={"action": "recall"},
)
assert response.status_code == 200
data = response.json()
for item in data["items"]:
assert item["action"] == "recall"
@pytest.mark.asyncio
async def test_audit_log_filter_by_transport(audit_api_client, bank_id):
"""Test filtering audit logs by transport type."""
await audit_api_client.put(
f"/v1/default/banks/{bank_id}",
json={"name": "Audit Test Bank"},
)
await asyncio.sleep(0.5)
# Filter by http transport
response = await audit_api_client.get(
f"/v1/default/banks/{bank_id}/audit-logs",
params={"transport": "http"},
)
assert response.status_code == 200
data = response.json()
for item in data["items"]:
assert item["transport"] == "http"
# Filter by mcp transport - should be empty (no MCP calls in this test)
response = await audit_api_client.get(
f"/v1/default/banks/{bank_id}/audit-logs",
params={"transport": "mcp"},
)
assert response.status_code == 200
data = response.json()
assert data["total"] == 0
@pytest.mark.asyncio
async def test_audit_log_filter_by_date_range(audit_api_client, bank_id):
"""Test filtering audit logs by date range."""
from datetime import datetime, timedelta, timezone
await audit_api_client.put(
f"/v1/default/banks/{bank_id}",
json={"name": "Audit Test Bank"},
)
await asyncio.sleep(0.5)
now = datetime.now(timezone.utc)
# Filter with start_date in the past - should include entries
past = (now - timedelta(hours=1)).isoformat()
response = await audit_api_client.get(
f"/v1/default/banks/{bank_id}/audit-logs",
params={"start_date": past},
)
assert response.status_code == 200
data = response.json()
assert data["total"] >= 1
# Filter with start_date in the future - should be empty
future = (now + timedelta(hours=1)).isoformat()
response = await audit_api_client.get(
f"/v1/default/banks/{bank_id}/audit-logs",
params={"start_date": future},
)
assert response.status_code == 200
data = response.json()
assert data["total"] == 0
@pytest.mark.asyncio
async def test_audit_log_pagination(audit_api_client, bank_id):
"""Test audit log pagination with limit and offset."""
await audit_api_client.put(
f"/v1/default/banks/{bank_id}",
json={"name": "Audit Test Bank"},
)
# Generate multiple audit entries
for i in range(5):
await audit_api_client.post(
f"/v1/default/banks/{bank_id}/memories/recall",
json={"query": f"test query {i}"},
)
await asyncio.sleep(1.5)
# Get first page
response = await audit_api_client.get(
f"/v1/default/banks/{bank_id}/audit-logs",
params={"limit": 2, "offset": 0},
)
assert response.status_code == 200
page1 = response.json()
assert len(page1["items"]) == 2
assert page1["limit"] == 2
assert page1["offset"] == 0
assert page1["total"] >= 5 # At least 5 recall + 1 create_bank
# Get second page
response = await audit_api_client.get(
f"/v1/default/banks/{bank_id}/audit-logs",
params={"limit": 2, "offset": 2},
)
assert response.status_code == 200
page2 = response.json()
assert len(page2["items"]) == 2
assert page2["offset"] == 2
# Entries should be different between pages
page1_ids = {item["id"] for item in page1["items"]}
page2_ids = {item["id"] for item in page2["items"]}
assert page1_ids.isdisjoint(page2_ids), "Pages should not overlap"
@pytest.mark.asyncio
async def test_audit_log_stats(audit_api_client, bank_id):
"""Test the audit log stats endpoint returns correct structure."""
await audit_api_client.put(
f"/v1/default/banks/{bank_id}",
json={"name": "Audit Test Bank"},
)
await audit_api_client.post(
f"/v1/default/banks/{bank_id}/memories/recall",
json={"query": "stats test"},
)
await asyncio.sleep(1.0)
# Get stats for last 24h
response = await audit_api_client.get(
f"/v1/default/banks/{bank_id}/audit-logs/stats",
params={"period": "1d"},
)
assert response.status_code == 200
data = response.json()
assert data["bank_id"] == bank_id
assert data["period"] == "1d"
assert data["trunc"] == "day"
assert "buckets" in data
assert isinstance(data["buckets"], list)
# Should have at least one bucket with our operations
assert len(data["buckets"]) >= 1
bucket = data["buckets"][0]
assert "time" in bucket
assert "actions" in bucket
assert "total" in bucket
assert bucket["total"] >= 1
@pytest.mark.asyncio
async def test_audit_log_stats_filter_by_action(audit_api_client, bank_id):
"""Test stats endpoint filters by action."""
await audit_api_client.put(
f"/v1/default/banks/{bank_id}",
json={"name": "Audit Test Bank"},
)
await audit_api_client.post(
f"/v1/default/banks/{bank_id}/memories/recall",
json={"query": "test"},
)
await asyncio.sleep(1.0)
# Stats filtered by recall
response = await audit_api_client.get(
f"/v1/default/banks/{bank_id}/audit-logs/stats",
params={"period": "1d", "action": "recall"},
)
assert response.status_code == 200
data = response.json()
for bucket in data["buckets"]:
# All actions in buckets should be "recall" only
for action_name in bucket["actions"]:
assert action_name == "recall"
@pytest.mark.asyncio
async def test_audit_log_stats_periods(audit_api_client, bank_id):
"""Test stats endpoint supports different periods."""
await audit_api_client.put(
f"/v1/default/banks/{bank_id}",
json={"name": "Audit Test Bank"},
)
await asyncio.sleep(0.5)
for period, expected_trunc in [("1d", "day"), ("7d", "day"), ("30d", "day")]:
response = await audit_api_client.get(
f"/v1/default/banks/{bank_id}/audit-logs/stats",
params={"period": period},
)
assert response.status_code == 200
data = response.json()
assert data["period"] == period
assert data["trunc"] == expected_trunc
@pytest.mark.asyncio
async def test_audit_log_disabled(memory):
"""Test that no audit logs are created when audit logging is disabled."""
# Ensure audit logging is disabled
memory._audit_logger._enabled = False
app = create_app(memory, initialize_memory=False)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
from datetime import datetime
bid = f"audit_disabled_test_{datetime.now().timestamp()}"
await client.put(f"/v1/default/banks/{bid}", json={"name": "No Audit"})
await client.post(
f"/v1/default/banks/{bid}/memories/recall",
json={"query": "test"},
)
await asyncio.sleep(0.5)
response = await client.get(f"/v1/default/banks/{bid}/audit-logs")
assert response.status_code == 200
data = response.json()
assert data["total"] == 0, "No audit entries should exist when audit logging is disabled"
@pytest.mark.asyncio
async def test_audit_log_action_allowlist(memory):
"""Test that only allowed actions are audited when allowlist is set."""
memory._audit_logger._enabled = True
memory._audit_logger._allowed_actions = frozenset({"recall"}) # Only audit recall
app = create_app(memory, initialize_memory=False)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
from datetime import datetime
bid = f"audit_allowlist_test_{datetime.now().timestamp()}"
# create_bank should NOT be audited
await client.put(f"/v1/default/banks/{bid}", json={"name": "Allowlist Test"})
# recall should be audited
await client.post(
f"/v1/default/banks/{bid}/memories/recall",
json={"query": "allowlist test"},
)
await asyncio.sleep(1.0)
response = await client.get(f"/v1/default/banks/{bid}/audit-logs")
assert response.status_code == 200
data = response.json()
actions = [item["action"] for item in data["items"]]
assert "recall" in actions, "recall should be audited"
assert "create_bank" not in actions, "create_bank should NOT be audited (not in allowlist)"
@pytest.mark.asyncio
async def test_audit_log_ordered_by_most_recent(audit_api_client, bank_id):
"""Test that audit logs are returned ordered by most recent first."""
await audit_api_client.put(
f"/v1/default/banks/{bank_id}",
json={"name": "Order Test Bank"},
)
for i in range(3):
await audit_api_client.post(
f"/v1/default/banks/{bank_id}/memories/recall",
json={"query": f"order test {i}"},
)
await asyncio.sleep(0.2) # Small gap between requests
await asyncio.sleep(1.0)
response = await audit_api_client.get(f"/v1/default/banks/{bank_id}/audit-logs")
assert response.status_code == 200
data = response.json()
# Check descending order by started_at
timestamps = [item["started_at"] for item in data["items"] if item["started_at"]]
assert timestamps == sorted(timestamps, reverse=True), "Audit logs should be ordered most recent first"
@@ -0,0 +1,600 @@
"""Integration tests for bank template import/export endpoints."""
import pytest
import pytest_asyncio
import httpx
from datetime import datetime
from hindsight_api.api import create_app
@pytest_asyncio.fixture
async def api_client(memory):
"""Create an async test client for the FastAPI app."""
app = create_app(memory, initialize_memory=False)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
yield client
@pytest.fixture
def bank_id():
return f"template_test_{datetime.now().timestamp()}"
@pytest.fixture
def sample_template():
return {
"version": "1",
"bank": {
"reflect_mission": "Test mission for reflect",
"retain_mission": "Extract test data carefully",
"retain_extraction_mode": "verbose",
"disposition_empathy": 5,
"disposition_skepticism": 2,
"enable_observations": True,
"observations_mission": "Track test patterns",
},
"mental_models": [
{
"id": "test-model-one",
"name": "Test Model One",
"source_query": "What are the key patterns?",
"tags": ["test"],
"max_tokens": 1024,
"trigger": {"refresh_after_consolidation": True},
},
{
"id": "test-model-two",
"name": "Test Model Two",
"source_query": "What are the common issues?",
},
],
"directives": [
{
"name": "Be concise",
"content": "Always respond concisely.",
"priority": 10,
},
{
"name": "Use examples",
"content": "Include examples when explaining concepts.",
"tags": ["style"],
},
],
}
class TestImportValidation:
"""Test template manifest validation."""
@pytest.mark.asyncio
async def test_import_dry_run_valid(self, api_client, bank_id, sample_template):
"""dry_run=true with a valid manifest returns what would happen."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import?dry_run=true",
json=sample_template,
)
assert resp.status_code == 200
data = resp.json()
assert data["dry_run"] is True
assert data["config_applied"] is True
assert set(data["mental_models_created"]) == {"test-model-one", "test-model-two"}
assert set(data["directives_created"]) == {"Be concise", "Use examples"}
@pytest.mark.asyncio
async def test_import_invalid_version(self, api_client, bank_id):
"""Reject manifest with unsupported version."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={"version": "999"},
)
assert resp.status_code == 400
@pytest.mark.asyncio
async def test_import_invalid_extraction_mode(self, api_client, bank_id):
"""Semantic validation catches bad extraction mode."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"bank": {"retain_extraction_mode": "invalid_mode"},
},
)
assert resp.status_code == 400
assert "retain_extraction_mode" in resp.json()["detail"]
@pytest.mark.asyncio
async def test_import_custom_instructions_without_custom_mode(self, api_client, bank_id):
"""Validate that custom_instructions requires extraction_mode=custom."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"bank": {
"retain_extraction_mode": "verbose",
"retain_custom_instructions": "some custom prompt",
},
},
)
assert resp.status_code == 400
assert "retain_custom_instructions" in resp.json()["detail"]
@pytest.mark.asyncio
async def test_import_duplicate_mental_model_ids(self, api_client, bank_id):
"""Reject manifest with duplicate mental model IDs."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"mental_models": [
{"id": "dup-id", "name": "First", "source_query": "q1"},
{"id": "dup-id", "name": "Second", "source_query": "q2"},
],
},
)
assert resp.status_code == 400
@pytest.mark.asyncio
async def test_import_duplicate_directive_names(self, api_client, bank_id):
"""Reject manifest with duplicate directive names."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"directives": [
{"name": "Same Name", "content": "First"},
{"name": "Same Name", "content": "Second"},
],
},
)
assert resp.status_code == 400
@pytest.mark.asyncio
async def test_import_missing_mental_model_id(self, api_client, bank_id):
"""Mental model without id is rejected."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"mental_models": [
{"name": "No ID Model", "source_query": "test query"},
],
},
)
assert resp.status_code == 400
@pytest.mark.asyncio
async def test_import_invalid_mental_model_id_format(self, api_client, bank_id):
"""Mental model with invalid ID format is rejected."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"mental_models": [
{"id": "UPPERCASE-NOT-ALLOWED", "name": "Bad", "source_query": "q"},
],
},
)
assert resp.status_code == 400
@pytest.mark.asyncio
async def test_import_empty_manifest(self, api_client, bank_id):
"""Import with no bank or mental_models is valid (no-op)."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={"version": "1"},
)
assert resp.status_code == 200
data = resp.json()
assert data["config_applied"] is False
assert data["mental_models_created"] == []
assert data["directives_created"] == []
@pytest.mark.asyncio
async def test_import_empty_mental_model_name(self, api_client, bank_id):
"""Semantic validation catches empty mental model name."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"mental_models": [
{"id": "test-mm", "name": " ", "source_query": "q"},
],
},
)
assert resp.status_code == 400
assert "name" in resp.json()["detail"]
@pytest.mark.asyncio
async def test_import_empty_directive_content(self, api_client, bank_id):
"""Semantic validation catches empty directive content."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"directives": [
{"name": "Bad Directive", "content": " "},
],
},
)
assert resp.status_code == 400
assert "content" in resp.json()["detail"]
class TestImportApply:
"""Test that import actually applies config, mental models, and directives."""
@pytest.mark.asyncio
async def test_import_applies_config(self, api_client, bank_id):
"""Import with bank config applies config overrides on a new bank."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"bank": {
"reflect_mission": "Imported mission",
"disposition_empathy": 4,
},
},
)
assert resp.status_code == 200
data = resp.json()
assert data["config_applied"] is True
assert data["dry_run"] is False
# Verify config was actually applied
config_resp = await api_client.get(f"/v1/default/banks/{bank_id}/config")
assert config_resp.status_code == 200
config = config_resp.json()
assert config["overrides"]["reflect_mission"] == "Imported mission"
assert config["overrides"]["disposition_empathy"] == 4
@pytest.mark.asyncio
async def test_import_into_existing_bank(self, api_client, bank_id):
"""Import into an already-existing bank applies config and creates resources."""
# Pre-create the bank
await api_client.put(f"/v1/default/banks/{bank_id}", json={})
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"bank": {"reflect_mission": "Existing bank mission"},
"mental_models": [
{"id": "existing-bank-mm", "name": "MM", "source_query": "q"},
],
"directives": [
{"name": "Existing Bank Directive", "content": "Be helpful"},
],
},
)
assert resp.status_code == 200
data = resp.json()
assert data["config_applied"] is True
assert "existing-bank-mm" in data["mental_models_created"]
assert "Existing Bank Directive" in data["directives_created"]
# Verify everything exists
config_resp = await api_client.get(f"/v1/default/banks/{bank_id}/config")
assert config_resp.json()["overrides"]["reflect_mission"] == "Existing bank mission"
mm_resp = await api_client.get(f"/v1/default/banks/{bank_id}/mental-models/existing-bank-mm")
assert mm_resp.status_code == 200
dir_resp = await api_client.get(f"/v1/default/banks/{bank_id}/directives")
assert dir_resp.status_code == 200
names = [d["name"] for d in dir_resp.json()["items"]]
assert "Existing Bank Directive" in names
@pytest.mark.asyncio
async def test_import_creates_mental_models(self, api_client, bank_id):
"""Import creates mental models and returns operation IDs."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"mental_models": [
{
"id": "import-mm-1",
"name": "Imported Model",
"source_query": "What patterns exist?",
"tags": ["imported"],
},
],
},
)
assert resp.status_code == 200
data = resp.json()
assert "import-mm-1" in data["mental_models_created"]
assert len(data["operation_ids"]) == 1
# Verify mental model exists
mm_resp = await api_client.get(f"/v1/default/banks/{bank_id}/mental-models/import-mm-1")
assert mm_resp.status_code == 200
mm = mm_resp.json()
assert mm["name"] == "Imported Model"
assert mm["source_query"] == "What patterns exist?"
assert mm["tags"] == ["imported"]
@pytest.mark.asyncio
async def test_import_updates_existing_mental_models(self, api_client, bank_id):
"""Re-importing updates existing mental models matched by ID."""
# First import
await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"mental_models": [
{
"id": "reusable-mm",
"name": "Original Name",
"source_query": "Original query",
},
],
},
)
# Second import with same ID but different content
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"mental_models": [
{
"id": "reusable-mm",
"name": "Updated Name",
"source_query": "Updated query",
},
],
},
)
assert resp.status_code == 200
data = resp.json()
assert "reusable-mm" in data["mental_models_updated"]
assert data["mental_models_created"] == []
# Verify update
mm_resp = await api_client.get(f"/v1/default/banks/{bank_id}/mental-models/reusable-mm")
assert mm_resp.status_code == 200
mm = mm_resp.json()
assert mm["name"] == "Updated Name"
assert mm["source_query"] == "Updated query"
@pytest.mark.asyncio
async def test_import_creates_directives(self, api_client, bank_id):
"""Import creates directives."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"directives": [
{
"name": "Test Directive",
"content": "Always be helpful and precise.",
"priority": 5,
"tags": ["test"],
},
],
},
)
assert resp.status_code == 200
data = resp.json()
assert "Test Directive" in data["directives_created"]
assert data["directives_updated"] == []
# Verify directive exists
dir_resp = await api_client.get(f"/v1/default/banks/{bank_id}/directives")
assert dir_resp.status_code == 200
items = dir_resp.json()["items"]
assert len(items) == 1
assert items[0]["name"] == "Test Directive"
assert items[0]["content"] == "Always be helpful and precise."
assert items[0]["priority"] == 5
assert items[0]["tags"] == ["test"]
@pytest.mark.asyncio
async def test_import_updates_existing_directives(self, api_client, bank_id):
"""Re-importing updates existing directives matched by name."""
# First import
await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"directives": [
{"name": "Reusable Directive", "content": "Original content", "priority": 1},
],
},
)
# Second import with same name but different content
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"directives": [
{"name": "Reusable Directive", "content": "Updated content", "priority": 10},
],
},
)
assert resp.status_code == 200
data = resp.json()
assert "Reusable Directive" in data["directives_updated"]
assert data["directives_created"] == []
# Verify update
dir_resp = await api_client.get(f"/v1/default/banks/{bank_id}/directives")
items = dir_resp.json()["items"]
directive = [d for d in items if d["name"] == "Reusable Directive"][0]
assert directive["content"] == "Updated content"
assert directive["priority"] == 10
@pytest.mark.asyncio
async def test_import_config_only(self, api_client, bank_id):
"""Import with only bank config (no mental_models or directives) works."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"bank": {"retain_extraction_mode": "verbose"},
},
)
assert resp.status_code == 200
data = resp.json()
assert data["config_applied"] is True
assert data["mental_models_created"] == []
assert data["directives_created"] == []
assert data["operation_ids"] == []
@pytest.mark.asyncio
async def test_import_mental_models_only(self, api_client, bank_id):
"""Import with only mental_models works."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"mental_models": [
{"id": "mm-only", "name": "MM Only", "source_query": "test"},
],
},
)
assert resp.status_code == 200
data = resp.json()
assert data["config_applied"] is False
assert "mm-only" in data["mental_models_created"]
assert data["directives_created"] == []
@pytest.mark.asyncio
async def test_import_directives_only(self, api_client, bank_id):
"""Import with only directives works."""
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"directives": [
{"name": "Dir Only", "content": "test directive"},
],
},
)
assert resp.status_code == 200
data = resp.json()
assert data["config_applied"] is False
assert data["mental_models_created"] == []
assert "Dir Only" in data["directives_created"]
class TestExport:
"""Test bank template export."""
@pytest.mark.asyncio
async def test_export_empty_bank(self, api_client, bank_id):
"""Export a bank with no overrides returns minimal manifest."""
# Create bank
await api_client.put(f"/v1/default/banks/{bank_id}", json={})
resp = await api_client.get(f"/v1/default/banks/{bank_id}/export")
assert resp.status_code == 200
data = resp.json()
assert data["version"] == "1"
assert data["bank"] is None
assert data["mental_models"] is None
assert data["directives"] is None
@pytest.mark.asyncio
async def test_export_after_import(self, api_client, bank_id):
"""Export after import returns the imported config, mental models, and directives."""
template = {
"version": "1",
"bank": {
"reflect_mission": "Roundtrip mission",
"disposition_empathy": 3,
},
"mental_models": [
{
"id": "roundtrip-mm",
"name": "Roundtrip Model",
"source_query": "What happened?",
"tags": ["roundtrip"],
"max_tokens": 512,
},
],
"directives": [
{
"name": "Roundtrip Directive",
"content": "Be thorough.",
"priority": 3,
"tags": ["roundtrip"],
},
],
}
# Import
import_resp = await api_client.post(f"/v1/default/banks/{bank_id}/import", json=template)
assert import_resp.status_code == 200
# Export
resp = await api_client.get(f"/v1/default/banks/{bank_id}/export")
assert resp.status_code == 200
data = resp.json()
assert data["version"] == "1"
assert data["bank"]["reflect_mission"] == "Roundtrip mission"
assert data["bank"]["disposition_empathy"] == 3
assert len(data["mental_models"]) == 1
mm = data["mental_models"][0]
assert mm["id"] == "roundtrip-mm"
assert mm["name"] == "Roundtrip Model"
assert mm["source_query"] == "What happened?"
assert mm["tags"] == ["roundtrip"]
assert mm["max_tokens"] == 512
assert len(data["directives"]) == 1
d = data["directives"][0]
assert d["name"] == "Roundtrip Directive"
assert d["content"] == "Be thorough."
assert d["priority"] == 3
assert d["tags"] == ["roundtrip"]
@pytest.mark.asyncio
async def test_export_reimport_roundtrip(self, api_client, bank_id):
"""Exported manifest can be re-imported into a new bank."""
# Set up source bank
await api_client.post(
f"/v1/default/banks/{bank_id}/import",
json={
"version": "1",
"bank": {"retain_mission": "Roundtrip test"},
"mental_models": [
{"id": "rt-mm", "name": "RT Model", "source_query": "test query"},
],
"directives": [
{"name": "RT Directive", "content": "test directive"},
],
},
)
# Export
export_resp = await api_client.get(f"/v1/default/banks/{bank_id}/export")
assert export_resp.status_code == 200
exported = export_resp.json()
# Import into a new bank
new_bank_id = f"{bank_id}_clone"
import_resp = await api_client.post(
f"/v1/default/banks/{new_bank_id}/import",
json=exported,
)
assert import_resp.status_code == 200
data = import_resp.json()
assert data["config_applied"] is True
assert "rt-mm" in data["mental_models_created"]
assert "RT Directive" in data["directives_created"]
@pytest.mark.asyncio
async def test_export_nonexistent_bank(self, api_client):
"""Export from a nonexistent bank returns the bank with defaults (auto-created)."""
resp = await api_client.get("/v1/default/banks/nonexistent-export-test/export")
# get_bank_profile auto-creates, so this returns a valid empty manifest
assert resp.status_code == 200
data = resp.json()
assert data["version"] == "1"
@@ -198,7 +198,7 @@ async def test_real_openai_batch_api(real_llm_config, test_contents_real, integr
for fact in facts:
assert hasattr(fact, "fact_text"), "Fact should have fact_text"
assert hasattr(fact, "fact_type"), "Fact should have fact_type"
assert fact.fact_type in ["world", "experience", "opinion"], f"Invalid fact_type: {fact.fact_type}"
assert fact.fact_type in ["world", "experience"], f"Invalid fact_type: {fact.fact_type}"
logger.info("\n✅ All assertions passed!")
@@ -36,7 +36,7 @@ class TestCausalRelationsValidation:
"""
context = "Personal life update"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
event_date = datetime(2024, 3, 15)
facts, _, usage = await extract_facts_from_text(
@@ -81,7 +81,7 @@ class TestCausalRelationsValidation:
"""
context = "Project update"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
event_date = datetime(2024, 6, 1)
facts, _, _ = await extract_facts_from_text(
@@ -118,7 +118,7 @@ class TestCausalRelationsValidation:
"""
context = "Personal achievement story"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
event_date = datetime(2024, 7, 15)
facts, _, _ = await extract_facts_from_text(
@@ -168,7 +168,7 @@ class TestCausalRelationsValidation:
"""
context = "Business impact analysis"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
event_date = datetime(2024, 4, 1)
facts, _, usage = await extract_facts_from_text(
@@ -205,7 +205,7 @@ class TestCausalRelationsValidation:
"""
context = "Career progression"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
event_date = datetime(2024, 5, 1)
facts, _, _ = await extract_facts_from_text(
@@ -35,7 +35,7 @@ After searching for weeks, I finally found a cheaper apartment in Brooklyn.
"""
context = "Personal story about housing change"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text, event_date=datetime(2024, 3, 15), context=context, llm_config=llm_config, agent_name="TestUser",
@@ -105,7 +105,7 @@ The renovation took three months and cost $15,000.
"""
context = "Home repair story"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text, event_date=datetime(2024, 6, 1), context=context, llm_config=llm_config, agent_name="TestUser",
@@ -136,7 +136,7 @@ Machine learning fascinated me so much that I changed my career to data science.
"""
context = "Career change story"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text, event_date=datetime(2024, 1, 1), context=context, llm_config=llm_config, agent_name="TestUser",
@@ -164,7 +164,7 @@ The new role enabled me to lead a team of engineers.
"""
context = "Work promotion story"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text, event_date=datetime(2024, 2, 15), context=context, llm_config=llm_config, agent_name="TestUser",
@@ -192,7 +192,7 @@ Reduced spending somewhat affected local businesses.
"""
context = "Economic impact story"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text, event_date=datetime(2024, 4, 1), context=context, llm_config=llm_config, agent_name="TestUser",
@@ -0,0 +1,88 @@
"""
Regression tests for Codex provider tool_choice normalization.
The reflect agent forces tool selection via named tool_choice dicts on early iterations:
{"type": "function", "function": {"name": "recall"}}
The Codex Responses API expects the function name at the top level instead:
{"type": "function", "name": "recall"}
Without normalization, Codex rejects the request with:
400 Unknown parameter: 'tool_choice.function'
"""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from hindsight_api.engine.providers.codex_llm import CodexLLM
TOOLS = [
{
"type": "function",
"function": {
"name": "recall",
"description": "Recall semantic memories",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
},
},
}
]
def build_llm() -> CodexLLM:
with patch.object(CodexLLM, "_load_codex_auth", return_value=("token", "account")):
return CodexLLM(
provider="openai-codex",
api_key="ignored",
base_url="https://chatgpt.com/backend-api",
model="gpt-5.4-mini",
)
@pytest.mark.asyncio
async def test_codex_normalizes_legacy_named_tool_choice_shape():
llm = build_llm()
response = MagicMock()
response.status_code = 200
response.raise_for_status.return_value = None
with patch.object(llm._client, "post", new_callable=AsyncMock) as mock_post:
mock_post.return_value = response
with patch.object(llm, "_parse_sse_tool_stream", new_callable=AsyncMock) as mock_parse:
mock_parse.return_value = (None, [])
await llm.call_with_tools(
messages=[{"role": "user", "content": "recall the memory"}],
tools=TOOLS,
tool_choice={"type": "function", "function": {"name": "recall"}},
max_retries=0,
)
sent_payload = mock_post.call_args.kwargs["json"]
assert sent_payload["tool_choice"] == {"type": "function", "name": "recall"}
@pytest.mark.asyncio
async def test_codex_forced_tool_choice_still_yields_tool_calls():
llm = build_llm()
response = MagicMock()
response.status_code = 200
response.raise_for_status.return_value = None
tool_call = {"id": "call-1", "name": "recall", "arguments": {"query": "memory"}}
with patch.object(llm._client, "post", new_callable=AsyncMock) as mock_post:
mock_post.return_value = response
with patch.object(llm, "_parse_sse_tool_stream", new_callable=AsyncMock) as mock_parse:
mock_parse.return_value = (None, [tool_call])
result = await llm.call_with_tools(
messages=[{"role": "user", "content": "recall the memory"}],
tools=TOOLS,
tool_choice={"type": "function", "function": {"name": "recall"}},
max_retries=0,
)
sent_payload = mock_post.call_args.kwargs["json"]
assert len(result.tool_calls) == 1
assert result.tool_calls[0].name == "recall"
assert sent_payload["tool_choice"] == {"type": "function", "name": "recall"}
@@ -0,0 +1,339 @@
"""
Tests for CohereCrossEncoder.
Tests the Cohere cross-encoder implementation, including Azure AI Foundry endpoint support.
"""
import os
from unittest.mock import MagicMock, patch
import httpx
import pytest
from hindsight_api.engine.cross_encoder import CohereCrossEncoder, create_cross_encoder_from_env
class TestCohereCrossEncoder:
"""Test suite for CohereCrossEncoder class."""
@pytest.mark.asyncio
async def test_initialization_native_cohere(self):
"""Test successful initialization with native Cohere API (no base_url)."""
encoder = CohereCrossEncoder(
api_key="test_key",
model="rerank-english-v3.0",
)
assert encoder.provider_name == "cohere"
assert encoder.api_key == "test_key"
assert encoder.model == "rerank-english-v3.0"
assert encoder._client is None
assert encoder._httpx_client is None
# Mock the cohere import
mock_cohere = MagicMock()
mock_cohere.Client = MagicMock()
with patch.dict("sys.modules", {"cohere": mock_cohere}):
await encoder.initialize()
assert encoder._client is not None
assert encoder._httpx_client is None
mock_cohere.Client.assert_called_once_with(api_key="test_key", timeout=60.0)
@pytest.mark.asyncio
async def test_initialization_azure_endpoint(self):
"""Test initialization with Azure AI Foundry endpoint (uses httpx)."""
encoder = CohereCrossEncoder(
api_key="test_key",
model="cohere-rerank-v3-english",
base_url="https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke",
)
assert encoder.base_url == "https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke"
await encoder.initialize()
assert encoder._httpx_client is not None
assert encoder._client is None
assert isinstance(encoder._httpx_client, httpx.Client)
@pytest.mark.asyncio
async def test_initialization_missing_package(self):
"""Test initialization fails when cohere package is missing (native API)."""
encoder = CohereCrossEncoder(
api_key="test_key",
model="rerank-english-v3.0",
)
with patch.dict("sys.modules", {"cohere": None}):
with pytest.raises(ImportError, match="cohere is required"):
await encoder.initialize()
@pytest.mark.asyncio
async def test_initialization_idempotent(self):
"""Test that calling initialize() multiple times is safe."""
encoder = CohereCrossEncoder(
api_key="test_key",
model="rerank-english-v3.0",
)
mock_cohere = MagicMock()
mock_cohere.Client = MagicMock()
with patch.dict("sys.modules", {"cohere": mock_cohere}):
await encoder.initialize()
assert encoder._client is not None
# Second call should be no-op
await encoder.initialize()
# Should only create client once
mock_cohere.Client.assert_called_once()
@pytest.mark.asyncio
async def test_predict_native_cohere_single_query(self):
"""Test prediction with native Cohere SDK."""
encoder = CohereCrossEncoder(
api_key="test_key",
model="rerank-english-v3.0",
)
# Create mock Cohere response
mock_result_1 = MagicMock()
mock_result_1.index = 0
mock_result_1.relevance_score = 0.9
mock_result_2 = MagicMock()
mock_result_2.index = 1
mock_result_2.relevance_score = 0.7
mock_result_3 = MagicMock()
mock_result_3.index = 2
mock_result_3.relevance_score = 0.5
mock_response = MagicMock()
mock_response.results = [mock_result_1, mock_result_2, mock_result_3]
mock_cohere_client = MagicMock()
mock_cohere_client.rerank = MagicMock(return_value=mock_response)
mock_cohere = MagicMock()
mock_cohere.Client = MagicMock(return_value=mock_cohere_client)
with patch.dict("sys.modules", {"cohere": mock_cohere}):
await encoder.initialize()
pairs = [
("What is Python?", "Python is a programming language"),
("What is Python?", "Python is a snake"),
("What is Python?", "Python is a British comedy group"),
]
scores = await encoder.predict(pairs)
assert len(scores) == 3
assert scores == [0.9, 0.7, 0.5]
# Verify rerank was called correctly
mock_cohere_client.rerank.assert_called_once()
call_args = mock_cohere_client.rerank.call_args
assert call_args.kwargs["model"] == "rerank-english-v3.0"
assert call_args.kwargs["query"] == "What is Python?"
assert len(call_args.kwargs["documents"]) == 3
assert call_args.kwargs["return_documents"] is False
@pytest.mark.asyncio
async def test_predict_azure_endpoint_single_query(self):
"""Test prediction with Azure AI Foundry endpoint (httpx direct call)."""
encoder = CohereCrossEncoder(
api_key="test_key",
model="cohere-rerank-v3-english",
base_url="https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke",
)
await encoder.initialize()
# Mock httpx response
mock_response = MagicMock()
mock_response.json.return_value = {
"results": [
{"index": 0, "relevance_score": 0.9},
{"index": 1, "relevance_score": 0.7},
{"index": 2, "relevance_score": 0.5},
]
}
encoder._httpx_client.post = MagicMock(return_value=mock_response)
pairs = [
("What is Python?", "Python is a programming language"),
("What is Python?", "Python is a snake"),
("What is Python?", "Python is a British comedy group"),
]
scores = await encoder.predict(pairs)
assert len(scores) == 3
assert scores == [0.9, 0.7, 0.5]
# Verify httpx.post was called with correct URL and payload
encoder._httpx_client.post.assert_called_once()
call_args = encoder._httpx_client.post.call_args
assert call_args[0][0] == "https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke"
assert call_args.kwargs["json"]["model"] == "cohere-rerank-v3-english"
assert call_args.kwargs["json"]["query"] == "What is Python?"
assert len(call_args.kwargs["json"]["documents"]) == 3
assert call_args.kwargs["json"]["return_documents"] is False
@pytest.mark.asyncio
async def test_predict_multiple_queries(self):
"""Test prediction with multiple different queries (grouped efficiently)."""
encoder = CohereCrossEncoder(
api_key="test_key",
model="rerank-english-v3.0",
)
# First query response
mock_result_1_1 = MagicMock()
mock_result_1_1.index = 0
mock_result_1_1.relevance_score = 0.9
mock_result_1_2 = MagicMock()
mock_result_1_2.index = 1
mock_result_1_2.relevance_score = 0.7
mock_response1 = MagicMock()
mock_response1.results = [mock_result_1_1, mock_result_1_2]
# Second query response
mock_result_2_1 = MagicMock()
mock_result_2_1.index = 0
mock_result_2_1.relevance_score = 0.8
mock_response2 = MagicMock()
mock_response2.results = [mock_result_2_1]
mock_cohere_client = MagicMock()
mock_cohere_client.rerank = MagicMock(side_effect=[mock_response1, mock_response2])
mock_cohere = MagicMock()
mock_cohere.Client = MagicMock(return_value=mock_cohere_client)
with patch.dict("sys.modules", {"cohere": mock_cohere}):
await encoder.initialize()
pairs = [
("What is Python?", "Python is a programming language"),
("What is Python?", "Python is a snake"),
("What is Java?", "Java is a programming language"),
]
scores = await encoder.predict(pairs)
assert len(scores) == 3
assert scores[0] == 0.9 # First query, first doc
assert scores[1] == 0.7 # First query, second doc
assert scores[2] == 0.8 # Second query, first doc
# Verify rerank was called twice (once per unique query)
assert mock_cohere_client.rerank.call_count == 2
@pytest.mark.asyncio
async def test_predict_empty_pairs(self):
"""Test prediction with empty input."""
encoder = CohereCrossEncoder(
api_key="test_key",
model="rerank-english-v3.0",
)
mock_cohere = MagicMock()
mock_cohere.Client = MagicMock()
with patch.dict("sys.modules", {"cohere": mock_cohere}):
await encoder.initialize()
scores = await encoder.predict([])
assert scores == []
@pytest.mark.asyncio
async def test_predict_not_initialized(self):
"""Test that predict fails if encoder not initialized."""
encoder = CohereCrossEncoder(
api_key="test_key",
model="rerank-english-v3.0",
)
pairs = [("query", "document")]
with pytest.raises(RuntimeError, match="not initialized"):
await encoder.predict(pairs)
@pytest.mark.asyncio
async def test_azure_endpoint_http_error(self):
"""Test that HTTP errors from Azure endpoint are raised."""
encoder = CohereCrossEncoder(
api_key="test_key",
model="cohere-rerank-v3-english",
base_url="https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke",
)
await encoder.initialize()
# Mock httpx to raise HTTP error
mock_response = MagicMock()
mock_response.raise_for_status.side_effect = httpx.HTTPStatusError(
"404 Not Found",
request=MagicMock(),
response=MagicMock(status_code=404),
)
encoder._httpx_client.post = MagicMock(return_value=mock_response)
pairs = [("What is Python?", "Python is a programming language")]
# Should raise the HTTP error
with pytest.raises(httpx.HTTPStatusError):
await encoder.predict(pairs)
class TestFactoryFunction:
"""Test suite for create_cross_encoder_from_env factory function."""
@pytest.mark.asyncio
async def test_create_cohere_from_env(self):
"""Test creating Cohere cross-encoder from environment variables."""
env_vars = {
"HINDSIGHT_API_RERANKER_PROVIDER": "cohere",
"HINDSIGHT_API_RERANKER_COHERE_API_KEY": "test_key",
"HINDSIGHT_API_RERANKER_COHERE_MODEL": "rerank-english-v3.0",
}
with patch.dict(os.environ, env_vars, clear=False):
from hindsight_api.config import HindsightConfig
config = HindsightConfig.from_env()
with patch("hindsight_api.config.get_config", return_value=config):
encoder = create_cross_encoder_from_env()
assert isinstance(encoder, CohereCrossEncoder)
assert encoder.api_key == "test_key"
assert encoder.model == "rerank-english-v3.0"
assert encoder.base_url is None
@pytest.mark.asyncio
async def test_create_cohere_with_azure_base_url_from_env(self):
"""Test creating Cohere cross-encoder with Azure base URL from environment."""
env_vars = {
"HINDSIGHT_API_RERANKER_PROVIDER": "cohere",
"HINDSIGHT_API_RERANKER_COHERE_API_KEY": "test_key",
"HINDSIGHT_API_RERANKER_COHERE_MODEL": "cohere-rerank-v3-english",
"HINDSIGHT_API_RERANKER_COHERE_BASE_URL": "https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke",
}
with patch.dict(os.environ, env_vars, clear=False):
from hindsight_api.config import HindsightConfig
config = HindsightConfig.from_env()
with patch("hindsight_api.config.get_config", return_value=config):
encoder = create_cross_encoder_from_env()
assert isinstance(encoder, CohereCrossEncoder)
assert encoder.api_key == "test_key"
assert encoder.model == "cohere-rerank-v3-english"
assert encoder.base_url == "https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke"
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,842 @@
"""
Tests for delta retain upsert optimization that only re-processes changed chunks.
"""
import logging
from datetime import datetime, timezone
import pytest
from hindsight_api import RequestContext
from hindsight_api.engine.memory_engine import Budget
logger = logging.getLogger(__name__)
def _ts():
return datetime.now(timezone.utc).timestamp()
# ============================================================
# Core Delta Retain Tests
# ============================================================
@pytest.mark.asyncio
async def test_delta_retain_unchanged_content_skips_llm(memory, request_context):
"""
When upserting a document with identical content, no new facts should be
extracted (LLM is not called for unchanged chunks). The existing facts
should be preserved.
"""
bank_id = f"test_delta_unchanged_{_ts()}"
document_id = "conversation-001"
try:
content = "Alice works at Google. Bob works at Microsoft."
# First retain — full processing
v1_units = await memory.retain_async(
bank_id=bank_id,
content=content,
context="team info",
document_id=document_id,
request_context=request_context,
)
assert len(v1_units) > 0, "v1 should create facts"
# Get v1 document state
doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context)
v1_unit_count = doc_v1["memory_unit_count"]
# Second retain — same content, should use delta path (no new facts)
v2_units = await memory.retain_async(
bank_id=bank_id,
content=content,
context="team info",
document_id=document_id,
request_context=request_context,
)
# No new units should be returned (nothing changed)
assert v2_units == [], "Delta retain with unchanged content should return empty unit list"
# Existing facts should still be there
doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context)
assert doc_v2["memory_unit_count"] == v1_unit_count, "Existing facts should be preserved"
# Verify recall still works
result = await memory.recall_async(
bank_id=bank_id,
query="Where does Alice work?",
budget=Budget.MID,
max_tokens=1000,
request_context=request_context,
)
assert len(result.results) > 0, "Should still recall facts after delta retain"
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_delta_retain_appended_content(memory, request_context):
"""
When a conversation grows (new content appended), only new chunks should
be processed. Facts from unchanged chunks should be preserved.
"""
bank_id = f"test_delta_append_{_ts()}"
document_id = "growing-conversation"
try:
# First version — short content (single chunk)
v1_content = "Alice is a software engineer at Google. She works on search infrastructure."
v1_units = await memory.retain_async(
bank_id=bank_id,
content=v1_content,
context="profile",
document_id=document_id,
request_context=request_context,
)
assert len(v1_units) > 0
# Get v1 facts via recall
v1_recall = await memory.recall_async(
bank_id=bank_id,
query="What does Alice do?",
budget=Budget.MID,
max_tokens=2000,
request_context=request_context,
)
v1_fact_texts = {r.text for r in v1_recall.results}
# Second version — original content + new content appended
# This should preserve facts from the first chunk and add new ones
v2_content = v1_content + "\n\nBob joined Google as a product manager in 2024. He previously worked at Meta on AR/VR products."
v2_units = await memory.retain_async(
bank_id=bank_id,
content=v2_content,
context="profile",
document_id=document_id,
request_context=request_context,
)
# Should have facts about Bob from the new content
v2_recall = await memory.recall_async(
bank_id=bank_id,
query="What does Bob do?",
budget=Budget.MID,
max_tokens=2000,
request_context=request_context,
)
bob_facts = [r for r in v2_recall.results if "bob" in r.text.lower()]
assert len(bob_facts) > 0, "Should have facts about Bob from appended content"
# Should still have facts about Alice from original content
alice_recall = await memory.recall_async(
bank_id=bank_id,
query="What does Alice do?",
budget=Budget.MID,
max_tokens=2000,
request_context=request_context,
)
assert len(alice_recall.results) > 0, "Should still have Alice facts from original content"
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_delta_retain_modified_chunk(memory, request_context):
"""
When content in the middle changes, that chunk should be re-processed
while other chunks are preserved.
"""
bank_id = f"test_delta_modified_{_ts()}"
document_id = "changing-doc"
try:
# v1: Alice works at Google
v1_content = "Alice works at Google as a senior engineer."
v1_units = await memory.retain_async(
bank_id=bank_id,
content=v1_content,
context="team",
document_id=document_id,
request_context=request_context,
)
assert len(v1_units) > 0
# v2: Alice works at Microsoft (changed)
v2_content = "Alice works at Microsoft as a principal engineer."
v2_units = await memory.retain_async(
bank_id=bank_id,
content=v2_content,
context="team",
document_id=document_id,
request_context=request_context,
)
# New facts should reflect the updated content
result = await memory.recall_async(
bank_id=bank_id,
query="Where does Alice work?",
budget=Budget.MID,
max_tokens=2000,
request_context=request_context,
)
all_texts = " ".join(r.text.lower() for r in result.results)
assert "microsoft" in all_texts, f"Should have updated fact about Microsoft, got: {all_texts}"
finally:
await memory.delete_bank(bank_id, request_context=request_context)
# ============================================================
# Entity & Link Tests
# ============================================================
@pytest.mark.asyncio
async def test_delta_retain_entities_preserved_for_unchanged_chunks(memory, request_context):
"""
Entities linked to unchanged chunks should be preserved after delta retain.
"""
bank_id = f"test_delta_entities_{_ts()}"
document_id = "entity-doc"
try:
v1_content = "Alice works at Google. She is a senior engineer in the Cloud division."
v1_units = await memory.retain_async(
bank_id=bank_id,
content=v1_content,
context="team",
document_id=document_id,
request_context=request_context,
)
assert len(v1_units) > 0
# Check entities exist
pool = await memory._get_pool()
async with pool.acquire() as conn:
v1_entities = await conn.fetch(
"SELECT canonical_name FROM entities WHERE bank_id = $1",
bank_id,
)
v1_entity_names = {e["canonical_name"].lower() for e in v1_entities}
assert len(v1_entity_names) > 0, "Should have entities after v1 retain"
# Upsert with same content — entities should persist
await memory.retain_async(
bank_id=bank_id,
content=v1_content,
context="team",
document_id=document_id,
request_context=request_context,
)
async with pool.acquire() as conn:
v2_entities = await conn.fetch(
"SELECT canonical_name FROM entities WHERE bank_id = $1",
bank_id,
)
v2_entity_names = {e["canonical_name"].lower() for e in v2_entities}
# All v1 entities should still exist
assert v1_entity_names.issubset(v2_entity_names), (
f"v1 entities {v1_entity_names} should be preserved, got {v2_entity_names}"
)
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_delta_retain_new_entities_created_for_new_chunks(memory, request_context):
"""
New entities should be created for newly added chunks during delta retain.
"""
bank_id = f"test_delta_new_entities_{_ts()}"
document_id = "entity-growth-doc"
try:
v1_content = "Alice works at Google."
await memory.retain_async(
bank_id=bank_id,
content=v1_content,
context="team",
document_id=document_id,
request_context=request_context,
)
pool = await memory._get_pool()
async with pool.acquire() as conn:
v1_entities = await conn.fetch(
"SELECT canonical_name FROM entities WHERE bank_id = $1",
bank_id,
)
v1_entity_names = {e["canonical_name"].lower() for e in v1_entities}
# Append content mentioning new entities
v2_content = v1_content + "\n\nBob joined Facebook. He works with Charlie on the Reality Labs project."
await memory.retain_async(
bank_id=bank_id,
content=v2_content,
context="team",
document_id=document_id,
request_context=request_context,
)
async with pool.acquire() as conn:
v2_entities = await conn.fetch(
"SELECT canonical_name FROM entities WHERE bank_id = $1",
bank_id,
)
v2_entity_names = {e["canonical_name"].lower() for e in v2_entities}
# Should have more entities after adding content with new people/orgs
assert len(v2_entity_names) > len(v1_entity_names), (
f"Should have more entities after append: v1={v1_entity_names}, v2={v2_entity_names}"
)
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_delta_retain_links_preserved_for_unchanged_chunks(memory, request_context):
"""
Memory links (temporal, semantic, entity) for unchanged chunks should be preserved.
"""
bank_id = f"test_delta_links_{_ts()}"
document_id = "links-doc"
try:
content = "Alice is a senior engineer at Google Cloud. She mentors junior engineers and reviews their code."
v1_units = await memory.retain_async(
bank_id=bank_id,
content=content,
context="team",
document_id=document_id,
request_context=request_context,
)
assert len(v1_units) > 0
# Count links after v1
pool = await memory._get_pool()
async with pool.acquire() as conn:
v1_link_count = await conn.fetchval(
"""SELECT COUNT(*) FROM memory_links ml
JOIN memory_units mu ON ml.from_unit_id = mu.id
WHERE mu.bank_id = $1 AND mu.document_id = $2""",
bank_id,
document_id,
)
# Upsert with same content
await memory.retain_async(
bank_id=bank_id,
content=content,
context="team",
document_id=document_id,
request_context=request_context,
)
# Links should be preserved
async with pool.acquire() as conn:
v2_link_count = await conn.fetchval(
"""SELECT COUNT(*) FROM memory_links ml
JOIN memory_units mu ON ml.from_unit_id = mu.id
WHERE mu.bank_id = $1 AND mu.document_id = $2""",
bank_id,
document_id,
)
assert v2_link_count == v1_link_count, (
f"Links should be preserved: v1={v1_link_count}, v2={v2_link_count}"
)
finally:
await memory.delete_bank(bank_id, request_context=request_context)
# ============================================================
# Document Metadata & Tags Tests
# ============================================================
@pytest.mark.asyncio
async def test_delta_retain_document_metadata_updated(memory, request_context):
"""
Document metadata (retain_params, tags) should be updated even when
chunk content hasn't changed.
"""
bank_id = f"test_delta_meta_{_ts()}"
document_id = "metadata-doc"
try:
content = "Alice works at Google."
# v1 with initial tags
await memory.retain_async(
bank_id=bank_id,
content=content,
context="initial context",
document_id=document_id,
request_context=request_context,
)
doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context)
assert doc_v1 is not None
# v2 with updated context (same content — triggers delta path)
await memory.retain_async(
bank_id=bank_id,
content=content,
context="updated context",
document_id=document_id,
request_context=request_context,
)
doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context)
assert doc_v2 is not None
assert doc_v2["updated_at"] >= doc_v1["updated_at"], "Document should have updated timestamp"
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_delta_retain_tags_propagated_to_existing_units(memory, request_context):
"""
When tags change during an upsert with unchanged content, the new tags
should be propagated to all existing memory units.
"""
bank_id = f"test_delta_tags_{_ts()}"
document_id = "tags-doc"
try:
content = "Alice works at Google."
# v1 with tag "team-a"
await memory.retain_batch_async(
bank_id=bank_id,
contents=[{
"content": content,
"document_id": document_id,
"tags": ["team-a"],
}],
request_context=request_context,
)
pool = await memory._get_pool()
async with pool.acquire() as conn:
v1_tags = await conn.fetch(
"SELECT tags FROM memory_units WHERE bank_id = $1 AND document_id = $2",
bank_id,
document_id,
)
assert all("team-a" in row["tags"] for row in v1_tags), "v1 units should have team-a tag"
# v2 with same content but different tags
await memory.retain_batch_async(
bank_id=bank_id,
contents=[{
"content": content,
"document_id": document_id,
"tags": ["team-b", "important"],
}],
request_context=request_context,
)
async with pool.acquire() as conn:
v2_tags = await conn.fetch(
"SELECT tags FROM memory_units WHERE bank_id = $1 AND document_id = $2",
bank_id,
document_id,
)
for row in v2_tags:
assert "team-b" in row["tags"], f"v2 units should have team-b tag, got {row['tags']}"
assert "important" in row["tags"], f"v2 units should have important tag, got {row['tags']}"
finally:
await memory.delete_bank(bank_id, request_context=request_context)
# ============================================================
# Chunk Management Tests
# ============================================================
@pytest.mark.asyncio
async def test_delta_retain_removed_chunks_delete_facts(memory, request_context):
"""
When content is shortened (chunks removed), facts from the removed
chunks should be deleted.
"""
bank_id = f"test_delta_removed_{_ts()}"
document_id = "shrinking-doc"
try:
# v1: longer content with facts about Alice and Bob
v1_content = (
"Alice is a senior engineer at Google Cloud. "
"She leads the infrastructure team and has been there for 5 years.\n\n"
"Bob is a product manager at Facebook Reality Labs. "
"He previously worked at Amazon on Alexa voice products."
)
v1_units = await memory.retain_async(
bank_id=bank_id,
content=v1_content,
context="profiles",
document_id=document_id,
request_context=request_context,
)
assert len(v1_units) > 0
doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context)
v1_count = doc_v1["memory_unit_count"]
# v2: Completely different content — all chunks change
v2_content = "Charlie works at Netflix as a data scientist."
v2_units = await memory.retain_async(
bank_id=bank_id,
content=v2_content,
context="profiles",
document_id=document_id,
request_context=request_context,
)
doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context)
assert doc_v2 is not None
# Should have facts about Charlie
result = await memory.recall_async(
bank_id=bank_id,
query="Who works at Netflix?",
budget=Budget.MID,
max_tokens=2000,
request_context=request_context,
)
all_texts = " ".join(r.text.lower() for r in result.results)
assert "charlie" in all_texts or "netflix" in all_texts, (
f"Should have facts about Charlie/Netflix after replacing content, got: {all_texts}"
)
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_delta_retain_chunks_have_content_hash(memory, request_context):
"""
After retain, chunks should have content_hash populated.
"""
bank_id = f"test_delta_hash_{_ts()}"
document_id = "hash-doc"
try:
content = "Alice works at Google as a software engineer."
await memory.retain_async(
bank_id=bank_id,
content=content,
document_id=document_id,
request_context=request_context,
)
pool = await memory._get_pool()
async with pool.acquire() as conn:
chunks = await conn.fetch(
"SELECT chunk_id, content_hash FROM chunks WHERE document_id = $1 AND bank_id = $2",
document_id,
bank_id,
)
assert len(chunks) > 0, "Should have stored chunks"
for chunk in chunks:
assert chunk["content_hash"] is not None, f"Chunk {chunk['chunk_id']} should have content_hash"
assert len(chunk["content_hash"]) == 64, "content_hash should be SHA256 hex (64 chars)"
finally:
await memory.delete_bank(bank_id, request_context=request_context)
# ============================================================
# Backward Compatibility Tests
# ============================================================
@pytest.mark.asyncio
async def test_retain_without_document_id_still_works(memory, request_context):
"""
Retain without document_id should still work normally (no delta path).
"""
bank_id = f"test_no_docid_{_ts()}"
try:
units = await memory.retain_async(
bank_id=bank_id,
content="Alice works at Google.",
context="test",
request_context=request_context,
)
assert len(units) > 0, "Should create facts without document_id"
result = await memory.recall_async(
bank_id=bank_id,
query="Where does Alice work?",
budget=Budget.MID,
max_tokens=1000,
request_context=request_context,
)
assert len(result.results) > 0
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_delta_retain_first_retain_full_path(memory, request_context):
"""
First retain of a new document should use the full path (no delta possible).
"""
bank_id = f"test_first_retain_{_ts()}"
document_id = "new-doc"
try:
units = await memory.retain_async(
bank_id=bank_id,
content="Alice works at Google.",
context="test",
document_id=document_id,
request_context=request_context,
)
assert len(units) > 0, "First retain should create facts via full path"
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
assert doc is not None
assert doc["memory_unit_count"] > 0
finally:
await memory.delete_bank(bank_id, request_context=request_context)
# ============================================================
# Edge Cases
# ============================================================
@pytest.mark.asyncio
async def test_delta_retain_empty_to_content(memory, request_context):
"""
Going from gibberish (zero facts) to real content should work.
"""
bank_id = f"test_delta_empty_{_ts()}"
document_id = "empty-to-content"
try:
# v1: content that probably produces zero facts
await memory.retain_async(
bank_id=bank_id,
content="!!!###$$$%%%",
document_id=document_id,
request_context=request_context,
)
doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context)
assert doc_v1 is not None
# v2: real content
v2_units = await memory.retain_async(
bank_id=bank_id,
content="Alice works at Google as a senior engineer.",
document_id=document_id,
request_context=request_context,
)
doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context)
assert doc_v2 is not None
assert doc_v2["memory_unit_count"] > 0 or len(v2_units) > 0, "Should have facts after updating with real content"
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_delta_retain_multiple_upserts(memory, request_context):
"""
Multiple sequential upserts should work correctly, with delta optimization
kicking in after the first retain.
"""
bank_id = f"test_delta_multi_{_ts()}"
document_id = "multi-upsert"
try:
# v1: initial
v1_content = "Alice works at Google."
await memory.retain_async(
bank_id=bank_id,
content=v1_content,
document_id=document_id,
request_context=request_context,
)
# v2: same content (delta: no changes)
await memory.retain_async(
bank_id=bank_id,
content=v1_content,
document_id=document_id,
request_context=request_context,
)
# v3: append
v3_content = v1_content + "\n\nBob works at Microsoft."
await memory.retain_async(
bank_id=bank_id,
content=v3_content,
document_id=document_id,
request_context=request_context,
)
# v4: same as v3 (delta: no changes again)
await memory.retain_async(
bank_id=bank_id,
content=v3_content,
document_id=document_id,
request_context=request_context,
)
# Final check: should have facts about both Alice and Bob
result = await memory.recall_async(
bank_id=bank_id,
query="Who works where?",
budget=Budget.MID,
max_tokens=2000,
request_context=request_context,
)
all_texts = " ".join(r.text.lower() for r in result.results)
assert "alice" in all_texts or "google" in all_texts, f"Should have Alice/Google facts, got: {all_texts}"
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
assert doc is not None
assert doc["memory_unit_count"] > 0
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_delta_retain_with_user_entities(memory, request_context):
"""
User-provided entities should work correctly with delta retain.
"""
bank_id = f"test_delta_user_entities_{_ts()}"
document_id = "user-entity-doc"
try:
content = "The project is going well."
# v1 with user entities
await memory.retain_batch_async(
bank_id=bank_id,
contents=[{
"content": content,
"document_id": document_id,
"entities": [{"text": "Project Alpha", "type": "PROJECT"}],
}],
request_context=request_context,
)
pool = await memory._get_pool()
async with pool.acquire() as conn:
v1_entities = await conn.fetch(
"SELECT canonical_name FROM entities WHERE bank_id = $1",
bank_id,
)
v1_names = {e["canonical_name"].lower() for e in v1_entities}
# v2 with additional entity, same content
# Note: same content = delta path (no re-extraction)
# The user entities for NEW chunks only get processed
v2_content = content + "\n\nThe timeline is on track for Q2 delivery."
await memory.retain_batch_async(
bank_id=bank_id,
contents=[{
"content": v2_content,
"document_id": document_id,
"entities": [
{"text": "Project Alpha", "type": "PROJECT"},
{"text": "Q2 Deadline", "type": "MILESTONE"},
],
}],
request_context=request_context,
)
# Should have entities from both v1 and v2
async with pool.acquire() as conn:
v2_entities = await conn.fetch(
"SELECT canonical_name FROM entities WHERE bank_id = $1",
bank_id,
)
v2_names = {e["canonical_name"].lower() for e in v2_entities}
# v1 entities should be preserved
assert v1_names.issubset(v2_names), f"v1 entities should be preserved: {v1_names} not in {v2_names}"
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_delta_retain_recall_with_chunks(memory, request_context):
"""
After delta retain, recall with include_chunks should return correct chunk data.
"""
bank_id = f"test_delta_recall_chunks_{_ts()}"
document_id = "recall-chunks-doc"
try:
content = "Alice is a senior engineer at Google Cloud. She designs distributed systems."
await memory.retain_async(
bank_id=bank_id,
content=content,
context="profile",
document_id=document_id,
request_context=request_context,
)
# Upsert with same content (delta: no changes)
await memory.retain_async(
bank_id=bank_id,
content=content,
context="profile",
document_id=document_id,
request_context=request_context,
)
# Recall with chunks
result = await memory.recall_async(
bank_id=bank_id,
query="What does Alice do?",
budget=Budget.MID,
max_tokens=2000,
include_chunks=True,
max_chunk_tokens=8192,
request_context=request_context,
)
assert len(result.results) > 0, "Should recall facts"
# Facts with chunk_ids should have corresponding chunks
facts_with_chunks = [r for r in result.results if r.chunk_id]
if facts_with_chunks and result.chunks:
for fact in facts_with_chunks:
assert fact.chunk_id in result.chunks, (
f"Chunk {fact.chunk_id} should be in returned chunks"
)
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@@ -141,6 +141,70 @@ async def test_memory_without_document(memory, request_context):
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_document_metadata_from_retain_params(memory, request_context):
"""Test that document_metadata is returned from retain_params.metadata in both get and list."""
bank_id = f"test_doc_meta_{datetime.now(timezone.utc).timestamp()}"
try:
document_id = "doc-with-metadata"
metadata = {"source": "slack", "channel": "#general"}
await memory.retain_batch_async(
bank_id=bank_id,
contents=[{"content": "Alice works at Google.", "context": "Team meeting", "metadata": metadata}],
document_id=document_id,
request_context=request_context,
)
# get_document should include document_metadata
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
assert doc is not None
assert doc["document_metadata"] == metadata
assert doc["retain_params"] is not None
assert doc["retain_params"]["metadata"] == metadata
# list_documents should also include document_metadata
docs_list = await memory.list_documents(
bank_id=bank_id, search_query=None, limit=100, offset=0, request_context=request_context
)
listed_doc = next(d for d in docs_list["items"] if d["id"] == document_id)
assert listed_doc["document_metadata"] == metadata
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_document_without_metadata(memory, request_context):
"""Test that document_metadata is None when no metadata was provided during retain."""
bank_id = f"test_doc_no_meta_{datetime.now(timezone.utc).timestamp()}"
try:
document_id = "doc-no-metadata"
await memory.retain_async(
bank_id=bank_id,
content="Bob works at Microsoft.",
context="Meeting",
document_id=document_id,
request_context=request_context,
)
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
assert doc is not None
assert doc["document_metadata"] is None
docs_list = await memory.list_documents(
bank_id=bank_id, search_query=None, limit=100, offset=0, request_context=request_context
)
listed_doc = next(d for d in docs_list["items"] if d["id"] == document_id)
assert listed_doc["document_metadata"] is None
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_document_persisted_with_zero_facts(memory, request_context):
"""
@@ -0,0 +1,59 @@
"""
Regression test for experience fact_type preservation.
The LLM extraction layer normalizes raw "assistant" "experience" early in parsing.
The subsequent conversion to ExtractedFactType must pass through the already-normalized
fact_type rather than re-checking for "assistant" (which would remap experience world).
See: https://github.com/vectorize-io/hindsight/pull/839
"""
from datetime import datetime, timezone
from unittest.mock import AsyncMock, patch
import pytest
from hindsight_api.config import _get_raw_config
from hindsight_api.engine.response_models import TokenUsage
from hindsight_api.engine.retain.fact_extraction import (
Fact,
RetainContent,
extract_facts_from_contents,
extract_facts_from_contents_batch_api,
)
@pytest.mark.asyncio
async def test_extract_facts_preserves_experience_type():
"""
When extract_facts_from_text returns a Fact with fact_type="experience",
extract_facts_from_contents must preserve it (not remap to "world").
"""
contents = [
RetainContent(
content="I fixed the failing tests after discovering they mocked the wrong interface.",
event_date=datetime(2026, 4, 1, tzinfo=timezone.utc),
context="assistant work log",
)
]
extracted_fact = Fact(
fact="Fixed the failing tests after discovering they mocked the wrong interface.",
fact_type="experience",
)
with patch(
"hindsight_api.engine.retain.fact_extraction.extract_facts_from_text",
new=AsyncMock(return_value=([extracted_fact], [(contents[0].content, 1)], TokenUsage())),
):
facts, _chunks, _usage = await extract_facts_from_contents(
contents=contents,
llm_config=None,
agent_name="TestAgent",
config=_get_raw_config(),
)
assert len(facts) == 1
assert facts[0].fact_type == "experience", (
f"Expected 'experience' but got '{facts[0].fact_type}'"
f"the conversion layer is remapping the already-normalized fact_type"
)
@@ -303,7 +303,6 @@ class TestOperationHooksParameters:
contents=contents,
document_id=document_id,
fact_type_override="world",
confidence_score=0.9,
request_context=ctx,
)
@@ -317,7 +316,6 @@ class TestOperationHooksParameters:
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
@@ -334,7 +332,6 @@ class TestOperationHooksParameters:
contents=contents,
document_id=document_id,
fact_type_override="experience",
confidence_score=0.8,
request_context=ctx,
)
@@ -345,7 +342,6 @@ class TestOperationHooksParameters:
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
@@ -0,0 +1,131 @@
"""
Test that first-person agent experiences are classified as 'experience' fact_type,
not 'world'. This is critical for AI agent systems that store their own operational
experiences (debugging, code changes, user interactions) separately from world knowledge.
"""
from datetime import datetime
import pytest
from hindsight_api import LLMConfig
from hindsight_api.config import _get_raw_config
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
class TestAgentExperienceClassification:
"""Tests that first-person coding agent experiences get classified as 'experience'."""
@pytest.mark.asyncio
async def test_code_changes_classified_as_experience(self):
"""First-person code change descriptions should be experience, not world."""
text = """
I changed the return type of the `process_request` function from `dict` to `ResponseModel`.
After that, I updated the three callers in `api/handlers.py` to destructure the new model fields.
The type checker was happy after the change but I noticed one test was still using the old dict keys.
"""
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
event_date=datetime(2025, 3, 28),
llm_config=llm_config,
agent_name="coding-agent",
context="agent work log",
config=_get_raw_config(),
)
assert len(facts) > 0, "Should extract at least one fact"
world_facts = [f for f in facts if f.fact_type == "world"]
experience_facts = [f for f in facts if f.fact_type == "experience"]
assert len(experience_facts) > len(world_facts), (
f"First-person code changes should be mostly 'experience', "
f"got {len(experience_facts)} experience vs {len(world_facts)} world. "
f"Facts: {[(f.fact, f.fact_type) for f in facts]}"
)
@pytest.mark.asyncio
async def test_debugging_session_classified_as_experience(self):
"""First-person debugging narrative should be experience, not world."""
text = """
The tests were failing with a ConnectionRefusedError on the Redis integration suite.
I traced it to the connection pool not being initialized before the first test ran.
I added a setup fixture that ensures the pool is warmed up, and all 47 tests pass now.
"""
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
event_date=datetime(2025, 3, 28),
llm_config=llm_config,
agent_name="coding-agent",
context="agent work log",
config=_get_raw_config(),
)
assert len(facts) > 0, "Should extract at least one fact"
world_facts = [f for f in facts if f.fact_type == "world"]
experience_facts = [f for f in facts if f.fact_type == "experience"]
assert len(experience_facts) > len(world_facts), (
f"First-person debugging should be mostly 'experience', "
f"got {len(experience_facts)} experience vs {len(world_facts)} world. "
f"Facts: {[(f.fact, f.fact_type) for f in facts]}"
)
@pytest.mark.asyncio
async def test_user_interaction_classified_as_experience(self):
"""Agent describing interactions with the user should be experience."""
text = """
The user asked me to refactor the authentication middleware to support JWT tokens.
I proposed splitting it into two modules: token_validation.py and session_management.py.
The user approved my approach and I started with the token validation logic.
I discovered that the existing tests were mocking the wrong interface, so I had to rewrite them first.
"""
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
event_date=datetime(2025, 3, 28),
llm_config=llm_config,
agent_name="coding-agent",
context="agent work log",
config=_get_raw_config(),
)
assert len(facts) > 0, "Should extract at least one fact"
world_facts = [f for f in facts if f.fact_type == "world"]
experience_facts = [f for f in facts if f.fact_type == "experience"]
assert len(experience_facts) > len(world_facts), (
f"Agent-user interactions should be mostly 'experience', "
f"got {len(experience_facts)} experience vs {len(world_facts)} world. "
f"Facts: {[(f.fact, f.fact_type) for f in facts]}"
)
@pytest.mark.asyncio
async def test_mixed_agent_and_world_facts(self):
"""Mix of agent experiences and world knowledge should be classified correctly."""
text = """
Python 3.12 introduced a new type parameter syntax for generic classes.
I migrated our codebase from the old TypeVar approach to the new syntax.
The migration touched 23 files but was mostly mechanical.
PEP 695 defines the new type statement that makes generics more readable.
"""
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
event_date=datetime(2025, 3, 28),
llm_config=llm_config,
agent_name="coding-agent",
context="agent work log",
config=_get_raw_config(),
)
assert len(facts) > 0, "Should extract at least one fact"
world_facts = [f for f in facts if f.fact_type == "world"]
experience_facts = [f for f in facts if f.fact_type == "experience"]
# Should have both types - world facts about Python 3.12/PEP 695,
# experience facts about the migration work
assert len(world_facts) >= 1, (
f"Should have at least 1 world fact about Python 3.12/PEP 695. "
f"Facts: {[(f.fact, f.fact_type) for f in facts]}"
)
assert len(experience_facts) >= 1, (
f"Should have at least 1 experience fact about the migration. "
f"Facts: {[(f.fact, f.fact_type) for f in facts]}"
)
@@ -38,7 +38,7 @@ I ran into my neighbor Sarah who mentioned she's planning a trip to Italy next m
"""
context = "Personal diary entry"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -83,7 +83,7 @@ User: Perfect, I'll make a reservation for Saturday at 7pm.
"""
context = "Restaurant recommendation conversation"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -140,7 +140,7 @@ I edited about 20 photos from my recent trip to the mountains.
"""
context = "Personal blog post"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -205,7 +205,7 @@ I edited about 20 photos from my recent trip to the mountains.
text = "\n".join([f"{turn['speaker']}: {turn['text']}" for turn in session])
context = f"Conversation between {data['conversation']['speaker_a']} and {data['conversation']['speaker_b']}"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -267,7 +267,7 @@ I'm planning to visit Japan next year.
"""
context = "Personal info"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -42,7 +42,7 @@ Marcus felt anxious about the upcoming interview.
"""
context = "Personal journal entry"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -75,7 +75,7 @@ The music was so loud I could barely hear myself think.
"""
context = "Personal experience"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -109,7 +109,7 @@ Maybe we should reconsider the timeline.
"""
context = "Team discussion"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -143,7 +143,7 @@ I'm unable to attend the conference due to scheduling conflicts.
"""
context = "Personal profile discussion"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -176,7 +176,7 @@ Unlike last year, we're ahead of schedule.
"""
context = "Project review"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -210,7 +210,7 @@ She's enthusiastic about the opportunity.
"""
context = "Team meeting"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -244,7 +244,7 @@ I'm planning to switch careers because I'm not fulfilled in my current role.
"""
context = "Personal goals discussion"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -282,7 +282,7 @@ Family is the most important thing to her.
"""
context = "Personal values discussion"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -315,7 +315,7 @@ I prefer presenting in person rather than virtually because I can read the room
"""
context = "Personal reflection"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
event_date = datetime(2024, 11, 13)
@@ -373,7 +373,7 @@ I'm planning to visit Tokyo next month.
"""
context = "Personal conversation"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
event_date = datetime(2024, 11, 13)
@@ -421,7 +421,7 @@ with a concert surrounded by music, joy and the warm summer breeze.
"""
context = "Conversation between Melanie and Caroline"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
event_date = datetime(2023, 8, 14, 14, 24)
last_error = None
@@ -496,7 +496,7 @@ It was a beautiful day and I plan to make this a regular habit.
"""
context = "Personal diary"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
event_date = datetime(2024, 11, 13)
@@ -547,7 +547,7 @@ It was a beautiful day and I plan to make this a regular habit.
"""Test that relative dates are converted to absolute dates."""
reference_date = datetime(2024, 3, 20, 14, 0, 0, tzinfo=UTC)
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
text = """
Yesterday I went hiking in Yosemite.
@@ -582,7 +582,7 @@ It was a beautiful day and I plan to make this a regular habit.
"""Test that facts without temporal info are still extracted."""
reference_date = datetime(2024, 3, 20, 14, 0, 0, tzinfo=UTC)
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
text = "Alice works at Google. She loves Python programming."
@@ -607,7 +607,7 @@ It was a beautiful day and I plan to make this a regular habit.
"""Test that absolute dates in text are preserved."""
reference_date = datetime(2024, 3, 20, 14, 0, 0, tzinfo=UTC)
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
text = """
On March 15, 2024, Alice joined Google.
@@ -662,7 +662,7 @@ great time! Every time I see it, I can't help but smile.
"""
context = "Conversation between Deborah and Jolene"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
event_date = datetime(2023, 2, 23)
@@ -715,7 +715,7 @@ I've learned so much from it.
"""
context = "Personal update"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -785,7 +785,7 @@ Jamie: Congratulations! I'd love to read it.
context = "Podcast episode between you (Marcus) and Jamie discussing AI research"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=transcript,
@@ -831,7 +831,7 @@ We presented our findings to the team yesterday.
context = "Personal work log"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=text,
@@ -867,7 +867,7 @@ Jamie: [teasing] We'll see who's right, my Niners pick is solid.
context = "podcast episode on match prediction of week 10 - Marcus (you) and Jamie - 14 nov"
agent_name = "Marcus"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
facts, _, _ = await extract_facts_from_text(
text=transcript,
@@ -929,7 +929,7 @@ so the algorithm learns to box out. See you next week!
context = "Podcast episode between you (Marcus) and Jamie about AI"
llm_config = LLMConfig.for_memory()
llm_config = LLMConfig.from_env()
max_retries = 3
last_error = None
@@ -139,3 +139,76 @@ async def test_retain_llm_max_retries_overrides_global():
assert facts == []
# Verify it retried exactly retain_llm_max_retries times
assert llm_config.call.call_count == 5
@pytest.mark.asyncio
async def test_none_event_date_with_empty_facts_no_crash():
"""
When event_date is None and the LLM returns an empty facts list,
the debug log should not crash with AttributeError on .isoformat().
Regression test for https://github.com/vectorize-io/hindsight/issues/874
"""
from hindsight_api.engine.retain.fact_extraction import _extract_facts_from_chunk
config = _make_config(llm_max_retries=1)
# LLM returns a valid dict but with no facts — triggers the debug log path
llm_config = _make_llm_config(mock_response={"facts": []})
with patch(
"hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
return_value=("system prompt", MagicMock()),
):
facts, usage = await _extract_facts_from_chunk(
chunk="A plain text document with no timestamp.",
chunk_index=0,
total_chunks=1,
event_date=None,
context="",
llm_config=llm_config,
config=config,
agent_name="test-agent",
)
assert facts == []
@pytest.mark.asyncio
async def test_none_event_date_with_valid_facts_no_crash():
"""
When event_date is None but the LLM returns valid facts,
extraction should succeed without errors.
"""
from hindsight_api.engine.retain.fact_extraction import _extract_facts_from_chunk
config = _make_config(llm_max_retries=1)
llm_config = _make_llm_config(mock_response={
"facts": [
{
"what": "Alice visited Paris",
"when": "2023",
"who": "Alice",
"why": "vacation",
}
]
})
with patch(
"hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
return_value=("system prompt", MagicMock()),
):
facts, usage = await _extract_facts_from_chunk(
chunk="Alice visited Paris in 2023.",
chunk_index=0,
total_chunks=1,
event_date=None,
context="",
llm_config=llm_config,
config=config,
agent_name="test-agent",
)
assert len(facts) == 1
assert "Alice visited Paris" in facts[0].fact
@@ -0,0 +1,336 @@
"""
Tests for Google embeddings implementation (Gemini API + Vertex AI).
These tests cover:
1. Initialization (Gemini API key, Vertex AI with ADC/service account)
2. Dimension detection via test embedding
3. Output dimensionality configuration
4. Encode (single text, multiple texts, batching, empty list, uninitialized)
5. Provider name and model name normalization
6. Factory function (create from env, validation errors)
"""
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from hindsight_api.config import (
ENV_EMBEDDINGS_GEMINI_API_KEY,
ENV_EMBEDDINGS_PROVIDER,
HindsightConfig,
)
from hindsight_api.engine.embeddings import GeminiEmbeddings, create_embeddings_from_env
def _make_mock_embedding(values: list[float]) -> MagicMock:
emb = MagicMock()
emb.values = values
return emb
def _make_mock_embed_result(embeddings_data: list[list[float]]) -> MagicMock:
result = MagicMock()
result.embeddings = [_make_mock_embedding(v) for v in embeddings_data]
return result
def _make_mock_genai(embed_result: Any = None) -> MagicMock:
if embed_result is None:
embed_result = _make_mock_embed_result([[0.1] * 768])
mock_genai = MagicMock()
mock_client = MagicMock()
mock_client.models.embed_content = MagicMock(return_value=embed_result)
mock_genai.Client = MagicMock(return_value=mock_client)
return mock_genai
def _make_mock_google_module(mock_genai: MagicMock) -> MagicMock:
mod = MagicMock()
mod.genai = mock_genai
mod.genai.types.EmbedContentConfig = MagicMock(side_effect=lambda **kw: MagicMock(**kw))
return mod
def _patch_google_import(mock_genai: MagicMock):
original_import = __import__
def mock_import(name, *args, **kwargs):
if name == "google":
return _make_mock_google_module(mock_genai)
if name == "google.genai":
return mock_genai
return original_import(name, *args, **kwargs)
return patch("builtins.__import__", side_effect=mock_import)
class TestGeminiEmbeddings:
"""Unit tests for GeminiEmbeddings with mocked google.genai."""
async def test_initialization_api_key_success(self):
"""Test successful Gemini API key initialization."""
mock_genai = _make_mock_genai()
emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
with _patch_google_import(mock_genai):
await emb.initialize()
assert emb._client is not None
assert emb.dimension == 768
assert emb.provider_name == "google"
assert emb._is_vertexai is False
mock_genai.Client.return_value.models.embed_content.assert_called_once()
async def test_initialization_vertexai_success(self):
"""Test successful Vertex AI initialization."""
mock_genai = _make_mock_genai()
emb = GeminiEmbeddings(
model="gemini-embedding-001",
vertexai_project_id="test-project",
vertexai_region="us-central1",
)
with _patch_google_import(mock_genai):
await emb.initialize()
assert emb._client is not None
assert emb.dimension == 768
assert emb.provider_name == "google"
assert emb._is_vertexai is True
mock_genai.Client.assert_called_once_with(
vertexai=True,
project="test-project",
location="us-central1",
)
async def test_initialization_missing_api_key(self):
"""Test that missing API key raises ValueError when no vertexai_project_id."""
mock_genai = _make_mock_genai()
emb = GeminiEmbeddings(model="gemini-embedding-001", api_key=None)
with _patch_google_import(mock_genai):
with pytest.raises(ValueError, match="requires an API key"):
await emb.initialize()
async def test_initialization_vertexai_missing_project_id(self):
"""Test that Vertex AI mode requires project_id."""
mock_genai = _make_mock_genai()
emb = GeminiEmbeddings(model="gemini-embedding-001", vertexai_project_id="temp")
emb.vertexai_project_id = None # Simulate misconfiguration
with _patch_google_import(mock_genai):
with pytest.raises(ValueError, match="is required for Vertex AI"):
await emb.initialize()
async def test_initialization_idempotent(self):
"""Test that calling initialize() twice is a no-op."""
mock_genai = _make_mock_genai()
emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
with _patch_google_import(mock_genai):
await emb.initialize()
first_client = emb._client
await emb.initialize()
assert emb._client is first_client
async def test_dimension_detection_via_test_embedding(self):
"""Test that dimension is detected via a test embedding call."""
test_embed = _make_mock_embed_result([[0.5] * 256])
mock_genai = _make_mock_genai(embed_result=test_embed)
emb = GeminiEmbeddings(model="some-new-model", api_key="test-key")
with _patch_google_import(mock_genai):
await emb.initialize()
assert emb.dimension == 256
async def test_output_dimensionality(self):
"""Test that output_dimensionality is passed via EmbedContentConfig."""
test_embed = _make_mock_embed_result([[0.1] * 256])
mock_genai = _make_mock_genai(embed_result=test_embed)
emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key", output_dimensionality=256)
with _patch_google_import(mock_genai):
await emb.initialize()
assert emb.dimension == 256
assert emb._embed_config is not None
call_kwargs = mock_genai.Client.return_value.models.embed_content.call_args
assert "config" in call_kwargs.kwargs
async def test_no_output_dimensionality(self):
"""Test that no EmbedContentConfig is built when output_dimensionality is None."""
mock_genai = _make_mock_genai()
emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key", output_dimensionality=None)
with _patch_google_import(mock_genai):
await emb.initialize()
assert emb._embed_config is None
call_kwargs = mock_genai.Client.return_value.models.embed_content.call_args
assert "config" not in call_kwargs.kwargs
def test_auto_detect_vertexai(self):
"""Test that _is_vertexai is auto-detected from vertexai_project_id."""
assert GeminiEmbeddings(model="m", api_key="k")._is_vertexai is False
assert GeminiEmbeddings(model="m", vertexai_project_id="p")._is_vertexai is True
def test_encode_single_text(self):
emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
mock_client = MagicMock()
mock_client.models.embed_content = MagicMock(return_value=_make_mock_embed_result([[0.1, 0.2, 0.3]]))
emb._client = mock_client
emb._dimension = 3
assert emb.encode(["hello"]) == [[0.1, 0.2, 0.3]]
def test_encode_multiple_texts(self):
emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
mock_client = MagicMock()
mock_client.models.embed_content = MagicMock(
return_value=_make_mock_embed_result([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])
)
emb._client = mock_client
emb._dimension = 2
result = emb.encode(["a", "b", "c"])
assert len(result) == 3
assert result[1] == [0.3, 0.4]
def test_encode_batching(self):
emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key", batch_size=2)
mock_client = MagicMock()
mock_client.models.embed_content = MagicMock(
side_effect=[_make_mock_embed_result([[0.1], [0.2]]), _make_mock_embed_result([[0.3]])]
)
emb._client = mock_client
emb._dimension = 1
assert emb.encode(["a", "b", "c"]) == [[0.1], [0.2], [0.3]]
assert mock_client.models.embed_content.call_count == 2
def test_encode_passes_config(self):
emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
mock_client = MagicMock()
mock_client.models.embed_content = MagicMock(return_value=_make_mock_embed_result([[0.1, 0.2]]))
emb._client = mock_client
emb._dimension = 2
emb._embed_config = MagicMock()
emb.encode(["hello"])
assert mock_client.models.embed_content.call_args.kwargs["config"] is emb._embed_config
def test_encode_empty_list(self):
emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
emb._client = MagicMock()
emb._dimension = 768
assert emb.encode([]) == []
def test_encode_before_initialization(self):
emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
with pytest.raises(RuntimeError, match="not initialized"):
emb.encode(["test"])
def test_dimension_before_initialization(self):
emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
with pytest.raises(RuntimeError, match="not initialized"):
_ = emb.dimension
def test_provider_name_always_google(self):
assert GeminiEmbeddings(model="m", api_key="k").provider_name == "google"
assert GeminiEmbeddings(model="m", vertexai_project_id="p").provider_name == "google"
def test_vertexai_strips_google_prefix(self):
mock_genai = _make_mock_genai()
emb = GeminiEmbeddings(model="google/gemini-embedding-001", vertexai_project_id="test-project")
emb._init_vertexai(mock_genai)
assert emb.model == "gemini-embedding-001"
def test_default_region(self):
emb = GeminiEmbeddings(model="m", vertexai_project_id="proj")
assert emb.vertexai_region == "us-central1"
def test_custom_region(self):
emb = GeminiEmbeddings(model="m", vertexai_project_id="proj", vertexai_region="europe-west1")
assert emb.vertexai_region == "europe-west1"
class TestGeminiEmbeddingsFactory:
"""Tests for create_embeddings_from_env() with 'google' provider."""
def _make_config(self, **overrides) -> HindsightConfig:
from dataclasses import fields
defaults = {}
for f in fields(HindsightConfig):
if f.type == "str":
defaults[f.name] = ""
elif f.type == "str | None":
defaults[f.name] = None
elif f.type == "int":
defaults[f.name] = 0
elif f.type == "int | None":
defaults[f.name] = None
elif f.type == "float":
defaults[f.name] = 0.0
elif f.type == "float | None":
defaults[f.name] = None
elif f.type == "bool":
defaults[f.name] = False
elif f.type == "list | None":
defaults[f.name] = None
else:
defaults[f.name] = None
defaults["embeddings_provider"] = "google"
defaults["embeddings_gemini_api_key"] = "test-key"
defaults["embeddings_gemini_model"] = "gemini-embedding-001"
defaults["embeddings_gemini_output_dimensionality"] = 768
defaults["embeddings_vertexai_project_id"] = None
defaults["embeddings_vertexai_region"] = None
defaults["embeddings_vertexai_service_account_key"] = None
defaults.update(overrides)
return HindsightConfig(**defaults)
def test_create_with_api_key(self):
config = self._make_config()
with patch("hindsight_api.config.get_config", return_value=config):
emb = create_embeddings_from_env()
assert isinstance(emb, GeminiEmbeddings)
assert emb.provider_name == "google"
assert emb.api_key == "test-key"
assert emb._is_vertexai is False
def test_create_with_vertexai(self):
config = self._make_config(
embeddings_gemini_api_key=None,
embeddings_vertexai_project_id="my-project",
embeddings_vertexai_region="us-east1",
)
with patch("hindsight_api.config.get_config", return_value=config):
emb = create_embeddings_from_env()
assert isinstance(emb, GeminiEmbeddings)
assert emb._is_vertexai is True
assert emb.api_key is None
assert emb.vertexai_project_id == "my-project"
def test_create_missing_all_credentials(self):
config = self._make_config(embeddings_gemini_api_key=None, embeddings_vertexai_project_id=None)
with patch("hindsight_api.config.get_config", return_value=config):
with pytest.raises(ValueError, match="is required"):
create_embeddings_from_env()
def test_vertexai_takes_priority(self):
config = self._make_config(embeddings_gemini_api_key="key", embeddings_vertexai_project_id="proj")
with patch("hindsight_api.config.get_config", return_value=config):
emb = create_embeddings_from_env()
assert emb._is_vertexai is True
assert emb.api_key is None
def test_create_with_custom_dimensionality(self):
config = self._make_config(embeddings_gemini_output_dimensionality=256)
with patch("hindsight_api.config.get_config", return_value=config):
emb = create_embeddings_from_env()
assert emb.output_dimensionality == 256
@@ -0,0 +1,275 @@
"""
Tests for Google Discovery Engine cross-encoder (Ranking REST API).
These tests cover:
1. Initialization (service account, ADC, missing project_id)
2. Predict (single query, multiple queries, batching, empty pairs, uninitialized)
3. Provider name
4. Factory function (create from env, validation errors)
"""
from unittest.mock import MagicMock, patch
import httpx
import pytest
from hindsight_api.config import (
ENV_RERANKER_GOOGLE_PROJECT_ID,
ENV_RERANKER_PROVIDER,
HindsightConfig,
)
from hindsight_api.engine.cross_encoder import GoogleCrossEncoder, create_cross_encoder_from_env
def _make_rank_response(records: list[tuple[str, float]]) -> dict:
"""Build a JSON response matching the Discovery Engine REST API format."""
return {"records": [{"id": rid, "score": score} for rid, score in records]}
def _make_mock_httpx_client(responses: list[dict] | None = None) -> MagicMock:
"""Create a mock httpx.Client that returns predefined responses."""
mock_client = MagicMock(spec=httpx.Client)
if responses:
side_effects = []
for resp_json in responses:
mock_resp = MagicMock(spec=httpx.Response)
mock_resp.json.return_value = resp_json
mock_resp.raise_for_status.return_value = None
side_effects.append(mock_resp)
mock_client.post.side_effect = side_effects
return mock_client
def _make_mock_credentials() -> MagicMock:
"""Create mock credentials with a valid token."""
creds = MagicMock()
creds.valid = True
creds.token = "mock-token"
return creds
class TestGoogleCrossEncoder:
"""Unit tests for GoogleCrossEncoder with mocked httpx + google-auth."""
async def test_initialization_adc_success(self):
"""Test successful initialization with ADC (no service account key)."""
mock_creds = _make_mock_credentials()
encoder = GoogleCrossEncoder(project_id="test-project")
with patch("google.auth.default", return_value=(mock_creds, "test-project")):
await encoder.initialize()
assert encoder._client is not None
assert encoder._credentials is mock_creds
assert encoder.provider_name == "google"
assert "test-project" in encoder._rank_url
async def test_initialization_service_account(self):
"""Test initialization with service account key."""
mock_creds = _make_mock_credentials()
encoder = GoogleCrossEncoder(
project_id="test-project",
service_account_key="/path/to/key.json",
)
with patch(
"google.oauth2.service_account.Credentials.from_service_account_file",
return_value=mock_creds,
):
await encoder.initialize()
assert encoder._client is not None
assert encoder._credentials is mock_creds
async def test_initialization_idempotent(self):
"""Test that calling initialize() twice is a no-op."""
mock_creds = _make_mock_credentials()
encoder = GoogleCrossEncoder(project_id="test-project")
with patch("google.auth.default", return_value=(mock_creds, "test-project")):
await encoder.initialize()
first_client = encoder._client
await encoder.initialize()
assert encoder._client is first_client
async def test_predict_single_query(self):
"""Test prediction with a single query and multiple documents."""
mock_creds = _make_mock_credentials()
mock_client = _make_mock_httpx_client([
_make_rank_response([("1", 0.95), ("0", 0.30)]),
])
encoder = GoogleCrossEncoder(project_id="test-project")
with patch("google.auth.default", return_value=(mock_creds, "p")):
await encoder.initialize()
encoder._client = mock_client
scores = await encoder.predict([
("What is AI?", "AI is artificial intelligence"),
("What is AI?", "The sky is blue"),
])
assert len(scores) == 2
assert scores[0] == 0.30 # id="0" -> index 0
assert scores[1] == 0.95 # id="1" -> index 1
mock_client.post.assert_called_once()
async def test_predict_multiple_queries(self):
"""Test prediction with multiple distinct queries."""
mock_creds = _make_mock_credentials()
mock_client = _make_mock_httpx_client([
_make_rank_response([("0", 0.9), ("1", 0.1)]),
_make_rank_response([("0", 0.8)]),
])
encoder = GoogleCrossEncoder(project_id="test-project")
with patch("google.auth.default", return_value=(mock_creds, "p")):
await encoder.initialize()
encoder._client = mock_client
scores = await encoder.predict([
("Query A", "Doc A1"),
("Query A", "Doc A2"),
("Query B", "Doc B1"),
])
assert len(scores) == 3
assert scores[0] == 0.9
assert scores[1] == 0.1
assert scores[2] == 0.8
assert mock_client.post.call_count == 2
async def test_predict_empty_pairs(self):
"""Test that empty pairs returns empty list."""
mock_creds = _make_mock_credentials()
encoder = GoogleCrossEncoder(project_id="test-project")
with patch("google.auth.default", return_value=(mock_creds, "p")):
await encoder.initialize()
scores = await encoder.predict([])
assert scores == []
async def test_predict_not_initialized(self):
"""Test that predict raises if not initialized."""
encoder = GoogleCrossEncoder(project_id="test-project")
with pytest.raises(RuntimeError, match="not initialized"):
await encoder.predict([("q", "d")])
async def test_predict_batching(self):
"""Test that >200 records are split into batches."""
mock_creds = _make_mock_credentials()
mock_client = _make_mock_httpx_client([
_make_rank_response([(str(i), 0.5) for i in range(200)]),
_make_rank_response([(str(i), 0.3) for i in range(50)]),
])
encoder = GoogleCrossEncoder(project_id="test-project")
with patch("google.auth.default", return_value=(mock_creds, "p")):
await encoder.initialize()
encoder._client = mock_client
pairs = [("same query", f"doc {i}") for i in range(250)]
scores = await encoder.predict(pairs)
assert len(scores) == 250
assert mock_client.post.call_count == 2
async def test_auth_header_sent(self):
"""Test that Authorization header is sent with requests."""
mock_creds = _make_mock_credentials()
mock_creds.token = "test-bearer-token"
mock_client = _make_mock_httpx_client([
_make_rank_response([("0", 0.9)]),
])
encoder = GoogleCrossEncoder(project_id="test-project")
with patch("google.auth.default", return_value=(mock_creds, "p")):
await encoder.initialize()
encoder._client = mock_client
await encoder.predict([("q", "d")])
call_kwargs = mock_client.post.call_args
assert call_kwargs.kwargs["headers"]["Authorization"] == "Bearer test-bearer-token"
def test_provider_name(self):
assert GoogleCrossEncoder(project_id="p").provider_name == "google"
def test_default_model(self):
encoder = GoogleCrossEncoder(project_id="p")
assert encoder.model == "semantic-ranker-default-004"
def test_custom_model(self):
encoder = GoogleCrossEncoder(project_id="p", model="semantic-ranker-fast-004")
assert encoder.model == "semantic-ranker-fast-004"
def test_default_location(self):
encoder = GoogleCrossEncoder(project_id="p")
assert encoder.location == "global"
class TestGoogleCrossEncoderFactory:
"""Tests for create_cross_encoder_from_env() with 'google' provider."""
def _make_config(self, **overrides) -> HindsightConfig:
from dataclasses import fields
defaults = {}
for f in fields(HindsightConfig):
if f.type == "str":
defaults[f.name] = ""
elif f.type == "str | None":
defaults[f.name] = None
elif f.type == "int":
defaults[f.name] = 0
elif f.type == "int | None":
defaults[f.name] = None
elif f.type == "float":
defaults[f.name] = 0.0
elif f.type == "float | None":
defaults[f.name] = None
elif f.type == "bool":
defaults[f.name] = False
elif f.type == "list | None":
defaults[f.name] = None
else:
defaults[f.name] = None
defaults["reranker_provider"] = "google"
defaults["reranker_google_model"] = "semantic-ranker-default-004"
defaults["reranker_google_project_id"] = "test-project"
defaults["reranker_google_service_account_key"] = None
defaults.update(overrides)
return HindsightConfig(**defaults)
def test_create_with_project_id(self):
config = self._make_config()
with patch("hindsight_api.config.get_config", return_value=config):
encoder = create_cross_encoder_from_env()
assert isinstance(encoder, GoogleCrossEncoder)
assert encoder.provider_name == "google"
assert encoder.project_id == "test-project"
assert encoder.service_account_key is None
def test_create_with_service_account(self):
config = self._make_config(reranker_google_service_account_key="/path/to/key.json")
with patch("hindsight_api.config.get_config", return_value=config):
encoder = create_cross_encoder_from_env()
assert isinstance(encoder, GoogleCrossEncoder)
assert encoder.service_account_key == "/path/to/key.json"
def test_create_missing_project_id(self):
config = self._make_config(reranker_google_project_id=None)
with patch("hindsight_api.config.get_config", return_value=config):
with pytest.raises(ValueError, match="is required"):
create_cross_encoder_from_env()
def test_create_with_custom_model(self):
config = self._make_config(reranker_google_model="semantic-ranker-fast-004")
with patch("hindsight_api.config.get_config", return_value=config):
encoder = create_cross_encoder_from_env()
assert encoder.model == "semantic-ranker-fast-004"
@@ -88,8 +88,17 @@ async def test_hierarchical_fields_categorization():
assert "entities_allow_free_form" in configurable
assert "entity_labels" in configurable
# Verify other configurable fields
assert "retain_default_strategy" in configurable
assert "retain_strategies" in configurable
assert "max_observations_per_scope" in configurable
assert "reflect_source_facts_max_tokens" in configurable
assert "llm_gemini_safety_settings" in configurable
assert "mcp_enabled_tools" in configurable
assert "retain_chunk_batch_size" in configurable
# Verify count is correct
assert len(configurable) == 20
assert len(configurable) == 22
# Verify credential fields (NEVER exposed)
assert "llm_api_key" in credentials
+24 -24
View File
@@ -1,10 +1,10 @@
"""
Tests for per-bank HNSW index lifecycle and UNION ALL retrieval.
Tests for per-bank vector index lifecycle and UNION ALL retrieval.
Covers:
- _hnsw_index_name deterministic naming
- Per-bank HNSW indexes created on bank creation (retain_async / ensure_bank_exists)
- Per-bank HNSW indexes dropped on bank deletion
- _bank_index_name deterministic naming
- Per-bank vector indexes created on bank creation (retain_async / ensure_bank_exists)
- Per-bank vector indexes dropped on bank deletion
- retrieve_semantic_bm25_combined groups results correctly by fact_type and source
"""
import uuid
@@ -12,7 +12,7 @@ from datetime import datetime, timezone
import pytest
from hindsight_api.engine.retain.bank_utils import _HNSW_FACT_TYPES, _hnsw_index_name
from hindsight_api.engine.retain.bank_utils import _BANK_INDEX_FACT_TYPES, _bank_index_name
# ---------------------------------------------------------------------------
@@ -20,36 +20,36 @@ from hindsight_api.engine.retain.bank_utils import _HNSW_FACT_TYPES, _hnsw_index
# ---------------------------------------------------------------------------
class TestHnswIndexName:
class TestBankIndexName:
def test_deterministic(self):
uid = "550e8400-e29b-41d4-a716-446655440000"
assert _hnsw_index_name("world", uid) == _hnsw_index_name("world", uid)
assert _bank_index_name("world", uid) == _bank_index_name("world", uid)
def test_strips_dashes(self):
uid = "550e8400-e29b-41d4-a716-446655440000"
name = _hnsw_index_name("world", uid)
name = _bank_index_name("world", uid)
# uid16 should be hex chars only
assert "-" not in name
def test_uses_first_16_hex_chars(self):
uid = "550e8400-e29b-41d4-a716-446655440000"
uid16 = uid.replace("-", "")[:16] # "550e8400e29b41d4"
assert name_ends_with(name=_hnsw_index_name("world", uid), suffix=uid16)
assert name_ends_with(name=_bank_index_name("world", uid), suffix=uid16)
def test_suffix_per_fact_type(self):
uid = "550e8400-e29b-41d4-a716-446655440000"
names = {ft: _hnsw_index_name(ft, uid) for ft in _HNSW_FACT_TYPES}
names = {ft: _bank_index_name(ft, uid) for ft in _BANK_INDEX_FACT_TYPES}
# All three names must be distinct
assert len(set(names.values())) == 3
def test_all_fact_types_covered(self):
assert set(_HNSW_FACT_TYPES) == {"world", "experience", "observation"}
assert set(_BANK_INDEX_FACT_TYPES) == {"world", "experience", "observation"}
def test_fits_pg_identifier_limit(self):
# PostgreSQL max identifier length is 63 chars
uid = "f" * 32 # simulated UUID without dashes
for ft in _HNSW_FACT_TYPES:
assert len(_hnsw_index_name(ft, uid)) <= 63
for ft in _BANK_INDEX_FACT_TYPES:
assert len(_bank_index_name(ft, uid)) <= 63
def name_ends_with(name: str, suffix: str) -> bool:
@@ -61,7 +61,7 @@ def name_ends_with(name: str, suffix: str) -> bool:
# ---------------------------------------------------------------------------
async def _get_bank_hnsw_indexes(pool, bank_id: str) -> list[str]:
async def _get_bank_vector_indexes(pool, bank_id: str) -> list[str]:
"""Return index names for memory_units that match the per-bank pattern."""
async with pool.acquire() as conn:
rows = await conn.fetch(
@@ -79,8 +79,8 @@ async def _get_bank_hnsw_indexes(pool, bank_id: str) -> list[str]:
@pytest.mark.asyncio
async def test_retain_creates_per_bank_hnsw_indexes(memory, request_context):
"""retain_async on a new bank must create 3 per-(bank, fact_type) HNSW indexes."""
async def test_retain_creates_per_bank_vector_indexes(memory, request_context):
"""retain_async on a new bank must create 3 per-(bank, fact_type) vector indexes."""
bank_id = f"test_hnsw_create_{uuid.uuid4().hex[:8]}"
try:
await memory.retain_async(
@@ -88,9 +88,9 @@ async def test_retain_creates_per_bank_hnsw_indexes(memory, request_context):
content="Alice is a software engineer.",
request_context=request_context,
)
indexes = await _get_bank_hnsw_indexes(memory._pool, bank_id)
assert len(indexes) == 3, f"Expected 3 per-bank HNSW indexes, got: {indexes}"
for ft_short in _HNSW_FACT_TYPES.values():
indexes = await _get_bank_vector_indexes(memory._pool, bank_id)
assert len(indexes) == 3, f"Expected 3 per-bank vector indexes, got: {indexes}"
for ft_short in _BANK_INDEX_FACT_TYPES.values():
assert any(ft_short in idx for idx in indexes), (
f"Missing index for fact_type short '{ft_short}' in {indexes}"
)
@@ -99,8 +99,8 @@ async def test_retain_creates_per_bank_hnsw_indexes(memory, request_context):
@pytest.mark.asyncio
async def test_delete_bank_drops_hnsw_indexes(memory, request_context):
"""delete_bank must drop all per-bank HNSW indexes."""
async def test_delete_bank_drops_vector_indexes(memory, request_context):
"""delete_bank must drop all per-bank vector indexes."""
bank_id = f"test_hnsw_drop_{uuid.uuid4().hex[:8]}"
await memory.retain_async(
@@ -109,12 +109,12 @@ async def test_delete_bank_drops_hnsw_indexes(memory, request_context):
request_context=request_context,
)
# Verify indexes exist before deletion
indexes_before = await _get_bank_hnsw_indexes(memory._pool, bank_id)
indexes_before = await _get_bank_vector_indexes(memory._pool, bank_id)
assert len(indexes_before) == 3
await memory.delete_bank(bank_id, request_context=request_context)
indexes_after = await _get_bank_hnsw_indexes(memory._pool, bank_id)
indexes_after = await _get_bank_vector_indexes(memory._pool, bank_id)
assert indexes_after == [], f"Indexes should be dropped after bank deletion, got: {indexes_after}"
@@ -133,7 +133,7 @@ async def test_retain_idempotent_bank_creation(memory, request_context):
content="Carol joined the company in 2022.",
request_context=request_context,
)
indexes = await _get_bank_hnsw_indexes(memory._pool, bank_id)
indexes = await _get_bank_vector_indexes(memory._pool, bank_id)
assert len(indexes) == 3
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@@ -0,0 +1,353 @@
"""Test observation tracking for a sequence of horse-related memories.
This test retains a series of facts about horses on a farm and inspects
how observations track the evolving state over time, with full prompt debugging.
"""
import json
import uuid
from dataclasses import dataclass, field
from typing import Any
import pytest
from hindsight_api.config import _get_raw_config
from hindsight_api.engine.consolidation import consolidator as consolidator_mod
from hindsight_api.engine.memory_engine import MemoryEngine
@pytest.fixture(autouse=True)
def enable_observations():
"""Enable observations for all tests in this module."""
config = _get_raw_config()
original_value = config.enable_observations
config.enable_observations = True
yield
config.enable_observations = original_value
@dataclass
class _ActionLog:
text: str
source_fact_ids: list[str] = field(default_factory=list)
observation_id: str = ""
@dataclass
class _ConsolidationResponse:
creates: list[_ActionLog] = field(default_factory=list)
updates: list[_ActionLog] = field(default_factory=list)
deletes: list[_ActionLog] = field(default_factory=list)
@dataclass
class _ConsolidationDebugEntry:
facts: str
observations_text: str
response: _ConsolidationResponse
# Store prompts/responses for debugging
_debug_log: list[_ConsolidationDebugEntry] = []
def _fact_line(m: dict[str, Any]) -> str:
text = f"[{m['id']}] {m['text']}"
temporal_parts = []
if m.get("occurred_start"):
temporal_parts.append(f"occurred_start={m['occurred_start']}")
if m.get("occurred_end"):
temporal_parts.append(f"occurred_end={m['occurred_end']}")
if m.get("mentioned_at"):
temporal_parts.append(f"mentioned_at={m['mentioned_at']}")
if temporal_parts:
text += f" ({', '.join(temporal_parts)})"
return text
async def _instrumented_consolidate(
original_fn: Any,
*,
llm_config: Any,
memories: list[dict[str, Any]],
union_observations: Any,
union_source_facts: Any,
config: Any = None,
remaining_observation_slots: int | None = None,
max_observations_per_scope: int = -1,
) -> Any:
"""Wrapper that captures the prompt and response for debugging."""
if union_observations:
obs_list = consolidator_mod._build_observations_for_llm(union_observations, union_source_facts)
observations_text = json.dumps(obs_list, indent=2)
else:
observations_text = "[]"
facts_lines = "\n".join(_fact_line(m) for m in memories)
result = await original_fn(
llm_config=llm_config,
memories=memories,
union_observations=union_observations,
union_source_facts=union_source_facts,
config=config,
remaining_observation_slots=remaining_observation_slots,
max_observations_per_scope=max_observations_per_scope,
)
_debug_log.append(_ConsolidationDebugEntry(
facts=facts_lines,
observations_text=observations_text,
response=_ConsolidationResponse(
creates=[_ActionLog(text=c.text, source_fact_ids=c.source_fact_ids) for c in result.creates],
updates=[
_ActionLog(text=u.text, observation_id=u.observation_id, source_fact_ids=u.source_fact_ids)
for u in result.updates
],
deletes=[_ActionLog(text="", observation_id=d.observation_id) for d in result.deletes],
),
))
return result
def _print_consolidation_debug(entry: _ConsolidationDebugEntry, index: int) -> None:
"""Print a single consolidation LLM call for debugging."""
print(f"\n --- LLM Call #{index} ---")
print(" FACTS sent to LLM:")
for line in entry.facts.split("\n"):
print(f" {line}")
print("\n EXISTING OBSERVATIONS sent to LLM:")
obs_data = json.loads(entry.observations_text)
if obs_data:
for obs in obs_data:
src_summary = ""
if obs.get("source_memories"):
src_texts = [sm["text"] for sm in obs["source_memories"]]
src_summary = f" (sources: {src_texts})"
print(f" [{obs['id'][:8]}..] proof={obs.get('proof_count', '?')}: {obs['text']}{src_summary}")
else:
print(" (none)")
resp = entry.response
print("\n LLM RESPONSE:")
if resp.creates:
for c in resp.creates:
print(f" CREATE: \"{c.text}\" (from facts: {[fid[:8] + '..' for fid in c.source_fact_ids]})")
if resp.updates:
for u in resp.updates:
print(
f" UPDATE [{u.observation_id[:8]}..]: \"{u.text}\""
f" (from facts: {[fid[:8] + '..' for fid in u.source_fact_ids]})"
)
if resp.deletes:
for d in resp.deletes:
print(f" DELETE [{d.observation_id[:8]}..]")
if not resp.creates and not resp.updates and not resp.deletes:
print(" (no actions)")
def _parse_history(hist: Any) -> list[str]:
"""Parse observation history from DB (may be list of dicts or JSON strings)."""
if not hist:
return []
parsed = hist if isinstance(hist, list) else json.loads(hist)
prev_texts = []
for h in parsed:
if isinstance(h, str):
h = json.loads(h)
prev_texts.append(h.get("previous_text", "?"))
return prev_texts
@pytest.mark.asyncio
@pytest.mark.flaky(reruns=2, reruns_delay=5)
async def test_horse_farm_observation_history(memory: MemoryEngine, request_context: Any) -> None:
"""Retain a sequence of horse facts and inspect how observations evolve."""
bank_id = f"test-horses-{uuid.uuid4().hex[:8]}"
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
messages = [
"I have a farm.",
"I have 2 horses.",
"I have a horse named Daisy.",
"I have a horse named Buttercup.",
"I sold Buttercup.",
"I now have 1 horse.",
"I have 5 horses on my farm.",
"I have a horse named Midnight.",
"I have horses named Midnight and Shadow.",
"I have horses named Shadow and Twister.",
"I am sad to report that Shadow has died.",
]
# Monkey-patch to intercept consolidation LLM calls
_original_consolidate = consolidator_mod._consolidate_batch_with_llm
async def _patched(**kwargs: Any) -> Any:
return await _instrumented_consolidate(_original_consolidate, **kwargs)
consolidator_mod._consolidate_batch_with_llm = _patched
_debug_log.clear()
try:
for i, content in enumerate(messages):
print(f"\n{'='*80}")
print(f"RETAIN #{i+1}: {content}")
print(f"{'='*80}")
log_start = len(_debug_log)
await memory.retain_async(
bank_id=bank_id,
content=content,
request_context=request_context,
)
await memory.wait_for_background_tasks()
for j, entry in enumerate(_debug_log[log_start:]):
_print_consolidation_debug(entry, j + 1)
# Dump current observations
pool = await memory._get_pool()
async with pool.acquire() as conn:
observations = await conn.fetch(
"""
SELECT id, text, proof_count, source_memory_ids, history
FROM memory_units
WHERE bank_id = $1 AND fact_type = 'observation'
ORDER BY created_at
""",
bank_id,
)
print(f"\n CURRENT OBSERVATIONS ({len(observations)}):")
for obs in observations:
prev_texts = _parse_history(obs["history"])
hist_str = f" (was: {' -> '.join(prev_texts)})" if prev_texts else ""
print(f" [{str(obs['id'])[:8]}..] proof={obs['proof_count']}: {obs['text']}{hist_str}")
finally:
consolidator_mod._consolidate_batch_with_llm = _original_consolidate
# Final summary
print(f"\n{'='*80}")
print("FINAL STATE")
print(f"{'='*80}")
pool = await memory._get_pool()
async with pool.acquire() as conn:
observations = await conn.fetch(
"""
SELECT id, text, proof_count, source_memory_ids, history
FROM memory_units
WHERE bank_id = $1 AND fact_type = 'observation'
ORDER BY created_at
""",
bank_id,
)
print(f"\nFinal observations ({len(observations)}):")
for obs in observations:
prev_texts = _parse_history(obs["history"])
if prev_texts:
chain = prev_texts + [obs["text"]]
print(f" - [proof={obs['proof_count']}] {obs['text']}")
print(f" evolution: {' -> '.join(chain)}")
else:
print(f" - [proof={obs['proof_count']}] {obs['text']}")
# Create a mental model to synthesize the observations
print(f"\n{'='*80}")
print("MENTAL MODEL")
print(f"{'='*80}")
# Patch reflect _execute_tool to log tool inputs/outputs
from hindsight_api.engine.reflect import agent as reflect_agent_mod
_original_execute = reflect_agent_mod._execute_tool
async def _logging_execute(tool_name: str, args: dict[str, Any], *a: Any, **kw: Any) -> dict[str, Any]:
result = await _original_execute(tool_name, args, *a, **kw)
normalized = reflect_agent_mod._normalize_tool_name(tool_name)
print(f"\n [REFLECT TOOL] {normalized}(args={args})")
if isinstance(result, dict):
if "observations" in result:
print(f" Observations returned ({result.get('count', '?')}, freshness={result.get('freshness', '?')}):")
for obs in result.get("observations", []):
print(f" - [proof={obs.get('proof_count', '?')}] {obs.get('text', '?')}")
if "memories" in result:
print(f" Memories returned ({result.get('count', '?')}):")
for mem in result.get("memories", []):
chunk = mem.get("chunk_text", "")
chunk_preview = f" | chunk: {chunk[:80]}..." if chunk else ""
print(f" - [{mem.get('fact_type', '?')}] {mem.get('text', '?')}{chunk_preview}")
if "mental_models" in result:
print(f" Mental models returned ({result.get('count', '?')}):")
for mm_item in result.get("mental_models", []):
print(f" - {mm_item.get('name', '?')}: {str(mm_item.get('content', '?'))[:120]}")
if "error" in result:
print(f" ERROR: {result['error']}")
return result
reflect_agent_mod._execute_tool = _logging_execute
source_query = (
"Produce a structured summary of all animals on the farm. Include:\n"
"1. A chronological timeline of events (acquisitions, sales, deaths) with dates\n"
"2. The list of all known horse names and their current status (alive, sold, died)\n"
"3. The current number of horses on the farm, accounting for all events\n"
"Reason step by step from the facts. If a horse died or was sold, subtract from the count."
)
try:
mm = await memory.create_mental_model(
bank_id=bank_id,
name="Farm Animals",
source_query=source_query,
content="(initial — awaiting refresh)",
request_context=request_context,
)
refreshed = await memory.refresh_mental_model(
bank_id=bank_id,
mental_model_id=mm["id"],
request_context=request_context,
)
content = refreshed["content"]
finally:
reflect_agent_mod._execute_tool = _original_execute
print(f"\nMental model content:\n{content}")
reflect_resp = refreshed.get("reflect_response")
if reflect_resp and isinstance(reflect_resp, str) and reflect_resp.strip():
try:
reflect_resp = json.loads(reflect_resp)
except json.JSONDecodeError:
reflect_resp = None
if isinstance(reflect_resp, dict):
based_on = reflect_resp.get("based_on", [])
if based_on:
print("\nBased on:")
for item in based_on:
if isinstance(item, str):
try:
item = json.loads(item)
except json.JSONDecodeError:
continue
print(f" - [{item.get('fact_type', '?')}] {item.get('text', '?')}")
# Verify the mental model captures key facts
content_lower = content.lower()
for name in ["daisy", "buttercup", "midnight", "shadow", "twister"]:
assert name in content_lower, f"Mental model should mention {name}. Got:\n{content}"
assert "sold" in content_lower or "sale" in content_lower, (
f"Mental model should mention Buttercup was sold. Got:\n{content}"
)
assert "died" in content_lower or "passed" in content_lower or "death" in content_lower, (
f"Mental model should mention Shadow's death. Got:\n{content}"
)
# Cleanup
await memory.delete_bank(bank_id, request_context=request_context)
@@ -584,6 +584,87 @@ async def test_delete_bank(api_client):
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
@pytest.mark.asyncio
async def test_delete_bank_nonexistent(api_client):
"""Test deleting a bank that doesn't exist returns success with zero counts."""
fake_bank_id = f"nonexistent_bank_{datetime.now().timestamp()}"
response = await api_client.delete(f"/v1/default/banks/{fake_bank_id}")
assert response.status_code == 200
result = response.json()
assert result["success"] is True
assert result["deleted_count"] == 0
@pytest.mark.asyncio
async def test_clear_memories_preserves_bank(api_client):
"""Test that clearing memories preserves the bank profile.
Workflow:
1. Create a bank with memories
2. Clear all memories via DELETE /memories
3. Verify the bank still exists with its profile intact
4. Verify all memories are gone
"""
test_bank_id = f"clear_memories_test_{datetime.now().timestamp()}"
try:
# 1. Create bank with memories
response = await api_client.post(
f"/v1/default/banks/{test_bank_id}/memories",
json={
"items": [
{"content": "Alice is a software engineer.", "context": "team info"},
{"content": "Bob works on infrastructure.", "context": "team info"},
]
},
)
assert response.status_code == 200
# Verify bank exists and has data
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/stats")
assert response.status_code == 200
assert response.json()["total_nodes"] > 0
response = await api_client.get("/v1/default/banks")
assert response.status_code == 200
bank_ids = [b["bank_id"] for b in response.json()["banks"]]
assert test_bank_id in bank_ids
# 2. Clear all memories
response = await api_client.delete(f"/v1/default/banks/{test_bank_id}/memories")
assert response.status_code == 200
assert response.json()["success"] is True
# 3. Bank should still exist in the list
response = await api_client.get("/v1/default/banks")
assert response.status_code == 200
bank_ids = [b["bank_id"] for b in response.json()["banks"]]
assert test_bank_id in bank_ids, "Bank should still exist after clearing memories"
# Profile should still be accessible
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/profile")
assert response.status_code == 200
# 4. Memories should be gone
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/stats")
assert response.status_code == 200
assert response.json()["total_nodes"] == 0
finally:
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
@pytest.mark.asyncio
async def test_clear_memories_nonexistent_bank(api_client):
"""Test clearing memories for a bank that doesn't exist returns success."""
fake_bank_id = f"nonexistent_clear_{datetime.now().timestamp()}"
response = await api_client.delete(f"/v1/default/banks/{fake_bank_id}/memories")
assert response.status_code == 200
assert response.json()["success"] is True
@pytest.mark.asyncio
async def test_async_retain(api_client):
"""Test asynchronous retain functionality.
@@ -1229,3 +1310,90 @@ async def test_retain_with_timestamp_async_complete_processing(api_client, test_
assert response.status_code == 200
items = response.json()["items"]
assert len(items) > 0, "Should have stored memories after async processing"
@pytest.mark.asyncio
async def test_http_recall_preserves_metadata(api_client, test_bank_id):
"""
Regression test for #797: HTTP recall must return metadata stored during retain.
The engine correctly preserves metadata, but _fact_to_result in http.py was
missing the metadata= kwarg, causing the HTTP endpoint to always return null.
"""
metadata = {"source": "slack", "channel": "engineering", "importance": "high"}
# Retain with metadata
response = await api_client.post(
f"/v1/default/banks/{test_bank_id}/memories",
json={
"items": [
{
"content": "The product launch is scheduled for March 1st.",
"metadata": metadata,
}
]
},
)
assert response.status_code == 200
# Recall via HTTP
response = await api_client.post(
f"/v1/default/banks/{test_bank_id}/memories/recall",
json={"query": "When is the product launch?", "budget": "low"},
)
assert response.status_code == 200
results = response.json()["results"]
assert len(results) > 0, "Should recall at least one fact"
# Find a result that has our metadata (LLM may extract multiple facts)
facts_with_metadata = [r for r in results if r.get("metadata")]
assert len(facts_with_metadata) > 0, "At least one fact must have metadata (regression #797)"
fact = facts_with_metadata[0]
assert fact["metadata"]["source"] == "slack"
assert fact["metadata"]["channel"] == "engineering"
assert fact["metadata"]["importance"] == "high"
@pytest.mark.asyncio
async def test_unknown_params_not_rejected(api_client):
"""Unknown query params and body fields should not cause a rejection (no 400).
The server should return 200 with an X-Ignored-Params header listing the
unknown parameters instead of rejecting the request. This ensures forward
compatibility when a newer client talks to an older server.
"""
test_bank_id = f"unknown_params_test_{datetime.now().timestamp()}"
# Ensure bank exists
await api_client.get(f"/v1/default/banks/{test_bank_id}/profile")
# Unknown query params on GET endpoint
response = await api_client.get(
f"/v1/default/banks/{test_bank_id}/memories/list",
params={"limit": 1, "tag": "source:slack", "created_after": "2026-01-01"},
)
assert response.status_code == 200
assert "X-Ignored-Params" in response.headers
ignored = response.headers["X-Ignored-Params"]
assert "tag" in ignored
assert "created_after" in ignored
# Unknown body fields on POST endpoint
response = await api_client.post(
f"/v1/default/banks/{test_bank_id}/memories",
json={
"items": [{"content": "test memory", "context": "test"}],
"unknown_future_field": True,
},
)
assert response.status_code == 200
assert "X-Ignored-Params" in response.headers
assert "unknown_future_field" in response.headers["X-Ignored-Params"]
# Known params only — no header
response = await api_client.get(
f"/v1/default/banks/{test_bank_id}/memories/list",
params={"limit": 1, "type": "world"},
)
assert response.status_code == 200
assert "X-Ignored-Params" not in response.headers
+135 -1
View File
@@ -1,11 +1,15 @@
"""Tests for link_utils datetime handling and temporal link computation."""
"""Tests for link_utils datetime handling, temporal link computation, and semantic link splitting."""
import numpy as np
import pytest
from datetime import datetime, timezone, timedelta
from hindsight_api.engine.retain.link_utils import (
_normalize_datetime,
_cap_links_per_unit,
compute_temporal_links,
compute_temporal_query_bounds,
compute_semantic_links_within_batch,
MAX_TEMPORAL_LINKS_PER_UNIT,
)
@@ -254,3 +258,133 @@ class TestComputeTemporalLinks:
assert len(links) == 1
assert links[0][3] >= 0.3
class TestCapLinksPerUnit:
"""Tests for the _cap_links_per_unit helper function."""
def test_empty_links(self):
assert _cap_links_per_unit([]) == []
def test_under_cap_unchanged(self):
links = [
("unit_a", "unit_x", "temporal", 0.9, None),
("unit_a", "unit_y", "temporal", 0.8, None),
]
result = _cap_links_per_unit(links, max_per_unit=5)
assert len(result) == 2
def test_caps_to_max_per_unit(self):
# Create 30 links from the same unit with descending weights
links = [("unit_a", f"unit_{i}", "temporal", 1.0 - i * 0.01, None) for i in range(30)]
result = _cap_links_per_unit(links, max_per_unit=10)
assert len(result) == 10
# Should keep the highest-weight links
weights = [lnk[3] for lnk in result]
assert weights == sorted(weights, reverse=True)
assert weights[0] == 1.0 # Highest weight kept
def test_caps_independently_per_unit(self):
links_a = [("unit_a", f"target_{i}", "temporal", 0.9 - i * 0.01, None) for i in range(10)]
links_b = [("unit_b", f"target_{i}", "temporal", 0.8 - i * 0.01, None) for i in range(10)]
result = _cap_links_per_unit(links_a + links_b, max_per_unit=5)
# 5 from unit_a + 5 from unit_b
assert len(result) == 10
from_a = [lnk for lnk in result if lnk[0] == "unit_a"]
from_b = [lnk for lnk in result if lnk[0] == "unit_b"]
assert len(from_a) == 5
assert len(from_b) == 5
def test_default_max_is_temporal_constant(self):
links = [("unit_a", f"target_{i}", "temporal", 1.0 - i * 0.01, None) for i in range(50)]
result = _cap_links_per_unit(links)
assert len(result) == MAX_TEMPORAL_LINKS_PER_UNIT
def test_preserves_tuple_structure(self):
links = [("from_id", "to_id", "temporal", 0.95, "entity_id")]
result = _cap_links_per_unit(links, max_per_unit=5)
assert result[0] == ("from_id", "to_id", "temporal", 0.95, "entity_id")
class TestComputeSemanticLinksWithinBatch:
"""Tests for compute_semantic_links_within_batch.
This function computes semantic links between units in the same batch
using numpy dot product (no DB access). It runs in Phase 2 (write
transaction) while the expensive ANN search against existing units runs
in Phase 1 on a separate connection to avoid TimeoutErrors from HNSW
index contention under concurrent load.
"""
def test_empty_returns_empty(self):
assert compute_semantic_links_within_batch([], []) == []
def test_single_unit_returns_empty(self):
emb = [np.random.randn(384).tolist()]
assert compute_semantic_links_within_batch(["u1"], emb) == []
def test_identical_embeddings_produce_links(self):
"""Two identical embeddings should have similarity=1.0 (above 0.7 threshold)."""
emb = [0.1] * 384
links = compute_semantic_links_within_batch(["u1", "u2"], [emb, emb])
assert len(links) == 2 # bidirectional: u1→u2, u2→u1
from_ids = {lnk[0] for lnk in links}
to_ids = {lnk[1] for lnk in links}
assert from_ids == {"u1", "u2"}
assert to_ids == {"u1", "u2"}
for lnk in links:
assert lnk[2] == "semantic"
assert lnk[3] >= 0.99 # near-1.0 similarity
assert lnk[4] is None # no entity_id
def test_orthogonal_embeddings_no_links(self):
"""Orthogonal embeddings should have similarity=0 (below 0.7 threshold)."""
emb1 = [1.0] + [0.0] * 383
emb2 = [0.0] + [1.0] + [0.0] * 382
links = compute_semantic_links_within_batch(["u1", "u2"], [emb1, emb2])
assert len(links) == 0
def test_respects_threshold(self):
"""Links below threshold should be excluded."""
emb1 = np.random.randn(384).tolist()
# Create a slightly similar embedding (add noise)
emb2 = [x + np.random.randn() * 0.5 for x in emb1]
# Normalize both
norm1 = np.linalg.norm(emb1)
norm2 = np.linalg.norm(emb2)
emb1 = [x / norm1 for x in emb1]
emb2 = [x / norm2 for x in emb2]
links_low = compute_semantic_links_within_batch(["u1", "u2"], [emb1, emb2], threshold=0.0)
links_high = compute_semantic_links_within_batch(["u1", "u2"], [emb1, emb2], threshold=0.99)
# Low threshold should have more links than high threshold
assert len(links_low) >= len(links_high)
def test_top_k_limits_per_unit(self):
"""Each unit should link to at most top_k other units."""
n = 10
# Create similar embeddings (all close to the same vector)
base = np.random.randn(384)
base = base / np.linalg.norm(base)
embs = [(base + np.random.randn(384) * 0.01).tolist() for _ in range(n)]
unit_ids = [f"u{i}" for i in range(n)]
links = compute_semantic_links_within_batch(unit_ids, embs, top_k=3, threshold=0.5)
# Each unit should have at most 3 outgoing links
from collections import Counter
from_counts = Counter(lnk[0] for lnk in links)
for count in from_counts.values():
assert count <= 3
def test_link_tuple_structure(self):
"""Verify the tuple format matches what _bulk_insert_links expects."""
emb = [0.1] * 384
links = compute_semantic_links_within_batch(["u1", "u2"], [emb, emb])
for lnk in links:
assert len(lnk) == 5
from_id, to_id, link_type, weight, entity_id = lnk
assert isinstance(from_id, str)
assert isinstance(to_id, str)
assert link_type == "semantic"
assert 0.0 <= weight <= 1.0
assert entity_id is None
@@ -277,6 +277,99 @@ class TestLiteLLMSDKEmbeddings:
call_args = mock_litellm.embedding.call_args
assert call_args.kwargs["api_base"] == "https://custom.api.com"
async def test_output_dimensions_passed_when_set(self, mock_litellm):
"""Test output dimensions are passed to LiteLLM when configured."""
with patch(
"builtins.__import__",
side_effect=lambda name, *args: mock_litellm if name == "litellm" else __import__(name, *args),
):
emb = LiteLLMSDKEmbeddings(
api_key="test_key",
model="cohere/embed-english-v3.0",
output_dimensions=768,
)
await emb.initialize()
init_call_args = mock_litellm.aembedding.call_args
assert init_call_args.kwargs["dimensions"] == 768
mock_litellm.embedding.return_value.data = [{"embedding": [0.1] * 768, "index": 0}]
emb.encode(["test"])
encode_call_args = mock_litellm.embedding.call_args
assert encode_call_args.kwargs["dimensions"] == 768
async def test_output_dimensions_omitted_when_unset(self, mock_litellm):
"""Test output dimensions are omitted when not configured."""
with patch(
"builtins.__import__",
side_effect=lambda name, *args: mock_litellm if name == "litellm" else __import__(name, *args),
):
emb = LiteLLMSDKEmbeddings(
api_key="test_key",
model="cohere/embed-english-v3.0",
)
await emb.initialize()
init_call_args = mock_litellm.aembedding.call_args
assert "dimensions" not in init_call_args.kwargs
mock_litellm.embedding.return_value.data = [{"embedding": [0.1] * 768, "index": 0}]
emb.encode(["test"])
encode_call_args = mock_litellm.embedding.call_args
assert "dimensions" not in encode_call_args.kwargs
async def test_output_dimensions_and_api_base_passed_when_both_set(self, mock_litellm):
"""Test both dimensions and api_base are forwarded when configured together."""
with patch(
"builtins.__import__",
side_effect=lambda name, *args: mock_litellm if name == "litellm" else __import__(name, *args),
):
emb = LiteLLMSDKEmbeddings(
api_key="test_key",
model="cohere/embed-english-v3.0",
api_base="https://custom.api.com",
output_dimensions=768,
)
await emb.initialize()
init_call_args = mock_litellm.aembedding.call_args
assert init_call_args.kwargs["api_base"] == "https://custom.api.com"
assert init_call_args.kwargs["dimensions"] == 768
mock_litellm.embedding.return_value.data = [{"embedding": [0.1] * 768, "index": 0}]
emb.encode(["test"])
encode_call_args = mock_litellm.embedding.call_args
assert encode_call_args.kwargs["api_base"] == "https://custom.api.com"
assert encode_call_args.kwargs["dimensions"] == 768
async def test_openai_invalid_output_dimensions_raises(self, mock_litellm):
"""Invalid dimensions fail during initialize() (probe call), not per HTTP request.
MemoryEngine runs this at app lifespan startup; the process typically fails to become
ready rather than returning a JSON error for a single API call. The RuntimeError
message should still chain the underlying provider/LiteLLM detail for logs.
"""
mock_litellm.aembedding.side_effect = Exception("invalid dimensions for model")
with patch(
"builtins.__import__",
side_effect=lambda name, *args: mock_litellm if name == "litellm" else __import__(name, *args),
):
emb = LiteLLMSDKEmbeddings(
api_key="test_key",
model="openai/text-embedding-3-small",
output_dimensions=9999,
)
with pytest.raises(
RuntimeError,
match="Failed to initialize LiteLLM SDK embeddings:.*invalid dimensions for model",
):
await emb.initialize()
class TestLiteLLMSDKEmbeddingsFactory:
"""Test the factory function for creating LiteLLM SDK embeddings."""
@@ -324,6 +417,21 @@ class TestLiteLLMSDKEmbeddingsFactory:
assert isinstance(embeddings, LiteLLMSDKEmbeddings)
assert embeddings.api_base == "https://custom.api.com"
def test_create_from_env_with_output_dimensions(self, monkeypatch):
"""Test creating embeddings with configured output dimensions."""
mock_config = MagicMock()
mock_config.embeddings_provider = "litellm-sdk"
mock_config.embeddings_litellm_sdk_api_key = "test_key"
mock_config.embeddings_litellm_sdk_model = "gemini/gemini-embedding-2-preview"
mock_config.embeddings_litellm_sdk_api_base = None
mock_config.embeddings_litellm_sdk_output_dimensions = 768
with patch("hindsight_api.config.get_config", return_value=mock_config):
embeddings = create_embeddings_from_env()
assert isinstance(embeddings, LiteLLMSDKEmbeddings)
assert embeddings.output_dimensions == 768
class TestLiteLLMSDKCohereEmbeddings:
"""Integration tests calling real Cohere API (matches CI pattern)."""
@@ -284,7 +284,7 @@ async def test_llm_provider_memory_operations(provider: str, model: str):
# Verify facts have required fields
for fact in facts:
assert fact.fact, f"{provider}/{model} fact missing text"
assert fact.fact_type in ["world", "experience", "opinion"], f"{provider}/{model} invalid fact_type: {fact.fact_type}"
assert fact.fact_type in ["world", "experience"], f"{provider}/{model} invalid fact_type: {fact.fact_type}"
# Test 2: Reflect (actual reflect function)
response = await reflect(

Some files were not shown because too many files have changed in this diff Show More