fix: tagged directives should be applied to tagged mental models (#303)

* fix: tagged directives should be applied to tagged mental models

* test: add unit test for based_on structure

Verify that reflect returns the correct based_on structure with:
- directives as dicts (id, name, content) in based_on.directives
- mental models as MemoryFact objects in based_on.mental-models
- memories separated properly

This ensures directives and mental models are not mixed together
in the API response.
This commit is contained in:
Nicolò Boschi 2026-02-05 13:22:56 +01:00 committed by GitHub
parent 093ecff48d
commit 278718dd84
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 298 additions and 75 deletions

View file

@ -1916,17 +1916,17 @@ def _register_routes(app: FastAPI):
directives = []
for fact_type, facts in core_result.based_on.items():
if fact_type == "directives":
# Directives have different structure (id, name, content)
# Directives are dicts with id, name, content (not MemoryFact objects)
for directive in facts:
directives.append(
ReflectDirective(
id=directive.id,
name=directive.name,
content=directive.content,
id=directive["id"],
name=directive["name"],
content=directive["content"],
)
)
elif fact_type == "mental_models":
# Mental models are MemoryFact with type "mental_models"
elif fact_type == "mental-models":
# Mental models are MemoryFact with type "mental-models" (note: hyphen, not underscore)
for fact in facts:
mental_models.append(
ReflectMentalModel(

View file

@ -3673,12 +3673,14 @@ class MemoryEngine(MemoryEngineInterface):
# Load directives from the dedicated directives table
# Directives are hard rules that must be followed in all responses
# Use isolation_mode=True to prevent tag-scoped directives from leaking into untagged operations
directives_raw = await self.list_directives(
bank_id=bank_id,
tags=tags,
tags_match=tags_match,
active_only=True,
request_context=request_context,
isolation_mode=True,
)
# Convert directive format to the expected format for reflect agent
# The agent expects: name, description (optional), observations (list of {title, content})
@ -3747,7 +3749,16 @@ class MemoryEngine(MemoryEngineInterface):
# Extract memories from recall tool outputs - only include memories the agent actually used
# agent_result.used_memory_ids contains validated IDs from the done action
used_memory_ids_set = set(agent_result.used_memory_ids) if agent_result.used_memory_ids else set()
based_on: dict[str, list[MemoryFact]] = {"world": [], "experience": [], "opinion": [], "observation": []}
# based_on stores facts, mental models, and directives
# Note: directives list stores raw directive dicts (not MemoryFact), which will be converted to Directive objects
based_on: dict[str, list[MemoryFact] | list[dict[str, Any]]] = {
"world": [],
"experience": [],
"opinion": [],
"observation": [],
"mental-models": [],
"directives": [],
}
seen_memory_ids: set[str] = set()
for tc in agent_result.tool_trace:
if tc.tool == "recall" and "memories" in tc.output:
@ -3849,38 +3860,15 @@ class MemoryEngine(MemoryEngineInterface):
)
# List all models lookup - don't add to based_on (too verbose, just a listing)
# Add directives to based_on["mental-models"] (they are mental models with subtype='directive')
for directive in directives:
# Extract summary from observations
summary_parts: list[str] = []
for obs in directive.get("observations", []):
# Support both Pydantic Observation objects and dicts
if hasattr(obs, "content"):
content = obs.content
title = obs.title
else:
content = obs.get("content", "")
title = obs.get("title", "")
if title and content:
summary_parts.append(f"{title}: {content}")
elif content:
summary_parts.append(content)
# Fallback to description if no observations
if not summary_parts and directive.get("description"):
summary_parts.append(directive["description"])
directive_name = directive.get("name", "")
directive_summary = "; ".join(summary_parts) if summary_parts else ""
based_on["mental-models"].append(
MemoryFact(
id=directive.get("id", ""),
text=f"{directive_name}: {directive_summary}",
fact_type="mental-models",
context="directive (directive)",
occurred_start=None,
occurred_end=None,
)
# Add directives to based_on["directives"]
# Store raw directive dicts (with id, name, content) for http.py to convert to ReflectDirective
for directive_raw in directives_raw:
based_on["directives"].append(
{
"id": directive_raw["id"],
"name": directive_raw["name"],
"content": directive_raw["content"],
}
)
# Build directives_applied from agent result
@ -4754,6 +4742,7 @@ class MemoryEngine(MemoryEngineInterface):
"id": str(fact.id),
"text": fact.text,
"type": fact_type,
"context": fact.context, # Include context to distinguish directives from mental models in UI
}
for fact in facts
]
@ -4942,6 +4931,7 @@ class MemoryEngine(MemoryEngineInterface):
limit: int = 100,
offset: int = 0,
request_context: "RequestContext",
isolation_mode: bool = False,
) -> list[dict[str, Any]]:
"""List directives for a bank.
@ -4953,6 +4943,9 @@ class MemoryEngine(MemoryEngineInterface):
limit: Maximum number of results
offset: Offset for pagination
request_context: Request context for authentication
isolation_mode: When True and tags=None, only return directives with no tags.
This prevents tag-scoped directives from leaking into untagged operations.
Default False (normal API behavior - returns all directives when tags=None)
Returns:
List of directive dicts
@ -4962,6 +4955,8 @@ class MemoryEngine(MemoryEngineInterface):
async with acquire_with_retry(pool) as conn:
# Build filters
from .search.tags import build_tags_where_clause
filters = ["bank_id = $1"]
params: list[Any] = [bank_id]
param_idx = 2
@ -4969,15 +4964,23 @@ class MemoryEngine(MemoryEngineInterface):
if active_only:
filters.append("is_active = TRUE")
# Apply tags filter:
# - If tags provided: use standard filtering (with strict modes support)
# - If tags=None and isolation_mode=True: only include directives with NO tags
# (prevents tag-scoped directives from leaking into untagged reflect/refresh)
# - If tags=None and isolation_mode=False: no filtering (normal API behavior)
if tags:
if tags_match == "all":
filters.append(f"tags @> ${param_idx}::varchar[]")
elif tags_match == "exact":
filters.append(f"tags = ${param_idx}::varchar[]")
else: # any
filters.append(f"tags && ${param_idx}::varchar[]")
params.append(tags)
param_idx += 1
tags_clause, tags_params, param_idx = build_tags_where_clause(
tags=tags, param_offset=param_idx, table_alias="", match=tags_match
)
if tags_clause:
# Remove leading "AND " from clause since we're building filters list
filters.append(tags_clause.replace("AND ", "", 1))
params.extend(tags_params)
elif isolation_mode:
# Isolation mode: only include directives with empty/null tags
# This ensures tag-scoped directives don't apply to untagged operations
filters.append("(tags IS NULL OR tags = '{}')")
params.extend([limit, offset])

View file

@ -312,6 +312,49 @@ class TestDirectiveTags:
# Cleanup
await memory.delete_bank(bank_id, request_context=request_context)
async def test_list_all_directives_without_filter(self, memory: MemoryEngine, request_context):
"""Test that listing directives without tags returns ALL directives (both tagged and untagged)."""
bank_id = f"test-directive-list-all-{uuid.uuid4().hex[:8]}"
# Ensure bank exists
await memory.get_bank_profile(bank_id, request_context=request_context)
# Create untagged directive
await memory.create_directive(
bank_id=bank_id,
name="Untagged Directive",
content="This has no tags",
request_context=request_context,
)
# Create tagged directive
await memory.create_directive(
bank_id=bank_id,
name="Tagged Directive",
content="This has tags",
tags=["project-x"],
request_context=request_context,
)
# List ALL directives (no tag filter, isolation_mode defaults to False)
all_directives = await memory.list_directives(
bank_id=bank_id,
request_context=request_context,
)
# Should return BOTH tagged and untagged directives
assert len(all_directives) == 2
directive_names = {d["name"] for d in all_directives}
assert "Untagged Directive" in directive_names
assert "Tagged Directive" in directive_names
# Verify the tagged directive has its tags
tagged = next(d for d in all_directives if d["name"] == "Tagged Directive")
assert tagged["tags"] == ["project-x"]
# Cleanup
await memory.delete_bank(bank_id, request_context=request_context)
class TestReflect:
"""Test reflect endpoint."""
@ -399,6 +442,161 @@ class TestDirectivesInReflect:
# Cleanup
await memory.delete_bank(bank_id, request_context=request_context)
async def test_tagged_directive_not_applied_without_tags(self, memory: MemoryEngine, request_context):
"""Test that directives with tags are NOT applied to untagged reflect operations."""
bank_id = f"test-directive-isolation-{uuid.uuid4().hex[:8]}"
# Ensure bank exists
await memory.get_bank_profile(bank_id, request_context=request_context)
# Add some untagged content
await memory.retain_batch_async(
bank_id=bank_id,
contents=[
{"content": "The sky is blue."},
{"content": "Water is wet."},
],
request_context=request_context,
)
# Add some tagged content for the project-x context
await memory.retain_batch_async(
bank_id=bank_id,
contents=[
{"content": "The sky is blue according to project X standards.", "tags": ["project-x"]},
{"content": "Project X color guidelines specify sky is blue.", "tags": ["project-x"]},
],
request_context=request_context,
)
await memory.wait_for_background_tasks()
# Create an untagged directive (should be applied)
await memory.create_directive(
bank_id=bank_id,
name="General Policy",
content="Always be polite and start responses with 'Hello!'",
request_context=request_context,
)
# Create a tagged directive (should NOT be applied to untagged reflect)
await memory.create_directive(
bank_id=bank_id,
name="Tagged Policy",
content="ALWAYS respond in ALL CAPS and end with 'PROJECT-X ONLY'",
tags=["project-x"],
request_context=request_context,
)
# Run reflect without tags - should only apply the untagged directive
result = await memory.reflect_async(
bank_id=bank_id,
query="What color is the sky?",
request_context=request_context,
)
response_lower = result.text.lower()
# Should follow the untagged directive (polite greeting)
assert "hello" in response_lower, f"Expected 'Hello' from untagged directive, but got: {result.text}"
# Should NOT follow the tagged directive (all caps and PROJECT-X)
# If it did follow, the entire response would be in caps
all_caps = result.text.replace(" ", "").replace("!", "").replace(".", "").isupper()
assert not all_caps, f"Tagged directive was incorrectly applied to untagged operation: {result.text}"
assert "project-x only" not in response_lower, f"Tagged directive was incorrectly applied: {result.text}"
# Now run reflect WITH the tag - should apply BOTH directives
result_tagged = await memory.reflect_async(
bank_id=bank_id,
query="What color is the sky?",
tags=["project-x"],
tags_match="all_strict",
request_context=request_context,
)
response_tagged_lower = result_tagged.text.lower()
# With strict matching and tags, should apply the tagged directive
assert "project-x only" in response_tagged_lower, f"Tagged directive should be applied with tags: {result_tagged.text}"
# Cleanup
await memory.delete_bank(bank_id, request_context=request_context)
async def test_reflect_based_on_structure(self, memory: MemoryEngine, request_context):
"""Test that reflect returns correct based_on structure with directives and memories separated."""
bank_id = f"test-reflect-based-on-{uuid.uuid4().hex[:8]}"
# Ensure bank exists
await memory.get_bank_profile(bank_id, request_context=request_context)
# Add some memories
await memory.retain_batch_async(
bank_id=bank_id,
contents=[
{"content": "Alice works at Google as a software engineer."},
{"content": "Bob is a product manager at Microsoft."},
{"content": "The team meets every Monday at 9am."},
],
request_context=request_context,
)
await memory.wait_for_background_tasks()
# Create a directive
directive = await memory.create_directive(
bank_id=bank_id,
name="Professional Tone",
content="Always maintain a professional and formal tone in responses.",
request_context=request_context,
)
directive_id = directive["id"]
# Run reflect which returns the core result
result = await memory.reflect_async(
bank_id=bank_id,
query="Who works at Google?",
request_context=request_context,
)
# Verify based_on structure exists
assert result.based_on is not None
# Verify directives key exists and contains our directive
assert "directives" in result.based_on
directives_list = result.based_on.get("directives", [])
# Verify directives are dicts with id, name, content (not MemoryFact objects)
assert len(directives_list) > 0, "Should have at least one directive"
directive_found = False
for d in directives_list:
assert isinstance(d, dict), f"Directive should be dict, got {type(d)}"
assert "id" in d, "Directive dict should have 'id'"
assert "name" in d, "Directive dict should have 'name'"
assert "content" in d, "Directive dict should have 'content'"
# Check if this is our directive
if d["id"] == directive_id:
directive_found = True
assert d["name"] == "Professional Tone"
assert "professional" in d["content"].lower()
assert directive_found, f"Our directive {directive_id} should be in based_on.directives"
# Verify memories (world/experience) are separate from directives
has_memories = "world" in result.based_on or "experience" in result.based_on
assert has_memories, "Should have world or experience memories"
# Verify that if mental-models key exists, it's separate from directives
if "mental-models" in result.based_on:
mental_models = result.based_on.get("mental-models", [])
# Verify mental models are MemoryFact objects, not dicts like directives
for mm in mental_models:
assert hasattr(mm, "fact_type"), "Mental model should be MemoryFact with fact_type"
assert mm.fact_type == "mental-models"
assert hasattr(mm, "context")
assert "mental model" in mm.context.lower()
# Cleanup
await memory.delete_bank(bank_id, request_context=request_context)
class TestDirectivesPromptInjection:
"""Test that directives are properly injected into the system prompt."""

View file

@ -55,6 +55,7 @@ interface ReflectResponseBasedOnFact {
id: string;
text: string;
type: string;
context?: string;
}
interface ReflectResponse {
@ -978,6 +979,30 @@ function MentalModelDetailPanel({
.flatMap(([factType, facts]) => facts.map((fact) => ({ ...fact, factType })))
: [];
// Helper to determine display label for fact type
const getFactTypeDisplay = (fact: any) => {
if (fact.factType === "mental-models") {
// Check context to distinguish directives from mental models
if (fact.context?.includes("directive")) {
return {
label: "directive",
color: "bg-purple-500/10 text-purple-600 dark:text-purple-400",
};
}
return {
label: "mental model",
color: "bg-indigo-500/10 text-indigo-600 dark:text-indigo-400",
};
}
if (fact.factType === "world") {
return { label: "world", color: "bg-blue-500/10 text-blue-600 dark:text-blue-400" };
}
if (fact.factType === "experience") {
return { label: "experience", color: "bg-green-500/10 text-green-600 dark:text-green-400" };
}
return { label: fact.factType, color: "bg-slate-500/10 text-slate-600 dark:text-slate-400" };
};
// Observations are now in based_on with type=observation
const observations = mentalModel.reflect_response?.based_on?.observation || [];
@ -1101,35 +1126,32 @@ function MentalModelDetailPanel({
Based On ({basedOnFacts.length} {basedOnFacts.length === 1 ? "fact" : "facts"})
</div>
<div className="space-y-3">
{basedOnFacts.map((fact, i) => (
<div
key={fact.id || i}
className="p-4 bg-muted/50 rounded-lg border border-border/50"
>
<div className="flex items-start justify-between gap-2 mb-2">
<span
className={`px-2 py-0.5 rounded text-xs font-medium ${
fact.factType === "world"
? "bg-blue-500/10 text-blue-600 dark:text-blue-400"
: fact.factType === "experience"
? "bg-green-500/10 text-green-600 dark:text-green-400"
: "bg-purple-500/10 text-purple-600 dark:text-purple-400"
}`}
>
{fact.factType}
</span>
<Button
variant="outline"
size="sm"
className="h-6 text-xs"
onClick={() => setViewMemoryId(fact.id)}
>
View
</Button>
{basedOnFacts.map((fact, i) => {
const display = getFactTypeDisplay(fact);
return (
<div
key={fact.id || i}
className="p-4 bg-muted/50 rounded-lg border border-border/50"
>
<div className="flex items-start justify-between gap-2 mb-2">
<span
className={`px-2 py-0.5 rounded text-xs font-medium ${display.color}`}
>
{display.label}
</span>
<Button
variant="outline"
size="sm"
className="h-6 text-xs"
onClick={() => setViewMemoryId(fact.id)}
>
View
</Button>
</div>
<p className="text-sm text-foreground leading-relaxed">{fact.text}</p>
</div>
<p className="text-sm text-foreground leading-relaxed">{fact.text}</p>
</div>
))}
);
})}
</div>
</div>
)}