# Spec: Memory-Aware Extraction

> From PRD: `.vibeflow/prds/memory-aware-extraction.md`
> Generated: 2026-03-30

## Objective

Substituir a extração cega (4 chamadas LLM) por extração informada pela memória existente (1 chamada LLM + 1 DB lookup), eliminando duplicatas e ruído na fonte e classificando importância semântica no momento da extração.

## Context

O write pipeline atual segue o fluxo: entity scan (LLM) → fact extraction (LLM) → relation extraction (LLM) → reconciliation (LLM). São 4 chamadas LLM por mensagem. A extração é cega — o LLM não sabe o que o sistema já sabe. Isso causa:

- 7% de fatos duplicados (LoCoMo: "Mel friend of Caroline" 8x)
- 10% de fatos conversacionais sem valor
- 44 tipos de relationship pra 56 relações (explosão sem vocabulário)
- Reconciliação como patch pós-extração que adiciona latência e nem sempre funciona

O código atual vive em `write/extract.py` (3 passes + dedup) e `write/reconcile.py` (embedding search + LLM). O pipeline é orquestrado por `write/pipeline.py`.

## Definition of Done

1. **Alias lookup funciona sem LLM** — Dado um message text e um user_id, o sistema identifica entidades conhecidas via string matching no cache de aliases (MemoryEntityAlias), retornando entity_keys sem chamada LLM.
2. **Pré-retrieval monta contexto existente** — Pra cada entidade conhecida encontrada, o sistema busca os top-K fatos existentes mais relevantes à mensagem via embedding similarity, respeitando o budget de tokens configurado.
3. **Extração informada produz output completo em 1 chamada** — O novo prompt recebe mensagem + speaker context + contexto existente e retorna entities + facts (com `action` e `importance_category`) + relations num único JSON. Fatos já conhecidos não são extraídos.
4. **Pipeline usa extração informada como default** — `run_write_pipeline()` chama a extração informada em vez da multi-pass. Reconciliação (`reconcile_facts`) é skippada quando a extração informada é usada — as decisões (NEW/UPDATE) já vêm no output.
5. **Fallback pra extração cega funciona** — Se a extração informada falhar (timeout, parse error, output malformado), o pipeline cai automaticamente pra extração cega (`extract_from_message`) + reconciliação (`reconcile_facts`), sem perda de dados.
6. **Cold start equivalente** — Com base vazia (0 aliases, 0 fatos), o comportamento é idêntico à extração cega atual. Sem lógica especial, sem branch de cold start.
7. **Segue patterns do projeto** — LLM calls via `_call_llm` helper com `temperature=0`, `response_format=json_object`, `asyncio.wait_for` timeout. Fail-safe retorna defaults. Prompts como constantes module-level. Pipeline não faz commit. Savepoints pra DB writes.

## Scope

### Arquivos novos

- **`src/arandu/write/informed_extract.py`** — Módulo principal. Contém:
  - `alias_lookup()` — Carrega aliases do user, faz string matching case-insensitive na mensagem, retorna dict de entity_keys encontradas.
  - `build_extraction_context()` — Pra cada entity_key, busca top-K fatos via embedding similarity (pgvector), formata como bloco de texto dentro do budget de tokens.
  - `informed_extract_from_message()` — Monta o prompt informado, chama LLM, parseia output com entities + facts (action, importance_category) + relations. Entry point público.
  - `INFORMED_EXTRACTION_PROMPT` — Prompt unificado como constante module-level.

### Arquivos modificados

- **`src/arandu/constants.py`** — `ExtractedFact` ganha campos `action: str = "NEW"` (default NEW pra backwards compat) e `importance_category: str = "specific_event"` (default). Novo `IMPORTANCE_CATEGORY_MAP: dict[str, float]` mapeando categorias pra valores numéricos.
- **`src/arandu/config.py`** — Novos campos em `MemoryConfig`: `informed_extraction_topk: int = 10`, `informed_extraction_context_budget_tokens: int = 800`, `enable_informed_extraction: bool = True`.
- **`src/arandu/write/pipeline.py`** — `run_write_pipeline()` tenta extração informada primeiro. Se sucesso: pula reconciliação, passa action/importance direto pro upsert. Se falha: fallback pra fluxo atual (extração cega + reconciliação).

### Arquivos inalterados (fallback)

- **`src/arandu/write/extract.py`** — Mantido intacto como fallback. `extract_from_message()` continua disponível.
- **`src/arandu/write/reconcile.py`** — Mantido intacto como fallback. `reconcile_facts()` continua disponível.

