fastmcp-feedback/tests/unit/test_embeddings.py

"""Tests for EmbeddingSink, OpenAIEmbedder and similarity search.

DB tests take the ``db_url`` fixture, so they run on SQLite and, with
FFB_TEST_DATABASE_URL set, on Postgres with pgvector.
"""

import asyncio
import contextlib
import hashlib
import json
import logging
import math
import os
import re
import time
import uuid
from datetime import UTC, datetime, timedelta

import pytest
from fastmcp import Client, FastMCP
from fastmcp.exceptions import ToolError
from sqlalchemy import func, select, text, update

from fastmcp_feedback import add_feedback_tools
from fastmcp_feedback.instrumentation import (
    REDACTED,
    DatabaseSink,
    EmbeddingSink,
    EventRecord,
    OpenAIEmbedder,
    Redactor,
    ToolCallRecord,
    build_metadata,
    extract_texts,
    instrument,
)
from fastmcp_feedback.instrumentation.embeddings import cosine_score

DIM = 256
SECRET = "sk-abcdefghijklmnopqrstuvwx"


def token_vector(text_: str, dim: int = DIM) -> list[float]:
    """Bag of hashed words, normalized: texts sharing words are close."""
    v = [0.0] * dim
    for tok in re.findall(r"[a-z0-9]+", text_.lower()):
        h = int.from_bytes(hashlib.sha256(tok.encode()).digest()[:4], "big")
        v[h % dim] += 1.0
    n = math.sqrt(sum(x * x for x in v)) or 1.0
    return [x / n for x in v]


class FakeEmbedder:
    def __init__(
        self,
        dim=DIM,
        model="fake-embed",
        *,
        fail=False,
        delay=0.0,
        gate=None,
        out_dim=None,
    ):
        self.dim = dim
        self.model = model
        self.fail = fail
        self.delay = delay
        self.gate = gate
        self.out_dim = out_dim or dim
        self.calls: list[list[str]] = []

    @property
    def seen(self) -> list[str]:
        return [t for batch in self.calls for t in batch]

    async def embed(self, texts):
        self.calls.append(list(texts))
        if self.gate is not None:
            await self.gate.wait()
        if self.delay:
            await asyncio.sleep(self.delay)
        if self.fail:
            raise RuntimeError("embedding backend down")
        return [token_vector(t, self.out_dim) for t in texts]


@pytest.fixture
async def db(db_url):
    made = []

    def factory(name="calls", **kwargs):
        kwargs.setdefault("create_tables", True)
        kwargs.setdefault("embedding_dim", DIM)
        sink = DatabaseSink(db_url(name), **kwargs)
        made.append(sink)
        return sink

    yield factory
    for sink in made:
        await sink.engine.dispose()


def call_record(tool="render", outcome="error", message="render timed out", **kw):
    return ToolCallRecord(
        tool=tool,
        started_at=datetime.now(UTC),
        duration_ms=1.0,
        ok=outcome == "ok",
        outcome=outcome,
        mode="meta",
        error_type=None if outcome == "ok" else kw.pop("error_type", "TimeoutError"),
        error_message=None if outcome == "ok" else message,
        **kw,
    )


def event(kind, key=None, **attrs):
    return EventRecord(kind=kind, occurred_at=datetime.now(UTC), key=key, attrs=attrs)


async def rows(sink, *where):
    t = sink.embeddings
    if not await sink._has_table(t.name):  # created on the first embedding
        return []
    async with sink.engine.connect() as conn:
        result = await conn.execute(select(t).where(*where).order_by(t.c.source_id))
        return [dict(r) for r in result.mappings()]


async def settle(mw, emb):
    await mw.flush()
    await emb.flush()


