Accept top-level fact arrays in retain parsing

hindsight-api-slim/hindsight_api/engine/retain/fact_extraction.py

@@ -1236,6 +1236,15 @@ def _build_request_body(llm_config, config, prompt: str, user_message: str, resp
     return request_body
 
 
+def _coerce_fact_response(response: Any) -> dict[str, Any] | None:
+    """Accept the schema wrapper, or a recoverable top-level facts array."""
+    if isinstance(response, dict):
+        return response
+    if isinstance(response, list) and all(isinstance(item, dict) for item in response):
+        return {"facts": response}
+    return None
+
+
 async def _extract_facts_from_chunk(
     chunk: str,
     chunk_index: int,
@@ -1341,7 +1350,8 @@ async def _extract_facts_from_chunk(
             has_malformed_facts = False
 
             # Handle malformed LLM responses
-            if not isinstance(extraction_response_json, dict):
+            coerced_response_json = _coerce_fact_response(extraction_response_json)
+            if coerced_response_json is None:
                 if attempt < llm_max_retries - 1:
                     logger.warning(
                         f"LLM returned non-dict JSON on attempt {attempt + 1}/{llm_max_retries}: {type(extraction_response_json).__name__}. Retrying..."
@@ -1356,6 +1366,7 @@ async def _extract_facts_from_chunk(
                         f"Fact extraction failed: LLM returned non-dict JSON after {llm_max_retries} attempts "
                         f"({type(extraction_response_json).__name__}). Raw: {str(extraction_response_json)[:500]}"
                     )
+            extraction_response_json = coerced_response_json
 
             raw_facts = extraction_response_json.get("facts", [])
 
@@ -2144,6 +2155,19 @@ async def extract_facts_from_contents_batch_api(
             )
             continue
 
+        response_type_name = type(extraction_response_json).__name__
+        extraction_response_json = _coerce_fact_response(extraction_response_json)
+        if extraction_response_json is None:
+            message = f"{custom_id}: LLM returned non-dict JSON ({response_type_name})"
+            logger.error(message)
+            extraction_errors.add(message)
+            chunks_metadata.append(
+                ChunkMetadata(
+                    chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
+                )
+            )
+            continue
+
         # Parse facts (reuse existing logic from _extract_facts_from_chunk)
         raw_facts = extraction_response_json.get("facts", [])
         chunk_facts = []

hindsight-api-slim/tests/test_batch_api.py

@@ -8,18 +8,15 @@
 - Worker recovery on restart
 """
 
-import asyncio
 import json
 import logging
 import uuid
 from datetime import datetime, timezone
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
-from hindsight_api import RequestContext
 from hindsight_api.config import HindsightConfig
-from hindsight_api.engine.llm_wrapper import create_llm_provider
 from hindsight_api.engine.retain.fact_extraction import (
     RetainContent,
     extract_facts_from_contents,
@@ -202,6 +199,109 @@ async def test_batch_api_normal_flow(mock_llm_config, test_contents, hindsight_c
             pass
 
 
+@pytest.mark.asyncio
+async def test_batch_api_accepts_top_level_fact_list(mock_llm_config, test_contents, hindsight_config):
+    """Batch extraction accepts a recoverable top-level facts array."""
+    batch_id = "batch_top_level_facts"
+    mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)
+    mock_llm_config._provider_impl.submit_batch = AsyncMock(
+        return_value={
+            "batch_id": batch_id,
+            "status": "validating",
+            "request_counts": {"total": 1, "completed": 0, "failed": 0},
+        }
+    )
+    mock_llm_config._provider_impl.get_batch_status = AsyncMock(
+        return_value={"status": "completed", "request_counts": {"total": 1, "completed": 1, "failed": 0}}
+    )
+    mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(
+        return_value=[
+            {
+                "custom_id": "chunk_0",
+                "response": {
+                    "body": {
+                        "choices": [
+                            {
+                                "message": {
+                                    "content": json.dumps(
+                                        [
+                                            {
+                                                "what": "Alice is a senior software engineer at TechCorp",
+                                                "when": "present",
+                                                "where": "TechCorp",
+                                                "who": "Alice",
+                                                "why": "Professional background information",
+                                                "fact_type": "world",
+                                                "fact_kind": "conversation",
+                                            }
+                                        ]
+                                    )
+                                }
+                            }
+                        ],
+                        "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
+                    }
+                },
+            }
+        ]
+    )
+
+    facts, chunks, usage = await extract_facts_from_contents_batch_api(
+        contents=[test_contents[0]],
+        llm_config=mock_llm_config,
+        agent_name="test_agent",
+        config=hindsight_config,
+        pool=None,
+        operation_id=None,
+        schema=None,
+    )
+
+    assert len(facts) == 1
+    assert "Alice" in facts[0].fact_text
+    assert len(chunks) == 1
+    assert chunks[0].fact_count == 1
+    assert usage.total_tokens == 150
+
+
+@pytest.mark.asyncio
+async def test_batch_api_rejects_top_level_non_fact_list(mock_llm_config, test_contents, hindsight_config):
+    """Batch extraction records malformed top-level lists instead of crashing."""
+    batch_id = "batch_malformed_list"
+    mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)
+    mock_llm_config._provider_impl.submit_batch = AsyncMock(return_value={"batch_id": batch_id})
+    mock_llm_config._provider_impl.get_batch_status = AsyncMock(
+        return_value={"status": "completed", "request_counts": {"total": 1, "completed": 1, "failed": 0}}
+    )
+    mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(
+        return_value=[
+            {
+                "custom_id": "chunk_0",
+                "response": {
+                    "body": {
+                        "choices": [{"message": {"content": json.dumps(["not a fact dict"])}}],
+                        "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
+                    }
+                },
+            }
+        ]
+    )
+
+    facts, chunks, usage = await extract_facts_from_contents_batch_api(
+        contents=[test_contents[0]],
+        llm_config=mock_llm_config,
+        agent_name="test_agent",
+        config=hindsight_config,
+        pool=None,
+        operation_id=None,
+        schema=None,
+    )
+
+    assert facts == []
+    assert len(chunks) == 1
+    assert chunks[0].fact_count == 0
+    assert usage.total_tokens == 0
+
+
 @pytest.mark.asyncio
 async def test_batch_api_crash_recovery(mock_llm_config, test_contents, hindsight_config, memory, request_context):
     """Test crash recovery: resume polling from existing batch_id."""

hindsight-api-slim/tests/test_fact_extraction_retry.py

@@ -60,8 +60,8 @@ async def test_non_dict_json_all_retries_raises():
 
     config = _make_config(llm_max_retries=3, retain_llm_max_retries=None)
 
-    # Mock: always returns a list (non-dict), which is invalid
-    llm_config = _make_llm_config(mock_response=[{"invalid": "response"}])
+    # Mock: always returns a list containing a non-dict item, which is invalid.
+    llm_config = _make_llm_config(mock_response=["invalid response"])
 
     with patch(
         "hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
@@ -82,6 +82,50 @@ async def test_non_dict_json_all_retries_raises():
     assert llm_config.call.call_count == 3
 
 
+@pytest.mark.asyncio
+async def test_top_level_fact_list_is_accepted_without_retry():
+    """
+    Some lax-JSON models return the facts array directly instead of wrapping it
+    in {"facts": [...]}. A top-level list of dict-shaped facts is recoverable
+    and should not burn retries.
+    """
+    from hindsight_api.engine.retain.fact_extraction import _extract_facts_from_chunk
+
+    config = _make_config(llm_max_retries=3, retain_llm_max_retries=None)
+    llm_config = _make_llm_config(
+        mock_response=[
+            {
+                "what": "Alice visited Paris",
+                "when": "2023",
+                "where": "Paris",
+                "who": "Alice",
+                "why": "vacation",
+                "fact_type": "world",
+                "fact_kind": "conversation",
+            }
+        ]
+    )
+
+    with patch(
+        "hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
+        return_value=("system prompt", MagicMock()),
+    ):
+        facts, _usage = await _extract_facts_from_chunk(
+            chunk="Alice visited Paris in 2023.",
+            chunk_index=0,
+            total_chunks=1,
+            event_date=datetime(2023, 1, 1, tzinfo=timezone.utc),
+            context="travel notes",
+            llm_config=llm_config,
+            config=config,
+            agent_name="test-agent",
+        )
+
+    assert llm_config.call.call_count == 1
+    assert len(facts) == 1
+    assert "Alice visited Paris" in facts[0].fact
+
+
 @pytest.mark.asyncio
 async def test_non_dict_json_with_default_max_retries_raises():
     """