second-brain/tests/test_smoke_embedding.py

"""Smoke test for the Postgres + embedding pipeline.

Runs only when SECOND_BRAIN_DATABASE_URL (or HERBYLAB_DATABASE_URL) is set
*and* an Ollama instance with nomic-embed-text is reachable. Otherwise the
test is skipped — there is no in-memory fixture for pgvector, and we
explicitly want the smoke check to exercise the real wire format.

It covers:
  1. SQLAlchemy can write a Source + Extraction into the `second_brain`
     schema.
  2. The embeddings module chunks the summary, calls Ollama, and writes
     vector rows into public.embeddings keyed correctly.
  3. Re-running embedding is idempotent (delete-before-insert wins).
"""

from __future__ import annotations

import os
import socket
import uuid

import httpx
import psycopg
import pytest


def _ollama_reachable() -> bool:
    url = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
    try:
        r = httpx.get(f"{url}/api/tags", timeout=2.0)
        if r.status_code != 200:
            return False
        models = [m["name"] for m in r.json().get("models", [])]
        return any(m.startswith("nomic-embed-text") for m in models)
    except (httpx.HTTPError, socket.error):
        return False


def _db_reachable() -> bool:
    url = os.environ.get("SECOND_BRAIN_DATABASE_URL") or os.environ.get(
        "HERBYLAB_DATABASE_URL"
    )
    if not url:
        return False
    raw = url.replace("postgresql+psycopg://", "postgresql://", 1)
    try:
        with psycopg.connect(raw, connect_timeout=2) as conn:  # noqa: F841
            return True
    except Exception:
        return False


pytestmark = pytest.mark.skipif(
    not (_db_reachable() and _ollama_reachable()),
    reason="Smoke test needs SECOND_BRAIN_DATABASE_URL + reachable Ollama (nomic-embed-text)",
)


def test_extraction_round_trip_with_embedding():
    from second_brain.database import Database, reset_database_singleton
    from second_brain.embeddings import close_pool, embed_extraction
    from second_brain.models import (
        Extraction,
        Source,
        SourceStatus,
        SourceType,
        utcnow,
    )

    reset_database_singleton()
    db = Database()

    # Unique URL so re-runs don't collide on the unique constraint.
    unique = uuid.uuid4().hex[:8]
    url = f"https://example.test/smoke/{unique}"

    summary_text = (
        "Smoke test extraction summary. "
        "The second-brain pipeline embeds this string into Postgres. "
        "Re-running the embed should be idempotent — delete-before-insert "
        "keys on (source_schema, source_table, source_id, embedding_model)."
    )

    with db.session() as sess:
        src = Source(
            url=url,
            title=f"smoke-test-{unique}",
            domain="development",
            source_type=SourceType.ARTICLE,
            status=SourceStatus.ANALYZED,
            ingested_at=utcnow(),
        )
        sess.add(src)
        sess.flush()
        ext = Extraction(source_id=src.id, summary=summary_text)
        sess.add(ext)
        sess.flush()
        extraction_id = ext.id
        source_id = src.id

    # Read the configured model so the test doesn't break if the default
    # is changed in settings.toml or via EMBEDDING_MODEL.
    from second_brain.config import load_config

    embedding_model = os.environ.get(
        "EMBEDDING_MODEL",
        load_config().embeddings.get("model", "nomic-embed-text"),
    )

    try:
        n1 = embed_extraction(extraction_id, summary_text)
        assert n1 is not None and n1 >= 1, f"first embed should write rows, got {n1}"

        # Idempotency: re-embed should produce the same chunk count and not
        # accumulate orphan rows.
        n2 = embed_extraction(extraction_id, summary_text)
        assert n2 == n1, f"re-embed should match first count: {n1} vs {n2}"

        raw_url = (
            os.environ.get("SECOND_BRAIN_DATABASE_URL")
            or os.environ.get("HERBYLAB_DATABASE_URL")
        ).replace("postgresql+psycopg://", "postgresql://", 1)
        with psycopg.connect(raw_url) as conn, conn.cursor() as cur:
            cur.execute(
                """
                SELECT COUNT(*)
                  FROM public.embeddings
                 WHERE source_schema = 'second_brain'
                   AND source_table = 'extractions'
                   AND source_id = %s
                   AND embedding_model = %s
                """,
                (extraction_id, embedding_model),
            )
            (count,) = cur.fetchone()
            assert count == n1, f"public.embeddings should hold {n1} rows, has {count}"
    finally:
        # Tidy up so the smoke test stays leak-free.
        with db.session() as sess:
            sess.query(Source).filter(Source.id == source_id).delete()
        raw_url = (
            os.environ.get("SECOND_BRAIN_DATABASE_URL")
            or os.environ.get("HERBYLAB_DATABASE_URL")
        ).replace("postgresql+psycopg://", "postgresql://", 1)
        with psycopg.connect(raw_url) as conn, conn.cursor() as cur:
            cur.execute(
                """
                DELETE FROM public.embeddings
                 WHERE source_schema = 'second_brain'
                   AND source_table = 'extractions'
                   AND source_id = %s
                """,
                (extraction_id,),
            )
            conn.commit()
        close_pool()
        reset_database_singleton()