class TestExtractTexts:
    def test_call_errors(self):
        ok = call_record(outcome="ok")
        err = call_record(message="render timed out", error_type="TimeoutError")
        soft = call_record(
            tool="find",
            outcome="soft_error",
            message="status=unknown_target",
            error_type="SoftError",
        )
        items = extract_texts([ok, err, soft])
        assert [(i.source_type, i.source_id, i.text) for i in items] == [
            ("call_error", err.id, "render: TimeoutError: render timed out"),
            ("call_error", soft.id, "find: SoftError: status=unknown_target"),
        ]

    def test_feedback(self):
        ev = event(
            "feedback.submitted",
            key="bug-7Q2X",
            title="Render hangs",
            description="the render tool hung on scene a",
            type="bug",
        )
        [item] = extract_texts([ev])
        assert (item.source_type, item.source_id, item.text) == (
            "feedback",
            "bug-7Q2X",
            "Render hangs\nthe render tool hung on scene a",
        )

    def test_event_text_keys(self):
        ev = event(
            "job.failed",
            key="job-1",
            error_excerpt="solver diverged",
            reason="residual NaN",
            nproc=8,
        )
        quiet = event("job.finished", key="job-2", duration_s=3)
        [item] = extract_texts([ev, quiet])
        assert (item.source_type, item.source_id, item.text) == (
            "event",
            ev.id,
            "solver diverged\nresidual NaN",
        )
        [custom] = extract_texts([quiet], event_text_keys=("duration_s",))
        assert custom.text == "3"

    def test_llm(self):
        ev = event("llm.call", model="m", prompt="why red?", completion="the light")
        [item] = extract_texts([ev])
        assert (item.source_type, item.source_id, item.text) == (
            "llm",
            ev.id,
            "prompt: why red?\ncompletion: the light",
        )

    def test_llm_without_text_is_a_plain_event(self):
        ev = event("llm.call", model="m", ok=False, error="upstream 529 overloaded")
        [item] = extract_texts([ev])
        assert (item.source_type, item.text) == ("event", "upstream 529 overloaded")

    def test_sources_filter(self):
        records = [
            call_record(),
            event("feedback.submitted", key="1", title="t", description="d"),
            event("job.failed", error="boom"),
            event("llm.call", prompt="p"),
        ]
        assert {i.source_type for i in extract_texts(records)} == {
            "call_error",
            "feedback",
            "event",
            "llm",
        }
        assert [
            i.source_type for i in extract_texts(records, sources=["feedback"])
        ] == ["feedback"]


