memory-os/scripts/test_ingestion.py

#!/usr/bin/env python3
"""End-to-end ingestion test for Memory OS.

Verifies the full pipeline: enqueue → ARQ worker → embedding → Qdrant upsert.

Usage:
    python3 scripts/test_ingestion.py

Environment:
    REDIS_PASSWORD        Redis password (required)
    QDRANT_API_KEY        Qdrant API key (default: "")
    REDIS_HOST            Redis host (default: localhost)
    REDIS_PORT            Redis port (default: 6379)
    QDRANT_HOST           Qdrant host (default: localhost)
    QDRANT_PORT           Qdrant port (default: 6333)
    COLLECTION_NAME       Qdrant collection (default: knowledge_base)

Returns exit code 0 on success, 1 on failure.
"""

import asyncio
import hashlib
import json
import os
import sys
import tempfile
import time
import uuid


# ── Config from env ──────────────────────────────────────────────────────────
REDIS_HOST = os.environ.get("REDIS_HOST", "localhost")
REDIS_PORT = int(os.environ.get("REDIS_PORT", "6379"))
REDIS_PASSWORD = os.environ.get("REDIS_PASSWORD", "")
QDRANT_HOST = os.environ.get("QDRANT_HOST", "localhost")
QDRANT_PORT = int(os.environ.get("QDRANT_PORT", "6333"))
QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY", "")
COLLECTION_NAME = os.environ.get("COLLECTION_NAME", "knowledge_base")
TIMEOUT = 60  # seconds for ARQ job completion
POLL_INTERVAL = 2  # seconds between polls


def fail(msg: str) -> None:
    print(f"❌ {msg}")
    sys.exit(1)


def ok(msg: str) -> None:
    print(f"✅ {msg}")


# ── Test document ────────────────────────────────────────────────────────────
TEST_ID = uuid.uuid4().hex[:8]
TEST_TEXT = (
    f"Memory OS ingestion test document {TEST_ID}. "
    "This file was automatically generated by test_ingestion.py "
    "to verify the end-to-end pipeline: enqueue, worker, embedding, Qdrant upsert."
)
TEST_PATH = f"/wiki/raw/test/ingestion-test-{TEST_ID}.md"


async def main() -> None:
    print(f"=== Memory OS Ingestion Test (doc: {TEST_ID}) ===")
    print()

    # ── 1. Create temp file for the worker to ingest ──────────────────────────
    print("1. Creating test document...")
    # The worker reads from the wiki volume mounted in Docker.
    # Write to /tmp as a fallback; the worker path depends on volume config.
    tmpdir = tempfile.mkdtemp(prefix="memoryos-test-")
    test_file = os.path.join(tmpdir, f"ingestion-test-{TEST_ID}.md")
    with open(test_file, "w") as f:
        f.write(f"# Test {TEST_ID}\n\n{TEST_TEXT}\n")
    ok(f"Created {test_file}")

    # ── 2. Enqueue ARQ job ───────────────────────────────────────────────────
    print("2. Enqueuing ARQ job...")
    try:
        from arq import create_pool
        from arq.connections import RedisSettings
    except ImportError:
        fail("arq not installed — run: pip install arq")

    redis_settings = RedisSettings(
        host=REDIS_HOST,
        port=REDIS_PORT,
        password=REDIS_PASSWORD or None,
    )

    redis = await create_pool(redis_settings)
    job = await redis.enqueue_job("process_wiki_file", TEST_PATH)
    job_id = job.job_id
    ok(f"Job {job_id} enqueued")

    # ── 3. Wait for completion ───────────────────────────────────────────────
    print(f"3. Waiting for worker (timeout={TIMEOUT}s)...")
    deadline = time.monotonic() + TIMEOUT
    result = None
    while time.monotonic() < deadline:
        job_info = await redis.get_job_result(job_id)
        if job_info is not None:
            result = job_info.result
            if job_info.success:
                ok(f"Job completed in {TIMEOUT - (deadline - time.monotonic()):.0f}s")
                break
            else:
                fail(f"Job failed: {job_info.result}")
        await asyncio.sleep(POLL_INTERVAL)

    if result is None:
        fail(f"Job timed out after {TIMEOUT}s — is the ARQ worker running?")

    # ── 4. Search Qdrant for the ingested point ──────────────────────────────
    print("4. Searching Qdrant for ingested point...")
    try:
        from qdrant_client import AsyncQdrantClient
        from qdrant_client.models import Filter, FieldCondition, MatchText
    except ImportError:
        fail("qdrant-client not installed — run: pip install qdrant-client")

    client = AsyncQdrantClient(
        host=QDRANT_HOST,
        port=QDRANT_PORT,
        api_key=QDRANT_API_KEY or None,
        https=False,
    )

    # Search by the unique test ID in the payload
    points, _ = await client.scroll(
        collection_name=COLLECTION_NAME,
        scroll_filter=Filter(
            must=[
                FieldCondition(
                    key="source_file",
                    match=MatchText(text=TEST_PATH),
                )
            ]
        ),
        with_vectors=True,
        limit=5,
    )

    if not points:
        fail(
            f"No points found for {TEST_PATH} in collection "
            f"'{COLLECTION_NAME}' — ingestion may have succeeded "
            "but the source_file path may differ from expected."
        )

    point = points[0]
    point_id = point.id
    ok(f"Found point {point_id}")

    # ── 5. Verify dense vector dimensions ────────────────────────────────────
    print("5. Verifying embedding dimensions...")
    if "dense" not in point.vector:
        fail("Point has no 'dense' vector — named vectors may not be configured")

    dims = len(point.vector["dense"])
    if dims != 4096:
        fail(f"Expected 4096 dimensions, got {dims}")

    ok(f"{dims} dimensions ✓")

    # ── 6. Cleanup ───────────────────────────────────────────────────────────
    print("6. Cleaning up test point...")
    await client.delete(
        collection_name=COLLECTION_NAME,
        points_selector=[point_id],
    )
    ok(f"Deleted point {point_id}")

    # Clean up local temp file
    os.remove(test_file)
    os.rmdir(tmpdir)

    await client.close()
    await redis.close()

    # ── 7. Summary ───────────────────────────────────────────────────────────
    print()
    print("═" * 60)
    ok("All checks passed — ingestion pipeline is operational")


if __name__ == "__main__":
    asyncio.run(main())