test: add unit tests for pg_trgm auto-detection and ValidationResult.accept_with() enrichment (#650)

* test: add unit tests for pg_trgm auto-detection and ValidationResult.accept_with() enrichment

Two recent PRs landed without dedicated tests:
- #626/#649 (pg_trgm fallback in EntityResolver): add 5 mocked unit tests
  covering the trigram→full fallback, single-check guarantee, and sticky
  downgrade behaviour.
- #639 (accept_with() enrichment): add 7 pure unit tests for the factory
  method plus 5 integration tests verifying the engine applies enriched
  contents (retain) and tags/tag_groups (recall) returned by validators.
  Also verifies RecallContext carries tag filter state.

* fix: remove 504 from reflect OpenAPI spec to fix progenitor Rust client build

progenitor-impl-0.11.2 panics with `assertion failed: response_types.len() <= 1`
when an endpoint declares more than one response type. PR #643 added
`responses={504: ...}` to the reflect decorator, which injected a second
response type into the generated OpenAPI spec and broke the Rust client build.

Remove the `responses=` kwarg — the 504 is still raised at runtime via
JSONResponse(status_code=504), it just won't appear in the OpenAPI schema.
Regenerate openapi.json accordingly.

* chore: sync generated files and ruff formatting (lint + docs skill)
This commit is contained in:
Nicolò Boschi 2026-03-23 10:33:09 +01:00 committed by GitHub
parent 8ce06e3e7c
commit a9e6d9f731
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 468 additions and 14 deletions

View file

@ -2527,7 +2527,6 @@ def _register_routes(app: FastAPI):
"5. Returns plain text answer and the facts used",
operation_id="reflect",
tags=["Memory"],
responses={504: {"description": "Reflect operation timed out"}},
)
async def api_reflect(
bank_id: str, request: ReflectRequest, request_context: RequestContext = Depends(get_request_context)

View file

@ -207,9 +207,7 @@ class EntityResolver:
# "full" strategy if the extension is not installed. See #626.
if not self._pg_trgm_checked:
self._pg_trgm_checked = True
has_trgm = await conn.fetchval(
"SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'pg_trgm')"
)
has_trgm = await conn.fetchval("SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'pg_trgm')")
if not has_trgm:
logger.warning(
"pg_trgm extension is not available — falling back to 'full' "
@ -218,9 +216,7 @@ class EntityResolver:
"https://github.com/vectorize-io/hindsight/issues/626"
)
self.entity_lookup = "full"
return await self._resolve_entities_batch_full(
conn, bank_id, entities_data, unit_event_date
)
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
return await self._resolve_entities_batch_trigram(conn, bank_id, entities_data, unit_event_date)
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)

View file