class TestEmbeddingSink:
    async def test_every_source_lands(self, db):
        sink = db()
        fake = FakeEmbedder()
        emb = EmbeddingSink(fake, sink)
        err = call_record()
        records = [
            err,
            event(
                "feedback.submitted",
                key="42",
                title="Render hangs",
                description="scene a",
            ),
            event("job.failed", key="job-1", error_excerpt="solver diverged"),
            event("llm.call", prompt="hello", completion="hi"),
        ]
        await emb.write(records)
        await emb.flush()
        stored = await rows(sink)
        assert {r["source_type"] for r in stored} == {
            "call_error",
            "feedback",
            "event",
            "llm",
        }
        for r in stored:
            assert r["dim"] == DIM and r["model"] == "fake-embed"
            assert len(r["embedding"]) == DIM
            assert r["text_hash"] == hashlib.sha256(r["text"].encode()).hexdigest()
            assert r["created_at"] is not None
        assert emb.embedded == 4 and emb.errors == 0
        await emb.aclose()

    async def test_submit_feedback_is_embedded(self, db, tmp_path):
        app = FastMCP("app")
        sink = db("fb")
        fake = FakeEmbedder()
        emb = EmbeddingSink(fake, sink)
        mw = instrument(app, [sink, emb])
        add_feedback_tools(
            app, database_url=f"sqlite:///{tmp_path}/fb.db", instrumentation=mw
        )
        async with Client(app) as client:
            result = await client.call_tool(
                "submit_feedback",
                {
                    "request": {
                        "type": "bug",
                        "title": "Render hangs",
                        "description": "the render tool hung on scene a",
                        "submitter": "model",
                    }
                },
            )
        ref = result.data["feedback_id"]
        await settle(mw, emb)
        [row] = await rows(sink, sink.embeddings.c.source_type == "feedback")
        assert row["source_id"] == ref
        assert row["text"] == "Render hangs\nthe render tool hung on scene a"
        [ev] = await mw.events_for(ref, kinds=["feedback.submitted"])
        assert ev["attrs"]["type"] == "bug" and ev["attrs"]["submitter"] == "model"
        await mw.aclose()

    async def test_redaction(self, db):
        sink = db()
        fake = FakeEmbedder()
        emb = EmbeddingSink(
            fake, sink, redactor=Redactor(extra_patterns=[r"acme_\w{8}"])
        )
        # Straight to the sink, so upstream redaction cannot hide a gap here.
        await emb.write(
            [
                call_record(message=f"401 with api_key={SECRET} and acme_12345678"),
                event("job.failed", error=f"Bearer {SECRET}"),
            ]
        )
        await emb.flush()
        assert fake.seen and all(
            SECRET not in t and "acme_12345678" not in t for t in fake.seen
        )
        stored = await rows(sink)
        assert len(stored) == 2
        for r in stored:
            assert SECRET not in r["text"] and REDACTED in r["text"]
        await emb.aclose()

    async def test_instrument_hands_over_its_redactor(self, db):
        app = FastMCP("app")
        sink = db()
        emb = EmbeddingSink(FakeEmbedder(), sink)
        redactor = Redactor(extra_patterns=[r"acme_\w{8}"])
        mw = instrument(app, [sink, emb], redactor=redactor)
        assert emb.redactor is redactor
        await mw.aclose()

    async def test_max_chars(self, db):
        sink = db()
        fake = FakeEmbedder()
        emb = EmbeddingSink(fake, sink, max_chars=50)
        await emb.write([event("job.failed", error="word " * 100)])
        await emb.flush()
        [row] = await rows(sink)
        assert len(row["text"]) == 50 and fake.seen == [row["text"]]
        await emb.aclose()

    async def test_dedupe_and_replace(self, db):
        sink = db()
        fake = FakeEmbedder()
        emb = EmbeddingSink(fake, sink)
        ev = event("feedback.submitted", key="7", title="Render hangs", description="a")
        await emb.write([ev])
        await emb.flush()
        await emb.write([ev, ev])
        await emb.flush()
        assert fake.seen == ["Render hangs\na"]
        assert emb.skipped == 2
        changed = event(
            "feedback.submitted", key="7", title="Render hangs", description="b"
        )
        await emb.write([changed])
        await emb.flush()
        [row] = await rows(sink)
        assert row["text"] == "Render hangs\nb"
        assert fake.seen == ["Render hangs\na", "Render hangs\nb"]
        # Another model's vectors are kept apart.
        other = EmbeddingSink(FakeEmbedder(model="other"), sink)
        await other.write([changed])
        await other.flush()
        assert len(await rows(sink)) == 2
        await emb.aclose()
        await other.aclose()

    async def test_queue_full_drops(self, db):
        sink = db()
        gate = asyncio.Event()
        fake = FakeEmbedder(gate=gate)
        emb = EmbeddingSink(fake, sink, max_queue=2, batch_size=1)
        start = time.perf_counter()
        await emb.write([event("job.failed", error=f"boom {i}") for i in range(6)])
        assert time.perf_counter() - start < 0.5
        await asyncio.sleep(0.05)  # the worker takes one and blocks on the gate
        await emb.write([event("job.failed", error=f"late {i}") for i in range(3)])
        assert emb.dropped >= 4
        gate.set()
        await emb.flush()
        assert emb.errors == 0
        assert len(await rows(sink)) == 9 - emb.dropped
        await emb.aclose()

    async def test_embedder_errors_never_reach_tool_calls(self, db):
        app = FastMCP("app")

        @app.tool
        def render(scene: str) -> str:
            raise RuntimeError(f"render of {scene} timed out")

        @app.tool
        def ping() -> str:
            return "pong"

        sink = db()
        emb = EmbeddingSink(FakeEmbedder(fail=True), sink)
        mw = instrument(app, [sink, emb])
        async with Client(app) as client:
            with contextlib.suppress(ToolError):
                await client.call_tool("render", {"scene": "a"})
            assert (await client.call_tool("ping", {})).data == "pong"
        await settle(mw, emb)
        assert emb.errors == 1
        assert mw.dispatcher.sink_errors == 0
        assert await rows(sink) == []
        async with sink.engine.connect() as conn:
            calls = (
                await conn.execute(select(func.count()).select_from(sink.table))
            ).scalar_one()
        assert calls == 2
        assert await mw.similar("render timed out") == []
        await mw.aclose()

    async def test_slow_embedder_stays_off_the_call_path(self, db):
        app = FastMCP("app")

        @app.tool
        def render() -> str:
            raise RuntimeError("render timed out")

        sink = db()
        emb = EmbeddingSink(FakeEmbedder(delay=2.0), sink, timeout=5.0)
        mw = instrument(app, [sink, emb], sink_timeout=0.5)
        async with Client(app) as client:
            start = time.perf_counter()
            for _ in range(3):
                with contextlib.suppress(ToolError):
                    await client.call_tool("render", {})
            elapsed = time.perf_counter() - start
        await mw.flush()
        assert elapsed < 1.5
        assert mw.dispatcher.sink_errors == 0
        await emb.flush()
        assert emb.embedded == 3
        await mw.aclose()

    async def test_wrong_vector_length_is_an_error(self, db):
        sink = db()
        emb = EmbeddingSink(FakeEmbedder(out_dim=DIM + 1), sink)
        await emb.write([event("job.failed", error="boom")])
        await emb.flush()
        assert emb.errors == 1 and await rows(sink) == []
        await emb.aclose()

    async def test_dim_mismatch_refused_at_construction(self, db):
        with pytest.raises(ValueError, match="embedding_dim=512"):
            EmbeddingSink(FakeEmbedder(dim=512), db())

    async def test_unknown_source_refused(self, db):
        with pytest.raises(ValueError, match="unknown sources"):
            EmbeddingSink(FakeEmbedder(), db(), sources=["feedback", "chat"])

    async def test_aclose_drains(self, db):
        sink = db()
        emb = EmbeddingSink(FakeEmbedder(delay=0.1), sink, batch_size=1)
        await emb.write([event("job.failed", error=f"e{i}") for i in range(3)])
        await emb.aclose()
        assert emb.embedded == 3 and emb.dropped == 0

    async def test_aclose_timeout_drops(self, db):
        sink = db()
        emb = EmbeddingSink(
            FakeEmbedder(gate=asyncio.Event()), sink, batch_size=1, timeout=0.2
        )
        await emb.write([event("job.failed", error=f"e{i}") for i in range(3)])
        await emb.aclose()
        assert emb.dropped == 2 and emb.embedded == 0