### Testes

- **`tests/test_informed_extract.py`** — Testes unitários pra alias_lookup, build_extraction_context, informed_extract_from_message, fallback behavior, cold start.

## Anti-scope

- **Entity profiles (living summaries)** — PRD separado. Este spec injeta fatos individuais, não perfis.
- **Remoção de `extract.py` ou `reconcile.py`** — Ficam como fallback. Não deletar.
- **Mudanças no read pipeline** — Zero alterações.
- **Mudanças no schema do banco** — Nenhuma migration. `action` e `importance_category` existem apenas no schema intermediário de extração (`ExtractedFact`), não no `MemoryFact` ORM model. O `importance` numérico do `MemoryFact` é setado via `IMPORTANCE_CATEGORY_MAP` no upsert.
- **Vocabulário controlado de relationships** — Não forçar tipos canônicos no prompt. A extração informada pode reduzir explosão naturalmente (LLM vê rels existentes no contexto).
- **Testes de integração com DB real** — Testes unitários com mocks, seguindo pattern de testing do projeto.
- **Otimização de performance do alias lookup** — v0 faz SELECT + in-memory matching. Sem trie ou cache distribuído.

## Technical Decisions

### 1. Alias lookup: SELECT + in-memory matching (não NER)

**Escolha:** Carregar todos os aliases do user em memória (`SELECT alias, canonical_entity_key FROM memory_entity_aliases WHERE ...`), normalizar, e fazer substring matching case-insensitive na mensagem.

**Trade-off:** Mais simples e previsível que NER ou embedding-based entity detection. Não detecta entidades novas — isso fica pro LLM na extração informada. Pode ter falsos positivos com aliases curtos (<3 chars) — mitigação: skip aliases com length < 3.

**Alternativa rejeitada:** Embedding-based entity detection na mensagem. Mais sofisticado mas adiciona uma chamada de embedding antes da extração, e o LLM já detecta entidades novas no prompt unificado.

### 2. Contexto: top-K por embedding similarity à mensagem

**Escolha:** Pra cada entidade encontrada, fazer `embedding_vec <=> message_embedding ORDER BY LIMIT K`. Usar a mesma query pattern de `reconcile.py` (pgvector cosine distance).

**Trade-off:** Injeta os fatos mais relevantes ao contexto atual, não os mais recentes ou importantes. Um fato antigo mas relevante à mensagem atual é mais útil como contexto do que um fato recente mas irrelevante.

**Budget:** K=10 fatos por entidade, budget total ~800 tokens. Com 2 entidades típicas: ~400 tokens de contexto. Mensagens com 5+ entidades: distribuir budget proporcionalmente (limit K = max(3, 800 / n_entities / avg_tokens_per_fact)).

### 3. Output schema: action como campo no ExtractedFact

**Escolha:** `ExtractedFact` ganha `action: str = "NEW"` com valores possíveis "NEW" e "UPDATE". UPDATE inclui campo `updates: str | None` com o texto do fato existente que está sendo substituído.

**Trade-off:** Default "NEW" garante backwards compatibility — o fluxo cego produz fatos sem action, que são tratados como NEW (comportamento atual). O pipeline detecta se a extração informada rodou pelo campo `extraction_meta["mode"] == "informed"` e decide se pula reconciliação.

**Nota:** Não há action "DELETE". Retratações ("I no longer work at X") produzem um UPDATE com o novo estado. O DELETE explícito fica pro fluxo de reconciliação (fallback).

### 4. Importance categories: LLM classifica, sistema mapeia pra float

**Escolha:** O prompt pede `importance_category` como string categórica. O sistema mapeia pra float via `IMPORTANCE_CATEGORY_MAP`:

```python
IMPORTANCE_CATEGORY_MAP: dict[str, float] = {
    "biographical_milestone": 0.90,   # marriage, birth, death, graduation
    "relationship_change": 0.85,      # new friend, breakup, new job
    "stable_preference": 0.80,        # likes art, interested in mental health
    "specific_event": 0.60,           # went to parade, saw meteor shower
    "routine_activity": 0.30,         # went to beach, had dinner
    "conversational": 0.10,           # greeting, thanking, agreeing
}
```

**Trade-off:** Categorias semânticas são mais estáveis e interpretáveis que pedir um float direto ao LLM (que tende a dar 0.8 pra tudo). O mapeamento é determinístico e tunable. Categorias desconhecidas mapeiam pro default 0.5.

