From 2a00df0bc0cd50e7c33d1d0abede38ab587dd2e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Wed, 7 Jan 2026 11:16:24 +0100 Subject: [PATCH] fix: improve causal links detection (#111) * fix: improve causal links detection * fix: improve causal links detection --- .../engine/retain/fact_extraction.py | 177 ++++++++++++-- .../tests/test_causal_relationships.py | 222 ++++++++++++++++++ 2 files changed, 384 insertions(+), 15 deletions(-) create mode 100644 hindsight-api/tests/test_causal_relationships.py diff --git a/hindsight-api/hindsight_api/engine/retain/fact_extraction.py b/hindsight-api/hindsight_api/engine/retain/fact_extraction.py index 63cbf88c..ab3494b6 100644 --- a/hindsight-api/hindsight_api/engine/retain/fact_extraction.py +++ b/hindsight-api/hindsight_api/engine/retain/fact_extraction.py @@ -110,7 +110,7 @@ class Fact(BaseModel): class CausalRelation(BaseModel): - """Causal relationship between facts.""" + """Causal relationship between facts (legacy - embedded in each fact).""" target_fact_index: int = Field( description="Index of the related fact in the facts array (0-based). " @@ -132,6 +132,36 @@ class CausalRelation(BaseModel): ) +class TopLevelCausalRelation(BaseModel): + """ + Causal relationship between two facts (top-level schema). + + This is the preferred format - defined AFTER all facts are extracted, + allowing the LLM to see the full list of facts before specifying relationships. + """ + + from_fact_index: int = Field( + description="Index of the source fact (0-based). The fact that causes/enables/prevents." + ) + to_fact_index: int = Field( + description="Index of the target fact (0-based). The fact that is caused/enabled/prevented." + ) + relation_type: Literal["causes", "caused_by", "enables", "prevents"] = Field( + description="Type of causal relationship: " + "'causes' = source fact directly causes the target fact, " + "'caused_by' = source fact was caused by the target fact, " + "'enables' = source fact enables/allows the target fact, " + "'prevents' = source fact prevents/blocks the target fact" + ) + strength: float = Field( + description="Strength of causal relationship (0.0 to 1.0). " + "1.0 = direct/strong causation, 0.5 = moderate, 0.3 = weak/indirect", + ge=0.0, + le=1.0, + default=1.0, + ) + + class ExtractedFact(BaseModel): """A single extracted fact with 5 required dimensions for comprehensive capture.""" @@ -254,9 +284,15 @@ class ExtractedFact(BaseModel): class FactExtractionResponse(BaseModel): - """Response containing all extracted facts.""" + """Response containing all extracted facts and their causal relationships.""" facts: list[ExtractedFact] = Field(description="List of extracted factual statements") + causal_relationships: list[TopLevelCausalRelation] | None = Field( + default=None, + description="Causal relationships between facts. Define these AFTER listing all facts. " + "Each relationship specifies from_fact_index -> to_fact_index with a relation type. " + "Indices must be valid (0 to N-1 where N is the number of facts).", + ) def chunk_text(text: str, max_chars: int) -> list[str]: @@ -573,7 +609,53 @@ WHAT TO EXTRACT vs SKIP ══════════════════════════════════════════════════════════════════════════ ✅ EXTRACT: User preferences (ALWAYS as separate facts!), feelings, plans, events, relationships, achievements -❌ SKIP: Greetings, filler ("thanks", "cool"), purely structural statements""" +❌ SKIP: Greetings, filler ("thanks", "cool"), purely structural statements + +══════════════════════════════════════════════════════════════════════════ +CAUSAL RELATIONSHIPS (CRITICAL - DEFINE AFTER ALL FACTS) +══════════════════════════════════════════════════════════════════════════ + +⚠️ IMPORTANT: Causal relationships are defined at the TOP LEVEL, AFTER listing all facts! + +The `causal_relationships` array goes at the root of your response (NOT inside each fact). +This allows you to see all facts first before defining how they relate. + +Format: +```json +{{ + "facts": [...all your extracted facts...], + "causal_relationships": [ + {{"from_fact_index": 0, "to_fact_index": 1, "relation_type": "causes", "strength": 0.9}}, + {{"from_fact_index": 1, "to_fact_index": 2, "relation_type": "enables", "strength": 0.7}} + ] +}} +``` + +Relationship types: +- "causes": Fact A directly causes Fact B (A → B) +- "caused_by": Fact A was caused by Fact B (A ← B) +- "enables": Fact A enables/allows Fact B to happen +- "prevents": Fact A prevents/blocks Fact B from happening + +⚠️ INDEX VALIDATION: If you extract N facts (indices 0 to N-1), both from_fact_index and to_fact_index MUST be in range [0, N-1]. + +Example (Event Date: March 15, 2024): +Input: "I lost my job in January. Because of that, I couldn't pay rent. So I had to move to a cheaper apartment." + +Facts extracted: +- Fact 0: "User lost their job in January due to layoffs" +- Fact 1: "User couldn't pay rent because of job loss" +- Fact 2: "User moved to a cheaper apartment" + +Causal relationships (at root level): +```json +"causal_relationships": [ + {{"from_fact_index": 0, "to_fact_index": 1, "relation_type": "causes", "strength": 1.0}}, + {{"from_fact_index": 1, "to_fact_index": 2, "relation_type": "causes", "strength": 0.9}} +] +``` + +This creates a chain: Job loss (0) → Can't pay rent (1) → Moved to cheaper apartment (2)""" import logging @@ -633,6 +715,9 @@ Text: return [] raw_facts = extraction_response_json.get("facts", []) + # Get top-level causal relationships (new schema) + top_level_causal_relations = extraction_response_json.get("causal_relationships", []) + if not raw_facts: logger.debug( f"LLM response missing 'facts' field or returned empty list. " @@ -643,6 +728,47 @@ Text: f"text: {chunk}" ) + # Build a map from fact index to causal relations (from top-level field) + # This converts from_fact_index -> [{target_fact_index, relation_type, strength}] + causal_relations_by_fact: dict[int, list[dict]] = {} + if top_level_causal_relations: + num_facts = len(raw_facts) + for rel in top_level_causal_relations: + if not isinstance(rel, dict): + continue + from_idx = rel.get("from_fact_index") + to_idx = rel.get("to_fact_index") + relation_type = rel.get("relation_type") + strength = rel.get("strength", 1.0) + + # Validate indices + if from_idx is None or to_idx is None or relation_type is None: + logger.warning(f"Skipping malformed top-level causal relation: {rel}") + continue + if from_idx < 0 or from_idx >= num_facts: + logger.warning( + f"Invalid from_fact_index {from_idx} in top-level causal relation " + f"(valid range: 0-{num_facts - 1}). Skipping." + ) + continue + if to_idx < 0 or to_idx >= num_facts: + logger.warning( + f"Invalid to_fact_index {to_idx} in top-level causal relation " + f"(valid range: 0-{num_facts - 1}). Skipping." + ) + continue + + # Add to the map for the from_fact_index + if from_idx not in causal_relations_by_fact: + causal_relations_by_fact[from_idx] = [] + causal_relations_by_fact[from_idx].append( + { + "target_fact_index": to_idx, + "relation_type": relation_type, + "strength": strength, + } + ) + for i, llm_fact in enumerate(raw_facts): # Skip non-dict entries but track them for retry if not isinstance(llm_fact, dict): @@ -747,19 +873,40 @@ Text: if validated_entities: fact_data["entities"] = validated_entities - # Add causal relations if present (validate as CausalRelation objects) - # Filter out invalid relations (missing required fields) - causal_relations = get_value("causal_relations") - if causal_relations: - validated_relations = [] - for rel in causal_relations: + # Add causal relations from both sources: + # 1. Top-level causal_relationships (preferred, new schema) + # 2. Per-fact causal_relations (legacy, for backward compatibility) + validated_relations = [] + + # First, add relations from top-level (already validated above) + if i in causal_relations_by_fact: + for rel in causal_relations_by_fact[i]: + try: + validated_relations.append(CausalRelation.model_validate(rel)) + except Exception as e: + logger.warning(f"Invalid top-level causal relation for fact {i}: {rel}: {e}") + + # Then, add any legacy per-fact relations (with index validation) + legacy_causal_relations = get_value("causal_relations") + if legacy_causal_relations: + num_facts = len(raw_facts) + for rel in legacy_causal_relations: if isinstance(rel, dict) and "target_fact_index" in rel and "relation_type" in rel: - try: - validated_relations.append(CausalRelation.model_validate(rel)) - except Exception as e: - logger.warning(f"Invalid causal relation {rel}: {e}") - if validated_relations: - fact_data["causal_relations"] = validated_relations + target_idx = rel.get("target_fact_index") + # Validate target index for legacy format too + if target_idx is not None and 0 <= target_idx < num_facts: + try: + validated_relations.append(CausalRelation.model_validate(rel)) + except Exception as e: + logger.warning(f"Invalid causal relation {rel}: {e}") + else: + logger.warning( + f"Invalid target_fact_index {target_idx} in per-fact causal relation " + f"from fact {i} (valid range: 0-{num_facts - 1}). Skipping." + ) + + if validated_relations: + fact_data["causal_relations"] = validated_relations # Always set mentioned_at to the event_date (when the conversation/document occurred) fact_data["mentioned_at"] = event_date.isoformat() diff --git a/hindsight-api/tests/test_causal_relationships.py b/hindsight-api/tests/test_causal_relationships.py new file mode 100644 index 00000000..d8c64ece --- /dev/null +++ b/hindsight-api/tests/test_causal_relationships.py @@ -0,0 +1,222 @@ +""" +Test suite for causal relationship extraction. + +Tests that the fact extraction system correctly identifies and validates +causal relationships between facts, with valid indices. +""" + +from datetime import datetime + +import pytest + +from hindsight_api import LLMConfig +from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text + + +class TestCausalRelationships: + """Tests for causal relationship extraction and validation.""" + + @pytest.mark.asyncio + async def test_causal_chain_extraction(self): + """ + Test that a clear causal chain is extracted with valid relationships. + + Story: Lost job -> couldn't pay rent -> had to move -> found new apartment + + This is a 4-fact causal chain where each fact causes the next. + The extracted causal relations should have valid indices (0-3). + """ + text = """ +I lost my job at the tech company in January because of layoffs. +Because I lost my job, I couldn't pay my rent anymore. +Since I couldn't afford rent, I had to move out of my apartment. +After searching for weeks, I finally found a cheaper apartment in Brooklyn. +""" + + context = "Personal story about housing change" + llm_config = LLMConfig.for_memory() + + facts, _ = await extract_facts_from_text( + text=text, + event_date=datetime(2024, 3, 15), + context=context, + llm_config=llm_config, + agent_name="TestUser" + ) + + assert len(facts) >= 3, f"Should extract at least 3 facts from the causal chain. Got {len(facts)}" + + # Collect all causal relations from all facts + all_causal_relations = [] + for i, fact in enumerate(facts): + if fact.causal_relations: + for rel in fact.causal_relations: + all_causal_relations.append({ + "from_fact_index": i, + "to_fact_index": rel.target_fact_index, + "relation_type": rel.relation_type, + "strength": rel.strength, + "from_fact_text": fact.fact[:50], + }) + + # Verify that ALL causal relation indices are valid + num_facts = len(facts) + invalid_relations = [] + for rel in all_causal_relations: + if rel["to_fact_index"] < 0 or rel["to_fact_index"] >= num_facts: + invalid_relations.append(rel) + + assert len(invalid_relations) == 0, ( + f"Found {len(invalid_relations)} causal relations with invalid indices! " + f"Valid range is 0-{num_facts - 1}. " + f"Invalid relations: {invalid_relations}" + ) + + # Should have at least some causal relations extracted + assert len(all_causal_relations) >= 2, ( + f"Should extract at least 2 causal relationships from this clear chain. " + f"Got {len(all_causal_relations)}: {all_causal_relations}" + ) + + # Verify relation types are valid + valid_types = {"causes", "caused_by", "enables", "prevents"} + for rel in all_causal_relations: + assert rel["relation_type"] in valid_types, ( + f"Invalid relation_type '{rel['relation_type']}'. Must be one of {valid_types}" + ) + + @pytest.mark.asyncio + async def test_complex_causal_web(self): + """ + Test a more complex scenario with multiple interconnected causes. + + This tests the LLM's ability to identify multiple causal links and + ensure all referenced indices exist. + """ + text = """ +The heavy rain caused flooding in the basement. +The flooding damaged the electrical system. +Because of the electrical damage, we had to call an electrician. +The electrician found that the wiring was old and needed replacement. +We decided to renovate the entire basement while fixing the wiring. +The renovation took three months and cost $15,000. +""" + + context = "Home repair story" + llm_config = LLMConfig.for_memory() + + facts, _ = await extract_facts_from_text( + text=text, + event_date=datetime(2024, 6, 1), + context=context, + llm_config=llm_config, + agent_name="TestUser" + ) + + assert len(facts) >= 4, f"Should extract at least 4 facts. Got {len(facts)}" + + # Validate all causal relation indices + num_facts = len(facts) + for i, fact in enumerate(facts): + if fact.causal_relations: + for rel in fact.causal_relations: + assert 0 <= rel.target_fact_index < num_facts, ( + f"Fact {i} has causal relation to invalid index {rel.target_fact_index}. " + f"Valid range is 0-{num_facts - 1}. " + f"Fact text: {fact.fact[:80]}..." + ) + + @pytest.mark.asyncio + async def test_no_self_referencing_causal_relations(self): + """ + Test that facts don't have causal relations pointing to themselves. + """ + text = """ +I started learning Python because I wanted to automate my work tasks. +Learning Python led me to discover machine learning. +Machine learning fascinated me so much that I changed my career to data science. +""" + + context = "Career change story" + llm_config = LLMConfig.for_memory() + + facts, _ = await extract_facts_from_text( + text=text, + event_date=datetime(2024, 1, 1), + context=context, + llm_config=llm_config, + agent_name="TestUser" + ) + + # Check no fact references itself + for i, fact in enumerate(facts): + if fact.causal_relations: + for rel in fact.causal_relations: + assert rel.target_fact_index != i, ( + f"Fact {i} has a self-referencing causal relation! " + f"Fact text: {fact.fact}" + ) + + @pytest.mark.asyncio + async def test_bidirectional_causal_relationships(self): + """ + Test that bidirectional causal relationships (causes and caused_by) + are handled correctly. + """ + text = """ +My promotion at work caused me to move to New York. +Moving to New York was caused by my promotion at work. +The new role enabled me to lead a team of engineers. +""" + + context = "Work promotion story" + llm_config = LLMConfig.for_memory() + + facts, _ = await extract_facts_from_text( + text=text, + event_date=datetime(2024, 2, 15), + context=context, + llm_config=llm_config, + agent_name="TestUser" + ) + + num_facts = len(facts) + + # Validate all indices + for i, fact in enumerate(facts): + if fact.causal_relations: + for rel in fact.causal_relations: + assert 0 <= rel.target_fact_index < num_facts, ( + f"Invalid target_fact_index {rel.target_fact_index} in fact {i}. " + f"Valid range: 0-{num_facts - 1}" + ) + + @pytest.mark.asyncio + async def test_causal_relation_strength_values(self): + """ + Test that causal relation strength values are within valid range [0.0, 1.0]. + """ + text = """ +The stock market crash directly caused the company to lay off employees. +The layoffs indirectly led to reduced consumer spending in the area. +Reduced spending somewhat affected local businesses. +""" + + context = "Economic impact story" + llm_config = LLMConfig.for_memory() + + facts, _ = await extract_facts_from_text( + text=text, + event_date=datetime(2024, 4, 1), + context=context, + llm_config=llm_config, + agent_name="TestUser" + ) + + for i, fact in enumerate(facts): + if fact.causal_relations: + for rel in fact.causal_relations: + assert 0.0 <= rel.strength <= 1.0, ( + f"Causal relation strength {rel.strength} is outside valid range [0.0, 1.0]. " + f"Fact {i}: {fact.fact[:50]}..." + )