class TestLlmText:
    def app(self, sinks, **kwargs):
        app = FastMCP("llm")
        holder = {}

        @app.tool
        async def ask(q: str) -> str:
            holder["mw"].record_llm_call(
                model="m",
                prompt=f"{q} token={SECRET}",
                completion="because",
                attrs={"prompt": "smuggled", "temperature": 0.2},
            )
            return "because"

        holder["mw"] = mw = instrument(app, sinks, **kwargs)
        return app, mw

    async def test_dropped_by_default(self, db):
        sink = db()
        fake = FakeEmbedder()
        emb = EmbeddingSink(fake, sink)
        app, mw = self.app([sink, emb])
        async with Client(app) as client:
            await client.call_tool("ask", {"q": "why is the sky blue"})
        await settle(mw, emb)
        [ev] = [e for e in mw.recent_events if e.kind == "llm.call"]
        assert "prompt" not in ev.attrs and "completion" not in ev.attrs
        assert ev.attrs["temperature"] == 0.2
        assert fake.seen == []
        assert await rows(sink) == []
        await mw.aclose()

    async def test_embedded_when_captured(self, db):
        sink = db()
        fake = FakeEmbedder()
        emb = EmbeddingSink(fake, sink)
        app, mw = self.app([sink, emb], capture_llm_text=True, max_error_chars=10)
        async with Client(app) as client:
            await client.call_tool("ask", {"q": "why is the sky blue " * 10})
        await settle(mw, emb)
        [ev] = [e for e in mw.recent_events if e.kind == "llm.call"]
        assert len(ev.attrs["prompt"]) == 80  # 8 * max_error_chars
        assert ev.attrs["completion"] == "because"
        [row] = await rows(sink)
        assert row["source_type"] == "llm" and row["source_id"] == ev.id
        assert row["text"].startswith("prompt: why is the sky blue")
        assert row["text"].endswith("completion: because")
        await mw.aclose()

    async def test_captured_text_is_redacted(self, db):
        sink = db()
        emb = EmbeddingSink(FakeEmbedder(), sink)
        app, mw = self.app([sink, emb], capture_llm_text=True)
        async with Client(app) as client:
            await client.call_tool("ask", {"q": "q"})
        await settle(mw, emb)
        [ev] = [e for e in mw.recent_events if e.kind == "llm.call"]
        assert SECRET not in ev.attrs["prompt"] and REDACTED in ev.attrs["prompt"]
        [row] = await rows(sink)
        assert SECRET not in row["text"]
        await mw.aclose()


