diff --git a/hindsight-api-slim/hindsight_api/engine/retain/fact_extraction.py b/hindsight-api-slim/hindsight_api/engine/retain/fact_extraction.py index eb9a50b7c8..dd0491927d 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/fact_extraction.py +++ b/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 = [] diff --git a/hindsight-api-slim/tests/test_batch_api.py b/hindsight-api-slim/tests/test_batch_api.py index 7770b5bba6..b58f655b3e 100644 --- a/hindsight-api-slim/tests/test_batch_api.py +++ b/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.""" diff --git a/hindsight-api-slim/tests/test_fact_extraction_retry.py b/hindsight-api-slim/tests/test_fact_extraction_retry.py index ba2ae09458..6ad0e15b22 100644 --- a/hindsight-api-slim/tests/test_fact_extraction_retry.py +++ b/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(): """