leocrm/tests/test_vector_performance.py

"""
Performance tests for pgvector HNSW/IVFFlat indexes.

B-VEC-TEST: Query latency benchmarks at 10k, 100k, 1M embeddings.
B-VEC-BATCH: Verify batch embedding reduces API calls.

These tests require a running PostgreSQL with pgvector extension.
They are skipped automatically when no TEST_DATABASE_URL is available.

Run manually:
    TEST_DATABASE_URL=postgresql+asyncpg://user:pass@localhost:5432/testdb \
        python -m pytest tests/test_vector_performance.py -v --tb=short
"""

from __future__ import annotations

import os
import time
import uuid
from typing import Any

import pytest

pytestmark = pytest.mark.skipif(
    not os.environ.get("TEST_DATABASE_URL"),
    reason="TEST_DATABASE_URL not set - pgvector performance tests require a real PostgreSQL with pgvector",
)


async def _ensure_pgvector(conn: Any) -> None:
    await conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
    await conn.execute("""
        CREATE TABLE IF NOT EXISTS vec_perf_test (
            id UUID PRIMARY KEY,
            tenant_id UUID NOT NULL,
            embedding vector(1536) NOT NULL,
            label TEXT DEFAULT ''
        )
    """)


async def _insert_batch(conn: Any, count: int, batch_size: int = 500) -> None:
    import random
    for offset in range(0, count, batch_size):
        n = min(batch_size, count - offset)
        rows = []
        for _ in range(n):
            vec = [random.uniform(-1, 1) for _ in range(1536)]
            rows.append((str(uuid.uuid4()), str(uuid.uuid4()), str(vec)))
        placeholders = ",".join(
            f"(${i*3+1}, ${i*3+2}, ${i*3+3}::vector)" for i in range(n)
        )
        params: list[str] = []
        for r in rows:
            params.extend([r[0], r[1], r[2]])
        await conn.execute(
            f"INSERT INTO vec_perf_test (id, tenant_id, embedding) VALUES {placeholders}",
            *params,
        )


async def _create_hnsw_index(conn: Any) -> None:
    await conn.execute("""
        CREATE INDEX IF NOT EXISTS idx_vec_perf_hnsw
        ON vec_perf_test USING hnsw (embedding vector_cosine_ops)
        WITH (m = 16, ef_construction = 128)
    """)


async def _create_ivfflat_index(conn: Any, lists: int = 100) -> None:
    await conn.execute(f"""
        CREATE INDEX IF NOT EXISTS idx_vec_perf_ivf
        ON vec_perf_test USING ivfflat (embedding vector_cosine_ops)
        WITH (lists = {lists})
    """)