### 5. Fallback: try/except no pipeline, não feature flag

**Escolha:** O pipeline tenta extração informada dentro de try/except. Se qualquer exceção (timeout, parse error, output malformado), loga warning e cai pro fluxo atual (extração cega + reconciliação). O fallback é transparente — o caller não sabe qual path rodou (exceto via trace).

**Trade-off:** Simples, sem feature flag. `enable_informed_extraction=False` em MemoryConfig desabilita completamente (vai direto pro fluxo antigo). Útil pra debugging e rollback.

### 6. Prompt unificado: entities + facts + relations em 1 call

**Escolha:** Um prompt que pede entities, facts (com action + importance_category), e relations no mesmo JSON. O entity scan separado é eliminado — o LLM identifica entidades como parte da extração informada.

**Trade-off:** Prompt mais complexo. Risco de degradação de qualidade em LLMs menores. Mitigação: fallback pro fluxo multi-pass se parse falhar. Testar com DeepSeek V3 e OpenAI.

**Formato do output esperado:**
```json
{
  "entities": [
    {"name": "Caroline", "type": "person", "aliases": []}
  ],
  "facts": [
    {
      "subject": "Caroline",
      "fact": "Caroline attended a pride parade in her city last weekend",
      "confidence": 0.95,
      "action": "NEW",
      "importance_category": "specific_event"
    },
    {
      "subject": "Caroline",
      "fact": "Caroline is an active member of the LGBTQ community",
      "confidence": 0.90,
      "action": "UPDATE",
      "updates": "Caroline supports LGBTQ rights",
      "importance_category": "stable_preference"
    }
  ],
  "relations": [
    {"source": "Caroline", "target": "pride parade", "rel_type": "attended"}
  ]
}
```

## Applicable Patterns

- **Pipeline Orchestration** (`patterns/pipeline-orchestration.md`) — `run_write_pipeline()` orquestra stages. A extração informada é um novo stage que substitui extract + reconcile quando bem-sucedido. Timing via `time.perf_counter()`. Trace via `trace.add_step()`.
- **LLM Interaction** (`patterns/llm-interaction.md`) — `_call_llm()` helper reutilizado. `temperature=0`, `response_format=json_object`. `strip_markdown_fences()` + `json.loads()` + Pydantic validation. Fail-safe retorna empty/defaults.
- **Data Access** (`patterns/data-access.md`) — Alias lookup e pré-retrieval usam `AsyncSession` com queries SQLAlchemy + raw `text()` pra pgvector. Savepoints pra isolamento.
- **Error Handling** (`patterns/error-handling.md`) — Fail-safe em todos os stages. Extração informada falha → fallback pra extração cega. Nunca propaga exceções ao caller.

## Risks

### 1. Qualidade do prompt unificado com LLMs menores
**Risco:** DeepSeek V3 pode não seguir instruções complexas (entities + facts com action/importance + relations) num único prompt. Parse errors frequentes ativariam fallback em toda mensagem, anulando o benefício.
**Mitigação:** Fallback transparente garante que a funcionalidade não degrada. Testar o prompt com 3+ providers (DeepSeek, OpenAI, Anthropic) antes de considerar estável. Se a taxa de fallback for >20%, simplificar o prompt (ex: separar relations como pass 2).

### 2. Alias lookup com falsos positivos
**Risco:** Alias curto como "AI" ou "Art" pode fazer match em contextos errados ("I went to the airport" contém "AI"). Injetar contexto da entidade errada polui a extração.
**Mitigação:** Skip aliases com length < 3. Match por word boundary (não substring puro). Testar com corpus real.

### 3. Pré-retrieval seleciona fatos irrelevantes
**Risco:** Embedding similarity da mensagem inteira pode surfar fatos não relacionados. "I went to the beach" pode puxar "Mel loves the beach" (relevante) mas também "Mel's address is Copacabana" (irrelevante se Copacabana associa com praia).
**Mitigação:** Min similarity threshold (0.3). Budget limita o dano (top-10, não top-100). A extração informada filtra naturalmente — contexto irrelevante é ignorado pelo LLM.

### 4. Latência do pré-retrieval em bases grandes
**Risco:** Se um user tem 1000+ fatos e 50+ entidades, o alias lookup + embedding search pode levar >100ms.
**Mitigação:** LIMIT no alias query. pgvector com IVFFlat index. Budget de tokens limita naturalmente o volume. Monitorar via trace.
