test_batch_api.py

"""
Test OpenAI Batch API integration for retain fact extraction.

Tests cover:
- Normal batch API flow (submit, poll, complete)
- Crash recovery (resume from existing batch_id)
- Hard error when provider doesn't support the batch API (no silent fallback)
- Worker recovery on restart
"""

import json
import logging
import uuid
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock

import pytest

from hindsight_api.config import HindsightConfig
from hindsight_api.engine.retain.fact_extraction import (
    RetainContent,
    extract_facts_from_contents,
    extract_facts_from_contents_batch_api,
)
from hindsight_api.worker.poller import WorkerPoller

logger = logging.getLogger(__name__)


@pytest.fixture
def mock_llm_config():
    """Create a mock LLM config with batch API support."""
    mock = MagicMock()
    mock.provider = "openai"
    mock.model = "gpt-4o-mini"
    mock._provider_impl = AsyncMock()
    return mock


@pytest.fixture
def test_contents():
    """Create test content for fact extraction."""
    return [
        RetainContent(
            content="Alice is a senior software engineer at TechCorp. She specializes in distributed systems.",
            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
            context="team overview",
        ),
        RetainContent(
            content="Bob joined the team last month as a junior developer. He is learning React.",
            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
            context="team overview",
        ),
    ]


@pytest.fixture
def hindsight_config():
    """Create test config with batch API enabled."""
    config = HindsightConfig.from_env()
    config.retain_batch_enabled = True
    config.retain_batch_poll_interval_seconds = 1  # Fast polling for tests
    config.retain_chunk_size = 4000
    config.retain_extraction_mode = "concise"
    config.retain_extract_causal_links = False
    return config