def _stats(latencies: list[float]) -> dict[str, float]:
    if not latencies:
        return {"p50": 0.0, "p95": 0.0, "p99": 0.0, "avg": 0.0}
    s = sorted(latencies)
    n = len(s)
    return {
        "p50": s[n // 2],
        "p95": s[int(n * 0.95)],
        "p99": s[int(n * 0.99)],
        "avg": sum(s) / n,
    }


async def _measure_hnsw_latency(conn: Any, num_queries: int = 100, k: int = 10, ef_search: int = 40) -> list[float]:
    import random
    await conn.execute(f"SET LOCAL hnsw.ef_search = {ef_search}")
    latencies: list[float] = []
    for _ in range(num_queries):
        query_vec = str([random.uniform(-1, 1) for _ in range(1536)])
        start = time.perf_counter()
        await conn.fetch(
            "SELECT id FROM vec_perf_test ORDER BY embedding <=> $1::vector LIMIT $2",
            query_vec, k,
        )
        latencies.append((time.perf_counter() - start) * 1000)
    return latencies


@pytest.mark.asyncio
@pytest.mark.parametrize("scale", [10_000, 100_000])
async def test_hnsw_query_latency(scale: int) -> None:
    from sqlalchemy.ext.asyncio import create_async_engine
    engine = create_async_engine(os.environ["TEST_DATABASE_URL"])
    try:
        async with engine.begin() as conn:
            await _ensure_pgvector(conn)
            await conn.execute("TRUNCATE vec_perf_test")
            await _insert_batch(conn, scale)
            await _create_hnsw_index(conn)
            await conn.execute("ANALYZE vec_perf_test")
        async with engine.connect() as conn:
            latencies = await _measure_hnsw_latency(conn, num_queries=100, k=10)
            stats = _stats(latencies)
            print(f"\nHNSW @ {scale} embeddings: {stats}")
            if scale == 10_000:
                assert stats["p95"] < 100, f"p95 too high: {stats['p95']:.1f}ms"
            elif scale == 100_000:
                assert stats["p95"] < 500, f"p95 too high: {stats['p95']:.1f}ms"
    finally:
        await engine.dispose()


@pytest.mark.asyncio
@pytest.mark.parametrize("scale", [10_000, 100_000])
async def test_ivfflat_query_latency(scale: int) -> None:
    from sqlalchemy.ext.asyncio import create_async_engine
    engine = create_async_engine(os.environ["TEST_DATABASE_URL"])
    try:
        async with engine.begin() as conn:
            await _ensure_pgvector(conn)
            await conn.execute("TRUNCATE vec_perf_test")
            await _insert_batch(conn, scale)
            lists = max(10, int(scale ** 0.5))
            await _create_ivfflat_index(conn, lists=lists)
            await conn.execute("ANALYZE vec_perf_test")
        async with engine.connect() as conn:
            await conn.execute("SET LOCAL ivfflat.probes = 10")
            import random
            latencies: list[float] = []
            for _ in range(100):
                query_vec = str([random.uniform(-1, 1) for _ in range(1536)])
                start = time.perf_counter()
                await conn.fetch(
                    "SELECT id FROM vec_perf_test ORDER BY embedding <=> $1::vector LIMIT 10",
                    query_vec,
                )
                latencies.append((time.perf_counter() - start) * 1000)
            stats = _stats(latencies)
            print(f"\nIVFFlat @ {scale} embeddings: {stats}")
            if scale == 10_000:
                assert stats["p95"] < 150, f"p95 too high: {stats['p95']:.1f}ms"
            elif scale == 100_000:
                assert stats["p95"] < 800, f"p95 too high: {stats['p95']:.1f}ms"
    finally:
        await engine.dispose()


@pytest.mark.asyncio
async def test_ef_search_tradeoff() -> None:
    from sqlalchemy.ext.asyncio import create_async_engine
    engine = create_async_engine(os.environ["TEST_DATABASE_URL"])
    try:
        async with engine.begin() as conn:
            await _ensure_pgvector(conn)
            await conn.execute("TRUNCATE vec_perf_test")
            await _insert_batch(conn, 10_000)
            await _create_hnsw_index(conn)
            await conn.execute("ANALYZE vec_perf_test")
        async with engine.connect() as conn:
            results: dict[int, dict[str, float]] = {}
            for ef in [10, 20, 40, 80, 120]:
                latencies = await _measure_hnsw_latency(conn, num_queries=50, k=10, ef_search=ef)
                results[ef] = _stats(latencies)
                print(f"  ef_search={ef}: p50={results[ef]['p50']:.1f}ms p95={results[ef]['p95']:.1f}ms")
            assert results[10]["p50"] <= results[120]["p50"] + 20, "ef_search=10 should be faster than ef_search=120"
    finally:
        await engine.dispose()


@pytest.mark.asyncio
async def test_batch_embedding_single_api_call() -> None:
    """B-VEC-BATCH: Verify batch embedding makes a single API call."""
    from unittest.mock import AsyncMock, MagicMock, patch
    with patch("app.ai.llm_client.litellm.aembedding", new_callable=AsyncMock) as mock_embed:
        mock_embed.return_value = MagicMock(
            data=[{"embedding": [0.1] * 1536} for _ in range(10)]
        )
        from app.ai.llm_client import llm_embed
        texts = [f"test text {i}" for i in range(10)]
        result = await llm_embed(texts, api_key="test-key")
        assert len(result) == 10
        assert mock_embed.call_count == 1, f"Expected 1 aembedding call for batch, got {mock_embed.call_count}"
        call_kwargs = mock_embed.call_args.kwargs
        assert isinstance(call_kwargs["input"], list)
        assert len(call_kwargs["input"]) == 10