FIXTURE = [
    ("call_error", "c1", "render: TimeoutError: render timed out after 30 s"),
    ("call_error", "c2", "upload: OSError: disk full on volume data"),
    ("event", "e1", "render job timed out waiting for the gpu"),
    ("feedback", "bug-1", "Render hangs\nthe render tool timed out on scene a"),
    ("feedback", "bug-2", "Dark mode please\nthe ui is too bright at night"),
    ("llm", "l1", "prompt: summarise the log\ncompletion: nothing notable"),
]


async def seed(sink, emb):
    from fastmcp_feedback.instrumentation import EmbeddingText

    for source_type, source_id, text_ in FIXTURE:
        emb._enqueue(EmbeddingText(source_type, source_id, text_))
    await emb.flush()
    assert len(await rows(sink)) == len(FIXTURE)


class TestSimilar:
    async def test_ranking_matches_brute_force(self, db, db_backend):
        sink = db()
        app = FastMCP("s")
        emb = EmbeddingSink(FakeEmbedder(), sink)
        mw = instrument(app, [sink, emb])
        await seed(sink, emb)

        query = "render timed out"
        q = token_vector(query)
        expected = sorted(
            ((cosine_score(q, token_vector(t)), sid) for _, sid, t in FIXTURE),
            reverse=True,
        )
        found = await mw.similar(query, k=len(FIXTURE))
        # Rows sharing no word with the query all score 0; their order is free.
        positive = [sid for score, sid in expected if score > 1e-6]
        assert len(positive) >= 3
        assert [r["source_id"] for r in found][: len(positive)] == positive
        assert {r["source_id"] for r in found} == {sid for _, sid in expected}
        want = {sid: score for score, sid in expected}
        for r in found:
            assert r["score"] == pytest.approx(want[r["source_id"]], abs=1e-4)
            assert 0.0 <= r["score"] <= 1.0
            assert r["created_at"].tzinfo is not None
            assert set(r) == {"source_type", "source_id", "text", "score", "created_at"}
        assert found[0]["source_id"] == "c1"

        top = await mw.similar(query, k=2)
        assert [r["source_id"] for r in top] == positive[:2]
        only = await mw.similar(query, source_types=["feedback"])
        assert [r["source_id"] for r in only] == ["bug-1", "bug-2"]
        strong = await mw.similar(query, min_score=0.4)
        assert strong and all(r["score"] >= 0.4 for r in strong)
        assert len(strong) < len(FIXTURE)
        await mw.aclose()

    async def test_similar_feedback_excludes_itself(self, db):
        sink = db()
        app = FastMCP("s")
        emb = EmbeddingSink(FakeEmbedder(), sink)
        mw = instrument(app, [sink, emb])
        await seed(sink, emb)
        found = await mw.similar_feedback("bug-1", k=3)
        assert "bug-1" not in [r["source_id"] for r in found]
        assert found[0]["source_id"] in ("c1", "e1")
        feedback_only = await mw.similar_feedback("bug-1", source_types=["feedback"])
        assert [r["source_id"] for r in feedback_only] == ["bug-2"]
        assert await mw.similar_feedback("bug-404") == []
        await mw.aclose()

    async def test_scan_limit_without_pgvector(self, db, db_backend, caplog):
        if db_backend == "postgres":
            pytest.skip("brute force only runs without pgvector")
        sink = db()
        app = FastMCP("s")
        emb = EmbeddingSink(FakeEmbedder(), sink, scan_limit=2)
        mw = instrument(app, [sink, emb])
        await seed(sink, emb)
        caplog.set_level(logging.INFO, "fastmcp_feedback.instrumentation.embeddings")
        assert len(await mw.similar("render", k=10)) <= 2
        await mw.similar("render", k=10)
        assert len([r for r in caplog.records if "scan_limit" in r.getMessage()]) == 1
        await mw.aclose()

    async def test_never_raises(self, db):
        app = FastMCP("s")
        mw = instrument(app, [])
        assert await mw.similar("x") == []
        assert await mw.similar_feedback("1") == []
        await mw.aclose()
        sink = db()
        emb = EmbeddingSink(FakeEmbedder(fail=True), sink)
        mw = instrument(FastMCP("t"), [sink, emb])
        assert await mw.similar("x") == []
        # No table yet: the lookup itself fails.
        assert await mw.similar_feedback("1") == []
        await mw.aclose()


