diff --git a/docker/standalone/start-all.sh b/docker/standalone/start-all.sh index e4e5d527..e1bf9443 100755 --- a/docker/standalone/start-all.sh +++ b/docker/standalone/start-all.sh @@ -97,7 +97,7 @@ fi if [ "$ENABLE_CP" = "true" ]; then echo "šŸŽ›ļø Starting Control Plane..." cd /app/control-plane - PORT=9999 node server.js & + PORT="${HINDSIGHT_CP_PORT:-9999}" node server.js & CP_PID=$! PIDS+=($CP_PID) else @@ -110,7 +110,7 @@ echo "āœ… Hindsight is running!" echo "" echo "šŸ“ Access:" if [ "$ENABLE_CP" = "true" ]; then - echo " Control Plane: http://localhost:9999" + echo " Control Plane: http://localhost:${HINDSIGHT_CP_PORT:-9999}" fi if [ "$ENABLE_API" = "true" ]; then echo " API: http://localhost:8888" diff --git a/hindsight-api/hindsight_api/alembic/versions/a2b3c4d5e6f7_add_text_signals_column.py b/hindsight-api/hindsight_api/alembic/versions/a2b3c4d5e6f7_add_text_signals_column.py new file mode 100644 index 00000000..bbda2b06 --- /dev/null +++ b/hindsight-api/hindsight_api/alembic/versions/a2b3c4d5e6f7_add_text_signals_column.py @@ -0,0 +1,88 @@ +"""Add text_signals column to memory_units for enriched BM25 indexing. + +text_signals stores a denormalized space-separated string of entity names +(and future signals) to improve full-text search recall without polluting +the stored fact text. + +- vchord: text_signals included in tokenize() at insert time +- native: search_vector GENERATED column regenerated to include text_signals +- pg_textsearch: no change (index only supports a single base column) + +Revision ID: a2b3c4d5e6f7 +Revises: z1u2v3w4x5y6 +Create Date: 2026-02-28 +""" + +import os +from collections.abc import Sequence + +from alembic import context, op + +revision: str = "a2b3c4d5e6f7" +down_revision: str | Sequence[str] | None = "aa2b3c4d5e6f" +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 _detect_text_search_extension() -> str: + return os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower() + + +def upgrade() -> None: + schema = _get_schema_prefix() + table = f"{schema}memory_units" + text_search_ext = _detect_text_search_extension() + + # Add text_signals column (nullable TEXT, populated at retain time) + op.execute(f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS text_signals TEXT") + + if text_search_ext == "native": + # Native PostgreSQL: drop and recreate the GENERATED tsvector column to include text_signals + op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS search_vector") + op.execute(f""" + ALTER TABLE {table} + ADD COLUMN search_vector tsvector + GENERATED ALWAYS AS ( + to_tsvector('english', + COALESCE(text, '') || ' ' || + COALESCE(context, '') || ' ' || + COALESCE(text_signals, '') + ) + ) STORED + """) + # Recreate GIN index (was dropped with the column) + op.execute(f""" + CREATE INDEX IF NOT EXISTS idx_memory_units_text_search + ON {table} USING gin(search_vector) + """) + + # vchord: tokenize() call in fact_storage.py is updated to include text_signals at insert time + # pg_textsearch: no change — index operates on the base `text` column only + + +def downgrade() -> None: + schema = _get_schema_prefix() + table = f"{schema}memory_units" + text_search_ext = _detect_text_search_extension() + + if text_search_ext == "native": + op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_text_search") + op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS search_vector") + op.execute(f""" + ALTER TABLE {table} + ADD COLUMN search_vector tsvector + GENERATED ALWAYS AS ( + to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, '')) + ) STORED + """) + op.execute(f""" + CREATE INDEX idx_memory_units_text_search + ON {table} USING gin(search_vector) + """) + + op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS text_signals") diff --git a/hindsight-api/hindsight_api/alembic/versions/b4c5d6e7f8a9_backfill_observation_scopes.py b/hindsight-api/hindsight_api/alembic/versions/b4c5d6e7f8a9_backfill_observation_scopes.py new file mode 100644 index 00000000..57784593 --- /dev/null +++ b/hindsight-api/hindsight_api/alembic/versions/b4c5d6e7f8a9_backfill_observation_scopes.py @@ -0,0 +1,34 @@ +"""Backfill observation_scopes column if missing. + +This migration ensures observation_scopes exists even on databases that had +revision z1u2v3w4x5y6 applied when it referred to the old text_signals migration +(before it was renamed to a2b3c4d5e6f7). The ADD COLUMN IF NOT EXISTS makes this +a no-op on databases that already have the column. + +Revision ID: b4c5d6e7f8a9 +Revises: a2b3c4d5e6f7 +Create Date: 2026-03-02 +""" + +from collections.abc import Sequence + +from alembic import context, op + +revision: str = "b4c5d6e7f8a9" +down_revision: str | Sequence[str] | None = "a2b3c4d5e6f7" +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() + op.execute(f"ALTER TABLE {schema}memory_units ADD COLUMN IF NOT EXISTS observation_scopes JSONB") + + +def downgrade() -> None: + pass # intentionally no-op — safe to leave the column in place diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index e705430b..409485b0 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -690,6 +690,13 @@ class HindsightConfig: consolidation_max_tokens: int observations_mission: str | None + # 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}]}] + entity_labels: list | None + # Whether to extract regular named entities alongside entity labels (default: True) + # When False: only label entities are extracted (or no entities at all if no labels configured) + entities_allow_free_form: bool + # Reflect agent settings reflect_mission: str | None @@ -770,6 +777,9 @@ class HindsightConfig: "retain_extraction_mode", "retain_mission", "retain_custom_instructions", + # Entity labels (controlled vocabulary for entity classification) + "entity_labels", + "entities_allow_free_form", # Consolidation settings "enable_observations", "observations_mission", @@ -1118,6 +1128,8 @@ class HindsightConfig: os.getenv(ENV_CONSOLIDATION_MAX_TOKENS, str(DEFAULT_CONSOLIDATION_MAX_TOKENS)) ), observations_mission=os.getenv(ENV_OBSERVATIONS_MISSION) or DEFAULT_OBSERVATIONS_MISSION, + entity_labels=None, + entities_allow_free_form=True, # Database migrations run_migrations_on_startup=os.getenv(ENV_RUN_MIGRATIONS_ON_STARTUP, "true").lower() == "true", # Database connection pool diff --git a/hindsight-api/hindsight_api/engine/entity_resolver.py b/hindsight-api/hindsight_api/engine/entity_resolver.py index 1d4feb50..770b72d1 100644 --- a/hindsight-api/hindsight_api/engine/entity_resolver.py +++ b/hindsight-api/hindsight_api/engine/entity_resolver.py @@ -12,6 +12,7 @@ import asyncpg from .db_utils import acquire_with_retry from .memory_engine import fq_table +from .retain.entity_labels import build_labels_lookup as _build_labels_lookup_from_config # Load spaCy model (singleton) _nlp = None @@ -31,6 +32,11 @@ class EntityResolver: """ self.pool = pool + @staticmethod + def _build_labels_lookup(entity_labels: list | None) -> set[str]: + """Build a set of valid 'key:value' entity label strings for fast lookup.""" + return _build_labels_lookup_from_config(entity_labels) + async def resolve_entities_batch( self, bank_id: str, @@ -38,6 +44,7 @@ class EntityResolver: context: str, unit_event_date, conn=None, + entity_labels: list | None = None, ) -> list[str]: """ Resolve multiple entities in batch (MUCH faster than sequential). @@ -58,14 +65,25 @@ class EntityResolver: if not entities_data: return [] + taxonomy_lookup = self._build_labels_lookup(entity_labels) if conn is None: async with acquire_with_retry(self.pool) as conn: - return await self._resolve_entities_batch_impl(conn, bank_id, entities_data, context, unit_event_date) + return await self._resolve_entities_batch_impl( + conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup + ) else: - return await self._resolve_entities_batch_impl(conn, bank_id, entities_data, context, unit_event_date) + return await self._resolve_entities_batch_impl( + conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup + ) async def _resolve_entities_batch_impl( - self, conn, bank_id: str, entities_data: list[dict], context: str, unit_event_date + self, + conn, + bank_id: str, + entities_data: list[dict], + context: str, + unit_event_date, + taxonomy_lookup: set[str] | None = None, ) -> list[str]: # Query ALL candidates for this bank all_entities = await conn.fetch( @@ -135,12 +153,19 @@ class EntityResolver: entities_to_update = [] # (entity_id, event_date) entities_to_create = [] # (idx, entity_data, event_date) + taxonomy_lookup = taxonomy_lookup or set() + for idx, entity_data in enumerate(entities_data): entity_text = entity_data["text"] nearby_entities = entity_data.get("nearby_entities", []) # Use per-entity date if available, otherwise fall back to batch-level date entity_event_date = entity_data.get("event_date", unit_event_date) + # Taxonomy entities: skip fuzzy matching, use exact canonical name + if taxonomy_lookup and entity_text.lower() in taxonomy_lookup: + entities_to_create.append((idx, entity_data, entity_event_date)) + continue + candidates = all_candidates.get(entity_text, []) if not candidates: diff --git a/hindsight-api/hindsight_api/engine/reflect/agent.py b/hindsight-api/hindsight_api/engine/reflect/agent.py index cdac3cdd..9874fdf8 100644 --- a/hindsight-api/hindsight_api/engine/reflect/agent.py +++ b/hindsight-api/hindsight_api/engine/reflect/agent.py @@ -431,7 +431,9 @@ async def run_reflect_agent( if is_last: # Force text response on last iteration - no tools - prompt = build_final_prompt(query, context_history, bank_profile, context, max_context_tokens=max_context_tokens) + prompt = build_final_prompt( + query, context_history, bank_profile, context, max_context_tokens=max_context_tokens + ) llm_start = time.time() response, usage = await llm_config.call( messages=[ @@ -486,7 +488,9 @@ async def run_reflect_agent( f"[REFLECT {reflect_id}] Context budget exceeded on iteration {iteration + 1}: " f"~{estimated_tokens} tokens >= {max_context_tokens} limit. Forcing final synthesis." ) - prompt = build_final_prompt(query, context_history, bank_profile, context, max_context_tokens=max_context_tokens) + prompt = build_final_prompt( + query, context_history, bank_profile, context, max_context_tokens=max_context_tokens + ) llm_start = time.time() response, usage = await llm_config.call( messages=[ @@ -588,7 +592,9 @@ async def run_reflect_agent( # For other errors: retry if no evidence yet (but cap consecutive errors to avoid long hangs) elif not has_gathered_evidence and iteration < max_iterations - 1 and consecutive_errors < 2: continue - prompt = build_final_prompt(query, context_history, bank_profile, context, max_context_tokens=max_context_tokens) + prompt = build_final_prompt( + query, context_history, bank_profile, context, max_context_tokens=max_context_tokens + ) llm_start = time.time() response, usage = await llm_config.call( messages=[ @@ -659,7 +665,9 @@ async def run_reflect_agent( directives_applied=directives_applied, ) # Empty response, force final - prompt = build_final_prompt(query, context_history, bank_profile, context, max_context_tokens=max_context_tokens) + prompt = build_final_prompt( + query, context_history, bank_profile, context, max_context_tokens=max_context_tokens + ) llm_start = time.time() response, usage = await llm_config.call( messages=[ diff --git a/hindsight-api/hindsight_api/engine/retain/embedding_processing.py b/hindsight-api/hindsight_api/engine/retain/embedding_processing.py index 7956184e..eccf874d 100644 --- a/hindsight-api/hindsight_api/engine/retain/embedding_processing.py +++ b/hindsight-api/hindsight_api/engine/retain/embedding_processing.py @@ -27,14 +27,21 @@ def augment_texts_with_dates(facts: list[ExtractedFact], format_date_fn) -> list """ augmented_texts = [] for fact in facts: - # Use occurred_start as the representative date + # Use occurred_start as the representative date, fall back to mentioned_at fact_date = fact.occurred_start or fact.mentioned_at + # Augment text with date and entity names for embedding (but store original text in DB) + # Entity names (including key:value labels) improve retrieval without polluting stored content if fact_date is not None: readable_date = format_date_fn(fact_date) - # Augment text with date for embedding (but store original text in DB) - augmented_text = f"{fact.fact_text} (happened in {readable_date})" + if fact.occurred_end and fact.occurred_end != fact.occurred_start: + readable_end = format_date_fn(fact.occurred_end) + augmented_text = f"{fact.fact_text} (happened from {readable_date} to {readable_end})" + else: + augmented_text = f"{fact.fact_text} (happened in {readable_date})" else: augmented_text = fact.fact_text + if fact.entities: + augmented_text = f"{augmented_text} [{', '.join(fact.entities)}]" augmented_texts.append(augmented_text) return augmented_texts diff --git a/hindsight-api/hindsight_api/engine/retain/entity_labels.py b/hindsight-api/hindsight_api/engine/retain/entity_labels.py new file mode 100644 index 00000000..d86a4890 --- /dev/null +++ b/hindsight-api/hindsight_api/engine/retain/entity_labels.py @@ -0,0 +1,194 @@ +""" +Entity labels models and helpers for retain pipeline. + +Defines a controlled vocabulary of key:value classification labels +(e.g., 'pedagogy:scaffolding', 'interest:active') that are extracted +at retain time and stored as entities. +""" + +from typing import Literal + +from pydantic import BaseModel, Field, create_model + + +class LabelValue(BaseModel): + """A single allowed value for a label group.""" + + value: str + description: str = "" + + +class LabelGroup(BaseModel): + """A label group (dimension) with its type and allowed values.""" + + key: str + description: str = "" + type: Literal["value", "multi-values", "text"] = "value" + optional: bool = True + tag: bool = False + values: list[LabelValue] = [] + + +class EntityLabelsConfig(BaseModel): + """Entity labels configuration for a bank (controlled vocabulary).""" + + attributes: list[LabelGroup] = [] + + +def parse_entity_labels(raw: dict | list | None) -> EntityLabelsConfig | None: + """ + Parse raw entity labels config into EntityLabelsConfig. + + Accepts: + - None → returns None + - list → list of attribute dicts (each may use legacy free_values/multi_value or new type field) + - dict → {attributes: [...]} + + Legacy migration (backward-compat): + - free_values=True → type="text" + - multi_value=True → type="multi-values" + - neither / free_values=False → type="value" + + Args: + raw: Raw entity labels config from bank config + + Returns: + EntityLabelsConfig or None if raw is None/empty + """ + if raw is None: + return None + + if isinstance(raw, list): + if not raw: + return None + attributes = [LabelGroup.model_validate(_migrate_label_group(a)) for a in raw] + return EntityLabelsConfig(attributes=attributes) + + if isinstance(raw, dict): + attrs_raw = raw.get("attributes", []) + if not attrs_raw: + return None + attributes = [LabelGroup.model_validate(_migrate_label_group(a)) for a in attrs_raw] + return EntityLabelsConfig(attributes=attributes) + + return None + + +def _migrate_label_group(raw: dict) -> dict: + """Migrate legacy free_values/multi_value fields to the new type field.""" + if not isinstance(raw, dict) or "type" in raw: + return raw + patched = dict(raw) + if patched.get("free_values"): + patched["type"] = "text" + elif patched.get("multi_value"): + patched["type"] = "multi-values" + else: + patched["type"] = "value" + # Remove legacy keys so Pydantic doesn't error on unknown fields + patched.pop("free_values", None) + patched.pop("multi_value", None) + return patched + + +def build_labels_model(labels_cfg: EntityLabelsConfig) -> type[BaseModel] | None: + """ + Build a dynamic Pydantic model for structured label extraction. + + Each LabelGroup becomes a typed field based on its type: + - type="text" → str | None (always optional) + - type="value", optional=True → Literal["v1","v2"] | None + - type="value", optional=False → Literal["v1","v2"] (required) + - type="multi-values" → list[Literal["v1","v2"]] + + Args: + labels_cfg: Parsed EntityLabelsConfig + + Returns: + Dynamic Pydantic model class, or None if no groups defined + """ + fields: dict = {} + for group in labels_cfg.attributes: + if not group.key: + continue + description = group.description or group.key + + if group.type == "text": + # Free-form: any string value accepted, always optional + fields[group.key] = (str | None, Field(default=None, description=description)) + else: + # Enum-constrained: must have defined values + if not group.values: + continue + values = tuple(v.value for v in group.values if v.value) + if not values: + continue + # Literal[("v1", "v2")] is equivalent to Literal["v1", "v2"] in Python 3.11+ + literal_type = Literal[values] # type: ignore[valid-type] + if group.type == "multi-values": + fields[group.key] = ( + list[literal_type], # type: ignore[valid-type] + Field(default_factory=list, description=description), + ) + elif group.optional: + fields[group.key] = ( + literal_type | None, # type: ignore[valid-type] + Field(default=None, description=description), + ) + else: + fields[group.key] = ( + literal_type, # type: ignore[valid-type] + Field(description=description), + ) + + if not fields: + return None + + return create_model("Labels", **fields) + + +def is_label_entity(text: str, labels_cfg: EntityLabelsConfig, labels_lookup: set[str]) -> bool: + """ + Return True if entity text belongs to any configured label group. + + For enum groups: checks the pre-built lookup set. + For text groups: checks that the text starts with a known key prefix. + """ + if text.lower() in labels_lookup: + return True + for group in labels_cfg.attributes: + if group.type == "text" and group.key and text.lower().startswith(f"{group.key.lower()}:"): + return True + return False + + +def build_labels_lookup(labels_cfg: EntityLabelsConfig | list | None) -> set[str]: + """ + Build a set of valid 'key:value' label strings (lowercase) for fast lookup. + + Accepts either EntityLabelsConfig or raw list/None for backwards compatibility. + + Args: + labels_cfg: EntityLabelsConfig, raw list of attribute dicts, or None + + Returns: + Set of lowercase 'key:value' strings + """ + if labels_cfg is None: + return set() + + # Accept raw list/dict for backwards compatibility + if not isinstance(labels_cfg, EntityLabelsConfig): + parsed = parse_entity_labels(labels_cfg) + if parsed is None: + return set() + labels_cfg = parsed + + valid = set() + for group in labels_cfg.attributes: + if group.type == "text": + continue # No fixed vocabulary — all values accepted in post-processing + for v in group.values: + if group.key and v.value: + valid.add(f"{group.key}:{v.value}".lower()) + return valid diff --git a/hindsight-api/hindsight_api/engine/retain/entity_processing.py b/hindsight-api/hindsight_api/engine/retain/entity_processing.py index d1801692..d4783c26 100644 --- a/hindsight-api/hindsight_api/engine/retain/entity_processing.py +++ b/hindsight-api/hindsight_api/engine/retain/entity_processing.py @@ -20,6 +20,7 @@ async def process_entities_batch( facts: list[ProcessedFact], log_buffer: list[str] = None, user_entities_per_content: dict[int, list[dict]] = None, + entity_labels: list | None = None, ) -> list[EntityLink]: """ Process entities for all facts and create entity links. @@ -90,6 +91,7 @@ async def process_entities_batch( fact_dates, entities_per_fact, log_buffer, # Pass log_buffer for detailed logging + entity_labels=entity_labels, ) return entity_links diff --git a/hindsight-api/hindsight_api/engine/retain/fact_extraction.py b/hindsight-api/hindsight_api/engine/retain/fact_extraction.py index af480129..3c885f02 100644 --- a/hindsight-api/hindsight_api/engine/retain/fact_extraction.py +++ b/hindsight-api/hindsight_api/engine/retain/fact_extraction.py @@ -10,13 +10,20 @@ import json import logging import re from datetime import datetime, timedelta -from typing import Literal +from typing import Literal, cast -from pydantic import BaseModel, ConfigDict, Field, field_validator +from pydantic import BaseModel, ConfigDict, Field, create_model, field_validator from ...config import get_config from ..llm_wrapper import LLMConfig, OutputTooLongError from ..response_models import TokenUsage +from .entity_labels import ( + EntityLabelsConfig, + build_labels_lookup, + build_labels_model, + is_label_entity, + parse_entity_labels, +) def _infer_temporal_date(fact_text: str, event_date: datetime | None) -> str | None: @@ -692,10 +699,62 @@ Example: "Lost job → couldn't pay rent → moved apartment" - Fact 2: Moved apartment, causal_relations: [{target_index: 1, relation_type: "caused_by"}]""" +def _build_labels_prompt_section(labels_cfg: EntityLabelsConfig | list | None, free_form_entities: bool = True) -> str: + """Build the entity labels classification section for the extraction prompt.""" + if labels_cfg is None: + return "" + + # Accept raw list for backwards compatibility + if isinstance(labels_cfg, list): + if not labels_cfg: + return "" + labels_cfg = parse_entity_labels(labels_cfg) + if labels_cfg is None: + return "" + + if not labels_cfg.attributes: + return "" + + if free_form_entities: + entities_instruction = "Classify each fact using the structured 'labels' field below. Continue extracting regular named entities in the 'entities' field." + else: + entities_instruction = "Classify each fact using the structured 'labels' field below. Do NOT add regular named entities — labels-only mode." + + lines = [ + "\n\n══════════════════════════════════════════════════════════════════════════", + "ENTITY LABELS - CLASSIFICATION ATTRIBUTES", + "══════════════════════════════════════════════════════════════════════════", + "", + entities_instruction, + "", + "For each fact, fill the 'labels' object. Each field is a label group:", + "", + ] + + for attr in labels_cfg.attributes: + if attr.type == "text": + # Free-text: no predefined values — LLM writes any relevant string or null + lines.append(f"- {attr.key} (free text or null): {attr.description}") + else: + mode = "multi-value (list)" if attr.type == "multi-values" else "single value or null" + lines.append(f"- {attr.key} ({mode}): {attr.description}") + for v in attr.values: + desc = f" — {v.description}" if v.description else "" + lines.append(f' • "{v.value}"{desc}') + lines.append("") + + lines.append("Only assign labels when clearly applicable. Leave null/empty if the fact does not match.") + return "\n".join(lines) + + def _build_extraction_prompt_and_schema(config) -> tuple[str, type]: """ Build extraction prompt and response schema based on config. + When a taxonomy is configured, dynamically builds a Pydantic model with a + typed `taxonomy_entities` field using an Enum built from valid taxonomy values. + This enables JSON schema enforcement for structured outputs. + Returns: Tuple of (prompt, response_schema) """ @@ -738,9 +797,53 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]: # Add causal relationships section if enabled if extract_causal_links: prompt = prompt + CAUSAL_RELATIONSHIPS_SECTION - response_schema = FactExtractionResponseVerbose if extraction_mode == "verbose" else FactExtractionResponse + base_fact_class = ExtractedFactVerbose if extraction_mode == "verbose" else ExtractedFact + base_response_class = FactExtractionResponseVerbose if extraction_mode == "verbose" else FactExtractionResponse else: - response_schema = FactExtractionResponseNoCausal + base_fact_class = ExtractedFactNoCausal + base_response_class = FactExtractionResponseNoCausal + + # Add entity labels section if configured and build dynamic schema + entity_labels_raw = getattr(config, "entity_labels", None) + labels_cfg = parse_entity_labels(entity_labels_raw) + free_form_entities = getattr(config, "entities_allow_free_form", True) + labels_section = _build_labels_prompt_section(labels_cfg, free_form_entities) + if labels_section: + prompt = prompt + labels_section + + response_schema = base_response_class + + if labels_cfg and labels_cfg.attributes: + LabelsModel = build_labels_model(labels_cfg) + if LabelsModel is not None: + dynamic_fields: dict = { + "labels": ( + LabelsModel, + Field( + description="Classification labels for this fact. Fill each applicable field; leave others null/empty." + ), + ) + } + if not free_form_entities: + dynamic_fields["entities"] = ( + list[Entity] | None, + Field(default=None, description="Leave empty — labels-only mode"), + ) + # Inherit parent's required fields and add 'labels' so it appears in the JSON schema + # required array (the base class json_schema_extra overrides required entirely) + base_extra = base_fact_class.model_config.get("json_schema_extra") + base_required = cast(dict, base_extra).get("required", []) if isinstance(base_extra, dict) else [] + DynamicFact = create_model( + "LabelsFact", + __base__=base_fact_class, + __config__=ConfigDict( + json_schema_mode="validation", + json_schema_extra={"required": [*base_required, "labels"]}, + ), + **dynamic_fields, + ) + DynamicResponse = create_model("LabelsResponse", facts=(list[DynamicFact], ...)) # type: ignore[valid-type] + response_schema = DynamicResponse return prompt, response_schema @@ -997,9 +1100,9 @@ async def _extract_facts_from_chunk( # Add entities if present (validate as Entity objects) # LLM sometimes returns strings instead of {"text": "..."} format entities = get_value("entities") + validated_entities = [] if entities: # Validate and normalize each entity - validated_entities = [] for ent in entities: if isinstance(ent, str): # Normalize string to Entity object @@ -1009,8 +1112,48 @@ async def _extract_facts_from_chunk( validated_entities.append(Entity.model_validate(ent)) except Exception as e: logger.warning(f"Invalid entity {ent}: {e}") - if validated_entities: - fact_data["entities"] = validated_entities + + # Post-process label entities from structured labels object + entity_labels_raw = getattr(config, "entity_labels", None) + labels_cfg = parse_entity_labels(entity_labels_raw) + free_form_entities = getattr(config, "entities_allow_free_form", True) + if labels_cfg and labels_cfg.attributes: + labels_lookup = build_labels_lookup(labels_cfg) + labels_data = llm_fact.get("labels") or {} + if isinstance(labels_data, dict): + existing_texts_lower = {e.text.lower() for e in validated_entities} + for group in labels_cfg.attributes: + value = labels_data.get(group.key) + if not value: + continue + values_list = value if isinstance(value, list) else [value] + for v in values_list: + if not isinstance(v, str) or not v.strip() or v.lower() in ("none", "null", "n/a"): + continue + label_str = f"{group.key}:{v.strip()}" + if group.type == "text": + if label_str.lower() not in existing_texts_lower: + validated_entities.append(Entity(text=label_str)) + existing_texts_lower.add(label_str.lower()) + elif ( + label_str.lower() in labels_lookup and label_str.lower() not in existing_texts_lower + ): + validated_entities.append(Entity(text=label_str)) + existing_texts_lower.add(label_str.lower()) + else: + logger.warning(f"Label '{label_str}' not in valid label values, skipping") + + # In labels-only mode, keep only label entities + if not free_form_entities: + validated_entities = [ + e for e in validated_entities if is_label_entity(e.text, labels_cfg, labels_lookup) + ] + elif not free_form_entities: + # No labels but free_form disabled: clear all entities + validated_entities = [] + + if validated_entities: + fact_data["entities"] = validated_entities # Add per-fact causal relations (only if enabled in config) if extract_causal_links: @@ -1606,8 +1749,8 @@ async def extract_facts_from_contents_batch_api( # Entities entities = get_value("entities") + validated_entities = [] if entities: - validated_entities = [] for ent in entities: if isinstance(ent, str): validated_entities.append(Entity(text=ent)) @@ -1616,8 +1759,45 @@ async def extract_facts_from_contents_batch_api( validated_entities.append(Entity.model_validate(ent)) except Exception: pass - if validated_entities: - fact_data["entities"] = validated_entities + + # Post-process label entities from structured labels object + entity_labels_raw = getattr(config, "entity_labels", None) + labels_cfg_batch = parse_entity_labels(entity_labels_raw) + free_form_entities_batch = getattr(config, "entities_allow_free_form", True) + if labels_cfg_batch and labels_cfg_batch.attributes: + labels_lookup_batch = build_labels_lookup(labels_cfg_batch) + labels_data = llm_fact.get("labels") or {} + if isinstance(labels_data, dict): + existing_texts_lower = {e.text.lower() for e in validated_entities} + for group in labels_cfg_batch.attributes: + value = labels_data.get(group.key) + if not value: + continue + values_list = value if isinstance(value, list) else [value] + for v in values_list: + if not isinstance(v, str) or not v.strip() or v.lower() in ("none", "null", "n/a"): + continue + label_str = f"{group.key}:{v.strip()}" + if group.type == "text": + if label_str.lower() not in existing_texts_lower: + validated_entities.append(Entity(text=label_str)) + existing_texts_lower.add(label_str.lower()) + elif ( + label_str.lower() in labels_lookup_batch + and label_str.lower() not in existing_texts_lower + ): + validated_entities.append(Entity(text=label_str)) + existing_texts_lower.add(label_str.lower()) + + if not free_form_entities_batch: + validated_entities = [ + e for e in validated_entities if is_label_entity(e.text, labels_cfg_batch, labels_lookup_batch) + ] + elif not free_form_entities_batch: + validated_entities = [] + + if validated_entities: + fact_data["entities"] = validated_entities # Causal relations if extract_causal_links: @@ -1718,6 +1898,9 @@ async def extract_facts_from_contents_batch_api( # Step 7: Add temporal offsets _add_temporal_offsets(extracted_facts, contents) + # Step 8: Auto-tag facts from label groups with tag=True + _inject_label_tags(extracted_facts, config) + logger.info(f"Batch API extracted {len(extracted_facts)} facts from {len(all_chunks_info)} chunks") return extracted_facts, chunks_metadata, total_usage @@ -1850,6 +2033,9 @@ async def extract_facts_from_contents( # Step 4: Add time offsets to preserve ordering within each content _add_temporal_offsets(extracted_facts, contents) + # Step 5: Auto-tag facts from label groups with tag=True + _inject_label_tags(extracted_facts, config) + return extracted_facts, chunks_metadata, total_usage @@ -1905,3 +2091,24 @@ def _add_temporal_offsets(facts: list[ExtractedFactType], contents: list[RetainC fact.occurred_end = parse_datetime_flexible(fact.occurred_end) + offset if fact.mentioned_at: fact.mentioned_at = parse_datetime_flexible(fact.mentioned_at) + offset + + +def _inject_label_tags(facts: list[ExtractedFactType], config) -> None: + """ + For label groups with tag=True, add extracted key:value label entities + to each fact's tags list. Modifies facts in place. + + This lets entity labels double as tags, enabling filtering via the + existing tags API without any extra query infrastructure. + """ + labels_cfg = parse_entity_labels(getattr(config, "entity_labels", None)) + if not labels_cfg: + return + tag_group_keys = {g.key.lower() for g in labels_cfg.attributes if g.tag} + if not tag_group_keys: + return + for fact in facts: + label_tags = [e for e in fact.entities if ":" in e and e.split(":", 1)[0].lower() in tag_group_keys] + if label_tags: + existing = set(fact.tags) + fact.tags = fact.tags + [t for t in label_tags if t not in existing] diff --git a/hindsight-api/hindsight_api/engine/retain/fact_storage.py b/hindsight-api/hindsight_api/engine/retain/fact_storage.py index 4d5cdc48..994b423f 100644 --- a/hindsight-api/hindsight_api/engine/retain/fact_storage.py +++ b/hindsight-api/hindsight_api/engine/retain/fact_storage.py @@ -48,6 +48,7 @@ async def insert_facts_batch( document_ids = [] tags_list = [] observation_scopes_list = [] + text_signals_list = [] for fact in facts: fact_texts.append(_sanitize_text(fact.fact_text)) @@ -73,6 +74,15 @@ async def insert_facts_batch( observation_scopes_list.append( json.dumps(fact.observation_scopes) if fact.observation_scopes is not None else None ) + # Build text_signals: entity names + date tokens for enriched BM25 indexing + signal_parts = [] + if fact.entities: + signal_parts.extend(e.name for e in fact.entities) + if fact.occurred_start: + signal_parts.append(fact.occurred_start.strftime("%B %-d %Y")) + if fact.occurred_end and fact.occurred_end != fact.occurred_start: + signal_parts.append(fact.occurred_end.strftime("%B %-d %Y")) + text_signals_list.append(" ".join(signal_parts) if signal_parts else None) # Batch insert all facts # Note: tags are passed as JSON strings and converted back to varchar[] via jsonb_array_elements_text + array_agg @@ -80,18 +90,19 @@ async def insert_facts_batch( config = get_config() if config.text_search_extension == "vchord": # VectorChord: manually tokenize and insert search_vector + # text_signals (entity names etc.) are included in the tokenize input for enriched BM25 query = f""" 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[] + $8::text[], $9::text[], $10::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[], $15::jsonb[], $16::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, - observation_scopes_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, - observation_scopes, search_vector) + observation_scopes, text_signals, search_vector) SELECT $1, text, embedding, event_date, occurred_start, occurred_end, mentioned_at, @@ -101,25 +112,29 @@ async def insert_facts_batch( '{{}}'::varchar[] ), observation_scopes_json, - tokenize(COALESCE(text, '') || ' ' || COALESCE(context, ''), 'llmlingua2')::bm25_catalog.bm25vector + text_signals, + tokenize( + COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, ''), + 'llmlingua2' + )::bm25_catalog.bm25vector FROM input_data RETURNING id """ else: # native or pg_textsearch - # Native PostgreSQL: search_vector is GENERATED ALWAYS, don't include it + # Native PostgreSQL: search_vector is GENERATED ALWAYS (expression includes text_signals), don't include it # pg_textsearch: indexes operate on base columns directly, don't populate search_vector query = f""" WITH input_data AS ( SELECT * FROM unnest( $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[] + $8::text[], $9::text[], $10::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[], $15::jsonb[], $16::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, - observation_scopes_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, - observation_scopes) + observation_scopes, text_signals) SELECT $1, text, embedding, event_date, occurred_start, occurred_end, mentioned_at, @@ -128,7 +143,8 @@ async def insert_facts_batch( (SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem), '{{}}'::varchar[] ), - observation_scopes_json + observation_scopes_json, + text_signals FROM input_data RETURNING id """ @@ -150,6 +166,7 @@ async def insert_facts_batch( document_ids, tags_list, observation_scopes_list, + text_signals_list, ) unit_ids = [str(row["id"]) for row in results] diff --git a/hindsight-api/hindsight_api/engine/retain/link_utils.py b/hindsight-api/hindsight_api/engine/retain/link_utils.py index 3985d80b..103dfc31 100644 --- a/hindsight-api/hindsight_api/engine/retain/link_utils.py +++ b/hindsight-api/hindsight_api/engine/retain/link_utils.py @@ -150,6 +150,7 @@ async def extract_entities_batch_optimized( fact_dates: list, llm_entities: list[list[dict]], log_buffer: list[str] = None, + entity_labels: list | None = None, ) -> list[tuple]: """ Process LLM-extracted entities for ALL facts in batch. @@ -239,6 +240,7 @@ async def extract_entities_batch_optimized( context=context, unit_event_date=None, # Not used when per-entity dates provided conn=conn, # Use main transaction connection + entity_labels=entity_labels, ) _log( diff --git a/hindsight-api/hindsight_api/engine/retain/orchestrator.py b/hindsight-api/hindsight_api/engine/retain/orchestrator.py index 223707d3..3d1829a7 100644 --- a/hindsight-api/hindsight_api/engine/retain/orchestrator.py +++ b/hindsight-api/hindsight_api/engine/retain/orchestrator.py @@ -472,6 +472,7 @@ async def retain_batch( non_duplicate_facts, log_buffer, user_entities_per_content=user_entities_per_content, + entity_labels=getattr(config, "entity_labels", None), ) log_buffer.append(f"[6] Process entities: {len(entity_links)} links in {time.time() - step_start:.3f}s") diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index 6ff01b31..d76766c0 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/hindsight-api/hindsight_api/main.py @@ -277,6 +277,8 @@ def main(): consolidation_llm_batch_size=config.consolidation_llm_batch_size, consolidation_max_tokens=config.consolidation_max_tokens, observations_mission=config.observations_mission, + entity_labels=config.entity_labels, + entities_allow_free_form=config.entities_allow_free_form, skip_llm_verification=config.skip_llm_verification, lazy_reranker=config.lazy_reranker, run_migrations_on_startup=config.run_migrations_on_startup, diff --git a/hindsight-api/tests/test_entity_labels.py b/hindsight-api/tests/test_entity_labels.py new file mode 100644 index 00000000..0328a437 --- /dev/null +++ b/hindsight-api/tests/test_entity_labels.py @@ -0,0 +1,1120 @@ +""" +Unit tests for entity labels models and helpers. + +Tests label parsing, enum building, prompt generation, lookup building, +entity post-processing, and embedding augmentation. + +Also includes LLM integration tests (require DB + LLM) that call retain +and verify label entities are extracted and stored correctly. +""" + +import uuid +from unittest.mock import MagicMock + +import pytest + +from hindsight_api.engine.retain.entity_labels import ( + EntityLabelsConfig, + LabelGroup, + LabelValue, + build_labels_lookup, + parse_entity_labels, +) + +# ─── parse_entity_labels ─────────────────────────────────────────────────────── + + +def test_parse_entity_labels_none(): + result = parse_entity_labels(None) + assert result is None + + +def test_parse_entity_labels_empty_list(): + result = parse_entity_labels([]) + assert result is None + + +def test_parse_entity_labels_list_format(): + """Legacy list format: just a list of attribute dicts (using new type field).""" + raw = [ + { + "key": "pedagogy", + "description": "Teaching strategy", + "type": "multi-values", + "values": [ + {"value": "scaffolding", "description": "Break down tasks"}, + {"value": "active_engagement", "description": "Group work"}, + ], + } + ] + result = parse_entity_labels(raw) + assert result is not None + assert isinstance(result, EntityLabelsConfig) + assert len(result.attributes) == 1 + attr = result.attributes[0] + assert attr.key == "pedagogy" + assert attr.type == "multi-values" + assert len(attr.values) == 2 + assert attr.values[0].value == "scaffolding" + + +def test_parse_entity_labels_dict_format(): + """New dict format (free_form_entities is now a separate config field, not in EntityLabelsConfig).""" + raw = { + "attributes": [ + { + "key": "interest", + "description": "User interest area", + "values": [{"value": "active", "description": "Active hobbies"}], + } + ], + } + result = parse_entity_labels(raw) + assert result is not None + assert len(result.attributes) == 1 + assert result.attributes[0].key == "interest" + + +def test_parse_entity_labels_dict_format_defaults(): + """Dict format parses attributes correctly.""" + raw = { + "attributes": [ + {"key": "topic", "values": [{"value": "math", "description": "Mathematics"}]} + ] + } + result = parse_entity_labels(raw) + assert result is not None + assert len(result.attributes) == 1 + + +# ─── build_labels_model ──────────────────────────────────────────────────────── + +# free_values schema variants + + +def test_build_labels_model_single_value(): + """Single-value group → Literal | None field (anyOf), defaults to None.""" + from hindsight_api.engine.retain.entity_labels import build_labels_model + + labels_cfg = EntityLabelsConfig( + attributes=[ + LabelGroup( + key="engagement", + values=[LabelValue(value="active"), LabelValue(value="passive")], + ) + ] + ) + Model = build_labels_model(labels_cfg) + assert Model is not None + + schema = Model.model_json_schema() + props = schema["properties"] + assert "engagement" in props + # Single-value: Pydantic emits anyOf[{enum: [...]}, {type: null}] + any_of = props["engagement"]["anyOf"] + enum_values = next(branch["enum"] for branch in any_of if "enum" in branch) + assert set(enum_values) == {"active", "passive"} + + # Defaults to None when omitted + instance = Model() + assert instance.engagement is None # type: ignore[attr-defined] + + +def test_build_labels_model_multi_value(): + """Multi-value group → list[Literal] field, defaults to empty list.""" + from hindsight_api.engine.retain.entity_labels import build_labels_model + + labels_cfg = EntityLabelsConfig( + attributes=[ + LabelGroup( + key="pedagogy", + type="multi-values", + values=[LabelValue(value="scaffolding"), LabelValue(value="active_engagement")], + ) + ] + ) + Model = build_labels_model(labels_cfg) + assert Model is not None + + schema = Model.model_json_schema() + props = schema["properties"] + assert "pedagogy" in props + assert props["pedagogy"]["type"] == "array" + assert set(props["pedagogy"]["items"]["enum"]) == {"scaffolding", "active_engagement"} + + instance = Model() + assert instance.pedagogy == [] # type: ignore[attr-defined] + + +def test_build_labels_model_mixed(): + """Mixed single + multi-value groups both present.""" + from hindsight_api.engine.retain.entity_labels import build_labels_model + + labels_cfg = EntityLabelsConfig( + attributes=[ + LabelGroup(key="engagement", values=[LabelValue(value="active")]), + LabelGroup(key="pedagogy", type="multi-values", values=[LabelValue(value="scaffolding")]), + ] + ) + Model = build_labels_model(labels_cfg) + assert Model is not None + schema = Model.model_json_schema() + assert "engagement" in schema["properties"] + assert "pedagogy" in schema["properties"] + + +def test_build_labels_model_none_when_no_values(): + """Returns None when no groups have values.""" + from hindsight_api.engine.retain.entity_labels import build_labels_model + + labels_cfg = EntityLabelsConfig(attributes=[LabelGroup(key="empty", values=[])]) + assert build_labels_model(labels_cfg) is None + + +def test_build_labels_model_free_values_optional(): + """type='text', optional=True → str | None field.""" + from hindsight_api.engine.retain.entity_labels import build_labels_model + + labels_cfg = EntityLabelsConfig( + attributes=[LabelGroup(key="topic", type="text", optional=True, values=[])] + ) + Model = build_labels_model(labels_cfg) + assert Model is not None + schema = Model.model_json_schema() + topic = schema["properties"]["topic"] + any_of_types = {branch.get("type") for branch in topic["anyOf"]} + assert "string" in any_of_types and "null" in any_of_types + assert Model().topic is None # type: ignore[attr-defined] + + +def test_build_labels_model_free_values_always_optional(): + """type='text' with optional=False is still treated as str | None — always optional.""" + from hindsight_api.engine.retain.entity_labels import build_labels_model + + labels_cfg = EntityLabelsConfig( + attributes=[LabelGroup(key="topic", type="text", optional=False, values=[])] + ) + Model = build_labels_model(labels_cfg) + assert Model is not None + schema = Model.model_json_schema() + # free_values groups are always optional (str | None), never in required + assert "topic" not in schema.get("required", []) + anyOf = schema["properties"]["topic"].get("anyOf", []) + assert any(b.get("type") == "string" for b in anyOf) + + +def test_build_labels_model_free_values_multi_still_optional(): + """type='text' is always str | None — multi-values only applies to enum types.""" + from hindsight_api.engine.retain.entity_labels import build_labels_model + + labels_cfg = EntityLabelsConfig( + attributes=[LabelGroup(key="tags", type="text", values=[])] + ) + Model = build_labels_model(labels_cfg) + assert Model is not None + schema = Model.model_json_schema() + # free_values groups are always str | None regardless of multi_value + assert "tags" not in schema.get("required", []) + anyOf = schema["properties"]["tags"].get("anyOf", []) + assert any(b.get("type") == "string" for b in anyOf) + + +def test_build_labels_model_free_values_no_values_still_creates_field(): + """type='text' group with no values still creates a field (description holds examples).""" + from hindsight_api.engine.retain.entity_labels import build_labels_model + + labels_cfg = EntityLabelsConfig( + attributes=[LabelGroup(key="mood", type="text", values=[])] + ) + Model = build_labels_model(labels_cfg) + assert Model is not None + assert "mood" in Model.model_json_schema()["properties"] + + +# ─── is_label_entity ────────────────────────────────────────────────────────── + + +def test_is_label_entity_enum_match(): + from hindsight_api.engine.retain.entity_labels import build_labels_lookup, is_label_entity, parse_entity_labels + + cfg = parse_entity_labels([{"key": "engagement", "values": [{"value": "active"}]}]) + lookup = build_labels_lookup(cfg) + assert is_label_entity("engagement:active", cfg, lookup) is True + + +def test_is_label_entity_enum_no_match(): + from hindsight_api.engine.retain.entity_labels import build_labels_lookup, is_label_entity, parse_entity_labels + + cfg = parse_entity_labels([{"key": "engagement", "values": [{"value": "active"}]}]) + lookup = build_labels_lookup(cfg) + assert is_label_entity("engagement:unknown", cfg, lookup) is False + assert is_label_entity("Alice", cfg, lookup) is False + + +def test_is_label_entity_free_values_prefix_match(): + from hindsight_api.engine.retain.entity_labels import build_labels_lookup, is_label_entity, parse_entity_labels + + cfg = parse_entity_labels([{"key": "topic", "type": "text", "values": []}]) + lookup = build_labels_lookup(cfg) + assert is_label_entity("topic:algebra", cfg, lookup) is True + assert is_label_entity("topic:anything at all", cfg, lookup) is True + + +def test_is_label_entity_free_values_no_match_other_key(): + from hindsight_api.engine.retain.entity_labels import build_labels_lookup, is_label_entity, parse_entity_labels + + cfg = parse_entity_labels([{"key": "topic", "type": "text", "values": []}]) + lookup = build_labels_lookup(cfg) + assert is_label_entity("Alice", cfg, lookup) is False + assert is_label_entity("engagement:active", cfg, lookup) is False + + +# ─── build_labels_lookup ─────────────────────────────────────────────────────── + + +def test_build_labels_lookup(): + labels_cfg = EntityLabelsConfig( + attributes=[ + LabelGroup( + key="Pedagogy", + values=[ + LabelValue(value="Scaffolding"), + LabelValue(value="Active_Engagement"), + ], + ) + ] + ) + lookup = build_labels_lookup(labels_cfg) + assert "pedagogy:scaffolding" in lookup + assert "pedagogy:active_engagement" in lookup + # Should be lowercase + assert all(v == v.lower() for v in lookup) + + +def test_build_labels_lookup_raw_list(): + """build_labels_lookup accepts raw list format for backwards compatibility.""" + raw = [ + { + "key": "interest", + "values": [{"value": "active", "description": "Active interest"}], + } + ] + lookup = build_labels_lookup(raw) + assert "interest:active" in lookup + + +def test_build_labels_lookup_none(): + lookup = build_labels_lookup(None) + assert lookup == set() + + +# ─── _build_labels_prompt_section ───────────────────────────────────────────── + + +def test_build_labels_prompt_section_none(): + from hindsight_api.engine.retain.fact_extraction import _build_labels_prompt_section + + result = _build_labels_prompt_section(None) + assert result == "" + + +def test_build_labels_prompt_section_empty_config(): + from hindsight_api.engine.retain.fact_extraction import _build_labels_prompt_section + + result = _build_labels_prompt_section(EntityLabelsConfig(attributes=[])) + assert result == "" + + +def test_build_labels_prompt_section_generates_key_values(): + from hindsight_api.engine.retain.fact_extraction import _build_labels_prompt_section + + labels_cfg = EntityLabelsConfig( + attributes=[ + LabelGroup( + key="pedagogy", + description="Teaching strategy", + type="multi-values", + values=[ + LabelValue(value="scaffolding", description="Break down tasks"), + LabelValue(value="active_engagement", description="Group work"), + ], + ) + ] + ) + result = _build_labels_prompt_section(labels_cfg) + # Structured format: values listed as "value" bullets under the key name + assert "scaffolding" in result + assert "active_engagement" in result + assert "pedagogy" in result + assert "Teaching strategy" in result + assert "multi" in result # prompt mentions multi-value nature + + +def test_build_labels_prompt_section_free_form_true(): + from hindsight_api.engine.retain.fact_extraction import _build_labels_prompt_section + + labels_cfg = EntityLabelsConfig( + attributes=[ + LabelGroup( + key="topic", + values=[LabelValue(value="math")], + ) + ], + ) + result = _build_labels_prompt_section(labels_cfg, free_form_entities=True) + # When free_form_entities=True: prompt says to also fill 'entities' field + assert "labels" in result + assert "entities" in result + + +def test_build_labels_prompt_section_free_form_false(): + from hindsight_api.engine.retain.fact_extraction import _build_labels_prompt_section + + labels_cfg = EntityLabelsConfig( + attributes=[ + LabelGroup( + key="topic", + values=[LabelValue(value="math")], + ) + ], + ) + result = _build_labels_prompt_section(labels_cfg, free_form_entities=False) + assert "labels-only mode" in result + + +# ─── augment_texts_with_entities ────────────────────────────────────────────── + + +def test_augment_texts_with_entities(): + """Entity names appear in embedding input but fact_text is unchanged.""" + from datetime import UTC, datetime + + from hindsight_api.engine.retain.embedding_processing import augment_texts_with_dates + from hindsight_api.engine.retain.types import ExtractedFact + + event_date = datetime(2024, 6, 1, tzinfo=UTC) + fact = ExtractedFact( + fact_text="User attended workshop", + fact_type="world", + entities=["pedagogy:scaffolding", "user"], + mentioned_at=event_date, + ) + + def fmt_date(dt): + return "June 2024" + + augmented = augment_texts_with_dates([fact], fmt_date) + assert len(augmented) == 1 + # Entity names should appear in augmented text + assert "pedagogy:scaffolding" in augmented[0] + assert "user" in augmented[0] + # Original fact text should be present + assert "User attended workshop" in augmented[0] + + +# ─── _inject_label_tags ─────────────────────────────────────────────────────── + + +def test_inject_label_tags_adds_tagged_entities(): + """tag=True group: extracted label entities are added to fact.tags.""" + from unittest.mock import MagicMock + + from hindsight_api.engine.retain.fact_extraction import _inject_label_tags + from hindsight_api.engine.retain.types import ExtractedFact + + config = MagicMock() + config.entity_labels = [ + {"key": "pedagogy", "type": "value", "tag": True, "values": [{"value": "scaffolding"}]}, + {"key": "engagement", "type": "value", "tag": False, "values": [{"value": "active"}]}, + ] + + fact = ExtractedFact( + fact_text="Teacher used scaffolding", + fact_type="world", + entities=["pedagogy:scaffolding", "engagement:active"], + tags=["session-1"], + ) + _inject_label_tags([fact], config) + + # pedagogy group has tag=True → added to tags + assert "pedagogy:scaffolding" in fact.tags + # engagement group has tag=False → NOT added + assert "engagement:active" not in fact.tags + # original tag preserved + assert "session-1" in fact.tags + + +def test_inject_label_tags_no_duplicate(): + """No duplicate if label entity already in tags.""" + from unittest.mock import MagicMock + + from hindsight_api.engine.retain.fact_extraction import _inject_label_tags + from hindsight_api.engine.retain.types import ExtractedFact + + config = MagicMock() + config.entity_labels = [ + {"key": "pedagogy", "type": "value", "tag": True, "values": [{"value": "scaffolding"}]}, + ] + + fact = ExtractedFact( + fact_text="...", + fact_type="world", + entities=["pedagogy:scaffolding"], + tags=["pedagogy:scaffolding"], + ) + _inject_label_tags([fact], config) + assert fact.tags.count("pedagogy:scaffolding") == 1 + + +def test_inject_label_tags_no_tag_groups_is_noop(): + """When no groups have tag=True, tags are unchanged.""" + from unittest.mock import MagicMock + + from hindsight_api.engine.retain.fact_extraction import _inject_label_tags + from hindsight_api.engine.retain.types import ExtractedFact + + config = MagicMock() + config.entity_labels = [ + {"key": "pedagogy", "type": "value", "tag": False, "values": [{"value": "scaffolding"}]}, + ] + + fact = ExtractedFact(fact_text="...", fact_type="world", entities=["pedagogy:scaffolding"]) + _inject_label_tags([fact], config) + assert fact.tags == [] + + +def test_inject_label_tags_no_labels_config_is_noop(): + """When entity_labels is None, tags are unchanged.""" + from unittest.mock import MagicMock + + from hindsight_api.engine.retain.fact_extraction import _inject_label_tags + from hindsight_api.engine.retain.types import ExtractedFact + + config = MagicMock() + config.entity_labels = None + + fact = ExtractedFact(fact_text="...", fact_type="world", entities=["pedagogy:scaffolding"]) + _inject_label_tags([fact], config) + assert fact.tags == [] + + +# ─── entity label post-processing ───────────────────────────────────────────── + + +def test_label_entity_post_processing(): + """Structured labels dict is parsed into key:value entity strings; invalid values filtered.""" + from hindsight_api.engine.retain.entity_labels import build_labels_lookup, parse_entity_labels + from hindsight_api.engine.retain.fact_extraction import Entity + + labels_cfg = parse_entity_labels( + [ + { + "key": "pedagogy", + "values": [ + {"value": "scaffolding", "description": ""}, + {"value": "active_engagement", "description": ""}, + ], + } + ] + ) + assert labels_cfg is not None + labels_lookup = build_labels_lookup(labels_cfg) + + # Simulated LLM response — structured dict, not a flat list + labels_data = {"pedagogy": "scaffolding"} # single-value field + + validated_entities: list[Entity] = [] + if isinstance(labels_data, dict) and labels_lookup: + existing_texts_lower: set[str] = set() + for group in labels_cfg.attributes: + value = labels_data.get(group.key) + if not value: + continue + values_list = value if isinstance(value, list) else [value] + for v in values_list: + label_str = f"{group.key}:{v}" + if label_str.lower() in labels_lookup and label_str.lower() not in existing_texts_lower: + validated_entities.append(Entity(text=label_str)) + existing_texts_lower.add(label_str.lower()) + + entity_texts = {e.text for e in validated_entities} + assert "pedagogy:scaffolding" in entity_texts + + +def test_label_entity_post_processing_invalid_value_ignored(): + """Values not in the lookup are silently dropped.""" + from hindsight_api.engine.retain.entity_labels import build_labels_lookup, parse_entity_labels + from hindsight_api.engine.retain.fact_extraction import Entity + + labels_cfg = parse_entity_labels( + [{"key": "pedagogy", "values": [{"value": "scaffolding", "description": ""}]}] + ) + labels_lookup = build_labels_lookup(labels_cfg) + + labels_data = {"pedagogy": "unknown_value"} + + validated_entities: list[Entity] = [] + existing_texts_lower: set[str] = set() + for group in labels_cfg.attributes: + value = labels_data.get(group.key) + if not value: + continue + values_list = value if isinstance(value, list) else [value] + for v in values_list: + label_str = f"{group.key}:{v}" + if label_str.lower() in labels_lookup and label_str.lower() not in existing_texts_lower: + validated_entities.append(Entity(text=label_str)) + + assert validated_entities == [] + + +def test_label_entity_post_processing_multi_value(): + """Multi-value list field produces one entity per value.""" + from hindsight_api.engine.retain.entity_labels import build_labels_lookup, parse_entity_labels + from hindsight_api.engine.retain.fact_extraction import Entity + + labels_cfg = parse_entity_labels( + [ + { + "key": "pedagogy", + "multi_value": True, + "values": [ + {"value": "scaffolding", "description": ""}, + {"value": "active_engagement", "description": ""}, + ], + } + ] + ) + labels_lookup = build_labels_lookup(labels_cfg) + + # Multi-value: LLM returns a list + labels_data = {"pedagogy": ["scaffolding", "active_engagement"]} + + validated_entities: list[Entity] = [] + existing_texts_lower: set[str] = set() + for group in labels_cfg.attributes: + value = labels_data.get(group.key) + if not value: + continue + values_list = value if isinstance(value, list) else [value] + for v in values_list: + label_str = f"{group.key}:{v}" + if label_str.lower() in labels_lookup and label_str.lower() not in existing_texts_lower: + validated_entities.append(Entity(text=label_str)) + existing_texts_lower.add(label_str.lower()) + + entity_texts = {e.text for e in validated_entities} + assert "pedagogy:scaffolding" in entity_texts + assert "pedagogy:active_engagement" in entity_texts + + +def _run_label_post_processing(labels_cfg, labels_data: dict) -> set[str]: + """Helper: mirrors the production label post-processing logic, returns entity text set.""" + from hindsight_api.engine.retain.entity_labels import build_labels_lookup + from hindsight_api.engine.retain.fact_extraction import Entity + + labels_lookup = build_labels_lookup(labels_cfg) + validated_entities: list[Entity] = [] + existing_texts_lower: set[str] = set() + + effective_data = labels_data or {} + if isinstance(effective_data, dict): + for group in labels_cfg.attributes: + value = effective_data.get(group.key) + if not value: + continue + values_list = value if isinstance(value, list) else [value] + for v in values_list: + if not isinstance(v, str) or not v.strip() or v.lower() in ("none", "null", "n/a"): + continue + label_str = f"{group.key}:{v.strip()}" + if group.type == "text": + if label_str.lower() not in existing_texts_lower: + validated_entities.append(Entity(text=label_str)) + existing_texts_lower.add(label_str.lower()) + elif label_str.lower() in labels_lookup and label_str.lower() not in existing_texts_lower: + validated_entities.append(Entity(text=label_str)) + existing_texts_lower.add(label_str.lower()) + + return {e.text for e in validated_entities} + + +def test_free_values_label_accepts_any_string(): + """type='text' group: any non-empty string produces a key:value entity.""" + from hindsight_api.engine.retain.entity_labels import parse_entity_labels + + labels_cfg = parse_entity_labels([{"key": "topic", "type": "text", "values": []}]) + entity_texts = _run_label_post_processing(labels_cfg, {"topic": "quadratic equations"}) + assert "topic:quadratic equations" in entity_texts + + +def test_free_values_label_rejects_none_sentinel(): + """type='text' group: string 'None' / 'null' / 'n/a' are rejected.""" + from hindsight_api.engine.retain.entity_labels import parse_entity_labels + + labels_cfg = parse_entity_labels([{"key": "topic", "type": "text", "values": []}]) + for sentinel in ("None", "null", "n/a", "NULL", "NONE"): + result = _run_label_post_processing(labels_cfg, {"topic": sentinel}) + assert result == set(), f"Sentinel '{sentinel}' should not produce an entity, got: {result}" + + +def test_free_values_label_is_single_value(): + """type='text' groups are always single-value (str | None).""" + from hindsight_api.engine.retain.entity_labels import build_labels_model, parse_entity_labels + + labels_cfg = parse_entity_labels( + [{"key": "topic", "type": "text", "values": []}] + ) + Model = build_labels_model(labels_cfg) + assert Model is not None + schema = Model.model_json_schema() + # Must be str | None, not list + assert schema["properties"]["topic"].get("type") != "array" + anyOf = schema["properties"]["topic"].get("anyOf", []) + assert any(b.get("type") == "string" for b in anyOf) + + +def test_free_values_label_not_in_lookup(): + """type='text' group values do NOT appear in the lookup set (no fixed vocabulary).""" + from hindsight_api.engine.retain.entity_labels import build_labels_lookup, parse_entity_labels + + labels_cfg = parse_entity_labels( + [{"key": "topic", "type": "text", "values": [{"value": "algebra"}]}] + ) + lookup = build_labels_lookup(labels_cfg) + assert "topic:algebra" not in lookup # example hints not added to lookup + assert len(lookup) == 0 + + +def test_optional_label_null_produces_no_entity(): + """JSON null (Python None) for an optional label → no entity created.""" + from hindsight_api.engine.retain.entity_labels import parse_entity_labels + + labels_cfg = parse_entity_labels( + [{"key": "engagement", "optional": True, "values": [{"value": "active"}, {"value": "passive"}]}] + ) + + # LLM returned null — content didn't match any value + entity_texts = _run_label_post_processing(labels_cfg, {"engagement": None}) + assert entity_texts == set(), f"Expected no entities for null optional label, got: {entity_texts}" + + +def test_optional_label_absent_key_produces_no_entity(): + """Missing key in labels dict for an optional label → no entity created.""" + from hindsight_api.engine.retain.entity_labels import parse_entity_labels + + labels_cfg = parse_entity_labels( + [{"key": "engagement", "optional": True, "values": [{"value": "active"}, {"value": "passive"}]}] + ) + + # LLM omitted the key entirely + entity_texts = _run_label_post_processing(labels_cfg, {}) + assert entity_texts == set(), f"Expected no entities for absent optional label, got: {entity_texts}" + + +def test_optional_label_string_none_produces_no_entity(): + """String 'None' from LLM for an optional label → no entity created (not in lookup).""" + from hindsight_api.engine.retain.entity_labels import parse_entity_labels + + labels_cfg = parse_entity_labels( + [{"key": "engagement", "optional": True, "values": [{"value": "active"}, {"value": "passive"}]}] + ) + + # LLM returned the string "None" instead of JSON null — must not be stored + entity_texts = _run_label_post_processing(labels_cfg, {"engagement": "None"}) + assert entity_texts == set(), ( + f"String 'None' must not produce engagement:None entity, got: {entity_texts}" + ) + + +def test_optional_label_null_does_not_affect_other_labels(): + """Null for one optional label doesn't suppress other valid labels on the same fact.""" + from hindsight_api.engine.retain.entity_labels import parse_entity_labels + + labels_cfg = parse_entity_labels( + [ + {"key": "engagement", "optional": True, "values": [{"value": "active"}, {"value": "passive"}]}, + {"key": "topic", "optional": True, "values": [{"value": "math"}, {"value": "science"}]}, + ] + ) + + # engagement is null, but topic is set + entity_texts = _run_label_post_processing(labels_cfg, {"engagement": None, "topic": "math"}) + assert "topic:math" in entity_texts, f"Expected topic:math entity, got: {entity_texts}" + assert not any("engagement" in t for t in entity_texts), ( + f"engagement should not appear, got: {entity_texts}" + ) + + +def test_free_form_entities_false_clears_entities(): + """When retain_free_form_entities=False, non-label entities are removed.""" + from hindsight_api.engine.retain.entity_labels import build_labels_lookup, parse_entity_labels + from hindsight_api.engine.retain.fact_extraction import Entity + + labels_cfg = parse_entity_labels( + { + "attributes": [ + { + "key": "pedagogy", + "values": [{"value": "scaffolding", "description": ""}], + } + ], + } + ) + labels_lookup = build_labels_lookup(labels_cfg) + free_form_entities = False # standalone config field + + # Mix of label and free-form entities + validated_entities = [ + Entity(text="pedagogy:scaffolding"), + Entity(text="Alice"), + Entity(text="Google"), + ] + + # Apply free_form filtering + if not free_form_entities and labels_lookup: + validated_entities = [e for e in validated_entities if e.text.lower() in labels_lookup] + + entity_texts = {e.text for e in validated_entities} + assert "pedagogy:scaffolding" in entity_texts + assert "Alice" not in entity_texts + assert "Google" not in entity_texts + + +def test_free_form_entities_true_keeps_all(): + """When retain_free_form_entities=True (default), all entities are kept.""" + from hindsight_api.engine.retain.entity_labels import build_labels_lookup, parse_entity_labels + from hindsight_api.engine.retain.fact_extraction import Entity + + labels_cfg = parse_entity_labels( + { + "attributes": [ + { + "key": "pedagogy", + "values": [{"value": "scaffolding", "description": ""}], + } + ], + } + ) + labels_lookup = build_labels_lookup(labels_cfg) + free_form_entities = True # default value + + validated_entities = [ + Entity(text="pedagogy:scaffolding"), + Entity(text="Alice"), + ] + + # With free_form_entities=True, should NOT filter + if not free_form_entities and labels_lookup: + validated_entities = [e for e in validated_entities if e.text.lower() in labels_lookup] + + entity_texts = {e.text for e in validated_entities} + assert "pedagogy:scaffolding" in entity_texts + assert "Alice" in entity_texts + + +# ─── _build_extraction_prompt_and_schema with labels ────────────────────────── + + +def test_extraction_schema_includes_labels_model(): + """When entity_labels configured, response schema has a structured Labels field.""" + from unittest.mock import MagicMock + + from hindsight_api.engine.retain.fact_extraction import _build_extraction_prompt_and_schema + + config = MagicMock() + config.entity_labels = [ + { + "key": "engagement", + "values": [{"value": "active"}, {"value": "passive"}], + }, + { + "key": "pedagogy", + "type": "multi-values", + "values": [{"value": "scaffolding"}, {"value": "active_engagement"}], + }, + ] + config.entities_allow_free_form = True + config.retain_extraction_mode = "concise" + config.retain_extract_causal_links = False + config.retain_mission = None + config.retain_custom_instructions = None + + prompt, schema = _build_extraction_prompt_and_schema(config) + + # Schema should be a dynamic response model + json_schema = schema.model_json_schema() + assert "facts" in json_schema["properties"] + + # Drill into the fact item schema + fact_schema = json_schema["$defs"]["LabelsFact"] + assert "labels" in fact_schema["properties"] + assert "labels" in fact_schema["required"] + + # Labels should be a nested object (not a flat array) + labels_ref = fact_schema["properties"]["labels"] + labels_def_key = labels_ref["$ref"].split("/")[-1] + labels_def = json_schema["$defs"][labels_def_key] + + assert "engagement" in labels_def["properties"] + assert "pedagogy" in labels_def["properties"] + + # engagement: single-value → anyOf[{enum: [...]}, {type: null}] + any_of = labels_def["properties"]["engagement"]["anyOf"] + engagement_enums = next(b["enum"] for b in any_of if "enum" in b) + assert set(engagement_enums) == {"active", "passive"} + + # pedagogy: multi-value → array of enum + assert labels_def["properties"]["pedagogy"]["type"] == "array" + assert set(labels_def["properties"]["pedagogy"]["items"]["enum"]) == {"scaffolding", "active_engagement"} + + # Prompt should reference the labels object + assert "labels" in prompt + assert "engagement" in prompt + assert "pedagogy" in prompt + + +def test_extraction_schema_labels_in_required(): + """labels field is in the required array so OpenAI structured outputs enforce it.""" + from unittest.mock import MagicMock + + from hindsight_api.engine.retain.fact_extraction import _build_extraction_prompt_and_schema + + config = MagicMock() + config.entity_labels = [{"key": "topic", "values": [{"value": "math"}]}] + config.entities_allow_free_form = True + config.retain_extraction_mode = "concise" + config.retain_extract_causal_links = False + config.retain_mission = None + config.retain_custom_instructions = None + + _, schema = _build_extraction_prompt_and_schema(config) + fact_schema = schema.model_json_schema()["$defs"]["LabelsFact"] + assert "labels" in fact_schema["required"] + + +def test_extraction_schema_no_labels_when_unconfigured(): + """Without entity_labels, schema falls back to a base FactExtraction class (no dynamic model).""" + from unittest.mock import MagicMock + + from hindsight_api.engine.retain.fact_extraction import ( + _build_extraction_prompt_and_schema, + ) + + config = MagicMock() + config.entity_labels = None + config.entities_allow_free_form = True + config.retain_extraction_mode = "concise" + config.retain_extract_causal_links = False + config.retain_mission = None + config.retain_custom_instructions = None + + _, schema = _build_extraction_prompt_and_schema(config) + # No labels field in schema — it's a plain base response model + json_schema = schema.model_json_schema() + # Verify 'labels' is NOT a required or present field in any fact definition + fact_defs = {k: v for k, v in json_schema.get("$defs", {}).items() if "facts" not in k.lower()} + for name, defn in fact_defs.items(): + assert "labels" not in defn.get("properties", {}), f"Found 'labels' in {name}" + + +# ─── LLM integration tests (require DB + LLM) ───────────────────────────────── + + +@pytest.mark.asyncio +async def test_retain_extracts_single_value_label(memory, request_context): + """ + End-to-end: retain content with entity_labels configured (single-value). + Verify that the LLM assigns the label and it ends up as a key:value entity on the memory unit. + """ + from hindsight_api.engine.memory_engine import fq_table + + bank_id = f"test-labels-single-{uuid.uuid4().hex[:8]}" + try: + await memory.get_bank_profile(bank_id=bank_id, request_context=request_context) + + # Configure entity_labels on the bank + await memory._config_resolver.update_bank_config( + bank_id=bank_id, + updates={ + "entity_labels": [ + { + "key": "engagement", + "description": "Student engagement level during the session", + "values": [ + {"value": "active", "description": "Student is actively participating"}, + {"value": "passive", "description": "Student is listening but not participating"}, + ], + } + ], + "entities_allow_free_form": False, # labels-only mode + }, + context=request_context, + ) + + unit_ids = await memory.retain_async( + bank_id=bank_id, + content=( + "During today's tutoring session, Maria asked many questions, " + "participated in every exercise, and solved the problems independently. " + "She was very engaged throughout." + ), + request_context=request_context, + ) + + assert len(unit_ids) > 0, "Should have extracted at least one fact" + + # Query entity names for the retained units + async with memory._pool.acquire() as conn: + rows = await conn.fetch( + f""" + SELECT e.canonical_name + FROM {fq_table("unit_entities")} ue + JOIN {fq_table("entities")} e ON e.id = ue.entity_id + WHERE ue.unit_id = ANY($1::uuid[]) + """, + [u for u in unit_ids], + ) + + entity_names = {r["canonical_name"].lower() for r in rows} + assert "engagement:active" in entity_names, ( + f"Expected 'engagement:active' label entity. Got: {entity_names}" + ) + # In labels-only mode, free-form entities like 'Maria' should be absent + assert not any("maria" in n for n in entity_names), ( + f"Free-form entity 'Maria' should not appear in labels-only mode. Got: {entity_names}" + ) + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_retain_extracts_multi_value_label(memory, request_context): + """ + End-to-end: retain content with a multi_value entity_labels group. + Verify that multiple label values can be assigned to a single fact. + """ + from hindsight_api.engine.memory_engine import fq_table + + bank_id = f"test-labels-multi-{uuid.uuid4().hex[:8]}" + try: + await memory.get_bank_profile(bank_id=bank_id, request_context=request_context) + + await memory._config_resolver.update_bank_config( + bank_id=bank_id, + updates={ + "entity_labels": [ + { + "key": "pedagogy", + "description": "Teaching strategies observed in the session", + "type": "multi-values", + "values": [ + {"value": "scaffolding", "description": "Teacher breaks tasks into smaller steps"}, + {"value": "direct_instruction", "description": "Teacher explains concepts directly"}, + {"value": "socratic_questioning", "description": "Teacher guides via questions"}, + ], + } + ], + "entities_allow_free_form": False, + }, + context=request_context, + ) + + unit_ids = await memory.retain_async( + bank_id=bank_id, + content=( + "The teacher broke the algebra problem into small steps and guided the student " + "through each one with questions like 'What do you notice about this equation?' " + "and 'What would happen if you moved this term to the other side?'. " + "The lesson was clearly structured with scaffolding and socratic questioning." + ), + request_context=request_context, + ) + + assert len(unit_ids) > 0 + + async with memory._pool.acquire() as conn: + rows = await conn.fetch( + f""" + SELECT e.canonical_name + FROM {fq_table("unit_entities")} ue + JOIN {fq_table("entities")} e ON e.id = ue.entity_id + WHERE ue.unit_id = ANY($1::uuid[]) + """, + [u for u in unit_ids], + ) + + entity_names = {r["canonical_name"].lower() for r in rows} + # At least one pedagogy label should be assigned + pedagogy_labels = {n for n in entity_names if n.startswith("pedagogy:")} + assert len(pedagogy_labels) > 0, ( + f"Expected at least one pedagogy:* label entity. Got: {entity_names}" + ) + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_retain_extracts_free_values_label(memory, request_context): + """ + End-to-end: retain content with a free_values entity_labels group. + Verify that the LLM produces a key:value entity with an open-ended value + (not constrained to a predefined enum list). + """ + from hindsight_api.engine.memory_engine import fq_table + + bank_id = f"test-labels-free-{uuid.uuid4().hex[:8]}" + try: + await memory.get_bank_profile(bank_id=bank_id, request_context=request_context) + + await memory._config_resolver.update_bank_config( + bank_id=bank_id, + updates={ + "entity_labels": [ + { + "key": "topic", + "description": "The specific subject being discussed in this session. Examples: algebra, geometry, quadratic equations.", + "type": "text", + "optional": True, + "values": [], + } + ], + "entities_allow_free_form": False, + }, + context=request_context, + ) + + unit_ids = await memory.retain_async( + bank_id=bank_id, + content=( + "The student and tutor spent the session working through quadratic equations. " + "They factored several expressions and practised the quadratic formula." + ), + request_context=request_context, + ) + + assert len(unit_ids) > 0 + + async with memory._pool.acquire() as conn: + rows = await conn.fetch( + f""" + SELECT e.canonical_name + FROM {fq_table("unit_entities")} ue + JOIN {fq_table("entities")} e ON e.id = ue.entity_id + WHERE ue.unit_id = ANY($1::uuid[]) + """, + [u for u in unit_ids], + ) + + entity_names = {r["canonical_name"].lower() for r in rows} + # A topic:* entity must exist — value is free-form so we only check the prefix + topic_entities = {n for n in entity_names if n.startswith("topic:")} + assert len(topic_entities) > 0, ( + f"Expected at least one topic:* free-value entity. Got: {entity_names}" + ) + # The value must not be the literal string "none" or "null" + assert not any(n in ("topic:none", "topic:null", "topic:n/a") for n in topic_entities), ( + f"topic entity should not be a null sentinel. Got: {topic_entities}" + ) + finally: + await memory.delete_bank(bank_id, request_context=request_context) diff --git a/hindsight-api/tests/test_hierarchical_config.py b/hindsight-api/tests/test_hierarchical_config.py index a04ecd61..e1597152 100644 --- a/hindsight-api/tests/test_hierarchical_config.py +++ b/hindsight-api/tests/test_hierarchical_config.py @@ -81,8 +81,12 @@ async def test_hierarchical_fields_categorization(): assert "disposition_literalism" in configurable assert "disposition_empathy" in configurable + # Verify entity labels fields are included + assert "entities_allow_free_form" in configurable + assert "entity_labels" in configurable + # Verify count is correct - assert len(configurable) == 11 + assert len(configurable) == 13 # Verify credential fields (NEVER exposed) assert "llm_api_key" in credentials diff --git a/hindsight-clients/go/integration_test.go b/hindsight-clients/go/integration_test.go index c4c3fb70..7037dfdc 100644 --- a/hindsight-clients/go/integration_test.go +++ b/hindsight-clients/go/integration_test.go @@ -67,7 +67,7 @@ func TestRetainWithContext(t *testing.T) { Items: []MemoryItem{ { Content: "Bob went hiking in the mountains", - Timestamp: *NewNullableTime(PtrTime(timestamp)), + Timestamp: *NewNullableTimestamp(&Timestamp{TimeTime: ×tamp}), Context: *NewNullableString(PtrString("outdoor activities")), }, }, diff --git a/hindsight-clients/python/hindsight_client/hindsight_client.py b/hindsight-clients/python/hindsight_client/hindsight_client.py index 63acf3ef..4039f166 100644 --- a/hindsight-clients/python/hindsight_client/hindsight_client.py +++ b/hindsight-clients/python/hindsight_client/hindsight_client.py @@ -180,16 +180,19 @@ class Hindsight: RetainResponse with success status and item count """ from hindsight_client_api.models.entity_input import EntityInput + from hindsight_client_api.models.timestamp import Timestamp memory_items = [] for item in items: entities = None if item.get("entities"): entities = [EntityInput(text=e["text"], type=e.get("type")) for e in item["entities"]] + raw_ts = item.get("timestamp") + timestamp_val = Timestamp(actual_instance=raw_ts) if raw_ts is not None else None memory_items.append( memory_item.MemoryItem( content=item["content"], - timestamp=item.get("timestamp"), + timestamp=timestamp_val, context=item.get("context"), metadata=item.get("metadata"), # Use item's document_id if provided, otherwise fall back to batch-level document_id @@ -591,16 +594,19 @@ class Hindsight: RetainResponse with success status and item count """ from hindsight_client_api.models.entity_input import EntityInput + from hindsight_client_api.models.timestamp import Timestamp memory_items = [] for item in items: entities = None if item.get("entities"): entities = [EntityInput(text=e["text"], type=e.get("type")) for e in item["entities"]] + raw_ts = item.get("timestamp") + timestamp_val = Timestamp(actual_instance=raw_ts) if raw_ts is not None else None memory_items.append( memory_item.MemoryItem( content=item["content"], - timestamp=item.get("timestamp"), + timestamp=timestamp_val, context=item.get("context"), metadata=item.get("metadata"), # Use item's document_id if provided, otherwise fall back to batch-level document_id diff --git a/hindsight-control-plane/package.json b/hindsight-control-plane/package.json index 783b8b45..c5b99593 100644 --- a/hindsight-control-plane/package.json +++ b/hindsight-control-plane/package.json @@ -11,7 +11,7 @@ "public" ], "scripts": { - "dev": "next dev --turbopack -p 9999", + "dev": "next dev --turbopack -p ${PORT:-9999}", "build": "next build && npm run build:standalone", "build:standalone": "rm -rf standalone && SERVER_JS=$(find .next/standalone -path '*/node_modules' -prune -o -name 'server.js' -print | head -1) && test -n \"$SERVER_JS\" || (echo 'Error: server.js not found in .next/standalone - standalone build failed' && exit 1) && STANDALONE_ROOT=$(dirname \"$SERVER_JS\") && cp -r \"$STANDALONE_ROOT\" standalone && cp -r .next/standalone/node_modules standalone/node_modules && mkdir -p standalone/.next && cp -r .next/static standalone/.next/static && mkdir -p standalone/public && (cp -r public/* standalone/public/ 2>/dev/null || true)", "start": "next start", diff --git a/hindsight-control-plane/src/components/bank-config-view.tsx b/hindsight-control-plane/src/components/bank-config-view.tsx index 0faefead..3f4ee894 100644 --- a/hindsight-control-plane/src/components/bank-config-view.tsx +++ b/hindsight-control-plane/src/components/bank-config-view.tsx @@ -14,7 +14,10 @@ import { SelectValue, } from "@/components/ui/select"; import { Alert, AlertDescription } from "@/components/ui/alert"; -import { Loader2, AlertCircle } from "lucide-react"; +import { Switch } from "@/components/ui/switch"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Label } from "@/components/ui/label"; +import { Loader2, AlertCircle, Plus, Trash2, ChevronDown, ChevronRight } from "lucide-react"; import { Card } from "@/components/ui/card"; // ─── Types ──────────────────────────────────────────────────────────────────── @@ -38,6 +41,21 @@ type ObservationsEdits = { observations_mission: string | null; }; +type LabelValue = { value: string; description: string }; +type LabelGroup = { + key: string; + description: string; + type: "value" | "multi-values" | "text"; + optional: boolean; + tag: boolean; + values: LabelValue[]; +}; + +type EntityLabelsEdits = { + entity_labels: LabelGroup[] | null; + entities_allow_free_form: boolean; +}; + type MCPEdits = { mcp_enabled_tools: string[] | null; }; @@ -96,6 +114,20 @@ function observationsSlice(config: Record): ObservationsEdits { }; } +function entityLabelsSlice(config: Record): EntityLabelsEdits { + const raw = config.entity_labels; + let attrs: LabelGroup[] | null = null; + if (Array.isArray(raw)) { + attrs = raw as LabelGroup[]; + } else if (raw && typeof raw === "object" && Array.isArray(raw.attributes)) { + attrs = raw.attributes as LabelGroup[]; + } + return { + entity_labels: attrs, + entities_allow_free_form: config.entities_allow_free_form ?? true, + }; +} + function mcpSlice(config: Record): MCPEdits { return { mcp_enabled_tools: config.mcp_enabled_tools ?? null, @@ -124,16 +156,21 @@ export function BankConfigView() { const [observationsEdits, setObservationsEdits] = useState( observationsSlice({}) ); + const [entityLabelsEdits, setEntityLabelsEdits] = useState( + entityLabelsSlice({}) + ); const [reflectEdits, setReflectEdits] = useState(DEFAULT_PROFILE); const [mcpEdits, setMcpEdits] = useState(mcpSlice({})); // Per-section saving/error state const [retainSaving, setRetainSaving] = useState(false); const [observationsSaving, setObservationsSaving] = useState(false); + const [entityLabelsSaving, setEntityLabelsSaving] = useState(false); const [reflectSaving, setReflectSaving] = useState(false); const [mcpSaving, setMcpSaving] = useState(false); const [retainError, setRetainError] = useState(null); const [observationsError, setObservationsError] = useState(null); + const [entityLabelsError, setEntityLabelsError] = useState(null); const [reflectError, setReflectError] = useState(null); const [mcpError, setMcpError] = useState(null); @@ -148,6 +185,10 @@ export function BankConfigView() { () => JSON.stringify(observationsEdits) !== JSON.stringify(observationsSlice(baseConfig)), [observationsEdits, baseConfig] ); + const entityLabelsDirty = useMemo( + () => JSON.stringify(entityLabelsEdits) !== JSON.stringify(entityLabelsSlice(baseConfig)), + [entityLabelsEdits, baseConfig] + ); const reflectDirty = useMemo( () => JSON.stringify(reflectEdits) !== JSON.stringify(baseProfile), [reflectEdits, baseProfile] @@ -182,6 +223,7 @@ export function BankConfigView() { setBaseProfile(prof); setRetainEdits(retainSlice(cfg)); setObservationsEdits(observationsSlice(cfg)); + setEntityLabelsEdits(entityLabelsSlice(cfg)); setReflectEdits(prof); setMcpEdits(mcpSlice(cfg)); } catch (err) { @@ -219,6 +261,24 @@ export function BankConfigView() { } }; + const saveEntityLabels = async () => { + if (!bankId) return; + setEntityLabelsSaving(true); + setEntityLabelsError(null); + try { + const payload = { + entity_labels: entityLabelsEdits.entity_labels, + entities_allow_free_form: entityLabelsEdits.entities_allow_free_form, + }; + await client.updateBankConfig(bankId, payload); + setBaseConfig((prev) => ({ ...prev, ...payload })); + } catch (err: any) { + setEntityLabelsError(err.message || "Failed to save entity labels settings"); + } finally { + setEntityLabelsSaving(false); + } + }; + const saveReflect = async () => { if (!bankId) return; setReflectSaving(true); @@ -340,6 +400,46 @@ export function BankConfigView() { )} + {/* Entity Labels Section */} + + +
+ + + setEntityLabelsEdits((prev) => ({ ...prev, entities_allow_free_form: v })) + } + /> +
+
+ + setEntityLabelsEdits((prev) => ({ + ...prev, + entity_labels: attrs.length > 0 ? attrs : null, + })) + } + /> +
+ {/* Observations Section */}
- + setObservationsEdits((prev) => ({ ...prev, enable_observations: v })) } /> @@ -430,15 +530,18 @@ export function BankConfigView() { label="Restrict tools" description="When off, all tools are available. When on, only the selected tools can be invoked for this bank." > -
- +
+ setMcpEdits({ mcp_enabled_tools: restricted ? [...ALL_TOOLS] : null, }) } /> +
{mcpEdits.mcp_enabled_tools !== null && ( @@ -713,22 +816,216 @@ function TraitRow({ ); } -// ─── Toggle ─────────────────────────────────────────────────────────────────── +// ─── EntityLabelsEditor ─────────────────────────────────────────────────────── + +function emptyAttribute(): LabelGroup { + return { + key: "", + description: "", + type: "value", + optional: true, + tag: false, + values: [], + }; +} + +function emptyValue(): LabelValue { + return { value: "", description: "" }; +} + +function EntityLabelsEditor({ + value, + onChange, +}: { + value: LabelGroup[]; + onChange: (attrs: LabelGroup[]) => void; +}) { + const [expanded, setExpanded] = useState>({}); + + const updateAttr = (i: number, patch: Partial) => { + const next = value.map((a, idx) => (idx === i ? { ...a, ...patch } : a)); + onChange(next); + }; + + const removeAttr = (i: number) => { + onChange(value.filter((_, idx) => idx !== i)); + setExpanded((prev) => { + const next = { ...prev }; + delete next[i]; + return next; + }); + }; + + const addAttr = () => { + const next = [...value, emptyAttribute()]; + onChange(next); + setExpanded((prev) => ({ ...prev, [next.length - 1]: true })); + }; + + const updateVal = (attrIdx: number, valIdx: number, patch: Partial) => { + const newValues = value[attrIdx].values.map((v, vi) => + vi === valIdx ? { ...v, ...patch } : v + ); + updateAttr(attrIdx, { values: newValues }); + }; + + const removeVal = (attrIdx: number, valIdx: number) => { + updateAttr(attrIdx, { values: value[attrIdx].values.filter((_, vi) => vi !== valIdx) }); + }; + + const addVal = (attrIdx: number) => { + updateAttr(attrIdx, { values: [...value[attrIdx].values, emptyValue()] }); + }; -function Toggle({ value, onChange }: { value: boolean; onChange: (v: boolean) => void }) { return ( - +
+
+
+

Label Groups

+

+ Classification labels extracted at retain time. Leave empty to disable. +

+
+ {value.length > 0 && ( + + {value.length} group{value.length !== 1 ? "s" : ""} + + )} +
+ + {value.length === 0 && ( +

No label groups defined.

+ )} + +
+ {value.map((attr, i) => { + const isOpen = expanded[i] ?? false; + const isText = attr.type === "text"; + const hasValues = !isText; + return ( +
+ {/* Attribute header */} +
+ + updateAttr(i, { key: e.target.value })} + className="h-8 text-xs font-mono w-36 shrink-0" + /> + updateAttr(i, { description: e.target.value })} + className="h-8 text-xs flex-1 min-w-0" + /> + {/* Type dropdown */} + + {/* Tag checkbox — also write extracted labels as tags */} + + +
+ + {/* Values list — enum and multi-values only */} + {isOpen && hasValues && ( +
+ {attr.values.length === 0 && ( +

No values yet.

+ )} + {attr.values.map((v, vi) => ( +
+ updateVal(i, vi, { value: e.target.value })} + className="h-8 text-xs font-mono w-32 shrink-0" + /> + updateVal(i, vi, { description: e.target.value })} + className="h-8 text-xs flex-1 min-w-0" + /> + +
+ ))} + +
+ )} +
+ ); + })} +
+ + +
); } diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index f003ae9b..4e642092 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -553,6 +553,8 @@ Controls the retain (memory ingestion) pipeline. | `HINDSIGHT_API_RETAIN_BATCH_ENABLED` | Use LLM Batch API for fact extraction (50% cost savings, only with async operations) | `false` | | `HINDSIGHT_API_RETAIN_BATCH_POLL_INTERVAL_SECONDS` | Batch API polling interval in seconds | `60` | +> **Entity labels** (`entity_labels`) and **free-form entity extraction** (`entities_allow_free_form`) are configured per bank via the [bank config API](/developer/api/memory-banks#retain-configuration), not as global environment variables — each bank can have its own controlled vocabulary. See [Entity Labels](/developer/retain#entity-labels) for details. + #### Customizing retain: when to use what There are three levels of customization for the retain pipeline. Start with the simplest that covers your needs: diff --git a/hindsight-docs/docs/developer/retain.md b/hindsight-docs/docs/developer/retain.md index fc85d5e0..b7676b95 100644 --- a/hindsight-docs/docs/developer/retain.md +++ b/hindsight-docs/docs/developer/retain.md @@ -199,6 +199,90 @@ Set `retain_mission` and `retain_extraction_mode` via the [bank config API](/dev --- +## Entity Labels + +**Entity labels** let you define a controlled vocabulary of classification labels that are extracted at retain time and stored as entities alongside regular named entities. Each label takes the form `key:value` (e.g. `pedagogy:scaffolding`, `engagement:active`). + +Because labels become entities, they automatically: +- Appear in the **knowledge graph** — two memories with `pedagogy:scaffolding` are linked +- Improve **semantic and BM25 retrieval** — label strings are included in both the dense embedding and the sparse `text_signals` field +- Support **labels-only mode** — optionally disable free-form entity extraction so only labels are stored + +Labels are configured per bank via `entity_labels` in the bank config. + +### Defining Label Groups + +Each label group defines one classification dimension: + +```json +{ + "entity_labels": [ + { + "key": "engagement", + "description": "Student engagement level during the session", + "type": "value", + "optional": true, + "values": [ + { "value": "active", "description": "Student is actively participating" }, + { "value": "passive", "description": "Student is listening but not participating" } + ] + }, + { + "key": "pedagogy", + "description": "Teaching strategies used", + "type": "multi-values", + "values": [ + { "value": "scaffolding", "description": "Breaking complex tasks into smaller steps" }, + { "value": "direct_instruction", "description": "Explicit explanation by the teacher" }, + { "value": "socratic_questioning", "description": "Guiding through questions rather than answers" } + ] + } + ] +} +``` + +| Field | Default | Description | +|-------|---------|-------------| +| `key` | — | Label group identifier. Becomes the prefix in `key:value` entities. | +| `description` | `""` | Shown to the LLM to help it assign the right label. | +| `type` | `"value"` | `"value"` → single enum value; `"multi-values"` → multiple enum values; `"text"` → free-form string. | +| `values` | `[]` | Allowed values for `"value"` and `"multi-values"` types. Ignored for `"text"` type. | +| `optional` | `true` | `true` → the LLM may skip this label if not applicable (default). `false` → LLM must always assign a value. Has no effect on `"multi-values"` groups (always optional). | +| `tag` | `false` | `true` → also write extracted `key:value` entities as tags on the memory unit, enabling filtering via the standard `tags`/`tags_match` API parameters. | + +### Enum vs Free-text Labels + +**Enum groups** (`type: "value"` or `type: "multi-values"`): the LLM must pick from the predefined `values` list. Values not in the list are silently dropped. This is the most reliable option — the vocabulary is stable and graph clustering is tight. Use `"multi-values"` when a single fact can match multiple values. + +**Free-text groups** (`type: "text"`): the LLM can write any string value. The `values` field is ignored — use the `description` to provide examples and guidance instead. + +```json +{ + "key": "topic", + "description": "The specific subject being discussed. Examples: algebra, geometry, quadratic equations.", + "type": "text", + "optional": true, + "values": [] +} +``` + +The trade-off with free-text: the LLM may use different phrasings for the same concept across sessions (`topic:fractions` vs `topic:fraction arithmetic`), so graph linking is less reliable than with enum groups. + +### Labels-only Mode + +By default, entity labels are extracted **alongside** regular named entities (people, places, concepts). Set `entities_allow_free_form: false` to disable free-form extraction and store only label entities: + +```json +{ + "entity_labels": [...], + "entities_allow_free_form": false +} +``` + +Configure both via the [bank config API](/developer/api/memory-banks#retain-configuration). + +--- + ## Observation Consolidation After `retain()` completes, Hindsight automatically triggers **observation consolidation** in the background. This process: diff --git a/hindsight-docs/examples/api/quickstart.go b/hindsight-docs/examples/api/quickstart.go index 06a3458d..a9da853c 100644 --- a/hindsight-docs/examples/api/quickstart.go +++ b/hindsight-docs/examples/api/quickstart.go @@ -62,7 +62,7 @@ func main() { { Content: "Alice got promoted", Context: *hindsight.NewNullableString(hindsight.PtrString("career update")), - Timestamp: *hindsight.NewNullableTime(hindsight.PtrTime(timestamp)), + Timestamp: *hindsight.NewNullableTimestamp(&hindsight.Timestamp{TimeTime: hindsight.PtrTime(timestamp)}), Tags: []string{"career"}, }, }, diff --git a/scripts/dev/start-control-plane.sh b/scripts/dev/start-control-plane.sh index 3f848285..cea65e59 100755 --- a/scripts/dev/start-control-plane.sh +++ b/scripts/dev/start-control-plane.sh @@ -28,6 +28,7 @@ fi # Map prefixed env vars to Next.js standard vars export HOSTNAME="${HINDSIGHT_CP_HOSTNAME:-0.0.0.0}" +export PORT="${HINDSIGHT_CP_PORT:-9999}" # Run dev server npm run dev -w @vectorize-io/hindsight-control-plane \ No newline at end of file diff --git a/scripts/dev/start.sh b/scripts/dev/start.sh index 78ec125f..36c932bc 100755 --- a/scripts/dev/start.sh +++ b/scripts/dev/start.sh @@ -11,6 +11,7 @@ if [ -f "$ROOT_DIR/.env" ]; then set +a fi API_PORT="${HINDSIGHT_API_PORT:-8888}" +CP_PORT="${HINDSIGHT_CP_PORT:-9999}" PIDS=() @@ -70,7 +71,7 @@ echo "" echo "Hindsight is running!" echo "" echo " API: http://localhost:${API_PORT}" -echo " Control Plane: http://localhost:9999" +echo " Control Plane: http://localhost:${CP_PORT}" echo "" echo "Press Ctrl+C to stop both services." echo ""