@pytest.mark.asyncio
async def test_batch_api_normal_flow(mock_llm_config, test_contents, hindsight_config, memory, request_context):
    """Test normal batch API flow: submit, poll, complete."""
    bank_id = f"test_batch_{datetime.now(timezone.utc).timestamp()}"

    try:
        # Mock batch API responses
        batch_id = "batch_test123"

        # Mock supports_batch_api
        mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)

        # Mock submit_batch - returns batch metadata
        mock_llm_config._provider_impl.submit_batch = AsyncMock(
            return_value={
                "batch_id": batch_id,
                "status": "validating",
                "request_counts": {"total": 2, "completed": 0, "failed": 0},
            }
        )

        # Mock get_batch_status - simulate polling sequence
        status_sequence = [
            {"status": "in_progress", "request_counts": {"total": 2, "completed": 1, "failed": 0}},
            {"status": "completed", "request_counts": {"total": 2, "completed": 2, "failed": 0}},
        ]
        mock_llm_config._provider_impl.get_batch_status = AsyncMock(side_effect=status_sequence)

        # Mock retrieve_batch_results - returns fact extraction results
        mock_results = [
            {
                "custom_id": "chunk_0",
                "response": {
                    "body": {
                        "choices": [
                            {
                                "message": {
                                    "content": json.dumps(
                                        {
                                            "facts": [
                                                {
                                                    "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},
                    }
                },
            },
            {
                "custom_id": "chunk_1",
                "response": {
                    "body": {
                        "choices": [
                            {
                                "message": {
                                    "content": json.dumps(
                                        {
                                            "facts": [
                                                {
                                                    "what": "Bob joined the team last month as a junior developer",
                                                    "when": "last month",
                                                    "where": "team",
                                                    "who": "Bob",
                                                    "why": "New team member information",
                                                    "fact_type": "world",
                                                    "fact_kind": "conversation",
                                                }
                                            ]
                                        }
                                    )
                                }
                            }
                        ],
                        "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
                    }
                },
            },
        ]
        mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(return_value=mock_results)

        # Call batch API extraction
        facts, chunks, usage = await extract_facts_from_contents_batch_api(
            contents=test_contents,
            llm_config=mock_llm_config,
            agent_name="test_agent",
            config=hindsight_config,
            pool=None,  # No DB pool for this test
            operation_id=None,
            schema=None,
        )

        # Verify results
        assert len(facts) == 2, "Should extract 2 facts (one per chunk)"
        # Facts are ExtractedFact objects with .fact_text field
        assert "Alice" in facts[0].fact_text and "senior software engineer" in facts[0].fact_text
        assert "Bob" in facts[1].fact_text and "junior developer" in facts[1].fact_text

        # Verify chunks metadata
        assert len(chunks) == 2, "Should have 2 chunks metadata"
        assert chunks[0].fact_count == 1
        assert chunks[1].fact_count == 1

        # Verify token usage
        assert usage.input_tokens == 200  # 100 per chunk
        assert usage.output_tokens == 100  # 50 per chunk
        assert usage.total_tokens == 300

        # Verify API calls
        mock_llm_config._provider_impl.submit_batch.assert_called_once()
        assert mock_llm_config._provider_impl.get_batch_status.call_count == 2
        mock_llm_config._provider_impl.retrieve_batch_results.assert_called_once_with(batch_id)

        logger.info("✅ Normal batch API flow test passed")

    finally:
        # Cleanup
        try:
            await memory.delete_bank(bank_id, request_context=request_context)
        except Exception:
            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_recovers_fenced_and_control_char_json(mock_llm_config, test_contents, hindsight_config):
    """#2701: batch content that bare json.loads can't parse but parse_llm_json can
    (markdown code fences + an embedded raw control character, e.g. a transient
    Gemini quirk) must still yield facts instead of dropping the whole chunk."""
    batch_id = "batch_recoverable_json"

    # Valid facts JSON, but wrapped in ```json fences AND containing a raw
    # control character (\x01) inside a string value. Bare json.loads fails on
    # both; parse_llm_json strips the fences and scrubs the control char.
    inner_json = json.dumps(
        {
            "facts": [
                {
                    "what": "Alice is a senior software engineer at TechCorp",
                    "when": "present",
                    "where": "TechCorp",
                    "who": "Alice",
                    "why": "Professional background\x01information",
                    "fact_type": "world",
                    "fact_kind": "conversation",
                }
            ]
        }
    )
    unparseable_content = f"```json\n{inner_json}\n```"

    # Sanity: the raw content is NOT parseable by the bare parser.
    with pytest.raises(json.JSONDecodeError):
        json.loads(unparseable_content)

    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": unparseable_content}}],
                        "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,
    )

    # The facts are recovered rather than lost.
    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_unparseable_json_still_records_error(mock_llm_config, test_contents, hindsight_config):
    """#2701: genuinely unparseable content (not recoverable by parse_llm_json)
    must preserve the existing behavior — record the error, fact_count=0, no crash."""
    batch_id = "batch_unparseable_json"

    # Not JSON at all, and not recoverable by fence-stripping or control-char scrubbing.
    unparseable_content = "this is not json {{{ ["

    with pytest.raises(json.JSONDecodeError):
        json.loads(unparseable_content)

    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": unparseable_content}}],
                        "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."""
    bank_id = f"test_crash_{datetime.now(timezone.utc).timestamp()}"
    operation_id = str(uuid.uuid4())  # Must be UUID for async_operations table

    try:
        # Ensure bank exists
        await memory.get_bank_profile(bank_id, request_context=request_context)

        # Setup: Store batch_id in async_operations table (simulates partial execution)
        batch_id = "batch_recovered_456"
        pool = memory._pool
        schema = request_context.tenant_id

        from hindsight_api.engine.task_backend import fq_table

        table = fq_table("async_operations", schema)

        # Create operation with batch_id already stored
        await pool.execute(
            f"""
            INSERT INTO {table} (operation_id, operation_type, bank_id, status, result_metadata)
            VALUES ($1, 'retain', $2, 'processing', $3::jsonb)
            """,
            operation_id,
            bank_id,
            json.dumps(
                {
                    "batch_id": batch_id,
                    "batch_provider": "openai",
                    "chunk_count": 2,
                }
            ),
        )

        # Mock batch API responses for resume scenario
        mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)

        # Mock get_batch_status - batch already in progress
        mock_llm_config._provider_impl.get_batch_status = AsyncMock(
            return_value={
                "status": "completed",
                "request_counts": {"total": 2, "completed": 2, "failed": 0},
            }
        )

        # Mock retrieve_batch_results
        mock_results = [
            {
                "custom_id": "chunk_0",
                "response": {
                    "body": {
                        "choices": [
                            {
                                "message": {
                                    "content": json.dumps(
                                        {
                                            "facts": [
                                                {
                                                    "what": "Alice is a senior software engineer",
                                                    "when": "present",
                                                    "where": "TechCorp",
                                                    "who": "Alice",
                                                    "why": "Background",
                                                    "fact_type": "world",
                                                    "fact_kind": "conversation",
                                                }
                                            ]
                                        }
                                    )
                                }
                            }
                        ],
                        "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
                    }
                },
            },
            {
                "custom_id": "chunk_1",
                "response": {
                    "body": {
                        "choices": [
                            {
                                "message": {
                                    "content": json.dumps(
                                        {
                                            "facts": [
                                                {
                                                    "what": "Bob is a junior developer",
                                                    "when": "last month",
                                                    "where": "team",
                                                    "who": "Bob",
                                                    "why": "New member",
                                                    "fact_type": "world",
                                                    "fact_kind": "conversation",
                                                }
                                            ]
                                        }
                                    )
                                }
                            }
                        ],
                        "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
                    }
                },
            },
        ]
        mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(return_value=mock_results)

        # Call batch API extraction with operation_id (crash recovery scenario)
        facts, chunks, usage = await extract_facts_from_contents_batch_api(
            contents=test_contents,
            llm_config=mock_llm_config,
            agent_name="test_agent",
            config=hindsight_config,
            pool=pool,
            operation_id=operation_id,  # Provides crash recovery context
            schema=schema,
        )

        # Verify results
        assert len(facts) == 2, "Should extract 2 facts after recovery"

        # CRITICAL: Verify submit_batch was NOT called (because batch_id already exists)
        mock_llm_config._provider_impl.submit_batch.assert_not_called()

        # Verify get_batch_status WAS called (polling resumed)
        mock_llm_config._provider_impl.get_batch_status.assert_called()

        # Verify retrieve_batch_results was called with the recovered batch_id
        mock_llm_config._provider_impl.retrieve_batch_results.assert_called_once_with(batch_id)

        logger.info("✅ Crash recovery test passed - resumed polling without re-submission")

    finally:
        # Cleanup
        try:
            await memory.delete_bank(bank_id, request_context=request_context)
        except Exception:
            pass


@pytest.mark.asyncio
async def test_batch_api_records_non_fatal_extraction_errors(
    mock_llm_config, test_contents, hindsight_config, memory, request_context
):
    """Batch API skipped chunks are surfaced in operation result_metadata."""
    bank_id = f"test_batch_errors_{datetime.now(timezone.utc).timestamp()}"
    operation_id = str(uuid.uuid4())

    try:
        await memory.get_bank_profile(bank_id, request_context=request_context)
        pool = memory._pool
        schema = request_context.tenant_id

        from hindsight_api.engine.task_backend import fq_table

        table = fq_table("async_operations", schema)
        await pool.execute(
            f"""
            INSERT INTO {table} (operation_id, operation_type, bank_id, status, result_metadata)
            VALUES ($1, 'retain', $2, 'processing', $3::jsonb)
            """,
            operation_id,
            bank_id,
            json.dumps({}),
        )

        batch_id = "batch_partial_errors"
        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": 2, "completed": 2, "failed": 0},
            }
        )
        mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(
            return_value=[
                {
                    "custom_id": "chunk_0",
                    "response": {
                        "body": {
                            "choices": [
                                {
                                    "message": {
                                        "content": json.dumps(
                                            {
                                                "facts": [
                                                    {
                                                        "what": "Alice is a senior software engineer",
                                                        "when": "present",
                                                        "where": "TechCorp",
                                                        "who": "Alice",
                                                        "why": "Background",
                                                        "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,
            llm_config=mock_llm_config,
            agent_name="test_agent",
            config=hindsight_config,
            pool=pool,
            operation_id=operation_id,
            schema=schema,
        )

        assert len(facts) == 1
        assert len(chunks) == 2
        assert chunks[1].fact_count == 0
        assert usage.total_tokens == 150

        row = await pool.fetchrow(f"SELECT result_metadata FROM {table} WHERE operation_id = $1", operation_id)
        metadata = (
            json.loads(row["result_metadata"]) if isinstance(row["result_metadata"], str) else row["result_metadata"]
        )
        assert metadata["batch_id"] == batch_id
        assert metadata["extraction_errors_count"] == 1
        assert metadata["extraction_errors_sample"] == ["chunk_1: missing batch result"]

    finally:
        try:
            await memory.delete_bank(bank_id, request_context=request_context)
        except Exception:
            pass


@pytest.mark.asyncio
async def test_batch_api_raises_for_unsupported_provider(mock_llm_config, test_contents, hindsight_config):
    """Batch extraction must surface a hard error (not silently fall back) when
    the configured provider doesn't support the batch API.

    The silent-fallback behavior was removed in #1463 because it created a
    mutual-recursion path between sync and batch extraction. Misconfiguration
    should fail loudly and be caught at startup; this test guards that
    contract.
    """
    mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=False)
    mock_llm_config.provider = "groq"

    with pytest.raises(RuntimeError, match="does not.*support the batch API"):
        await extract_facts_from_contents_batch_api(
            contents=test_contents,
            llm_config=mock_llm_config,
            agent_name="test_agent",
            config=hindsight_config,
            pool=None,
            operation_id=None,
            schema=None,
        )

    mock_llm_config._provider_impl.submit_batch.assert_not_called()


@pytest.mark.asyncio
async def test_worker_batch_recovery(memory, request_context):
    """Test that WorkerPoller._recover_batch_operations finds and resets orphaned batches."""
    bank_id = f"test_worker_recovery_{datetime.now(timezone.utc).timestamp()}"
    operation_id = str(uuid.uuid4())  # Must be UUID for async_operations table

    try:
        # Ensure bank exists
        await memory.get_bank_profile(bank_id, request_context=request_context)

        pool = memory._pool
        schema = request_context.tenant_id

        from hindsight_api.engine.task_backend import fq_table

        table = fq_table("async_operations", schema)

        # Create orphaned batch operation (simulates worker crash during polling)
        batch_id = "batch_orphaned_999"
        task_payload = {
            "operation_type": "retain",
            "bank_id": bank_id,
            "contents": [{"content": "test", "event_date": "2024-01-15T00:00:00Z"}],
        }

        await pool.execute(
            f"""
            INSERT INTO {table} (operation_id, operation_type, bank_id, status, worker_id, result_metadata, task_payload)
            VALUES ($1, 'retain', $2, 'processing', 'worker_crashed', $3::jsonb, $4::jsonb)
            """,
            operation_id,
            bank_id,
            json.dumps(
                {
                    "batch_id": batch_id,
                    "batch_provider": "openai",
                    "chunk_count": 1,
                }
            ),
            json.dumps(task_payload),
        )

        # Create WorkerPoller
        from hindsight_api.extensions.builtin.tenant import DefaultTenantExtension

        tenant_extension = DefaultTenantExtension(config={"schema": schema} if schema else {})

        poller = WorkerPoller(
            backend=pool,
            worker_id="test_worker_recovery",
            executor=memory,
            poll_interval_ms=100,
            schema=schema,
            tenant_extension=tenant_extension,
            max_slots=5,
            slot_reservations={"consolidation": 2},
        )

        # Run recovery
        recovered_count = await poller._recover_batch_operations(schema)

        # Verify recovery
        assert recovered_count == 1, "Should recover 1 batch operation"

        # Verify operation was reset to pending
        row = await pool.fetchrow(
            f"SELECT status, worker_id FROM {table} WHERE operation_id = $1",
            operation_id,
        )

        assert row["status"] == "pending", "Operation should be reset to pending"
        assert row["worker_id"] is None, "Worker ID should be cleared"

        logger.info("✅ Worker batch recovery test passed")

    finally:
        # Cleanup
        try:
            await memory.delete_bank(bank_id, request_context=request_context)
        except Exception:
            pass


@pytest.mark.asyncio
async def test_batch_api_via_extract_facts_from_contents(
    mock_llm_config, test_contents, hindsight_config, memory, request_context
):
    """Test that extract_facts_from_contents routes to batch API when enabled."""
    bank_id = f"test_routing_{datetime.now(timezone.utc).timestamp()}"

    try:
        # Enable batch API in config
        hindsight_config.retain_batch_enabled = True

        # Mock batch API support
        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_123", "status": "validating", "request_counts": {}}
        )
        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({"facts": []})}}],
                            "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
                        }
                    },
                }
            ]
        )

        # Call main extract_facts_from_contents (should route to batch API)
        facts, chunks, usage = await extract_facts_from_contents(
            contents=test_contents,
            llm_config=mock_llm_config,
            agent_name="test_agent",
            config=hindsight_config,
            pool=None,
            operation_id=None,
            schema=None,
        )

        # Verify batch API was called
        mock_llm_config._provider_impl.submit_batch.assert_called_once()

        logger.info("✅ Routing to batch API test passed")

    finally:
        # Cleanup
        try:
            await memory.delete_bank(bank_id, request_context=request_context)
        except Exception:
            pass