class TestPrune:
    async def test_keeps_feedback_embeddings(self, db):
        sink = db(prune_batch=2)
        emb = EmbeddingSink(FakeEmbedder(), sink)
        await seed(sink, emb)
        t = sink.embeddings
        old = datetime.now(UTC) - timedelta(days=40)
        async with sink.engine.begin() as conn:
            await conn.execute(
                update(t).where(t.c.source_id != "c2").values(created_at=old)
            )
        deleted = await sink.prune(older_than=datetime.now(UTC) - timedelta(days=30))
        assert deleted == 3
        left = sorted(r["source_id"] for r in await rows(sink))
        assert left == ["bug-1", "bug-2", "c2"]
        await emb.aclose()

    async def test_prune_without_embeddings_table(self, db):
        sink = db()
        await sink.write([call_record()])
        assert not await sink._has_table(sink.embeddings.name)
        assert await sink.prune(older_than=datetime.now(UTC) + timedelta(days=1)) == 1


class TestCreateTables:
    async def test_embeddings_table_only_with_an_embedding_sink(self, db):
        sink = db()
        await sink.write([call_record()])
        assert not await sink._has_table(sink.embeddings.name)
        emb = EmbeddingSink(FakeEmbedder(), sink)
        await emb.write([call_record()])
        await emb.flush()
        assert await sink._has_table(sink.embeddings.name)
        await emb.aclose()

    async def test_create_tables_false_creates_nothing(self, db):
        sink = db(create_tables=False)
        assert await sink.ensure_embeddings() is True
        assert not await sink._has_table(sink.embeddings.name)

    async def test_failure_is_logged_not_raised(self, db, monkeypatch, caplog):
        sink = db()

        async def boom():
            raise RuntimeError("no room")

        monkeypatch.setattr(sink, "_create_embeddings_table", boom)
        caplog.set_level(logging.WARNING)
        emb = EmbeddingSink(FakeEmbedder(), sink)
        assert await sink.ensure_embeddings() is False
        assert any("Could not create" in r.getMessage() for r in caplog.records)
        await emb.write([call_record()])
        await emb.flush()
        assert emb.errors == 1  # the insert fails; counted, not raised
        await emb.aclose()

    @pytest.mark.postgres
    async def test_postgres_extension_table_and_index(self, db, db_backend):
        if db_backend != "postgres":
            pytest.skip("Postgres only")
        sink = db()
        await sink._ready()
        async with sink.engine.begin() as conn:
            await conn.execute(text("DROP EXTENSION IF EXISTS vector"))
        assert await sink.ensure_embeddings() is True
        async with sink.engine.connect() as conn:
            ext = (
                await conn.execute(
                    text("SELECT count(*) FROM pg_extension WHERE extname = 'vector'")
                )
            ).scalar_one()
            index = (
                await conn.execute(
                    text(
                        "SELECT indexdef FROM pg_indexes "
                        "WHERE indexname = 'ffb_embeddings_hnsw'"
                    )
                )
            ).scalar_one()
            coltype = (
                await conn.execute(
                    text(
                        "SELECT format_type(atttypid, atttypmod) FROM pg_attribute "
                        "WHERE attrelid = 'ffb_embeddings'::regclass "
                        "AND attname = 'embedding'"
                    )
                )
            ).scalar_one()
        assert ext == 1
        assert "hnsw" in index and "vector_cosine_ops" in index
        assert coltype == f"vector({DIM})"

    @pytest.mark.postgres
    async def test_postgres_without_privilege(self, db, db_backend, db_url, caplog):
        if db_backend != "postgres":
            pytest.skip("Postgres only")
        from sqlalchemy.engine import make_url

        admin = db("priv", embedding_dim=DIM)
        role = f"ffb_limited_{uuid.uuid4().hex[:8]}"
        password = uuid.uuid4().hex
        url = make_url(db_url("priv"))
        async with admin.engine.begin() as conn:
            await conn.execute(text("DROP EXTENSION IF EXISTS vector"))
            await conn.execute(text(f"CREATE ROLE {role} LOGIN PASSWORD '{password}'"))
            await conn.execute(text(f"GRANT ALL ON SCHEMA public TO {role}"))
        limited = DatabaseSink(
            url.set(username=role, password=password).render_as_string(
                hide_password=False
            ),
            create_tables=True,
            embedding_dim=DIM,
        )
        try:
            caplog.set_level(logging.WARNING)
            emb = EmbeddingSink(FakeEmbedder(), limited)
            assert await limited.ensure_embeddings() is False
            messages = [r.getMessage() for r in caplog.records]
            assert any("CREATE EXTENSION vector failed" in m for m in messages)
            assert any("Could not create" in m for m in messages)
            await limited.write([call_record()])  # the rest still works
            await emb.write([call_record()])
            await emb.flush()
            assert emb.errors == 1
            await emb.aclose()
        finally:
            await limited.engine.dispose()
            async with admin.engine.begin() as conn:
                await conn.execute(text(f"DROP OWNED BY {role}"))
                await conn.execute(text(f"DROP ROLE {role}"))