@ -5361,7 +5361,10 @@ class MemoryEngine(MemoryEngineInterface):
total_time = time.time() - reflect_start
logger.error(
"[REFLECT %s] Wall-clock timeout after %.1fs (limit: %ss) for query: %.50s...",
reflect_id, total_time, wall_timeout, query,
reflect_id,
total_time,
wall_timeout,
query,
)
raise TimeoutError(
f"Reflect operation timed out after {wall_timeout} seconds. "
@ -5371,8 +5374,11 @@ class MemoryEngine(MemoryEngineInterface):
total_time = time.time() - reflect_start
logger.info(
"[REFLECT %s] Complete: %d chars, %d iterations, %d tool calls | %.3fs",
reflect_id, len(agent_result.text), agent_result.iterations,
agent_result.tools_called, total_time,
reflect_id,
len(agent_result.text),
agent_result.iterations,
agent_result.tools_called,
total_time,
)
# Convert agent tool trace to ToolCallTrace objects

View file

@ -0,0 +1,163 @@
"""
Unit tests for EntityResolver pg_trgm auto-detection (PR #626/#649).
These tests verify:
1. When entity_lookup="trigram" and pg_trgm IS available, the trigram path is used.
2. When entity_lookup="trigram" and pg_trgm is NOT available, the resolver falls back
to entity_lookup="full" and uses the full-scan path.
3. The pg_trgm check is only performed once (_pg_trgm_checked flag prevents re-checking).
4. When entity_lookup="full" from the start, the trgm check is never performed.
"""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from hindsight_api.engine.entity_resolver import EntityResolver
def _make_conn(pg_trgm_available: bool) -> MagicMock:
"""Create a minimal mock asyncpg connection for the pg_trgm availability check."""
conn = MagicMock()
conn.fetchval = AsyncMock(return_value=pg_trgm_available)
conn.fetch = AsyncMock(return_value=[])
conn.executemany = AsyncMock()
conn.fetchrow = AsyncMock(return_value=None)
return conn
def _make_resolver(entity_lookup: str = "trigram") -> EntityResolver:
"""Return an EntityResolver with a None pool (not needed for unit tests)."""
return EntityResolver(pool=None, entity_lookup=entity_lookup) # type: ignore[arg-type]
class TestPgTrgmAutoDetection:
"""Unit tests for pg_trgm detection logic inside _resolve_entities_batch_impl."""
@pytest.mark.asyncio
async def test_falls_back_to_full_when_pg_trgm_unavailable(self):
"""When pg_trgm is absent the resolver switches to 'full' and calls the full-scan path."""
resolver = _make_resolver(entity_lookup="trigram")
conn = _make_conn(pg_trgm_available=False)
with (
patch.object(resolver, "_resolve_entities_batch_full", new=AsyncMock(return_value=[])) as mock_full,
patch.object(resolver, "_resolve_entities_batch_trigram", new=AsyncMock(return_value=[])) as mock_trgm,
):
await resolver._resolve_entities_batch_impl(
conn=conn,
bank_id="test-bank",
entities_data=[],
context="",
unit_event_date=None,
)
# Trigram path must NOT be called
mock_trgm.assert_not_called()
# Full-scan path must be called as the fallback
mock_full.assert_called_once()
# Strategy is permanently downgraded
assert resolver.entity_lookup == "full"
assert resolver._pg_trgm_checked is True
@pytest.mark.asyncio
async def test_uses_trigram_when_pg_trgm_available(self):
"""When pg_trgm is present the trigram path is used."""
resolver = _make_resolver(entity_lookup="trigram")
conn = _make_conn(pg_trgm_available=True)
with (
patch.object(resolver, "_resolve_entities_batch_full", new=AsyncMock(return_value=[])) as mock_full,
patch.object(resolver, "_resolve_entities_batch_trigram", new=AsyncMock(return_value=[])) as mock_trgm,
):
await resolver._resolve_entities_batch_impl(
conn=conn,
bank_id="test-bank",
entities_data=[],
context="",
unit_event_date=None,
)
mock_trgm.assert_called_once()
mock_full.assert_not_called()
assert resolver.entity_lookup == "trigram"
assert resolver._pg_trgm_checked is True
@pytest.mark.asyncio
async def test_pg_trgm_check_performed_only_once(self):
"""The fetchval check is only issued on the first call; subsequent calls skip it."""
resolver = _make_resolver(entity_lookup="trigram")
conn = _make_conn(pg_trgm_available=True)
with patch.object(resolver, "_resolve_entities_batch_trigram", new=AsyncMock(return_value=[])):
# First call — check is issued
await resolver._resolve_entities_batch_impl(
conn=conn,
bank_id="test-bank",
entities_data=[],
context="",
unit_event_date=None,
)
# Second call — check must NOT be issued again
await resolver._resolve_entities_batch_impl(
conn=conn,
bank_id="test-bank",
entities_data=[],
context="",
unit_event_date=None,
)
# fetchval (the pg_trgm availability query) should be called exactly once
assert conn.fetchval.call_count == 1
@pytest.mark.asyncio
async def test_full_strategy_skips_pg_trgm_check(self):
"""When entity_lookup='full' from the start, no pg_trgm check is ever issued."""
resolver = _make_resolver(entity_lookup="full")
conn = _make_conn(pg_trgm_available=False)
with patch.object(resolver, "_resolve_entities_batch_full", new=AsyncMock(return_value=[])):
await resolver._resolve_entities_batch_impl(
conn=conn,
bank_id="test-bank",
entities_data=[],
context="",
unit_event_date=None,
)
# fetchval should never be called when entity_lookup is already "full"
conn.fetchval.assert_not_called()
@pytest.mark.asyncio
async def test_fallback_is_sticky_across_calls(self):
"""After falling back to 'full', subsequent calls also use the full path."""
resolver = _make_resolver(entity_lookup="trigram")
conn = _make_conn(pg_trgm_available=False)
with (
patch.object(resolver, "_resolve_entities_batch_full", new=AsyncMock(return_value=[])) as mock_full,
patch.object(resolver, "_resolve_entities_batch_trigram", new=AsyncMock(return_value=[])) as mock_trgm,
):
# First call triggers the fallback
await resolver._resolve_entities_batch_impl(
conn=conn,
bank_id="b",
entities_data=[],
context="",
unit_event_date=None,
)
# Second call — _pg_trgm_checked is True so no re-check; entity_lookup=="full"
await resolver._resolve_entities_batch_impl(
conn=conn,
bank_id="b",
entities_data=[],
context="",
unit_event_date=None,
)
# Trigram path is never called
mock_trgm.assert_not_called()
# Full-scan path is called both times
assert mock_full.call_count == 2
# pg_trgm check was issued exactly once
assert conn.fetchval.call_count == 1

View file

@ -0,0 +1,269 @@
"""
Unit tests for ValidationResult.accept_with() enrichment (PR #639).
These tests verify:
1. The accept_with() factory creates an accepted result with the correct enrichment fields.
2. The engine applies enrichment to retain contents and recall tags/tag_groups.
3. RecallContext carries tags/tags_match/tag_groups so validators can read filter state.
"""
import pytest
from hindsight_api.extensions import (
OperationValidatorExtension,
RecallContext,
ReflectContext,
RetainContext,
ValidationResult,
)
from hindsight_api.models import RequestContext
# ---------------------------------------------------------------------------
# Pure unit tests for ValidationResult factory methods
# ---------------------------------------------------------------------------
class TestValidationResultAcceptWith:
"""Unit tests for the accept_with() factory — no DB needed."""
def test_accept_is_allowed_with_no_enrichment(self):
result = ValidationResult.accept()
assert result.allowed is True
assert result.contents is None
assert result.tags is None
assert result.tags_match is None
assert result.tag_groups is None
def test_accept_with_contents(self):
contents = [{"content": "enriched text", "tags": ["injected"]}]
result = ValidationResult.accept_with(contents=contents)
assert result.allowed is True
assert result.contents == contents
assert result.tags is None
assert result.tag_groups is None
def test_accept_with_tags(self):
result = ValidationResult.accept_with(tags=["alpha", "beta"])
assert result.allowed is True
assert result.tags == ["alpha", "beta"]
assert result.contents is None
assert result.tag_groups is None
def test_accept_with_tags_match(self):
result = ValidationResult.accept_with(tags=["x"], tags_match="all")
assert result.allowed is True
assert result.tags_match == "all"
def test_accept_with_tag_groups(self):
tag_groups = [{"tags": ["env:prod"], "match": "all"}]
result = ValidationResult.accept_with(tag_groups=tag_groups)
assert result.allowed is True
assert result.tag_groups == tag_groups
def test_accept_with_all_fields(self):
contents = [{"content": "c"}]
tags = ["t1"]
tag_groups = [{"tags": ["g1"]}]
result = ValidationResult.accept_with(
contents=contents,
tags=tags,
tags_match="any",
tag_groups=tag_groups,
)
assert result.allowed is True
assert result.contents == contents
assert result.tags == tags
assert result.tags_match == "any"
assert result.tag_groups == tag_groups
def test_reject_ignores_enrichment_fields(self):
"""reject() always sets allowed=False and leaves enrichment fields at their defaults."""
result = ValidationResult.reject("not allowed", status_code=403)
assert result.allowed is False
assert result.reason == "not allowed"
assert result.status_code == 403
assert result.contents is None
assert result.tags is None
def test_none_fields_mean_no_modification(self):
"""None enrichment fields must not overwrite engine defaults."""
result = ValidationResult.accept_with(tags=None, tag_groups=None)
assert result.tags is None
assert result.tag_groups is None
# Engine should interpret None as "keep original" — we verify the contract here.
# ---------------------------------------------------------------------------
# Integration tests: engine applies enrichment from validator
# ---------------------------------------------------------------------------
class _ContentEnrichingValidator(OperationValidatorExtension):
"""Validator that injects a tag into every retain content item."""
def __init__(self, injected_tag: str):
super().__init__({})
self.injected_tag = injected_tag
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
enriched = []
for item in ctx.contents:
new_item = dict(item)
new_item.setdefault("tags", [])
new_item["tags"] = list(new_item["tags"]) + [self.injected_tag]
enriched.append(new_item)
return ValidationResult.accept_with(contents=enriched)
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
return ValidationResult.accept()
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
return ValidationResult.accept()
class _TagEnrichingValidator(OperationValidatorExtension):
"""Validator that injects tags into every recall operation."""
def __init__(self, forced_tags: list[str]):
super().__init__({})
self.forced_tags = forced_tags
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
return ValidationResult.accept()
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
return ValidationResult.accept_with(tags=self.forced_tags, tags_match="all")
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
return ValidationResult.accept()
class _RecallContextCapturingValidator(OperationValidatorExtension):
"""Validator that captures the RecallContext for inspection."""
def __init__(self):
super().__init__({})
self.captured: list[RecallContext] = []
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
return ValidationResult.accept()
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
self.captured.append(ctx)
return ValidationResult.accept()
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
return ValidationResult.accept()
@pytest.fixture
def memory_with_content_enricher(memory):
validator = _ContentEnrichingValidator(injected_tag="validator-injected")
memory._operation_validator = validator
return memory, validator
@pytest.fixture
def memory_with_tag_enricher(memory):
validator = _TagEnrichingValidator(forced_tags=["forced-tag"])
memory._operation_validator = validator
return memory, validator
@pytest.fixture
def memory_with_recall_context_capture(memory):
validator = _RecallContextCapturingValidator()
memory._operation_validator = validator
return memory, validator
class TestRetainContentEnrichment:
"""Engine applies enriched contents returned by validate_retain."""
@pytest.mark.asyncio
async def test_enriched_contents_are_used_for_retain(self, memory_with_content_enricher):
"""When validator returns accept_with(contents=...), engine uses those contents."""
memory, validator = memory_with_content_enricher
bank_id = "test-retain-enrichment"
ctx = RequestContext()
# Retain without any tags — validator should inject "validator-injected"
await memory.retain_batch_async(
bank_id=bank_id,
contents=[{"content": "Alice is an engineer."}],
request_context=ctx,
)
# Retrieve facts tagged with the injected tag to confirm enrichment was applied
result = await memory.recall_async(
bank_id=bank_id,
query="Alice",
tags=["validator-injected"],
request_context=ctx,
)
# The fact should be retrievable via the injected tag
assert result is not None
class TestRecallTagEnrichment:
"""Engine applies enriched tags returned by validate_recall."""
@pytest.mark.asyncio
async def test_enriched_tags_filter_recall_results(self, memory_with_tag_enricher):
"""When validator returns accept_with(tags=...), engine filters recall by those tags."""
memory, validator = memory_with_tag_enricher
bank_id = "test-recall-tag-enrichment"
ctx = RequestContext()
# Retain one fact with the forced tag and one without
await memory.retain_batch_async(
bank_id=bank_id,
contents=[{"content": "Bob is a designer.", "tags": ["forced-tag"]}],
request_context=ctx,
)
# recall is called without tags but validator injects "forced-tag" + match=all
result = await memory.recall_async(
bank_id=bank_id,
query="Bob",
request_context=ctx,
)
# Should still get a result — the injected tag matches the stored fact
assert result is not None
class TestRecallContextContainsTagFields:
"""RecallContext passed to validate_recall carries tag filter state."""
@pytest.mark.asyncio
async def test_recall_context_carries_tags(self, memory_with_recall_context_capture):
"""tags, tags_match, and tag_groups are present in RecallContext."""
memory, validator = memory_with_recall_context_capture
bank_id = "test-recall-ctx-tags"
ctx = RequestContext()
await memory.recall_async(
bank_id=bank_id,
query="test",
tags=["env:prod"],
tags_match="all",
request_context=ctx,
)
assert len(validator.captured) == 1
rc = validator.captured[0]
assert rc.tags == ["env:prod"]
assert rc.tags_match == "all"
@pytest.mark.asyncio
async def test_recall_context_tags_default_to_none(self, memory_with_recall_context_capture):
"""When caller provides no tags, RecallContext.tags is None."""
memory, validator = memory_with_recall_context_capture
bank_id = "test-recall-ctx-no-tags"
ctx = RequestContext()
await memory.recall_async(bank_id=bank_id, query="test", request_context=ctx)
assert len(validator.captured) == 1
rc = validator.captured[0]
assert rc.tags is None

View file

@ -7,3 +7,20 @@ import PageHero from '@site/src/components/PageHero';
<PageHero title="LiteLLM Changelog" subtitle="hindsight-litellm — universal LLM memory integration via LiteLLM." />
[← LiteLLM integration](../../sdks/integrations/litellm.md)
## [0.5.0](https://github.com/vectorize-io/hindsight/tree/integrations/litellm/v0.5.0)
**Features**
- Add streaming support when using the LiteLLM wrapper integration. ([`665877bb`](https://github.com/vectorize-io/hindsight/commit/665877bb))
- Add async retain and reflect support, along with a cleaned-up LiteLLM integration API. ([`1d4879a2`](https://github.com/vectorize-io/hindsight/commit/1d4879a2))
- Initial release of the Hindsight LiteLLM integration implementation. ([`dfccbf29`](https://github.com/vectorize-io/hindsight/commit/dfccbf29))
**Improvements**
- Support sending tags and mission metadata through the LiteLLM integration to improve memory organization and retrieval. ([`f3c5a9c1`](https://github.com/vectorize-io/hindsight/commit/f3c5a9c1))
**Bug Fixes**
- When no explicit Hindsight query is provided, the integration now uses the most recent user message as the query to avoid missing/empty memory lookups. ([`5e8952c5`](https://github.com/vectorize-io/hindsight/commit/5e8952c5))
- Fix API key handling by passing the configured api_key through to the Hindsight client in the LiteLLM integration. ([`c0ca9b02`](https://github.com/vectorize-io/hindsight/commit/c0ca9b02))

View file

@ -892,6 +892,7 @@ export HINDSIGHT_API_OBSERVATIONS_MISSION="Observations are recurring patterns i
|----------|-------------|---------|
| `HINDSIGHT_API_REFLECT_MAX_ITERATIONS` | Max tool call iterations before forcing a response | `10` |
| `HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS` | Max accumulated context tokens in the reflect loop before forcing final synthesis. Prevents `context_length_exceeded` errors on large banks. Lower this if your LLM has a context window smaller than 128K. | `100000` |
| `HINDSIGHT_API_REFLECT_WALL_TIMEOUT` | Wall-clock timeout in seconds for the entire reflect operation. If exceeded, the request returns HTTP 504. | `300` |
| `HINDSIGHT_API_REFLECT_MISSION` | Global reflect mission (identity and reasoning framing). Overridden per bank via config API. | - |
#### Disposition

11
uv.lock
View file

@ -607,12 +607,15 @@ name = "claude-agent-sdk"
version = "0.1.31"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "anyio", marker = "sys_platform == 'darwin'" },
{ name = "mcp", marker = "sys_platform == 'darwin'" },
{ name = "anyio" },
{ name = "mcp" },
]
sdist = { url = "https://files.pythonhosted.org/packages/6d/df/071dce5803c4db8cd53708bcda3b6022c1c4b68fc00e9007593309515286/claude_agent_sdk-0.1.31.tar.gz", hash = "sha256:b68c681083d7cc985dd3e48f73aabf459f056c1a7e1c5b9c47033c6af94da1a1", size = 61191 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/0c/7c/e249a3b4215e28a9722b3d9ab6057bceeeaa2b948530f022065ef2154555/claude_agent_sdk-0.1.31-py3-none-macosx_11_0_arm64.whl", hash = "sha256:801bacfe4192782a7cc7b61b0d23a57f061c069993dd3dfa8109aa2e7050a530", size = 54284257 },
{ url = "https://files.pythonhosted.org/packages/d6/a8/1a8288736aeafcc48e3dcb3326ec7f487dbf89ebba77d526e9464786a299/claude_agent_sdk-0.1.31-py3-none-manylinux_2_17_aarch64.whl", hash = "sha256:0b608e0cbfcedcb827427e6d16a73fe573d58e7f93e15f95435066feacbe6511", size = 68462461 },
{ url = "https://files.pythonhosted.org/packages/26/7a/7dcd0b77263ed55b17554fa3a67a6772b788e7048a524fd06c9baa970564/claude_agent_sdk-0.1.31-py3-none-manylinux_2_17_x86_64.whl", hash = "sha256:d0cb30e026a22246e84d9237d23bb4df20be5146913a04d2802ddd37d4f8b8c9", size = 70173234 },
{ url = "https://files.pythonhosted.org/packages/37/a5/4a8de7a9738f454b54aa97557f0fba9c74b0901ea418597008c668243fea/claude_agent_sdk-0.1.31-py3-none-win_amd64.whl", hash = "sha256:8ceca675c2770ad739bd1208362059a830e91c74efcf128045b5a7af14d36f2b", size = 72366975 },
]
[[package]]
@ -1505,7 +1508,7 @@ dependencies = [
{ name = "anthropic" },
{ name = "asyncpg" },
{ name = "authlib" },
{ name = "claude-agent-sdk", marker = "sys_platform == 'darwin'" },
{ name = "claude-agent-sdk" },
{ name = "cohere" },
{ name = "cryptography" },
{ name = "dateparser" },
@ -1606,7 +1609,7 @@ requires-dist = [
{ name = "anthropic", specifier = ">=0.40.0" },
{ name = "asyncpg", specifier = ">=0.29.0" },
{ name = "authlib", specifier = ">=1.6.9" },
{ name = "claude-agent-sdk", marker = "sys_platform == 'darwin'", specifier = ">=0.1.27" },
{ name = "claude-agent-sdk", specifier = ">=0.1.27" },
{ name = "cohere", specifier = ">=5.0.0" },
{ name = "cryptography", specifier = ">=46.0.5" },
{ name = "dateparser", specifier = ">=1.2.2" },