class TestMetadata:
    def test_build_metadata_tables(self):
        md = build_metadata("x_", embedding_dim=768)
        t = md.tables["x_embeddings"]
        assert t.c.embedding.type.dim == 768
        assert {c.name for c in t.c} == {
            "id",
            "source_type",
            "source_id",
            "model",
            "dim",
            "text_hash",
            "text",
            "created_at",
            "embedding",
        }
        without = build_metadata("y_", embedding_dim=None)
        assert "y_embeddings" not in without.tables
        assert "y_tool_calls" in without.tables


class TestOpenAIEmbedder:
    def transport(self, handler):
        from fastmcp_feedback.instrumentation.embeddings import _http_module

        return _http_module().MockTransport(handler)

    async def test_request_batches_and_order(self):
        seen = []

        def handler(request):
            body = json.loads(request.content)
            seen.append((request.url.path, request.headers.get("authorization"), body))
            data = [
                {"index": i, "embedding": [float(len(t)), 0.0, 1.0]}
                for i, t in enumerate(body["input"])
            ]
            from fastmcp_feedback.instrumentation.embeddings import _http_module

            return _http_module().Response(200, json={"data": list(reversed(data))})

        e = OpenAIEmbedder(
            "https://gw.example/v1/",
            api_key="k",
            model="m",
            dim=3,
            batch_size=2,
            transport=self.transport(handler),
        )
        out = await e.embed(["a", "bb", "ccc"])
        assert out == [[1.0, 0.0, 1.0], [2.0, 0.0, 1.0], [3.0, 0.0, 1.0]]
        assert [s[0] for s in seen] == ["/v1/embeddings", "/v1/embeddings"]
        assert seen[0][1] == "Bearer k"
        assert [s[2]["input"] for s in seen] == [["a", "bb"], ["ccc"]]
        assert all(s[2]["model"] == "m" for s in seen)
        assert await e.embed([]) == []

    async def test_dim_mismatch_raises(self):
        from fastmcp_feedback.instrumentation.embeddings import _http_module

        def handler(request):
            return _http_module().Response(
                200, json={"data": [{"index": 0, "embedding": [0.1, 0.2]}]}
            )

        e = OpenAIEmbedder(
            "https://gw.example/v1", dim=3, transport=self.transport(handler)
        )
        with pytest.raises(ValueError, match="2-dim vector, expected 3"):
            await e.embed(["a"])

    async def test_http_error_raises(self):
        from fastmcp_feedback.instrumentation.embeddings import _http_module

        e = OpenAIEmbedder(
            "https://gw.example/v1",
            dim=3,
            transport=self.transport(
                lambda r: _http_module().Response(403, json={"error": "nope"})
            ),
        )
        with pytest.raises(Exception, match="403"):
            await e.embed(["a"])


LIVE_URL = os.environ.get("FFB_TEST_EMBED_URL")


@pytest.mark.skipif(not LIVE_URL, reason="FFB_TEST_EMBED_URL is not set")
async def test_live_gateway():
    """FFB_TEST_EMBED_URL=https://.../v1 FFB_TEST_EMBED_KEY=... [FFB_TEST_EMBED_MODEL]."""
    e = OpenAIEmbedder(
        LIVE_URL,
        api_key=os.environ.get("FFB_TEST_EMBED_KEY"),
        model=os.environ.get("FFB_TEST_EMBED_MODEL", "mxbai-embed-large"),
        dim=int(os.environ.get("FFB_TEST_EMBED_DIM", "1024")),
    )
    a, b, c = await e.embed(
        ["render timed out", "the render took too long and timed out", "dark mode"]
    )
    assert len(a) == e.dim
    assert cosine_score(a, b) > cosine_score(a, c)